Update main.py

main.py

@@ -10,7 +10,7 @@ import torch
 import numpy as np
 from PIL import Image
 from ultralytics import YOLO
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoProcessor, AutoModelForCausalLM,Blip2ForConditionalGeneration
 
 # Type definitions
 class ProcessResponse(BaseModel):
@@ -40,10 +40,9 @@ class ModelManager:
                 trust_remote_code=True
             )
             
-            self.model = AutoModelForCausalLM.from_pretrained(
-                "banao-tech/OmniParse",
-                torch_dtype=torch.float16,
-                trust_remote_code=True
+            self.model = Blip2ForConditionalGeneration.from_pretrained("banao-tech/OmniParser",torch_dtype=torch.float16,
+                trust_remote_code=True)
+                
             ).to(self.device)
             
             return True