Dockerfile

Dockerfile

@@ -0,0 +1,34 @@
+FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04
+
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+ENV FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    python3.10 \
+    python3-pip \
+    python3.10-dev \
+    build-essential \
+    ninja-build \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create and set working directory
+WORKDIR /app
+
+# Copy requirements first to leverage Docker cache
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Install flash-attention
+RUN pip3 install --no-cache-dir flash-attn --no-build-isolation
+
+# Copy the rest of the application
+COPY . .
+
+# Set the default command
+CMD ["python3", "-m", "src.video_processor.processor"] 

src/video_processor/processor.py

@@ -1,10 +1,7 @@
 import torch
 from transformers import AutoProcessor, AutoModelForImageTextToText
 from typing import List, Dict
-import numpy as np
 import logging
-import subprocess
-import os
 
 logger = logging.getLogger(__name__)
 
@@ -17,25 +14,6 @@ class VideoAnalyzer:
             raise RuntimeError("CUDA is required but not available!")
             
         logger.info("Initializing VideoAnalyzer")
-        
-        # Try to install flash-attention at runtime
-        logger.info("Attempting to install flash-attention...")
-        try:
-            env = os.environ.copy()
-            env['FLASH_ATTENTION_SKIP_CUDA_BUILD'] = "TRUE"
-            subprocess.run(
-                'pip install flash-attn --no-build-isolation',
-                env=env,
-                shell=True,
-                check=True,
-                capture_output=True
-            )
-            logger.info("Successfully installed flash-attention")
-            use_flash_attn = True
-        except subprocess.CalledProcessError as e:
-            logger.warning(f"Failed to install flash-attention: {e.stderr.decode()}")
-            use_flash_attn = False
-            
         self.model_path = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
         logger.info(f"Loading model from {self.model_path}")
         
@@ -44,13 +22,14 @@ class VideoAnalyzer:
             self.model_path, 
             torch_dtype=torch.bfloat16
         )
+
         
         self.model = AutoModelForImageTextToText.from_pretrained(
             self.model_path,
             torch_dtype=torch.bfloat16,
-            _attn_implementation="flash_attention_2" if use_flash_attn else "sdpa"
+            _attn_implementation="flash_attention_2"
         ).to(DEVICE)
-        logger.info(f"Model loaded on device: {self.model.device} using attention implementation: {'flash_attention_2' if use_flash_attn else 'sdpa'}")
+        logger.info(f"Model loaded on device: {self.model.device} using attention implementation: flash_attention_2")
         
     def process_video(self, video_path: str, frame_interval: int = 30) -> List[Dict]:
         logger.info(f"Processing video: {video_path} with frame_interval={frame_interval}")