Upload 43 files

src/vision/florence.py

@@ -21,16 +21,14 @@ except ImportError:
 
 def get_device() -> str:
     """Select best available device."""
-    if IS_SPACES:
-        return "cuda"
     if torch.cuda.is_available():
         return "cuda"
     elif torch.backends.mps.is_available():
         return "mps"
     return "cpu"
 
-DEVICE: str = get_device()
-DTYPE: torch.dtype = torch.float16 if DEVICE == "cuda" else torch.float32
+DEVICE: str = "cpu" if IS_SPACES else get_device()
+DTYPE: torch.dtype = torch.float16 if IS_SPACES or DEVICE == "cuda" else torch.float32
 
 class FlorenceVisionEngine(VisionEngine):
     def __init__(self):
@@ -44,7 +42,7 @@ class FlorenceVisionEngine(VisionEngine):
             return
             
         try:
-            print(f"Loading Florence-2 on {DEVICE.upper()}...")
+            print(f"Loading Florence-2 on {DEVICE.upper()} (will move to GPU during inference if on Spaces)...")
             
             # Hotfix for Florence-2 in newer transformers versions
             import transformers
@@ -60,7 +58,13 @@ class FlorenceVisionEngine(VisionEngine):
                 trust_remote_code=True,
                 torch_dtype=DTYPE,
                 attn_implementation="eager"
-            ).to(DEVICE).eval()
+            ).eval()
+
+            # Hotfix for Florence-2 processor tokenizer compatibility
+            if not hasattr(transformers.PreTrainedTokenizerBase, "additional_special_tokens"):
+                transformers.PreTrainedTokenizerBase.additional_special_tokens = property(
+                    lambda self: getattr(self, "_additional_special_tokens", [])
+                )
 
             self.processor = AutoProcessor.from_pretrained(
                 CONFIG.MODEL_NAME,
@@ -99,6 +103,11 @@ class FlorenceVisionEngine(VisionEngine):
         if self.model is None or self.processor is None:
             raise RuntimeError("Model not loaded. Call load() first.")
 
+        # Ensure model is on the right device when inference runs
+        target_device = "cuda" if IS_SPACES else DEVICE
+        if next(self.model.parameters()).device.type != target_device:
+            self.model.to(target_device)
+
         image = preprocess_image(image)
         image = auto_enhance(image)
         max_tokens = CONFIG.MAX_NEW_TOKENS.get(task_token, 64)
@@ -107,7 +116,7 @@ class FlorenceVisionEngine(VisionEngine):
             text=task_token,
             images=image,
             return_tensors="pt",
-        ).to(DEVICE)
+        ).to(target_device)
         
         if "pixel_values" in inputs:
             inputs["pixel_values"] = inputs["pixel_values"].to(DTYPE)