Update app.py

app.py

@@ -1,395 +1,369 @@
-import os
-import re
-import time
-from typing import Tuple, Optional, List, Dict
-
 import gradio as gr
-import numpy as np
-import torch
-import spaces
-from PIL import Image, ImageDraw, ImageFont
-
-# Transformers & Qwen Utils
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
-from qwen_vl_utils import process_vision_info
-
-# -----------------------------------------------------------------------------
-# 1. PROMPT DEFINITIONS (from prompt.py)
-# -----------------------------------------------------------------------------
-
-OS_ACTIONS = """
-def final_answer(answer: any) -> any:
-    \"\"\"
-    Provides a final answer to the given problem.
-    Args:
-        answer: The final answer to the problem
-    \"\"\"
-
-def move_mouse(self, x: float, y: float) -> str:
-    \"\"\"
-    Moves the mouse cursor to the specified coordinates
-    Args:
-        x: The x coordinate (horizontal position)
-        y: The y coordinate (vertical position)
-    \"\"\"
-
-def click(x: Optional[float] = None, y: Optional[float] = None) -> str:
-    \"\"\"
-    Performs a left-click at the specified normalized coordinates
-    Args:
-        x: The x coordinate (horizontal position)
-        y: The y coordinate (vertical position)
-    \"\"\"
-
-def double_click(x: Optional[float] = None, y: Optional[float] = None) -> str:
-    \"\"\"
-    Performs a double-click at the specified normalized coordinates
-    Args:
-        x: The x coordinate (horizontal position)
-        y: The y coordinate (vertical position)
-    \"\"\"
-
-def type(text: str) -> str:
-    \"\"\"
-    Types the specified text at the current cursor position.
-    Args:
-        text: The text to type
-    \"\"\"
-
-def press(keys: str | list[str]) -> str:
-    \"\"\"
-    Presses a keyboard key
-    Args:
-        keys: The key or list of keys to press (e.g. "enter", "space", "backspace", "ctrl", etc.).
-    \"\"\"
-
-def navigate_back() -> str:
-    \"\"\"
-    Goes back to the previous page in the browser. If using this tool doesn't work, just click the button directly.
-    \"\"\"
-
-def drag(from_coord: list[float], to_coord: list[float]) -> str:
-    \"\"\"
-    Clicks [x1, y1], drags mouse to [x2, y2], then release click.
-    Args:
-        x1: origin x coordinate
-        y1: origin y coordinate
-        x2: end x coordinate
-        y2: end y coordinate
-    \"\"\"
-
-def scroll(direction: Literal["up", "down"] = "down", amount: int = 1) -> str:
-    \"\"\"
-    Moves the mouse to selected coordinates, then uses the scroll button: this could scroll the page or zoom, depending on the app. DO NOT use scroll to move through linux desktop menus.
-    Args:
-        x: The x coordinate (horizontal position) of the element to scroll/zoom, defaults to None to not focus on specific coordinates
-        y: The y coordinate (vertical position) of the element to scroll/zoom, defaults to None to not focus on specific coordinates
-        direction: The direction to scroll ("up" or "down"), defaults to "down". For zoom, "up" zooms in, "down" zooms out.
-        amount: The amount to scroll. A good amount is 1 or 2.
-    \"\"\"
-
-def wait(seconds: float) -> str:
-    \"\"\"
-    Waits for the specified number of seconds. Very useful in case the prior order is still executing (for example starting very heavy applications like browsers or office apps)
-    Args:
-        seconds: Number of seconds to wait, generally 2 is enough.
-    \"\"\"
-"""
-
-OS_SYSTEM_PROMPT = f"""You are a helpful GUI agent. You’ll be given a task and a screenshot of the screen. Complete the task using Python function calls.
+import json, os, re, traceback, contextlib
+from typing import Any, List, Dict
 
-For each step:
-	•	First, <think></think> to express the thought process guiding your next action and the reasoning behind it.
-	•	Then, use <code></code> to perform the action. it will be executed in a stateful environment.
-
-The following functions are exposed to the Python interpreter:
-<code>
-{OS_ACTIONS}
-</code>
-
-The state persists between code executions: so if in one step you've created variables or imported modules, these will all persist.
-"""
-
-# -----------------------------------------------------------------------------
-# 2. MODEL DEFINITION (Adapted for Fara-7B / Qwen2.5-VL)
-# -----------------------------------------------------------------------------
+import spaces
+import torch
+from PIL import Image, ImageDraw
+import requests
+from transformers import AutoModelForImageTextToText, AutoProcessor
+from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
 
