Updated app.py

app.py

@@ -1,58 +1,85 @@
-import ast
-import base64
-from io import BytesIO
-
 import spaces
+import ast
 import torch
 from PIL import Image, ImageDraw
 import gradio as gr
+import base64
+from io import BytesIO
+
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
-from qwen_vl_utils import process_vision_info     # keep this file in repo
+from qwen_vl_utils import process_vision_info
 
-# ─── Model & Processor ──────────────────────────────────────────────────────
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+_MODEL = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     "ByteDance-Seed/UI-TARS-1.5-7B",
     device_map="auto",
     torch_dtype=torch.float16
 )
-processor = AutoProcessor.from_pretrained(
+
+_PROCESSOR = AutoProcessor.from_pretrained(
     "ByteDance-Seed/UI-TARS-1.5-7B",
     size={"shortest_edge": 100 * 28 * 28, "longest_edge": 16384 * 28 * 28},
     use_fast=True,
 )
 
-# ─── Helpers ────────────────────────────────────────────────────────────────
-def draw_point(img: Image.Image, pos=None, r: int = 5):
-    if pos:
-        x, y = pos[0] * img.width, pos[1] * img.height
-        ImageDraw.Draw(img).ellipse((x - r, y - r, x + r, y + r), fill="red")
+model = _MODEL
+processor = _PROCESSOR
+
+
+def draw_point(image: Image.Image, point=None, radius: int = 5):
+    """Overlay a red dot on the screenshot where the model clicked."""
+    img = image.copy()
+    if point:
+        x, y = point[0] * img.width, point[1] * img.height
+        ImageDraw.Draw(img).ellipse(
+            (x - radius, y - radius, x + radius, y + radius), fill="red"
+        )
     return img
 
-# ─── Core Inference ─────────────────────────────────────────────────────────
+
 @spaces.GPU
 def navigate(screenshot, task: str):
+    """Run one inference step on the GUI‑reasoning model."""
+    messages = []
+
     prompt_header = (
-        "You are an expert GUI agent.\n"
-        "Given the task and previous screenshots, decide only the next action.\n\n"
-        "=== TASK ===\n"
-        f"{task}\n"
-        "==============\n\n"
-        "Respond strictly as JSON: {\"Thought\": \"…\", \"Action\": {…}}\n"
-        "Valid actions: click, left_double, right_single, drag, hotkey, type, scroll, wait, finished.\n"
-        "Use English for Thought. Summarize your plan in one sentence."
+        "🤖 **GUI Agent Instructions**\n\n"
+        "You're an intelligent agent solving UI tasks through:\n"
+        "1. Visual understanding of screenshots\n"
+        "2. Natural language task interpretation\n"
+        "3. Action sequence generation\n\n"
+        "## Action Reference\n"
+        "```
+        "| Action Type       | Syntax                  |\n"
+        "|--------------------|-------------------------|\n"
+        "| Click Button       | click(start_box='(x,y)') |\n"
+        "| Double-Click       | left_double(start_box='(x,y)') |\n"
+        "| Drag Element       | drag(start_box='(x1,y1)', end_box='(x2,y2)') |\n"
+        "| Hotkey Input       | hotkey(key='Ctrl+A')    |\n"
+        "| Text Input         | type(content='Search term') |\n"
+        "| Scroll Action      | scroll(start_box='(x,y)', direction='down') |\n"
+        "| Wait & Recheck     | wait()                  |\n"
+        "| Task Completion    | finished(content='Result') |\n"
+        "```"
+        "\n**Note:**\n"
+        "1. Use 'win' key instead of 'meta' in hotkey commands\n"
+        "2. Include position coordinates in all spatial actions\n"
+        "3. Keep 'Thought' concise - max 3 sentence strategy\n"
+        f"\n**Task:**\n{task}"
     )
 
-    messages = [{
+    current = {
         "role": "user",
         "content": [
             {"type": "text", "text": prompt_header},
             {"type": "image_url", "image_url": screenshot}
         ]
-    }]
+    }
+    messages.append(current)
 
     images, videos = process_vision_info(messages)
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
     inputs = processor(
         text=[text],
         images=images,
@@ -61,45 +88,45 @@ def navigate(screenshot, task: str):
         return_tensors="pt",
     ).to("cuda")
 
-    gen_ids = model.generate(**inputs, max_new_tokens=128)
+    generated = model.generate(**inputs, max_new_tokens=128)
+    trimmed = [
+        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated)
+    ]
     raw_out = processor.batch_decode(
-        [out[len(inp):] for inp, out in zip(inputs.input_ids, gen_ids)],
-        skip_special_tokens=True,
-        clean_up_tokenization_spaces=False
+        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )[0]
 
     try:
-        for act in ast.literal_eval(raw_out) if isinstance(ast.literal_eval(raw_out), list) else [ast.literal_eval(raw_out)]:
-            p = act.get("position")
-            if p and len(p) == 2:
-                screenshot = draw_point(screenshot, p)
+        actions = ast.literal_eval(raw_out)
+        for act in (actions if isinstance(actions, list) else [actions]):
+            pos = act.get("position")
+            if pos and isinstance(pos, list) and len(pos) == 2:
+                screenshot = draw_point(screenshot, pos)
     except Exception:
         pass
 
     return screenshot, raw_out, messages
 
-# ─── UI ─────────────────────────────────────────────────────────────────────
-with gr.Blocks(title="UI-Tars Navigation Demo", theme=gr.themes.Soft()) as demo:
-    with gr.Row():
-        with gr.Column(scale=3):
-            in_img = gr.Image(type="pil", height=380, show_label=False)
-            in_task = gr.Textbox(
-                placeholder="Describe what you want to do (e.g. “Open Gmail and compose a message”)",
-                lines=1,
-                label="Task"
-            )
-            run_btn = gr.Button("Run", variant="primary")
-        with gr.Column(scale=4):
-            out_img = gr.Image(label="With Click Point", height=380)
-            out_raw = gr.Textbox(label="Raw Action JSON", interactive=False)
-            with gr.Accordion("Conversation History", open=False):
-                out_hist = gr.JSON()
-
-    run_btn.click(navigate, inputs=[in_img, in_task], outputs=[out_img, out_raw, out_hist])
+
+demo = gr.Blocks()
+with gr.Row():
+    gr.Image(type="pil", label="🖼️ Screenshot Input").style(width=400)
+    gr.Textbox(
+        lines=1,
+        placeholder="e.g., 'Book a flight to Paris'",
+        label="🔍 Task Description"
+    ).style(width=400)
+
+with gr.Row():
+    gr.Image(label="📍 Click Point Visualization").style(width=400)
+    gr.Textbox(label="📝 Action Response").style(width=400)
+
+with gr.Row():
+    gr.JSON(label="📜 Conversation History").style(width=800)
 
 demo.launch(
     server_name="0.0.0.0",
     server_port=7860,
     share=False,
-    ssr_mode=False,
+    ssr_mode=False
 )