update

README.md

@@ -5,7 +5,7 @@ colorFrom: green
 colorTo: yellow
 sdk: gradio
 sdk_version: 5.49.1
-app_file: app.py
+app_file: demo/gradio_demo.py
 pinned: false
 license: apache-2.0
 short_description: VLM-FO1-3B-Demo

demo/gradio_demo.py

@@ -12,6 +12,7 @@ from vlm_fo1.mm_utils import (
 )
 from vlm_fo1.task_templates import *
 import torch
+import spaces
 
 
 TASK_TYPES = {
@@ -128,7 +129,7 @@ def extract_bbox_and_original_image(edited_image: dict):
 
     return original_image, bbox_list
 
-
+@spaces.GPU
 def process(image, prompt, threshold):
     image, bbox_list = extract_bbox_and_original_image(image)
     image = image.convert('RGB')
@@ -237,18 +238,6 @@ def launch_demo():
     
     return demo
 
-import subprocess
-import sys
-
-def run_step(description, command):
-    """Prints description, runs command, and exits on failure."""
-    print(f"--- {description} ---")
-    result = subprocess.run(command, shell=True)
-    if result.returncode != 0:
-        print(f"{description} failed, exit.")
-        sys.exit(1)
-    print(f"--- {description} successfully ---")
-
 if __name__ == "__main__":
     model_path = 'omlab/VLM-FO1_Qwen2.5-VL-3B-v01' 
     upn_ckpt_path = "./resources/upn_large.pth"