incl. — a GPU Duration slider that dynamically allocates ZeroGPU time.

app.py

@@ -5,45 +5,46 @@ import json
 import ast
 import re
 from io import BytesIO
+from typing import Iterable
+from threading import Thread
+
+import numpy as np
 import torch
 import spaces
-import numpy as np
 from PIL import Image, ImageDraw, ImageFont
 import supervision as sv
-from typing import Iterable
-import gradio as gr
-from gradio import Server
-from fastapi.responses import HTMLResponse
-from threading import Thread
 from transformers import (
     Qwen3_5ForConditionalGeneration,
     AutoProcessor,
     TextIteratorStreamer,
 )
 
-# ------------------------------------------------------------------
-# Config & Constants
-# ------------------------------------------------------------------
+import gradio as gr
+from gradio import Server
+from fastapi.responses import HTMLResponse
+
+# --- Config ---
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
 MODEL_NAME = "Qwen/Qwen3.8-27B"
 
+GPU_DURATIONS = [60, 90, 120, 150, 180, 250, 300]
+DEFAULT_GPU_DURATION_IDX = 1
+
+# --- Colors for Annotations ---
 BRIGHT_YELLOW = sv.Color(r=255, g=230, b=0)
 DARK_OUTLINE = sv.Color(r=40, g=40, b=40)
 BLACK = sv.Color(r=0, g=0, b=0)
 WHITE = sv.Color(r=255, g=255, b=255)
 
-# Spatial path colors
 SPATIAL_LINE = (255, 69, 0)       # OrangeRed
 SPATIAL_DOT = (255, 69, 0)
 SPATIAL_RING = (255, 255, 255)
-SPATIAL_LABEL_BG = (80, 20, 0)
+SPATIAL_LABEL_BG = (80, 25, 0)
 SPATIAL_LABEL_TXT = (255, 255, 255)
-SPATIAL_ARROW = (255, 165, 0)     # Orange for direction
+SPATIAL_ARROW = (230, 149, 0)
 
