CPU only setup

app.py

@@ -1,6 +1,5 @@
 import os
 import gradio as gr
-import spaces
 from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
 from qwen_vl_utils import process_vision_info
 import torch
@@ -9,47 +8,56 @@ from PIL import Image, ImageDraw
 from io import BytesIO
 import re
 
+# -------- Runtime / device --------
+# Force CPU usage
+device = "cpu"
 
-# ---- HF Spaces: ensure we read the platform port ----
+# Hugging Face Spaces port
+PORT = int(os.getenv("PORT", "7860"))
+
+# -------- Model / Processor --------
+# NOTE: device_map=None + .to(device) keeps everything on CPU
 models = {
     "OS-Copilot/OS-Atlas-Base-7B": Qwen2VLForConditionalGeneration.from_pretrained(
         "OS-Copilot/OS-Atlas-Base-7B",
-        torch_dtype="auto",
-        device_map="auto",
-    ),
+        dtype="auto",          # use 'dtype' (new) rather than deprecated 'torch_dtype'
+        device_map=None
+    ).to(device)
 }
 
 processors = {
     "OS-Copilot/OS-Atlas-Base-7B": AutoProcessor.from_pretrained("OS-Copilot/OS-Atlas-Base-7B")
 }
 
-
+# -------- Helpers --------
 def image_to_base64(image: Image.Image) -> str:
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-
 def draw_bounding_boxes(image: Image.Image, bounding_boxes, outline_color="red", line_width=2):
     draw = ImageDraw.Draw(image)
-    for box in bounding_boxes:
+    for box in bounding_boxes or []:
         xmin, ymin, xmax, ymax = box
         draw.rectangle([xmin, ymin, xmax, ymax], outline=outline_color, width=line_width)
     return image
 
-
 def rescale_bounding_boxes(bounding_boxes, original_width, original_height, scaled_width=1000, scaled_height=1000):
+    if not bounding_boxes:
+        return []
     x_scale = original_width / scaled_width
     y_scale = original_height / scaled_height
-    rescaled_boxes = []
-    for box in bounding_boxes:
-        xmin, ymin, xmax, ymax = box
-        rescaled_boxes.append([xmin * x_scale, ymin * y_scale, xmax * x_scale, ymax * y_scale])
-    return rescaled_boxes
-
+    return [
+        [xmin * x_scale, ymin * y_scale, xmax * x_scale, ymax * y_scale]
+        for (xmin, ymin, xmax, ymax) in bounding_boxes
+    ]
 
-@spaces.GPU
+# -------- Inference --------
 def run_example(image, text_input, model_id="OS-Copilot/OS-Atlas-Base-7B"):
+    # Basic validation so the Space doesn't 500
+    if image is None or (text_input is None or str(text_input).strip() == ""):
+        return "", [], image
+
     model = models[model_id].eval()
     processor = processors[model_id]
 
@@ -64,6 +72,7 @@ def run_example(image, text_input, model_id="OS-Copilot/OS-Atlas-Base-7B"):
         }
     ]
 
+    # Build inputs
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
     inputs = processor(
@@ -72,32 +81,43 @@ def run_example(image, text_input, model_id="OS-Copilot/OS-Atlas-Base-7B"):
         videos=video_inputs,
         padding=True,
         return_tensors="pt",
-    ).to("cuda")
+    )
+
+    # Move tensors to CPU explicitly
+    inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
 
-    generated_ids = model.generate(**inputs, max_new_tokens=128)
-    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
-    output_text = processor.batch_decode(
+    # Generate
+    with torch.no_grad():
+        generated_ids = model.generate(**inputs, max_new_tokens=128)
+
+    # Post-process
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)]
+    output_texts = processor.batch_decode(
         generated_ids_trimmed, skip_special_tokens=False, clean_up_tokenization_spaces=False
     )
-    text = output_text[0]
+    text = output_texts[0] if output_texts else ""
 
-    # ---- simple, defensive parsing so the Space doesn't 500 if pattern not found ----
+    # Parse object_ref and bbox defensively
     object_ref_pattern = r"<\|object_ref_start\|>(.*?)<\|object_ref_end\|>"
     box_pattern = r"<\|box_start\|>(.*?)<\|box_end\|>"
 
     object_match = re.search(object_ref_pattern, text or "")
     box_match = re.search(box_pattern, text or "")
 
-    object_ref = object_match.group(1) if object_match else ""
-    box_content = box_match.group(1) if box_match else ""
+    object_ref = object_match.group(1).strip() if object_match else ""
+    box_content = box_match.group(1).strip() if box_match else ""
 
     boxes = []
     if box_content:
         try:
-            parsed = [tuple(map(int, pair.strip("()").split(","))) for pair in box_content.split("),(")]
-            # expecting two points -> convert to [xmin, ymin, xmax, ymax]
-            if len(parsed) >= 2:
-                boxes = [[parsed[0][0], parsed[0][1], parsed[1][0], parsed[1][1]]]
+            # Expecting "(x1,y1),(x2,y2)" -> convert to [xmin, ymin, xmax, ymax]
+            parts = [p.strip() for p in box_content.split("),(")]
+            parts[0] = parts[0].lstrip("(")
+            parts[-1] = parts[-1].rstrip(")")
+            coords = [tuple(map(int, p.split(","))) for p in parts]
+            if len(coords) >= 2:
+                (x1, y1), (x2, y2) = coords[0], coords[1]
+                boxes = [[x1, y1, x2, y2]]
         except Exception:
             boxes = []
 
@@ -106,7 +126,7 @@ def run_example(image, text_input, model_id="OS-Copilot/OS-Atlas-Base-7B"):
 
     return object_ref, scaled_boxes, annotated
 
-
+# -------- UI --------
 css = """
   #output {
     height: 500px;
@@ -141,8 +161,8 @@ with gr.Blocks(css=css) as demo:
         inputs=[input_img, text_input],
         outputs=[model_output_text, model_output_box, annotated_image],
         fn=run_example,
-        cache_examples=False,
-        label="Try examples",
+        cache_examples=False,   # IMPORTANT: don't run inference at startup
+        label="Try examples"
     )
 
     submit_btn.click(
@@ -151,6 +171,11 @@ with gr.Blocks(css=css) as demo:
         [model_output_text, model_output_box, annotated_image],
     )
 
-# ---- HF Spaces: bind to all interfaces + use provided port; disable API schema to avoid json-schema bug ----
-PORT = int(os.getenv("PORT", "7860"))
-demo.queue(api_open=False).launch(server_name="0.0.0.0", server_port=PORT, show_error=True, debug=True)
+# -------- Launch (Spaces-friendly) --------
+# api_open=False avoids known JSON-schema crashes on some versions
+demo.queue(api_open=False).launch(
+    server_name="0.0.0.0",
+    server_port=PORT,
+    show_error=True,
+    debug=True
+)