Update main.py

main.py

@@ -30,7 +30,6 @@ from ultralytics import YOLO
 
 # if not os.path.exists("/data/icon_detect"):
 #     os.makedirs("/data/icon_detect")
-
 try:
     yolo_model = torch.load("weights/icon_detect/best.pt", map_location="cuda", weights_only=False)["model"]
     yolo_model = yolo_model.to("cuda")
@@ -43,13 +42,18 @@ processor = AutoProcessor.from_pretrained(
     "microsoft/Florence-2-base", trust_remote_code=True
 )
 
-
-model = AutoModelForCausalLM.from_pretrained(
-    "weights/icon_caption_florence",
-    torch_dtype=torch.float16,
-    trust_remote_code=True,
-    device_map='auto',
-).to("cuda")
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        "weights/icon_caption_florence",
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+    ).to("cuda")
+except:
+    model = AutoModelForCausalLM.from_pretrained(
+        "weights/icon_caption_florence",
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+    )
 caption_model_processor = {"processor": processor, "model": model}
 print("finish loading model!!!")
 
@@ -66,6 +70,7 @@ def process(
     image_input: Image.Image, box_threshold: float, iou_threshold: float
 ) -> ProcessResponse:
     image_save_path = "imgs/saved_image_demo.png"
+    os.makedirs(os.path.dirname(image_save_path), exist_ok=True)
     image_input.save(image_save_path)
     image = Image.open(image_save_path)
     box_overlay_ratio = image.size[0] / 3200
@@ -121,8 +126,20 @@ async def process_image(
     try:
         contents = await image_file.read()
         image_input = Image.open(io.BytesIO(contents)).convert("RGB")
+        
+        # Add debug logging
+        print(f"Processing image: {image_file.filename}")
+        print(f"Image size: {image_input.size}")
+        
+        response = process(image_input, box_threshold, iou_threshold)
+        
+        # Validate response
+        if not response.image:
+            raise ValueError("Empty image in response")
+            
+        return response
+        
     except Exception as e:
-        raise HTTPException(status_code=400, detail="Invalid image file")
-
-    response = process(image_input, box_threshold, iou_threshold)
-    return response
+        import traceback
+        traceback.print_exc()  # This will show full error in logs
+        raise HTTPException(status_code=500, detail=str(e))