-# ------------------------------------------------------------------
-# Model Loading
-# ------------------------------------------------------------------
+# --- Model Loading ---
 print(f"Loading model: {MODEL_NAME} ...")
 qwen_model = Qwen3_5ForConditionalGeneration.from_pretrained(
     MODEL_NAME, torch_dtype=DTYPE, device_map=DEVICE, attn_implementation="kernels-community/flash-attn2@v3",
@@ -51,78 +52,22 @@ qwen_model = Qwen3_5ForConditionalGeneration.from_pretrained(
 qwen_processor = AutoProcessor.from_pretrained(MODEL_NAME)
 print("Model loaded.")
 
-# ------------------------------------------------------------------
-# Examples Config
-# ------------------------------------------------------------------
-EXAMPLES_CONFIG = [
-    {"image": "examples/1.jpg", "prompt": "Detect the yellow car that is parked.", "mode": "Detect"},
-    {"image": "examples/2.jpg", "prompt": "Point to all the red cars.", "mode": "Point"},
-    {"image": "examples/3.jpg", "prompt": "Map a path from the door to the lamp.", "mode": "Spatial"},
-]
-
-def make_thumb_b64(path, max_dim=220):
-    if not os.path.exists(path):
-        return ""
-    try:
-        img = Image.open(path).convert("RGB")
-        img.thumbnail((max_dim, max_dim), Image.LANCZOS)
-        buf = BytesIO()
-        img.save(buf, format="JPEG", quality=65)
-        return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
-    except Exception as e:
-        return ""
-
-def encode_full_image(path):
-    if not os.path.exists(path):
-        return ""
-    try:
-        with open(path, "rb") as f:
-            data = f.read()
-        ext = path.rsplit(".", 1)[-1].lower()
-        mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
-        return f"data:{mime};base64,{base64.b64encode(data).decode()}"
-    except Exception as e:
-        return ""
-
-def build_client_config():
-    examples = []
-    for i, ex in enumerate(EXAMPLES_CONFIG):
-        examples.append({
-            "idx": i,
-            "thumb": make_thumb_b64(ex["image"]),
-            "prompt": ex["prompt"],
-            "mode": ex["mode"],
-        })
-    return {"examples": examples, "modes": ["Detect", "Point", "Spatial"], "default_mode": "Detect"}
-
-print("Building client config…")
-CLIENT_CONFIG = build_client_config()
-
-# ------------------------------------------------------------------
-# Helpers
-# ------------------------------------------------------------------
+# --- JSON Parsing & Extraction ---
 def safe_parse_json(text: str):
     text = re.sub(r"```(json)?", "", text).strip()
     match = re.search(r'(\[.*\]|\{.*\})', text, re.DOTALL)
     if match:
         json_str = match.group(1)
         json_str_clean = re.sub(r',\s*([}\]])', r'\1', json_str)
-        try:
-            return json.loads(json_str_clean)
+        try: return json.loads(json_str_clean)
         except json.JSONDecodeError:
-            try:
-                return ast.literal_eval(json_str_clean)
-            except Exception:
-                pass
+            try: return ast.literal_eval(json_str_clean)
+            except Exception: pass
     text_clean = re.sub(r',\s*([}\]])', r'\1', text)
-    try:
-        return json.loads(text_clean)
-    except json.JSONDecodeError:
-        pass
-    try:
-        return ast.literal_eval(text_clean)
-    except Exception:
-        pass
+    try: return json.loads(text_clean)
+    except json.JSONDecodeError: pass
+    try: return ast.literal_eval(text_clean)
+    except Exception: pass
     return []
 
 def _extract_point(item: dict):
@@ -141,23 +86,17 @@ def _extract_bbox(item: dict):
 
 def _load_font(size: int = 16):
     size = max(6, int(size))
-    try:
-        return ImageFont.truetype("arial.ttf", size)
+    try: return ImageFont.truetype("arial.ttf", size)
     except (IOError, OSError):
-        try:
-            return ImageFont.truetype("DejaVuSans.ttf", size)
-        except (IOError, OSError):
-            return ImageFont.load_default()
-
-def pil_to_b64_png(image: Image.Image) -> str:
-    buf = BytesIO()
-    image.save(buf, format="PNG")
-    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
+        try: return ImageFont.truetype("DejaVuSans.ttf", size)
+        except (IOError, OSError): return ImageFont.load_default()
 
+# --- Annotation Functions ---
 def annotate_image(image: Image.Image, result: dict, point_radius: int = 6, box_thickness: int = 2, text_scale: float = 0.5):
     if not isinstance(image, Image.Image) or not isinstance(result, dict): return image
     image = image.convert("RGB")
     ow, oh = image.size
+
     point_radius = max(1, int(point_radius))
     box_thickness = max(1, int(box_thickness))
     text_scale = max(0.1, float(text_scale))
@@ -260,8 +199,10 @@ def annotate_spatial_path(image: Image.Image, result: dict, dot_radius: int = 6,
         halo_r = dot_radius + 8
         ring_r = dot_radius + 3
         draw.ellipse((cx - halo_r, cy - halo_r, cx + halo_r, cy + halo_r), fill=SPATIAL_LINE + (50,))
-        draw.ellipse((cx - ring_r, cy - ring_r, cx + ring_r, cy + ring_r), outline=SPATIAL_RING, width=max(1, round(3 * scale_ratio)))
-        draw.ellipse((cx - dot_radius, cy - dot_radius, cx + dot_radius, cy + dot_radius), fill=SPATIAL_DOT, outline=SPATIAL_DOT)
+        draw.ellipse((cx - ring_r, cy - ring_r, cx + ring_r, cy + ring_r),
+                     outline=SPATIAL_RING, width=max(1, round(3 * scale_ratio)))
+        draw.ellipse((cx - dot_radius, cy - dot_radius, cx + dot_radius, cy + dot_radius),
+                     fill=SPATIAL_DOT, outline=SPATIAL_DOT)
         num_text = str(i + 1)
         nbbox = draw.textbbox((0, 0), num_text, font=font_num)
         nw = nbbox[2] - nbbox[0]
@@ -275,7 +216,11 @@ def annotate_spatial_path(image: Image.Image, result: dict, dot_radius: int = 6,
         tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
         lx, ly = cx + dot_radius + 10, cy - th - 8
         pad = 5
-        draw.rectangle((lx - pad, ly - pad, lx + tw + pad, ly + th + pad), fill=SPATIAL_LABEL_BG, outline=SPATIAL_LINE, width=1)
+        draw.rectangle(
+            (lx - pad, ly - pad, lx + tw + pad, ly + th + pad),
+            fill=SPATIAL_LABEL_BG,
+            outline=SPATIAL_LINE, width=1,
+        )
         draw.text((lx, ly), label, fill=SPATIAL_LABEL_TXT, font=font_label)
 
     n_pts = len(pts)
@@ -284,35 +229,86 @@ def annotate_spatial_path(image: Image.Image, result: dict, dot_radius: int = 6,
     bbox = draw.textbbox((0, 0), legend_text, font=legend_font)
     tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
     fx, fy = 10, h - th - 22
-    draw.rectangle((fx - 8, fy - 6, fx + tw + 16, fy + th + 10), fill=SPATIAL_LABEL_BG + (220,))
+    draw.rectangle((fx - 8, fy - 6, fx + tw + 16, fy + th + 10),
+                   fill=SPATIAL_LABEL_BG + (220,))
     draw.text((fx, fy), legend_text, fill=SPATIAL_LABEL_TXT, font=legend_font)
 
     return image
 
-# ------------------------------------------------------------------
-# Gradio Server (Server mode): FastAPI + Gradio queue/API engine
-# ------------------------------------------------------------------
+def pil_to_b64_png(image: Image.Image) -> str:
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
+
+def get_gpu_duration(image_b64, prompt, task, point_radius, box_thickness, text_scale, gpu_duration_seconds):
+    """Dynamically returns GPU duration based on the UI slider value."""
+    try:
+        return int(gpu_duration_seconds)
+    except (TypeError, ValueError):
+        return GPU_DURATIONS[DEFAULT_GPU_DURATION_IDX]
+
+# --- Config Examples ---
+EXAMPLES_CONFIG = [
+    {"image": "examples/1.jpg", "prompt": "Detect the yellow car that is parked.", "task": "Detect"},
+    {"image": "examples/2.jpg", "prompt": "Point to all the red cars.", "task": "Point"},
+    {"image": "examples/3.jpg", "prompt": "Map a path from the door to the lamp.", "task": "Spatial"},
+]
+
+def make_thumb_b64(path, max_dim=220):
+    if not os.path.exists(path): return ""
+    try:
+        img = Image.open(path).convert("RGB")
+        img.thumbnail((max_dim, max_dim), Image.LANCZOS)
+        buf = BytesIO()
+        img.save(buf, format="JPEG", quality=65)
+        return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
+    except Exception as e:
+        return ""
+
+def encode_full_image(path):
+    if not os.path.exists(path): return ""
+    try:
+        with open(path, "rb") as f: data = f.read()
+        ext = path.rsplit(".", 1)[-1].lower()
+        mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
+        return f"data:{mime};base64,{base64.b64encode(data).decode()}"
+    except Exception as e:
+        return ""
+
+def build_client_config():
+    examples = []
+    for i, ex in enumerate(EXAMPLES_CONFIG):
+        examples.append({
+            "idx": i,
+            "thumb": make_thumb_b64(ex["image"]),
+            "prompt": ex["prompt"],
+            "task": ex["task"],
+        })
+    return {"examples": examples}
+
+CLIENT_CONFIG = build_client_config()
+
+# --- Gradio Server ---
 app = Server(title="Qwen3.8-27B-Object-Detection")
 
 @app.mcp.tool(name="run_inference")
 @app.api(name="run_inference")
-@spaces.GPU(size="xlarge", duration=90)
+@spaces.GPU(size="xlarge", duration=get_gpu_duration)
 def infer(
     image_b64: str,
-    mode: str,
     prompt: str,
+    task: str,
     point_radius: int,
     box_thickness: int,
     text_scale: float,
+    gpu_duration_seconds: int,
 ) -> dict:
-    """Runs object detection, point localization, or spatial mapping."""
+    """Runs Qwen3.8 vision model for detection, point localization, or spatial mapping."""
     gc.collect()
     torch.cuda.empty_cache()
 
-    if not image_b64:
-        raise gr.Error("Please upload an image.")
-    if not prompt or prompt.strip() == "":
-        raise gr.Error("Please provide a prompt.")
+    if not image_b64: raise gr.Error("Please upload an image.")
+    if not prompt or not prompt.strip(): raise gr.Error("Please provide a prompt.")
 
     try:
         header, data = image_b64.split(",", 1)
@@ -321,21 +317,20 @@ def infer(
         raise gr.Error(f"Invalid image data: {e}")
 
     pil_image.thumbnail((512, 512))
-    category = mode
 
-    if category == "Detect":
+    if task == "Detect":
         full_prompt = (
             f"Provide bounding box coordinates for {prompt}. "
             f"Report strictly in JSON format as a list of objects with 'label' and "
             f"'bbox_2d' (xmin, ymin, xmax, ymax in 0-1000 scale)."
         )
-    elif category == "Point":
+    elif task == "Point":
         full_prompt = (
             f"Provide 2d point coordinates for {prompt}. "
             f"Report strictly in JSON format as a list of objects with 'label' and "
             f"'point_2d' (x, y in 0-1000 scale)."
         )
-    elif category == "Spatial":
+    elif task == "Spatial":
         full_prompt = (
             f"Identify the key spatial waypoints to map a path/route for: {prompt}. "
             f"Return the points in the order they should be connected along the path, "
@@ -353,6 +348,7 @@ def infer(
             {"type": "text", "text": full_prompt},
         ],
     }]
+    
     text = qwen_processor.apply_chat_template(
         messages, tokenize=False, add_generation_prompt=True
     )
@@ -381,10 +377,8 @@ def infer(
         full_text += tok
     thread.join()
 
-    result_text = full_text
-    result_image = pil_image.copy()
-
-    if category == "Point":
+    # --- Post-process ---
+    if task == "Point":
         parsed = safe_parse_json(full_text)
         if isinstance(parsed, dict):
             for k in ["points", "keypoints", "point"]:
@@ -392,9 +386,7 @@ def infer(
                     parsed = parsed[k]
                     break
             else:
-                for v in parsed.values():
-                    if isinstance(v, list): parsed = v; break
-                else: parsed = []
+                parsed = next((v for v in parsed.values() if isinstance(v, list)), [])
 
         result = {"points": []}
         if isinstance(parsed, list):
@@ -405,12 +397,12 @@ def infer(
                     result["points"].append({"label": item.get("label", ""), "x": x / 1000.0, "y": y / 1000.0})
 
         if result["points"]:
-            result_image = annotate_image(pil_image.copy(), result, point_radius, box_thickness, text_scale)
-            result_text = json.dumps(result, indent=2)
+            annotated_img = annotate_image(pil_image.copy(), result, point_radius=point_radius, box_thickness=box_thickness, text_scale=text_scale)
+            return {"image": pil_to_b64_png(annotated_img), "text": json.dumps(result, indent=2)}
         else:
-            result_text = f"Could not extract any points.\nRaw model output:\n{full_text}"
+            return {"image": pil_to_b64_png(pil_image), "text": f"Could not extract any points.\nRaw model output:\n{full_text}"}
 
-    elif category == "Detect":
+    elif task == "Detect":
         parsed = safe_parse_json(full_text)
         if isinstance(parsed, dict):
             for k in ["objects", "detections", "bboxes", "boxes", "results"]:
@@ -418,9 +410,7 @@ def infer(
                     parsed = parsed[k]
                     break
             else:
-                for v in parsed.values():
-                    if isinstance(v, list): parsed = v; break
-                else: parsed = []
+                parsed = next((v for v in parsed.values() if isinstance(v, list)), [])
 
         result = {"objects": []}
         if isinstance(parsed, list):
@@ -435,12 +425,12 @@ def infer(
                     })
 
         if result["objects"]:
-            result_image = annotate_image(pil_image.copy(), result, point_radius, box_thickness, text_scale)
-            result_text = json.dumps(result, indent=2)
+            annotated_img = annotate_image(pil_image.copy(), result, point_radius=point_radius, box_thickness=box_thickness, text_scale=text_scale)
+            return {"image": pil_to_b64_png(annotated_img), "text": json.dumps(result, indent=2)}
         else:
-            result_text = f"Could not extract any objects.\nRaw model output:\n{full_text}"
+            return {"image": pil_to_b64_png(pil_image), "text": f"Could not extract any objects.\nRaw model output:\n{full_text}"}
 
-    elif category == "Spatial":
+    elif task == "Spatial":
         parsed = safe_parse_json(full_text)
         if isinstance(parsed, dict):
             for k in ["points", "waypoints", "path", "route", "nodes", "map"]:
@@ -448,9 +438,7 @@ def infer(
                     parsed = parsed[k]
                     break
             else:
-                for v in parsed.values():
-                    if isinstance(v, list): parsed = v; break
-                else: parsed = []
+                parsed = next((v for v in parsed.values() if isinstance(v, list)), [])
 
         result = {"points": []}
         if isinstance(parsed, list):
@@ -461,17 +449,13 @@ def infer(
                     result["points"].append({"label": item.get("label", "waypoint"), "x": x / 1000.0, "y": y / 1000.0})
 
         if result["points"]:
-            wp_lines = "\n".join(f"  {i+1}. {p['label']}  →  ({p['x']:.3f}, {p['y']:.3f})" for i, p in enumerate(result["points"]))
-            result_text = (
-                f"Spatial map generated.\n"
-                f"Waypoints ({len(result['points'])}):\n{wp_lines}\n"
-                f"Path segments: {max(0, len(result['points']) - 1)}"
-            )
-            result_image = annotate_spatial_path(pil_image.copy(), result, point_radius, box_thickness * 2, text_scale)
+            annotated_img = annotate_spatial_path(pil_image.copy(), result, dot_radius=point_radius, line_width=box_thickness * 2, text_scale=text_scale)
+            return {"image": pil_to_b64_png(annotated_img), "text": json.dumps(result, indent=2)}
         else:
-            result_text = f"Could not extract any spatial waypoints.\nRaw model output:\n{full_text}"
+            return {"image": pil_to_b64_png(pil_image), "text": f"Could not extract any spatial waypoints.\nRaw model output:\n{full_text}"}
+
+    return {"image": pil_to_b64_png(pil_image), "text": full_text}
 
-    return {"image": pil_to_b64_png(result_image), "text": result_text}
 
 @app.api(name="load_example", queue=False)
 def load_example(idx: float) -> dict:
@@ -480,13 +464,13 @@ def load_example(idx: float) -> dict:
     except (ValueError, TypeError):
         i = -1
     if i < 0 or i >= len(EXAMPLES_CONFIG):
-        return {"image": "", "prompt": "", "mode": "Detect", "name": "", "status": "error"}
+        return {"image": "", "prompt": "", "task": "", "name": "", "status": "error"}
     ex = EXAMPLES_CONFIG[i]
     b64 = encode_full_image(ex["image"])
     return {
         "image": b64,
         "prompt": ex["prompt"],
-        "mode": ex["mode"],
+        "task": ex["task"],
         "name": os.path.basename(ex["image"]),
         "status": "ok" if b64 else "error"
     }