+# --- Configuration ---
 MODEL_ID = "microsoft/Fara-7B"
 
-class FaraTransformersModel:
-    def __init__(self, model_id: str, to_device: str = "cuda"):
-        print(f"Loading {model_id}...")
-        self.model_id = model_id
-        
-        # Load Processor
-        self.processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-        
-        # Load Model
-        # Fara is based on Qwen2.5-VL architecture
-        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model_id,
-            trust_remote_code=True, 
-            torch_dtype=torch.bfloat16,
-            device_map="auto" if to_device == "cuda" else None
-        )
-        
-        if to_device == "cpu":
-            self.model.to("cpu")
-            
-        self.model.eval()
-        print("Model loaded successfully.")
-
-    def generate(self, messages: list[dict], **kwargs):
-        # 1. Prepare Text Prompts
-        text = self.processor.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        
-        # 2. Process Images (Qwen-VL specific utility)
-        image_inputs, video_inputs = process_vision_info(messages)
-        
-        # 3. Create Inputs
-        inputs = self.processor(
-            text=[text],
-            images=image_inputs,
-            videos=video_inputs,
-            padding=True,
-            return_tensors="pt",
-        )
-        
-        inputs = inputs.to(self.model.device)
-
-        # 4. Generate
-        with torch.no_grad():
-            generated_ids = self.model.generate(**inputs, **kwargs)
-        
-        # 5. Decode
-        generated_ids_trimmed = [
-            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-        ]
-        
-        output_text = self.processor.batch_decode(
-            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-        )[0]
-        
-        return output_text
-
-# Initialize Model Globally (Lazy loading handled by Gradio usually, but here we init for Spaces)
-# We use a global variable that is loaded on first run or at startup
-print(f"Initializing model class for {MODEL_ID}...")
-# Actual loading happens on GPU decorator or first call usually, 
-# but for the class structure we initialize it here.
-# Note: Actual torch.load happens inside the class init.
-fara_model = FaraTransformersModel(MODEL_ID, to_device="cuda" if torch.cuda.is_available() else "cpu")
-
-
-# -----------------------------------------------------------------------------
-# 3. HELPER FUNCTIONS (Parsing & Visualization)
-# -----------------------------------------------------------------------------
-
-def array_to_image(image_array: np.ndarray) -> Image.Image:
-    if image_array is None:
-        raise ValueError("No image provided. Please upload an image before submitting.")
-    return Image.fromarray(np.uint8(image_array))
-
-def get_navigation_prompt(task, image):
-    """Constructs the chat messages for Fara."""
+# ---------------- Device / DType helpers ----------------
+
+def pick_device() -> str:
+    """
+    On HF Spaces (ZeroGPU), CUDA is only available inside @spaces.GPU calls.
+    We still honor FORCE_DEVICE for local testing.
+    """
+    forced = os.getenv("FORCE_DEVICE", "").lower().strip()
+    if forced in {"cpu", "cuda", "mps"}:
+        return forced
+    if torch.cuda.is_available():
+        return "cuda"
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+def pick_dtype(device: str) -> torch.dtype:
+    if device == "cuda":
+        major, _ = torch.cuda.get_device_capability()
+        return torch.bfloat16 if major >= 8 else torch.float16  # Ampere+ -> bf16
+    if device == "mps":
+        return torch.float16
+    return torch.float32  # CPU: FP32 is usually fastest & most stable
+
+def move_to_device(batch, device: str):
+    if isinstance(batch, dict):
+        return {k: (v.to(device, non_blocking=True) if hasattr(v, "to") else v) for k, v in batch.items()}
+    if hasattr(batch, "to"):
+        return batch.to(device, non_blocking=True)
+    return batch
+
+# --- Chat/template helpers ---
+def apply_chat_template_compat(processor, messages: List[Dict[str, Any]]) -> str:
+    tok = getattr(processor, "tokenizer", None)
+    if hasattr(processor, "apply_chat_template"):
+        return processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    if tok is not None and hasattr(tok, "apply_chat_template"):
+        return tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    texts = []
+    for m in messages:
+        for c in m.get("content", []):
+            if isinstance(c, dict) and c.get("type") == "text":
+                texts.append(c.get("text", ""))
+    return "\n".join(texts)
+
+def batch_decode_compat(processor, token_id_batches, **kw):
+    tok = getattr(processor, "tokenizer", None)
+    if tok is not None and hasattr(tok, "batch_decode"):
+        return tok.batch_decode(token_id_batches, **kw)
+    if hasattr(processor, "batch_decode"):
+        return processor.batch_decode(token_id_batches, **kw)
+    raise AttributeError("No batch_decode available on processor or tokenizer.")
+
+def get_image_proc_params(processor) -> Dict[str, int]:
+    ip = getattr(processor, "image_processor", None)
+    return {
+        "patch_size": getattr(ip, "patch_size", 14),
+        "merge_size": getattr(ip, "merge_size", 1),
+        "min_pixels": getattr(ip, "min_pixels", 256 * 256),
+        "max_pixels": getattr(ip, "max_pixels", 1280 * 1280),
+    }
+
+def trim_generated(generated_ids, inputs):
+    in_ids = getattr(inputs, "input_ids", None)
+    if in_ids is None and isinstance(inputs, dict):
+        in_ids = inputs.get("input_ids", None)
+    if in_ids is None:
+        return [out_ids for out_ids in generated_ids]
+    return [out_ids[len(in_seq):] for in_seq, out_ids in zip(in_ids, generated_ids)]
+
+# --- Parsing helper: normalize various UI-TARS click formats to (x, y) ---
+def parse_click_coordinates(text: str, img_w: int, img_h: int):
+    """
+    Returns (x, y) in image coordinates, clamped to bounds, or None.
+    Handles:
+      - Click(start_box='(x,y)') / Click(end_box='(x,y)')
+      - Click(box='(x1,y1,x2,y2)') -> center
+      - Click(x, y)
+      - Click({'x':..., 'y':...}) / Click({"x":...,"y":...})
+    Preference: start_box > end_box when both exist.
+    """
+    s = str(text)
+
+    # 1) start_box / end_box
+    pairs = re.findall(r"(start_box|end_box)\s*=\s*['\"]\(\s*(\d+)\s*,\s*(\d+)\s*\)['\"]", s)
+    if pairs:
+        start = next(((int(x), int(y)) for k, x, y in pairs if k == "start_box"), None)
+        if start:
+            x, y = start
+            return max(0, min(x, img_w - 1)), max(0, min(y, img_h - 1))
+        end = next(((int(x), int(y)) for k, x, y in pairs if k == "end_box"), None)
+        if end:
+            x, y = end
+            return max(0, min(x, img_w - 1)), max(0, min(y, img_h - 1))
+
+    # 2) box='(x1,y1,x2,y2)' -> center
+    m = re.search(r"box\s*=\s*['\"]\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)['\"]", s)
+    if m:
+        x1, y1, x2, y2 = map(int, m.groups())
+        cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
+        return max(0, min(cx, img_w - 1)), max(0, min(cy, img_h - 1))
+
+    # 3) Direct Click(x, y)
+    m = re.search(r"Click\s*\(\s*(\d+)\s*,\s*(\d+)\s*\)", s)
+    if m:
+        x, y = int(m.group(1)), int(m.group(2))
+        return max(0, min(x, img_w - 1)), max(0, min(y, img_h - 1))
+
+    # 4) JSON-ish dicts
+    m = re.search(r"Click\s*\(\s*[{[][^)}]*['\"]?x['\"]?\s*:\s*(\d+)\s*,\s*['\"]?y['\"]?\s*:\s*(\d+)[^)}]*\)\s*", s)
+    if m:
+        x, y = int(m.group(1)), int(m.group(2))
+        return max(0, min(x, img_w - 1)), max(0, min(y, img_h - 1))
+
+    return None
+
+# --- Load model/processor ON CPU at import time (ZeroGPU safe) ---
+print(f"Loading model and processor for {MODEL_ID} on CPU startup (ZeroGPU safe)...")
+model = None
+processor = None
+model_loaded = False
+load_error_message = ""
+
+try:
+    model = AutoModelForImageTextToText.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float32,  # CPU-safe dtype at import
+        trust_remote_code=True,
+    )
+    # IMPORTANT: use_fast=False to avoid the breaking change error you hit
+    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=False)
+    model.eval()
+    model_loaded = True
+    print("Model and processor loaded on CPU.")
+except Exception as e:
+    load_error_message = (
+        f"Error loading model/processor: {e}\n"
+        "This might be due to network/model ID/library versions.\n"
+        "Check the full traceback in the logs."
+    )
+    print(load_error_message)
+    traceback.print_exc()
+
+# --- Prompt builder ---
+def get_localization_prompt(pil_image: Image.Image, instruction: str) -> List[dict]:
+    guidelines: str = (
+        "Localize an element on the GUI image according to my instructions and "
+        "output a click position as Click(x, y) with x num pixels from the left edge "
+        "and y num pixels from the top edge."
+    )
     return [
-        {
-            "role": "system",
-            "content": [{"type": "text", "text": OS_SYSTEM_PROMPT}],
-        },
         {
             "role": "user",
             "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": f"Instruction: {task}\n\nPrevious actions:\nNone"},
+                {"type": "image", "image": pil_image},
+                {"type": "text", "text": f"{guidelines}\n{instruction}"}
             ],
