Update main.py

main.py

@@ -37,23 +37,33 @@ except:
     yolo_model = torch.load("weights/icon_detect/best.pt", map_location="cpu", weights_only=False)["model"]
 
 from transformers import AutoProcessor, AutoModelForCausalLM
+import torch
 
-processor = AutoProcessor.from_pretrained(
-    "microsoft/Florence-2-base", trust_remote_code=True
-)
+# Check if CUDA is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32  # Use float32 on CPU
 
 try:
     model = AutoModelForCausalLM.from_pretrained(
         "weights/icon_caption_florence",
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-    ).to("cuda")
-except:
+        torch_dtype=dtype,  # Dynamic dtype based on device
+        trust_remote_code=True
+    ).to(device)
+except Exception as e:
+    print(f"Error loading model: {str(e)}")
+    # Fallback to CPU with float32
     model = AutoModelForCausalLM.from_pretrained(
         "weights/icon_caption_florence",
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-    )
+        torch_dtype=torch.float32,
+        trust_remote_code=True
+    ).to("cpu")
+
+# Force config for DaViT vision tower
+if not hasattr(model.config, 'vision_config'):
+    model.config.vision_config = {}
+if 'model_type' not in model.config.vision_config:
+    model.config.vision_config['model_type'] = 'davit'
+
 caption_model_processor = {"processor": processor, "model": model}
 print("finish loading model!!!")