Upload 2 files

vision_tools.py

@@ -39,7 +39,7 @@ from scenedetect.detectors import ContentDetector
 
 import os, base64, requests, subprocess, contextlib, time
 
-from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration
+from transformers import AutoProcessor, LlavaOneForConditionalGeneration
 from PIL import Image
 
 from libs.audio_tools_ana_2 import process_audio_for_video
@@ -261,11 +261,13 @@ def describe_montage_sequence(
     path_model = "BSC-LT/salamandra-7b-vision"
 
     processor = AutoProcessor.from_pretrained(path_model)
-    model = LlavaOnevisionForConditionalGeneration.from_pretrained(
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float16 if device == "cuda" else torch.float32
+    model = LlavaOneForConditionalGeneration.from_pretrained(
         path_model,
-        torch_dtype=torch.float16,
+        torch_dtype=dtype,
         low_cpu_mem_usage=True
-    ).to("cuda")
+    ).to(device)
 
     img = cv2.imread(montage_path, cv2.IMREAD_COLOR)
     if img is None:
@@ -305,9 +307,9 @@ def describe_montage_sequence(
         inputs = processor(images=batch, text=prompt_batch, return_tensors="pt")
         for k, v in inputs.items():
             if v.dtype.is_floating_point:
-                inputs[k] = v.to("cuda", torch.float16)
+                inputs[k] = v.to(device, dtype)
             else:
-                inputs[k] = v.to("cuda")
+                inputs[k] = v.to(device)
         
         output = model.generate(**inputs, max_new_tokens=1024)
         text = processor.decode(output[0], skip_special_tokens=True)