-        },
+        }
     ]
 
-def parse_actions_from_response(response: str) -> list[str]:
-    """Parse actions from model response using regex pattern."""
-    # Look for code blocks
-    pattern = r"<code>(.*?)</code>"
-    matches = re.findall(pattern, response, re.DOTALL)
-    if not matches:
-        # Fallback: if model forgets code tags but writes function calls
-        if "click(" in response or "type(" in response:
-            return [response]
-    return matches
-
-def extract_coordinates_from_action(action_code: str) -> list[dict]:
-    """Extract coordinates from action code for localization actions."""
-    localization_actions = []
-    
-    # Patterns for different action types
-    patterns = {
-        'click': r'click\((?:x=)?([0-9.]+)(?:,\s*(?:y=)?([0-9.]+))?\)',
-        'double_click': r'double_click\((?:x=)?([0-9.]+)(?:,\s*(?:y=)?([0-9.]+))?\)',
-        'move_mouse': r'move_mouse\((?:self,\s*)?(?:x=)?([0-9.]+)(?:,\s*(?:y=)?([0-9.]+))\)',
-        'drag': r'drag\(\[([0-9.]+),\s*([0-9.]+)\],\s*\[([0-9.]+),\s*([0-9.]+)\]\)'
-    }
-    
-    for action_type, pattern in patterns.items():
-        matches = re.finditer(pattern, action_code)
-        for match in matches:
-            if action_type == 'drag':
-                from_x, from_y, to_x, to_y = match.groups()
-                localization_actions.append({
-                    'type': 'drag_from', 'x': float(from_x), 'y': float(from_y), 'action': action_type
-                })
-                localization_actions.append({
-                    'type': 'drag_to', 'x': float(to_x), 'y': float(to_y), 'action': action_type
-                })
-            else:
-                x_val = match.group(1)
-                y_val = match.group(2) if match.group(2) else x_val
-                if x_val and y_val:
-                    localization_actions.append({
-                        'type': action_type, 'x': float(x_val), 'y': float(y_val), 'action': action_type
-                    })
-    
-    return localization_actions
-
-def create_localized_image(original_image: Image.Image, coordinates: list[dict]) -> Optional[Image.Image]:
-    """Create an image with localization markers drawn on it."""
-    if not coordinates:
-        return None
-    
-    img_copy = original_image.copy()
-    draw = ImageDraw.Draw(img_copy)
-    width, height = img_copy.size
-    
+# --- Inference core (device passed in; AMP used when suitable) ---
+@torch.inference_mode()
+def run_inference_localization(
+    messages_for_template: List[dict[str, Any]],
+    pil_image_for_processing: Image.Image,
+    device: str,
+    dtype: torch.dtype,
+) -> str:
+    text_prompt = apply_chat_template_compat(processor, messages, ) if False else apply_chat_template_compat(processor, messages_for_template)
+
+    inputs = processor(
+        text=[text_prompt],
+        images=[pil_image_for_processing],
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = move_to_device(inputs, device)
+
+    # AMP contexts
+    if device == "cuda":
+        amp_ctx = torch.autocast(device_type="cuda", dtype=dtype)
+    elif device == "mps":
+        amp_ctx = torch.autocast(device_type="mps", dtype=torch.float16)
+    else:
+        amp_ctx = contextlib.nullcontext()
+
+    with amp_ctx:
+        generated_ids = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            do_sample=False,
+        )
+
+    generated_ids_trimmed = trim_generated(generated_ids, inputs)
+    decoded_output = batch_decode_compat(
+        processor,
+        generated_ids_trimmed,
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=False
+    )
+    return decoded_output[0] if decoded_output else ""
+
+# --- Gradio processing function (ZeroGPU-visible) ---
+@spaces.GPU(duration=120)  # keep GPU attached briefly between calls (seconds)
+def predict_click_location(input_pil_image: Image.Image, instruction: str):
+    if not model_loaded or not processor or not model:
+        return f"Model not loaded. Error: {load_error_message}", None, "device: n/a | dtype: n/a"
+    if not input_pil_image:
+        return "No image provided. Please upload an image.", None, "device: n/a | dtype: n/a"
+    if not instruction or instruction.strip() == "":
+        return "No instruction provided. Please type an instruction.", input_pil_image.copy().convert("RGB"), "device: n/a | dtype: n/a"
+
+    # Decide device/dtype *inside* the GPU-decorated call
+    device = pick_device()
+    dtype = pick_dtype(device)
+
+    # Optional perf knobs for CUDA
+    if device == "cuda":
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.set_float32_matmul_precision("high")
+
+    # If needed, move model now that GPU is available
     try:
-        font = ImageFont.load_default()
-    except:
-        font = None
-    
-    colors = {'click': 'red', 'double_click': 'blue', 'move_mouse': 'green', 'drag_from': 'orange', 'drag_to': 'purple'}
-    
-    for i, coord in enumerate(coordinates):
-        # Normalize if model outputs 0-1 range (Fara usually does)
-        # If model outputs pixels, we need to handle that. 
-        # Fara/SmolVLM usually output normalized coordinates 0-1000 or 0-1.
-        # Assuming Fara outputs 0-1 floats based on the System Prompt definition.
-        
-        pixel_x = int(coord['x'] * width) if coord['x'] <= 1.0 else int(coord['x'])
-        pixel_y = int(coord['y'] * height) if coord['y'] <= 1.0 else int(coord['y'])
-        
-        color = colors.get(coord['type'], 'red')
-        
-        r = 10
-        draw.ellipse([pixel_x - r, pixel_y - r, pixel_x + r, pixel_y + r], outline=color, width=3)
-        
-        label = f"{coord['type']}"
-        draw.text((pixel_x + 12, pixel_y - 10), label, fill=color, font=font)
-        
-        # Draw drag arrows
-        if coord['type'] == 'drag_from' and i + 1 < len(coordinates) and coordinates[i + 1]['type'] == 'drag_to':
-            next_coord = coordinates[i + 1]
-            end_x = int(next_coord['x'] * width) if next_coord['x'] <= 1.0 else int(next_coord['x'])
-            end_y = int(next_coord['y'] * height) if next_coord['y'] <= 1.0 else int(next_coord['y'])
-            draw.line([pixel_x, pixel_y, end_x, end_y], fill='orange', width=3)
-
-    return img_copy
-
-# -----------------------------------------------------------------------------
-# 4. APP LOGIC (ZeroGPU)
-# -----------------------------------------------------------------------------
-
-@spaces.GPU(duration=60)
-def navigate(input_numpy_image: np.ndarray, task: str) -> Tuple[str, Optional[Image.Image]]:
-    if input_numpy_image is None:
-        return "Please upload an image.", None
-
-    input_pil_image = array_to_image(input_numpy_image)
-    
-    # 1. Build Prompt
-    prompt = get_navigation_prompt(task, input_pil_image)
-
-    # 2. Generate
-    if fara_model is None:
-        raise ValueError("Model not loaded")
-    
-    navigation_str = fara_model.generate(prompt, max_new_tokens=500)
-    print(f"Raw Output: {navigation_str}")
-    
-    # 3. Parse
-    navigation_str = navigation_str.strip()
-    actions = parse_actions_from_response(navigation_str)
-    
-    all_coordinates = []
-    for action_code in actions:
-        coordinates = extract_coordinates_from_action(action_code)
-        all_coordinates.extend(coordinates)
-    
-    # 4. Visualize
-    localized_image = input_pil_image
-    if all_coordinates:
-        visualized = create_localized_image(input_pil_image, all_coordinates)
-        if visualized:
-            localized_image = visualized
-            
-    return navigation_str, localized_image
-
-# -----------------------------------------------------------------------------
-# 5. GRADIO UI
-# -----------------------------------------------------------------------------
-
-title = "Fara-7B GUI Operator 🖥️"
-description = """
-This demo uses **microsoft/Fara-7B** to understand GUI screenshots and generate navigation actions.
-Upload a screenshot, define a task, and see the model's planned actions.
-"""
+        p = next(model.parameters())
+        cur_dev = p.device.type
+        cur_dtype = p.dtype
+    except StopIteration:
+        cur_dev, cur_dtype = "cpu", torch.float32
 
-# Load examples safely
-examples = []
-example_paths = [
-    ("Search for UK Prime Minister", "./assets/google.png"),
-    ("Find trending models", "./assets/huggingface.png")
-]
-
-# We skip checking file existence to allow script to run, 
-# but in a real space, ensure ./assets/ folder exists or remove examples
-safe_examples = []
-for label, path in example_paths:
-    if os.path.exists(path):
-        safe_examples.append([path, label])
-
-
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
-    gr.Markdown(description)
-
-    with gr.Row():
-        with gr.Column():
-            input_image_component = gr.Image(label="Upload Interface Screenshot", height=500)
-            task_component = gr.Textbox(
-                label="Task Instruction",
-                placeholder="e.g., Click the Search bar and type 'Hello World'",
-                lines=2
-            )
-            submit_button = gr.Button("Generate Action", variant="primary")
+    if cur_dev != device or cur_dtype != dtype:
+        model.to(device=device, dtype=dtype)
+        model.eval()
+
+    # 1) Resize according to image processor params (safe defaults if missing)
+    try:
+        ip = get_image_proc_params(processor)
+        resized_height, resized_width = smart_resize(
+            input_pil_image.height,
+            input_pil_image.width,
+            factor=ip["patch_size"] * ip["merge_size"],
+            min_pixels=ip["min_pixels"],
+            max_pixels=ip["max_pixels"],
+        )
+        resized_image = input_pil_image.resize(
+            size=(resized_width, resized_height),
+            resample=Image.Resampling.LANCZOS
+        )
+    except Exception as e:
+        traceback.print_exc()
+        return f"Error resizing image: {e}", input_pil_image.copy().convert("RGB"), f"device: {device} | dtype: {dtype}"
 
-        with gr.Column():
-            output_image_component = gr.Image(label="Visualized Action", height=500)
-            output_code_component = gr.Textbox(label="Model Output (Code)", lines=10, show_copy_button=True)
+    # 2) Build messages with image + instruction
+    messages = get_localization_prompt(resized_image, instruction)
 
-    submit_button.click(
-        fn=navigate, 
-        inputs=[input_image_component, task_component], 
-        outputs=[output_code_component, output_image_component]
-    )
+    # 3) Run inference
+    try:
+        coordinates_str = run_inference_localization(messages, resized_image, device, dtype)
+    except Exception as e:
+        traceback.print_exc()
+        return f"Error during model inference: {e}", resized_image.copy().convert("RGB"), f"device: {device} | dtype: {dtype}"
+
+    # 4) Parse coordinates and draw marker
+    output_image_with_click = resized_image.copy().convert("RGB")
+    coords = parse_click_coordinates(coordinates_str, resized_width, resized_height)
+
+    if coords is not None:
+        x, y = coords
+        draw = ImageDraw.Draw(output_image_with_click)
+        radius = max(5, min(resized_width // 100, resized_height // 100, 15))
+        bbox = (x - radius, y - radius, x + radius, y + radius)
+        draw.ellipse(bbox, outline="red", width=max(2, radius // 4))
+        print(f"Predicted and drawn click at: ({x}, {y}) on resized image ({resized_width}x{resized_height})")
+    else:
+        print(f"Could not parse a click from model output: {coordinates_str}")
+
+    return coordinates_str, output_image_with_click, f"device: {device} | dtype: {str(dtype).replace('torch.', '')}"
+
+# --- Load Example Data ---
+example_image = None
+example_instruction = "Enter the server address readyforquantum.com to check its security"
+try:
+    example_image_url = "https://readyforquantum.com/img/screentest.jpg"
+    example_image = Image.open(requests.get(example_image_url, stream=True).raw)
+except Exception as e:
+    print(f"Could not load example image from URL: {e}")
+    traceback.print_exc()
+    try:
+        example_image = Image.new("RGB", (200, 150), color="lightgray")
+        draw = ImageDraw.Draw(example_image)
+        draw.text((10, 10), "Example image\nfailed to load", fill="black")
+    except Exception:
+        pass
+
+# --- Gradio UI ---
+title = "GUI Nav Demo"
+article = f"""
+<p style='text-align: center'>
+Model: <a href='https://huggingface.co/{MODEL_ID}' target='_blank'>{MODEL_ID}</a>
+</p>
+"""
+
+if not model_loaded:
+    with gr.Blocks() as demo:
+        gr.Markdown(f"# <center>⚠️ Error: Model Failed to Load ⚠️</center>")
+        gr.Markdown(f"<center>{load_error_message}</center>")
+        gr.Markdown("<center>See logs for the full traceback.</center>")
+else:
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
+        gr.Markdown(article)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image_component = gr.Image(type="pil", label="Input UI Image", height=400)
+                instruction_component = gr.Textbox(
+                    label="Instruction",
+                    placeholder="e.g., Click the 'Login' button",
+                    info="Type the action you want the model to localize on the image."
+                )
+                submit_button = gr.Button("Localize Click", variant="primary")
+
+            with gr.Column(scale=1):
+                output_coords_component = gr.Textbox(
+                    label="Predicted Coordinates / Action",
+                    interactive=False
+                )
+                output_image_component = gr.Image(
+                    type="pil",
+                    label="Image with Predicted Click Point",
+                    height=400,
+                    interactive=False
+                )
+                runtime_info = gr.Textbox(
+                    label="Runtime Info",
+                    value="device: n/a | dtype: n/a",
+                    interactive=False
+                )
+
+        if example_image:
+            gr.Examples(
+                examples=[[example_image, example_instruction]],
+                inputs=[input_image_component, instruction_component],
+                outputs=[output_coords_component, output_image_component, runtime_info],
+                fn=predict_click_location,
+                cache_examples="lazy",
+            )
 
-    if safe_examples:
-        gr.Examples(
-            examples=safe_examples,
-            inputs=[input_image_component, task_component],
-            outputs=[output_code_component, output_image_component],
-            fn=navigate,
-            cache_examples=True,
+        submit_button.click(
+            fn=predict_click_location,
+            inputs=[input_image_component, instruction_component],
+            outputs=[output_coords_component, output_image_component, runtime_info]
         )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch(debug=True)