Update main.py

main.py

@@ -32,9 +32,9 @@ from ultralytics import YOLO
 #     os.makedirs("/data/icon_detect")
 
 try:
-    yolo_model = YOLO("microsoft/OmniParser/icon_detect/model_v1_5.pt").to("cuda")
+    yolo_model = YOLO("best.pt").to("cuda")
 except:
-    yolo_model = YOLO("microsoft/OmniParser/icon_detect/model_v1_5.pt")
+    yolo_model = YOLO("best.pt")
 
 from transformers import AutoProcessor, AutoModelForCausalLM
 
@@ -44,7 +44,7 @@ processor = AutoProcessor.from_pretrained(
 
 try:
     model = AutoModelForCausalLM.from_pretrained(
-        "microsoft/OmniParser/icon_caption_florence",
+        "microsoft/OmniParser",
         torch_dtype=torch.float16,
         trust_remote_code=True,
     ).to("cuda")