use device

src/video_processor/processor.py

@@ -18,25 +18,17 @@ class VideoAnalyzer:
         self.model_path = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
         logger.info(f"Loading model from {self.model_path}")
         
-        cache_dir = "/models"
-        logger.info(f"Using cache directory: {cache_dir}")
-        
         # Load processor and model
         self.processor = AutoProcessor.from_pretrained(
             self.model_path, 
-            cache_dir=cache_dir,
             torch_dtype=torch.bfloat16
         )
         
-        # Load model directly to CUDA
-        device_map = {"": 0}  # Force model to GPU 0
         self.model = AutoModelForImageTextToText.from_pretrained(
             self.model_path,
             torch_dtype=torch.bfloat16,
-            device_map=device_map,
-            _attn_implementation="flash_attention_2",
-            cache_dir=cache_dir
-        )
+            _attn_implementation="flash_attention_2"
+        ).to(DEVICE)
         logger.info(f"Model loaded on device: {self.model.device}")
         
     def process_video(self, video_path: str, frame_interval: int = 30) -> List[Dict]: