Updated app.py

app.py

@@ -1,141 +1,95 @@
-# app.py
 import spaces
-import ast
 import torch
 from PIL import Image, ImageDraw
 import gradio as gr
+import re
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
-from qwen_vl_utils import process_vision_info # Make sure this file is in your repository
 
-# --- Model and Processor Initialization ---
-_MODEL = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    "ByteDance-Seed/UI-TARS-1.5-7B",
-    device_map="auto",
-    torch_dtype=torch.float16
-)
-_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,
-)
-model = _MODEL
-processor = _PROCESSOR
+@spaces.GPU
+def load_model_and_processor():
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        "ByteDance-Seed/UI-TARS-1.5-7B",
+        device_map="auto",
+        torch_dtype=torch.float16
+    )
+    processor = AutoProcessor.from_pretrained(
+        "ByteDance-Seed/UI-TARS-1.5-7B",
+        use_fast=True,
+    )
+    return model, processor
+
+model, processor = load_model_and_processor()
 
-def draw_point(image: Image.Image, point=None, radius: int = 15):
-    """Overlays a larger, more visible red dot on the screenshot."""
+def draw_point(image: Image.Image, point_str: str, radius: int = 10):
     img = image.copy()
-    if point and isinstance(point, list) and len(point) == 2:
-        x, y = point[0] * img.width, point[1] * img.height
-        draw = ImageDraw.Draw(img)
-        draw.ellipse(
-            (x - radius, y - radius, x + radius, y + radius), fill="rgba(255, 0, 0, 180)", outline="white", width=2
-        )
+    try:
+        coord_regex = r'click\(.*?<\|box_start\|>\s*\((\s*[\d.]+)\s*,\s*([\d.]+)\s*\).*?\)'
+        match = re.search(coord_regex, point_str)
+        
+        if match:
+            x_norm, y_norm = float(match.group(1)), float(match.group(2))
+            x = x_norm * img.width
+            y = y_norm * img.height
+            
+            draw = ImageDraw.Draw(img)
+            draw.ellipse(
+                (x - radius, y - radius, x + radius, y + radius), 
+                fill="red", 
+                outline="white", 
+                width=2
+            )
+    except Exception:
+        pass
     return img
 
 @spaces.GPU
-def navigate(screenshot, task: str):
-    """Runs a single inference step of the GUI reasoning model."""
+def navigate(screenshot: Image.Image, task: str):
     if not screenshot or not task:
-        return None, "Please provide both a screenshot and a task.", []
+        raise gr.Error("Please provide both a screenshot and a task.")
 
-    messages = []
-    
     prompt_header = (
-        "You are a precise GUI agent. Your task is to analyze a screenshot and a user instruction, then output your thought process and the next action. "
-        "You MUST ONLY output a 'Thought' and an 'Action'. Do not add any other text.\n\n"
-        "## Output Format\n"
-        "```\n"
-        "Thought: [Your brief, single-sentence reasoning here.]\n"
-        "Action: [The specific action to take, e.g., click(...) or type(...)]\n"
-        "```\n\n"
-        "## Action Space\n"
-        "click(start_box='<|box_start|>(x1, y1)<|box_end|>')\n"
-        "type(content='...')\n"
-        "scroll(start_box='<|box_start|>(x1, y1)<|box_end|>', direction='...')\n"
-        "finished(content='...')\n\n"
+        "You are a GUI agent. You are given a task and your action history, with screenshots."
+        "You need to perform the next action to complete the task. \n\n## Output Format\n```\nThought: ...\nAction: ...\n```\n\n## Action Space\n\nclick(start_box='<|box_start|>(x1, y1)<|box_end|>')\nleft_double(start_box='<|box_start|>(x1, y1)<|box_end|>')\nright_single(start_box='<|box_start|<(x1, y1)>|box_end|>')\ndrag(start_box='<|box_start|>(x1, y1)<|box_end|>', end_box='<|box_start|>(x3, y3)<|box_end|>')\nhotkey(key='')\ntype(content='') #If you want to submit your input, use \"\\n\" at the end of `content`.\nscroll(start_box='<|box_start|>(x1, y1)<|box_end|>', direction='down or up or right or left')\nwait() #Sleep for 5s and take a screenshot to check for any changes.\nfinished(content='xxx') # Use escape characters \\', \\\", and \\n in content part to ensure we can parse the content in normal python string format.\n\n\n## Note\n- Use English in `Thought` part.\n- Write a small plan and finally summarize your next action (with its target element) in one sentence in `Thought` part. Always use 'win' instead of 'meta' key\n\n"
         f"## User Instruction\n{task}"
     )
     
