Update app.py

app.py

@@ -35,6 +35,13 @@ model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     MODEL_ID_M, trust_remote_code=True,
     torch_dtype=torch.float16).to(device).eval()
 
+# Load ViGoRL-MCTS-SFT-3b-Spatial
+MODEL_ID_P = "gsarch/ViGoRL-MCTS-SFT-3b-Spatial"
+processor_p = AutoProcessor.from_pretrained(MODEL_ID_P, trust_remote_code=True)
+model_p = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_P, trust_remote_code=True,
+    torch_dtype=torch.float16).to(device).eval()
+
 # Load OCRFlux-3B
 MODEL_ID_X = "ChatDOC/OCRFlux-3B"
 processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)
@@ -105,7 +112,7 @@ def generate_image(model_name: str,
     elif model_name == "OCRFlux-3B":
         processor = processor_x
         model = model_x
-    elif model_name == "owlet-safety-3b-1":
+    elif model_name == "Owlet-Safety-3B-1":
         processor = processor_o
         model = model_o
     elif model_name == "MonkeyOCR-pro-1.2B":
@@ -114,6 +121,9 @@ def generate_image(model_name: str,
     elif model_name == "OpenVLThinker-7B":
         processor = processor_a
         model = model_a
+    elif model_name == "ViGoRL-MCTS-SFT-3B":
+        processor = processor_p
+        model = model_p
     else:
         yield "Invalid model selected.", "Invalid model selected."
         return
@@ -172,7 +182,7 @@ def generate_video(model_name: str,
     elif model_name == "OCRFlux-3B":
         processor = processor_x
         model = model_x
-    elif model_name == "owlet-safety-3b-1":
+    elif model_name == "Owlet-Safety-3B-1":
         processor = processor_o
         model = model_o
     elif model_name == "MonkeyOCR-pro-1.2B":
@@ -181,6 +191,9 @@ def generate_video(model_name: str,
     elif model_name == "OpenVLThinker-7B":
         processor = processor_a
         model = model_a
+    elif model_name == "ViGoRL-MCTS-SFT-3B":
+        processor = processor_p
+        model = model_p
     else:
         yield "Invalid model selected.", "Invalid model selected."
         return
@@ -326,8 +339,9 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 #download_btn = gr.Button("Download Result.md")
 
             model_choice = gr.Radio(choices=[
-                 "Camel-Doc-OCR-080125(v2)", "OCRFlux-3B", "OpenVLThinker-7B",
-                 "owlet-safety-3b-1", "MonkeyOCR-pro-1.2B"
+                 "Camel-Doc-OCR-080125(v2)", "OCRFlux-3B",
+                 "OpenVLThinker-7B", "ViGoRL-MCTS-SFT-3B",
+                 "Owlet-Safety-3B-1", "MonkeyOCR-pro-1.2B"
             ],
                                     label="Select Model",
                                     value="Camel-Doc-OCR-080125(v2)")