add new models and fix sessions

aana_chat_with_video/alembic/env.py

@@ -5,6 +5,9 @@
 
 from aana.configs.settings import settings
 from aana.storage.models.base import BaseEntity
+# Import all models to be included in the migration
+import aana.storage.models  # noqa: F401
+import aana_chat_with_video.storage.models  # noqa: F401
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
@@ -20,10 +23,6 @@
 # from myapp import mymodel
 # target_metadata = mymodel.Base.metadata
 
-# Import all models to be included in the migration
-import aana.storage.models  # noqa: F401
-import aana_chat_with_video.storage.models  # noqa: F401
-
 target_metadata = BaseEntity.metadata
 
 # other values from the config, defined by the needs of env.py,

aana_chat_with_video/alembic/versions/b9860676dd49_set_server_default_for_task_completed_.py

@@ -0,0 +1,47 @@
+"""Set server default for task.completed_at and task.assigned_at to none and add num_retries.
+Revision ID: b9860676dd49
+Revises: 5ad873484aa3
+Create Date: 2024-08-22 07:54:55.921710
+"""
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "b9860676dd49"
+down_revision: str | None = "5ad873484aa3"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    """Upgrade database to this revision from previous."""
+    with op.batch_alter_table("tasks", schema=None) as batch_op:
+        batch_op.alter_column(
+            "completed_at",
+            server_default=None,
+        )
+        batch_op.alter_column(
+            "assigned_at",
+            server_default=None,
+        )
+        batch_op.add_column(
+            sa.Column(
+                "num_retries",
+                sa.Integer(),
+                nullable=False,
+                comment="Number of retries",
+                server_default=sa.text("0"),
+            )
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database from this revision to previous."""
+    with op.batch_alter_table("tasks", schema=None) as batch_op:
+        batch_op.drop_column("num_retries")
+
+    # ### end Alembic commands ###

aana_chat_with_video/alembic/versions/d93a90261ee5_added_extended_video.py

@@ -1,7 +1,7 @@
 """added extended video.
 
 Revision ID: d93a90261ee5
-Revises: 5ad873484aa3
+Revises: b9860676dd49
 Create Date: 2024-07-29 12:41:04.976640
 
 """
@@ -12,7 +12,7 @@
 
 # revision identifiers, used by Alembic.
 revision: str = 'd93a90261ee5'
-down_revision: str | None = '5ad873484aa3'
+down_revision: str | None = 'b9860676dd49'
 branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 

aana_chat_with_video/configs/deployments.py

@@ -1,6 +1,5 @@
 from aana.core.models.sampling import SamplingParams
 from aana.core.models.types import Dtype
-from aana.deployments.hf_blip2_deployment import HFBlip2Config, HFBlip2Deployment
 from aana.deployments.vad_deployment import VadConfig, VadDeployment
 from aana.deployments.vllm_deployment import VLLMConfig, VLLMDeployment
 from aana.deployments.whisper_deployment import (
@@ -18,7 +17,7 @@
             max_ongoing_requests=1000,
             ray_actor_options={"num_gpus": 0.25},
             user_config=WhisperConfig(
-                model_size=WhisperModelSize.MEDIUM,
+                model_size=WhisperModelSize.TURBO,
                 compute_type=WhisperComputeType.FLOAT16,
             ).model_dump(mode="json"),
         ),
@@ -41,15 +40,19 @@
     },
     {
         "name": "captioning_deployment",
-        "instance": HFBlip2Deployment.options(
+        "instance": VLLMDeployment.options(
             num_replicas=1,
-            max_ongoing_requests=1000,
             ray_actor_options={"num_gpus": 0.25},
-            user_config=HFBlip2Config(
-                model="Salesforce/blip2-opt-2.7b",
-                dtype=Dtype.FLOAT16,
-                batch_size=2,
-                num_processing_threads=2,
+            user_config=VLLMConfig(
+                model="Qwen/Qwen2-VL-2B-Instruct",
+                dtype=Dtype.AUTO,
+                gpu_memory_reserved=12000,
+                max_model_len=32768,
+                enforce_eager=True,
+                default_sampling_params=SamplingParams(
+                    temperature=0.0, top_p=1.0, top_k=-1, max_tokens=512
+                ),
+                engine_args={"trust_remote_code": True},
             ).model_dump(mode="json"),
         ),
     },

aana_chat_with_video/configs/settings.py

@@ -5,7 +5,7 @@ class Settings(AanaSettings):
     """A pydantic model for App settings."""
 
     asr_model_name: str = "whisper_medium"
-    captioning_model_name: str = "hf_blip2_opt_2_7b"
+    captioning_model_name: str = "qwen2-vl-2b-instruct"
     max_video_len: int = 60 * 20  # 20 minutes
 
 

aana_chat_with_video/endpoints/delete_video.py

@@ -2,6 +2,7 @@
 
 from aana.api.api_generation import Endpoint
 from aana.core.models.media import MediaId
+from aana.storage.session import get_session
 from aana_chat_with_video.storage.repository.extended_video import (
     ExtendedVideoRepository,
 )
@@ -16,12 +17,8 @@ class DeleteVideoOutput(TypedDict):
 class DeleteVideoEndpoint(Endpoint):
     """Delete video endpoint."""
 
-    async def initialize(self):
-        """Initialize the endpoint."""
-        await super().initialize()
-        self.video_repo = ExtendedVideoRepository(self.session)
-
     async def run(self, media_id: MediaId) -> DeleteVideoOutput:
         """Delete video."""
-        self.video_repo.delete(media_id)
+        with get_session() as session:
+            ExtendedVideoRepository(session).delete(media_id)
         return DeleteVideoOutput(media_id=media_id)

aana_chat_with_video/endpoints/get_video_status.py

@@ -2,6 +2,7 @@
 
 from aana.api.api_generation import Endpoint
 from aana.core.models.media import MediaId
+from aana.storage.session import get_session
 from aana_chat_with_video.core.models.video_status import VideoStatus
 from aana_chat_with_video.storage.repository.extended_video import (
     ExtendedVideoRepository,
@@ -17,12 +18,8 @@ class VideoStatusOutput(TypedDict):
 class GetVideoStatusEndpoint(Endpoint):
     """Get video status endpoint."""
 
-    async def initialize(self):
-        """Initialize the endpoint."""
-        await super().initialize()
-        self.video_repo = ExtendedVideoRepository(self.session)
-
     async def run(self, media_id: MediaId) -> VideoStatusOutput:
         """Load video metadata."""
-        video_status = self.video_repo.get_status(media_id)
+        with get_session() as session:
+            video_status = ExtendedVideoRepository(session).get_status(media_id)
         return VideoStatusOutput(status=video_status)

aana_chat_with_video/endpoints/index_video.py

@@ -1,6 +1,7 @@
 from collections.abc import AsyncGenerator
 from typing import TYPE_CHECKING, Annotated, TypedDict
-
+import asyncio
+from aana.storage.session import get_session
 from pydantic import Field
 
 from aana.api.api_generation import Endpoint
@@ -13,6 +14,8 @@
 from aana.core.models.vad import VadParams
 from aana.core.models.video import VideoInput, VideoMetadata, VideoParams
 from aana.core.models.whisper import BatchedWhisperParams
+from aana.core.models.image_chat import ImageChatDialog
+
 from aana.deployments.aana_deployment_handle import AanaDeploymentHandle
 from aana.exceptions.db import MediaIdAlreadyExistsException
 from aana.exceptions.io import VideoTooLongException
@@ -66,9 +69,6 @@ async def initialize(self):
         self.captioning_handle = await AanaDeploymentHandle.create(
             "captioning_deployment"
         )
-        self.extended_video_repo = ExtendedVideoRepository(self.session)
-        self.transcript_repo = ExtendedVideoTranscriptRepository(self.session)
-        self.caption_repo = ExtendedVideoCaptionRepository(self.session)
 
     async def run(  # noqa: C901
         self,
@@ -79,9 +79,10 @@ async def run(  # noqa: C901
     ) -> AsyncGenerator[IndexVideoOutput, None]:
         """Transcribe video in chunks."""
         media_id = video.media_id
-        if self.extended_video_repo.check_media_exists(media_id):
-            raise MediaIdAlreadyExistsException(table_name="media", media_id=video)
-
+        with get_session() as session:
+            if ExtendedVideoRepository(session).check_media_exists(media_id):
+                raise MediaIdAlreadyExistsException(table_name="media", media_id=video)
+
         video_duration = None
         if video.url is not None:
             video_metadata = get_video_metadata(video.url)
@@ -106,7 +107,9 @@ async def run(  # noqa: C901
                 max_len=settings.max_video_len,
             )
 
-        self.extended_video_repo.save(video=video_obj, duration=video_duration)
+        with get_session() as session:
+            ExtendedVideoRepository(session).save(video=video_obj, duration=video_duration)
+
         yield {
             "media_id": media_id,
             "metadata": VideoMetadata(
@@ -117,9 +120,10 @@ async def run(  # noqa: C901
         }
 
         try:
-            self.extended_video_repo.update_status(
-                media_id, VideoProcessingStatus.RUNNING
-            )
+            with get_session() as session:
+                ExtendedVideoRepository(session).update_status(
+                    media_id, VideoProcessingStatus.RUNNING
+                )
             audio: Audio = extract_audio(video=video_obj)
 
             # TODO: Update once batched whisper PR is merged
@@ -158,43 +162,51 @@ async def run(  # noqa: C901
 
                 timestamps.extend(frames_dict["timestamps"])
                 frame_ids.extend(frames_dict["frame_ids"])
+                chat_prompt = "Describe the content of the following image in a single sentence:"
+                dialogs = [
+                    ImageChatDialog.from_prompt(prompt=chat_prompt, images=[frame]) for frame in frames_dict["frames"]
+                ]
 
-                captioning_output = await self.captioning_handle.generate_batch(
-                    images=frames_dict["frames"]
-                )
-                captions.extend(captioning_output["captions"])
+                # Collect the tasks to run concurrently and wait for them to finish
+                tasks = [self.captioning_handle.chat(dialog) for dialog in dialogs]
+                captioning_output = await asyncio.gather(*tasks)
+                captioning_output = [caption["message"].content for caption in captioning_output]
+                captions.extend(captioning_output)
 
                 yield {
-                    "captions": captioning_output["captions"],
+                    "captions": captioning_output,
                     "timestamps": frames_dict["timestamps"],
                 }
 
-            transcription_entity = self.transcript_repo.save(
-                model_name=settings.asr_model_name,
-                media_id=video_obj.media_id,
-                transcription=transcription,
-                segments=segments,
-                transcription_info=transcription_info,
-            )
+            with get_session() as session:
+                transcription_entity = ExtendedVideoTranscriptRepository(session).save(
+                    model_name=settings.asr_model_name,
+                    media_id=video_obj.media_id,
+                    transcription=transcription,
+                    segments=segments,
+                    transcription_info=transcription_info,
+                )
 
-            caption_entities = self.caption_repo.save_all(
-                model_name=settings.captioning_model_name,
-                media_id=video_obj.media_id,
-                captions=captions,
-                timestamps=timestamps,
-                frame_ids=frame_ids,
-            )
+                caption_entities = ExtendedVideoCaptionRepository(session).save_all(
+                    model_name=settings.captioning_model_name,
+                    media_id=video_obj.media_id,
+                    captions=captions,
+                    timestamps=timestamps,
+                    frame_ids=frame_ids,
+                )
 
-            yield {
-                "transcription_id": transcription_entity.id,
-                "caption_ids": [c.id for c in caption_entities],
-            }
+                yield {
+                    "transcription_id": transcription_entity.id,
+                    "caption_ids": [c.id for c in caption_entities],
+                }
         except BaseException:
-            self.extended_video_repo.update_status(
-                media_id, VideoProcessingStatus.FAILED
-            )
+            with get_session() as session:
+                ExtendedVideoRepository(session).update_status(
+                    media_id, VideoProcessingStatus.FAILED
+                )
             raise
         else:
-            self.extended_video_repo.update_status(
-                media_id, VideoProcessingStatus.COMPLETED
-            )
+            with get_session() as session:
+                ExtendedVideoRepository(session).update_status(
+                    media_id, VideoProcessingStatus.COMPLETED
+                )

aana_chat_with_video/endpoints/load_video_metadata.py

@@ -3,6 +3,7 @@
 from aana.api.api_generation import Endpoint
 from aana.core.models.media import MediaId
 from aana.core.models.video import VideoMetadata
+from aana.storage.session import get_session
 from aana_chat_with_video.storage.repository.extended_video import (
     ExtendedVideoRepository,
 )
@@ -17,12 +18,8 @@ class LoadVideoMetadataOutput(TypedDict):
 class LoadVideoMetadataEndpoint(Endpoint):
     """Load video metadata endpoint."""
 
-    async def initialize(self):
-        """Initialize the endpoint."""
-        await super().initialize()
-        self.video_repo = ExtendedVideoRepository(self.session)
-
     async def run(self, media_id: MediaId) -> LoadVideoMetadataOutput:
         """Load video metadata."""
-        video_metadata = self.video_repo.get_metadata(media_id)
+        with get_session() as session:
+            video_metadata = ExtendedVideoRepository(session).get_metadata(media_id)
         return LoadVideoMetadataOutput(metadata=video_metadata)