Update app.py

app.py

@@ -80,7 +80,7 @@ class VideoHighlightDetector:
             return_tensors="pt"
         ).to(self.device)
         
-        outputs = self.model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
+        outputs = self._generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
         return self.processor.decode(outputs[0], skip_special_tokens=True).lower().split("assistant: ")[1]
 
     def determine_highlights(self, video_description: str, prompt_num: int = 1) -> str:
@@ -117,7 +117,7 @@ class VideoHighlightDetector:
             return_tensors="pt"
         ).to(self.device)
         
-        outputs = self.model.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
+        outputs = self._generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
         return self.processor.decode(outputs[0], skip_special_tokens=True).split("Assistant: ")[1]
 
     def process_segment(self, video_path: str, highlight_types: str) -> bool:
@@ -147,7 +147,7 @@ class VideoHighlightDetector:
             return_tensors="pt"
         ).to(self.device)
         
-        outputs = self.model.generate(**inputs, max_new_tokens=64, do_sample=False)
+        outputs = self._generate(**inputs, max_new_tokens=64, do_sample=False)
         response = self.processor.decode(outputs[0], skip_special_tokens=True).lower().split("assistant: ")[1]
         print(f"Segment response {response}")
         return "yes" in response
@@ -194,6 +194,10 @@ class VideoHighlightDetector:
         logger.info(f"Running ffmpeg command: {' '.join(cmd)}")
         subprocess.run(cmd, check=True)
 
+    def _generate(self, **kwargs):
+        with autocast(device_type="cuda", dtype=torch.float16):
+            return self.model.generate(**kwargs)
+
 def create_ui(examples_path: str, model_path: str):
     examples_data = load_examples(examples_path)