Implementation of batched whisper and updates on Audio pipeline

.vscode/settings.json

@@ -5,7 +5,6 @@
         "editor.formatOnSave": true,
     },
     "python.testing.pytestArgs": [
-        // "--import-mode=importlib",
         "aana"
     ],
     "python.testing.unittestEnabled": false,

aana/configs/deployments.py

@@ -3,6 +3,7 @@
     StableDiffusion2Config,
     StableDiffusion2Deployment,
 )
+from aana.deployments.vad_deployment import VadConfig, VadDeployment
 from aana.deployments.vllm_deployment import VLLMConfig, VLLMDeployment
 from aana.deployments.whisper_deployment import (
     WhisperComputeType,
@@ -59,4 +60,17 @@
             dtype=Dtype.FLOAT16,
         ).model_dump(),
     ),
+    "vad_deployment": VadDeployment.options(
+        num_replicas=1,
+        max_concurrent_queries=1000,
+        ray_actor_options={"num_gpus": 0.05},
+        user_config=VadConfig(
+            model=(
+                "https://whisperx.s3.eu-west-2.amazonaws.com/model_weights/segmentation/"
+                "0b5b3216d60a2d32fc086b47ea8c67589aaeb26b7e07fcbe620d6d0b83e209ea/pytorch_model.bin"
+            ),
+            onset=0.5,
+            sample_rate=16000,
+        ).model_dump(),
+    ),
 }

aana/configs/endpoints.py

@@ -113,6 +113,32 @@
                 ),
             ],
         ),
+        Endpoint(
+            name="whisper_transcribe_in_chunks",
+            path="/video/transcribe_in_chunks",
+            summary="Transcribe a video using Whisper Medium by segmenting it into chunks",
+            outputs=[
+                EndpointOutput(
+                    name="transcription",
+                    output="video_transcriptions_batched_whisper_medium",
+                    streaming=True,
+                ),
+                EndpointOutput(
+                    name="segments",
+                    output="video_transcriptions_segments_batched_whisper_medium",
+                    streaming=True,
+                ),
+                EndpointOutput(
+                    name="info",
+                    output="video_transcriptions_info_batched_whisper_medium",
+                    streaming=True,
+                ),
+                EndpointOutput(
+                    name="transcription_id", output="transcription_id_batched"
+                ),
+            ],
+            streaming=True,
+        ),
         Endpoint(
             name="delete_media_id",
             path="/video/delete",
@@ -130,17 +156,17 @@
             outputs=[
                 EndpointOutput(
                     name="transcription",
-                    output="video_transcriptions_whisper_medium",
+                    output="video_transcriptions_batched_whisper_medium",
                     streaming=True,
                 ),
                 EndpointOutput(
                     name="segments",
-                    output="video_transcriptions_segments_whisper_medium",
+                    output="video_transcriptions_segments_batched_whisper_medium",
                     streaming=True,
                 ),
                 EndpointOutput(
                     name="info",
-                    output="video_transcriptions_info_whisper_medium",
+                    output="video_transcriptions_info_batched_whisper_medium",
                     streaming=True,
                 ),
                 EndpointOutput(
@@ -152,7 +178,9 @@
                     name="timestamps", output="video_timestamps", streaming=True
                 ),
                 EndpointOutput(name="caption_ids", output="caption_ids"),
-                EndpointOutput(name="transcription_id", output="transcription_id"),
+                EndpointOutput(
+                    name="transcription_id", output="transcription_id_batched"
+                ),
             ],
             streaming=True,
         ),

aana/configs/pipeline.py

@@ -19,6 +19,8 @@
 from aana.models.pydantic.prompt import Prompt
 from aana.models.pydantic.question import Question
 from aana.models.pydantic.sampling_params import SamplingParams
+from aana.models.pydantic.vad_output import VadSegments
+from aana.models.pydantic.vad_params import VadParams
 from aana.models.pydantic.video_input import VideoInput, VideoInputList
 from aana.models.pydantic.video_metadata import VideoMetadata
 from aana.models.pydantic.video_params import VideoParams
@@ -261,6 +263,28 @@
             },
         ],
     },
+    {
+        "name": "extract_audios",
+        "type": "ray_task",
+        "function": "aana.utils.video.extract_audio",
+        "batched": True,
+        "flatten_by": "video_batch.videos.[*]",
+        "dict_output": False,
+        "inputs": [
+            {
+                "name": "video_objects",
+                "key": "video",
+                "path": "video_batch.videos.[*].video",
+            },
+        ],
+        "outputs": [
+            {
+                "name": "audio_objects",
+                "key": "output",
+                "path": "video_batch.videos.[*].audio",
+            },
+        ],
+    },
     {
         "name": "video_params",
         "type": "input",
@@ -353,9 +377,9 @@
         "method": "transcribe_batch",
         "inputs": [
             {
-                "name": "video_objects",
-                "key": "media_batch",
-                "path": "video_batch.videos.[*].video",
+                "name": "audio_objects",
+                "key": "audio_batch",
+                "path": "video_batch.videos.[*].audio",
             },
             {
                 "name": "whisper_params",
@@ -418,6 +442,26 @@
             },
         ],
     },
+    {
+        "name": "extract_audio",
+        "type": "ray_task",
+        "function": "aana.utils.video.extract_audio",
+        "dict_output": False,
+        "inputs": [
+            {
+                "name": "video_object",
+                "key": "video",
+                "path": "video.video",
+            },
+        ],
+        "outputs": [
+            {
+                "name": "audio_object",
+                "key": "output",
+                "path": "video.audio",
+            },
+        ],
+    },
     {
         "name": "generate_frames_for_video",
         "type": "ray_task",
@@ -487,9 +531,9 @@
         "method": "transcribe_stream",
         "inputs": [
             {
-                "name": "video_object",
-                "key": "media",
-                "path": "video.video",
+                "name": "audio_object",
+                "key": "audio",
+                "path": "video.audio",
             },
             {
                 "name": "whisper_params",
@@ -519,6 +563,91 @@
             },
         ],
     },
+    {
+        "name": "vad_params",
+        "type": "input",
+        "inputs": [],
+        "outputs": [
+            {
+                "name": "vad_params",
+                "key": "vad_params",
+                "path": "video.vad_params",
+                "data_model": VadParams,
+            }
+        ],
+    },
+    {
+        "name": "vad_transcribe_in_chunks_audio",
+        "type": "ray_deployment",
+        "deployment_name": "vad_deployment",
+        "method": "asr_preprocess_vad",
+        "inputs": [
+            {
+                "name": "audio_object",
+                "key": "audio",
+                "path": "video.audio",
+            },
+            {
+                "name": "vad_params",
+                "key": "params",
+                "path": "video.vad_params",
+            },
+        ],
+        "outputs": [
+            {
+                "name": "video_transcriptions_vad_segments",
+                "key": "segments",
+                "path": "video.vad_segments",
+                "data_model": VadSegments,
+            },
+        ],
+    },
+    {
+        "name": "whisper_medium_transcribe_in_chunks_video",
+        "type": "ray_deployment",
+        "deployment_name": "whisper_deployment_medium",
+        "data_type": "generator",
+        "generator_path": "video",
+        "method": "transcribe_in_chunks",
+        "inputs": [
+            {
+                "name": "audio_object",
+                "key": "audio",
+                "path": "video.audio",
+            },
+            {
+                "name": "video_transcriptions_vad_segments",
+                "key": "segments",
+                "path": "video.vad_segments",
+            },
+            {
+                "name": "whisper_params",
+                "key": "params",
+                "path": "video_batch.whisper_params",
+                "data_model": WhisperParams,
+            },
+        ],
+        "outputs": [
+            {
+                "name": "video_transcriptions_segments_batched_whisper_medium",
+                "key": "segments",
+                "path": "video.segments_batched",
+                "data_model": AsrSegments,
+            },
+            {
+                "name": "video_transcriptions_info_batched_whisper_medium",
+                "key": "transcription_info",
+                "path": "video.transcription_info_batched",
+                "data_model": AsrTranscriptionInfo,
+            },
+            {
+                "name": "video_transcriptions_batched_whisper_medium",
+                "key": "transcription",
+                "path": "video.transcription_batched",
+                "data_model": AsrTranscription,
+            },
+        ],
+    },
     {
         "name": "media_id",
         "type": "input",
@@ -865,6 +994,44 @@
             }
         ],
     },
+    {
+        "name": "save_video_transcription_batched",
+        "type": "function",
+        "function": "aana.utils.db.save_video_transcription",
+        "kwargs": {
+            "model_name": "whisper_medium",
+        },
+        "dict_output": True,
+        "inputs": [
+            {
+                "name": "video_media_id",
+                "key": "media_id",
+                "path": "video.media_id",
+            },
+            {
+                "name": "video_transcriptions_info_batched_whisper_medium",
+                "key": "transcription_info",
+                "path": "video.transcription_info_batched",
+            },
+            {
+                "name": "video_transcriptions_segments_batched_whisper_medium",
+                "key": "segments",
+                "path": "video.segments_batched",
+            },
+            {
+                "name": "video_transcriptions_batched_whisper_medium",
+                "key": "transcription",
+                "path": "video.transcription_batched",
+            },
+        ],
+        "outputs": [
+            {
+                "name": "transcription_id_batched",
+                "key": "transcription_id",
+                "path": "video.transcription_id_batched",
+            }
+        ],
+    },
     {
         "name": "save_transcripts_batch_medium",
         "type": "function",

aana/configs/settings.py

@@ -19,6 +19,8 @@ class Settings(BaseSettings):
     tmp_data_dir: Path = Path("/tmp/aana_data")  # noqa: S108
     image_dir: Path = tmp_data_dir / "images"
     video_dir: Path = tmp_data_dir / "videos"
+    audio_dir: Path = tmp_data_dir / "audios"
+    model_dir: Path = tmp_data_dir / "models"
     num_workers: int = 2
 
     db_config: DBConfig = {