-    content = [
-        {"type": "text", "text": prompt_header},
-        {"type": "image_url", "image_url": screenshot}
-    ]
-    messages.append({"role": "user", "content": content})
-
-    images, videos = process_vision_info(messages)
+    messages = [{"role": "user", "content": [{"type": "text", "text": prompt_header}, {"type": "image", "image": screenshot}]}]
+    
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = processor(
-        text=[text], images=images, videos=videos, padding=True, return_tensors="pt"
-    ).to("cuda")
-
-    generated = model.generate(**inputs, max_new_tokens=256)
-    trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated)]
-    raw_out = processor.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    
+    inputs = processor(text=[text], images=[screenshot], return_tensors="pt").to(model.device)
 
+    generated_ids = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
+    
+    response = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    
     try:
-        if "Action:" in raw_out:
-            action_part = raw_out.split("Action:")[1].strip()
-            if action_part.startswith("```") and action_part.endswith("```"):
-                action_part = action_part[3:-3].strip()
-            
-            action_dict = ast.literal_eval(action_part)
-            
-            box_str = action_dict.get("start_box")
-            if box_str and isinstance(box_str, str) and "( " in box_str:
-                coords_part = box_str.split('( ')[1].split(' )')[0]
-                x_str, y_str = coords_part.split(', ')
-                pos = [float(x_str), float(y_str)]
-                screenshot = draw_point(screenshot, pos)
-    except (Exception, SyntaxError) as e:
-        print(f"Could not parse action or draw point: {e}")
-        pass
-
-    return screenshot, raw_out, messages
-
-# --- UI Definition with Fixes ---
-with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 90% !important;}") as demo:
-    gr.Markdown(
-        """
-        # ✨ Enhanced UI-Tars Navigation Demo
-        **Upload a screenshot and provide a task to see how the AI plans its next action.**
-        The model will analyze the image and your instruction, then output its thought process and the specific action it would take. A red dot will indicate the target location for clicks or scrolls.
-        """
-    )
-    with gr.Row():
-        with gr.Column(scale=1):
-            screenshot_in = gr.Image(type="pil", label="Screenshot")
-            task_in = gr.Textbox(
-                lines=2,
-                placeholder="e.g., Click on the 'Sign In' button.",
-                label="Task Instruction",
-            )
-            submit_btn = gr.Button("Analyze Action", variant="primary")
-            
-        with gr.Column(scale=2):
-            # --- FIX APPLIED HERE ---
-            # The 'interactive' argument has been removed from all output components
-            # to ensure compatibility with the execution environment's Gradio version.
-            screenshot_out = gr.Image(label="Result: Screenshot with Click Point")
-            with gr.Accordion("Model Output Details", open=False):
-                raw_out = gr.Textbox(label="Full Model Output (Thought & Action)")
-                history_out = gr.JSON(label="Conversation History for Debugging")
+        action_text = response.split('[/INST]')[-1].strip()
+    except IndexError:
+        action_text = response
 
-    submit_btn.click(
-        fn=navigate,
-        inputs=[screenshot_in, task_in],
-        outputs=[screenshot_out, raw_out, history_out],
-    )
+    output_image = draw_point(screenshot, action_text)
     
-    gr.Markdown(
-        """
-        ---
-        *Model: [ByteDance-Seed/UI-TARS-1.5-7B](https://huggingface.co/ByteDance-Seed/UI-TARS-1.5-7B)*
-        """
-    )
+    return output_image, action_text
+
+demo = gr.Interface(
+    fn=navigate,
+    inputs=[
+        gr.Image(type="pil", label="Screenshot"),
+        gr.Textbox(
+            lines=1,
+            placeholder="e.g. Search the weather for New York",
+            label="Task",
+        )
+    ],
+    outputs=[
+        gr.Image(label="With Click Point"),
+        gr.Textbox(label="Raw Action Output"),
+    ],
+    title="UI-Tars Navigation Demo",
+    description="Upload a UI screenshot, describe a task, and see the AI-predicted next action. This model helps automate GUI interactions.",
+    allow_flagging="never",
+)
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-    )
+    demo.launch()