Update main.py

main.py

@@ -32,7 +32,7 @@ from ultralytics import YOLO
 #     os.makedirs("/data/icon_detect")
 
 try:
-    yolo_model = YOLO("weights/icon_detect/best.pt").to("cuda")
+    yolo_model = YOLO("weights/icon_detect/best.pt").to("cpu")
 except:
     yolo_model = YOLO("weights/icon_detect/best.pt")
 
@@ -47,7 +47,7 @@ try:
         "banao-tech/OmniParse",
         torch_dtype=torch.float16,
         trust_remote_code=True,
-    ).to("cuda")
+    ).to("cpu")
 except:
     model = AutoModelForCausalLM.from_pretrained(
         "banao-tech/OmniParse",