app.py

import os
import re
import json
import time
import unicodedata
import gc
from io import BytesIO
from typing import Iterable, Tuple, Optional, List, Dict, Any

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,
    AutoModelForImageTextToText
)
from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
from qwen_vl_utils import process_vision_info

# Gradio Theme
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

# -----------------------------------------------------------------------------
# 1. THEME CONFIGURATION
# -----------------------------------------------------------------------------

colors.steel_blue = colors.Color(
    name="steel_blue",
    c50="#EBF3F8",
    c100="#D3E5F0",
    c200="#A8CCE1",
    c300="#7DB3D2",
    c400="#529AC3",
    c500="#4682B4",
    c600="#3E72A0",
    c700="#36638C",
    c800="#2E5378",
    c900="#264364",
    c950="#1E3450",
)

class SteelBlueTheme(Soft):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.steel_blue,
        neutral_hue: colors.Color | str = colors.slate,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_800)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_500)",
            block_title_text_weight="600",
            block_border_width="3px",
            block_shadow="*shadow_drop_lg",
            button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px",
        )

steel_blue_theme = SteelBlueTheme()
css = "#main-title h1 { font-size: 2.3em !important; } #out_img { height: 600px; object-fit: contain; }"

# -----------------------------------------------------------------------------
# 2. GLOBAL MODEL LOADING
# -----------------------------------------------------------------------------

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Running on device: {device}")

# --- Load Fara-7B ---
print("🔄 Loading Fara-7B...")
MODEL_ID_V = "microsoft/Fara-7B"
try:
    processor_v = AutoProcessor.from_pretrained(MODEL_ID_V, trust_remote_code=True)
    model_v = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_ID_V,
        trust_remote_code=True,
        torch_dtype=torch.float16
    ).to(device).eval()
except Exception as e:
    print(f"Failed to load Fara: {e}")
    model_v = None
    processor_v = None

# --- Load UI-TARS-1.5-7B ---
print("🔄 Loading UI-TARS-1.5-7B...")
MODEL_ID_X = "ByteDance-Seed/UI-TARS-1.5-7B" 
try:
    processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True, use_fast=False)
    model_x = AutoModelForImageTextToText.from_pretrained(
        MODEL_ID_X,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    ).to(device).eval()
except Exception as e:
    print(f"Failed to load UI-TARS: {e}")
    model_x = None
    processor_x = None

# --- Load Holo2-8B ---
print("🔄 Loading Holo2-8B...")
MODEL_ID_H = "Hcompany/Holo2-8B"
try:
    processor_h = AutoProcessor.from_pretrained(MODEL_ID_H, trust_remote_code=True)
    model_h = AutoModelForImageTextToText.from_pretrained(
        MODEL_ID_H,
        trust_remote_code=True,
        torch_dtype=torch.float16
    ).to(device).eval()
except Exception as e:
    print(f"Failed to load Holo2: {e}")
    model_h = None
    processor_h = None

print("✅ Models loading sequence complete.")

# -----------------------------------------------------------------------------
# 3. UTILS & HELPERS
# -----------------------------------------------------------------------------

def array_to_image(image_array: np.ndarray) -> Image.Image:
    if image_array is None: raise ValueError("No image provided.")
    return Image.fromarray(np.uint8(image_array))

# --- Compatibility Helpers ---
def apply_chat_template_compat(processor, messages: List[Dict[str, Any]]) -> str:
    """Helper to handle chat template application across different processors"""
    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)
    
    # Fallback if no template method found
    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):
    """Helper to handle batch decoding"""
    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 trim_generated(generated_ids, inputs):
    """Removes input tokens from output if necessary"""
    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)]

def get_image_proc_params(processor) -> Dict[str, int]:
    """Extracts resizing parameters from the processor configuration"""
    ip = getattr(processor, "image_processor", None)
    return {
        "patch_size": getattr(ip, "patch_size", 14),
        "merge_size": getattr(ip, "merge_size", 2), # Default to 2, Holo2 might differ
        "min_pixels": getattr(ip, "min_pixels", 256 * 256),
        "max_pixels": getattr(ip, "max_pixels", 1280 * 1280),
    }

# -----------------------------------------------------------------------------
# 4. PROMPTS
# -----------------------------------------------------------------------------

# --- Fara Prompt ---
def get_fara_prompt(task, image):
    OS_SYSTEM_PROMPT = """You are a GUI agent. You are given a task and a screenshot of the current status. 
    You need to generate the next action to complete the task.
    Output your action inside a <tool_call> block using JSON format.
    Include "coordinate": [x, y] in pixels for interactions.
    Examples:
    <tool_call>{"name": "User", "arguments": {"action": "click", "coordinate": [400, 300]}}</tool_call>
    <tool_call>{"name": "User", "arguments": {"action": "type", "coordinate": [100, 200], "text": "hello"}}</tool_call>
    """
    return [
        {"role": "system", "content": [{"type": "text", "text": OS_SYSTEM_PROMPT}]},
        {"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": f"Instruction: {task}"}]},
    ]

# --- UI-TARS Prompt ---
def get_uitars_prompt(task, image):
    guidelines = (
        "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": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": f"{guidelines}\n{task}"}
            ],
        }
    ]

# --- Holo2 Prompt ---
def get_holo_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": "user",
            "content": [
                {"type": "image", "image": pil_image},
                {"type": "text", "text": f"{guidelines}\n{instruction}"}
            ],
        }
    ]

# -----------------------------------------------------------------------------
# 5. PARSING LOGIC
# -----------------------------------------------------------------------------

def parse_coordinate_response(text: str) -> List[Dict]:
    """
    Parses UI-TARS and Holo2 output formats.
    Targets formats like: Click(x, y), point=[x, y], etc.
    """
    actions = []
    text = text.strip()
    
    print(f"Parsing Coordinate output: {text}")

    # Regex 1: Click(x, y) - Standard prompt output for UI-TARS & Holo2
    matches_click = re.findall(r"Click\s*\(\s*(\d+)\s*,\s*(\d+)\s*\)", text, re.IGNORECASE)
    for m in matches_click:
        actions.append({"type": "click", "x": int(m[0]), "y": int(m[1]), "text": ""})

    # Regex 2: point=[x, y] - Common model internal format
    matches_point = re.findall(r"point=\[\s*(\d+)\s*,\s*(\d+)\s*\]", text, re.IGNORECASE)
    for m in matches_point:
        actions.append({"type": "click", "x": int(m[0]), "y": int(m[1]), "text": ""})

    # Regex 3: start_box='(x, y)' - Another variant
    matches_box = re.findall(r"start_box=['\"]?\(\s*(\d+)\s*,\s*(\d+)\s*\)['\"]?", text, re.IGNORECASE)
    for m in matches_box:
        actions.append({"type": "click", "x": int(m[0]), "y": int(m[1]), "text": ""})

    # Remove duplicates
    unique_actions = []
    seen = set()
    for a in actions:
        key = (a['type'], a['x'], a['y'])
        if key not in seen:
            seen.add(key)
            unique_actions.append(a)
            
    return unique_actions

def parse_fara_response(response: str) -> List[Dict]:
    """Parse Fara <tool_call> JSON format"""
    actions = []
    matches = re.findall(r"<tool_call>(.*?)</tool_call>", response, re.DOTALL)
    for match in matches:
        try:
            data = json.loads(match.strip())
            args = data.get("arguments", {})
            coords = args.get("coordinate", [])
            action_type = args.get("action", "unknown")
            text_content = args.get("text", "")
            if coords and len(coords) == 2:
                actions.append({
                    "type": action_type, "x": float(coords[0]), "y": float(coords[1]), "text": text_content
                })
        except: pass
    return actions

def create_localized_image(original_image: Image.Image, actions: list[dict]) -> Optional[Image.Image]:
    if not actions: return None
    img_copy = original_image.copy()
    draw = ImageDraw.Draw(img_copy)
    
    try: font = ImageFont.load_default()
    except: font = None
    
    for act in actions:
        pixel_x = int(act['x'])
        pixel_y = int(act['y'])
            
        color = 'red' if 'click' in act['type'].lower() else 'blue'
        
        # Draw Target Crosshair/Circle
        r = 15
        line_width = 4
        
        # Circle
        draw.ellipse([pixel_x - r, pixel_y - r, pixel_x + r, pixel_y + r], outline=color, width=line_width)
        # Center dot
        draw.ellipse([pixel_x - 3, pixel_y - 3, pixel_x + 3, pixel_y + 3], fill=color)
        
        # Label
        label = f"{act['type']}"
        if act['text']: label += f": {act['text']}"
        
        text_pos = (pixel_x + 20, pixel_y - 10)
        # Draw text background
        bbox = draw.textbbox(text_pos, label, font=font)
        draw.rectangle((bbox[0]-4, bbox[1]-2, bbox[2]+4, bbox[3]+2), fill="black")
        draw.text(text_pos, label, fill="white", font=font)

    return img_copy

# -----------------------------------------------------------------------------
# 6. CORE LOGIC
# -----------------------------------------------------------------------------

@spaces.GPU(duration=120)
def process_screenshot(input_numpy_image: np.ndarray, task: str, model_choice: str):
    if input_numpy_image is None: return "⚠️ Please upload an image.", None

    input_pil_image = array_to_image(input_numpy_image)
    orig_w, orig_h = input_pil_image.size
    
    actions = []
    raw_response = ""

    # -----------------------
    # MODEL: UI-TARS-1.5-7B
    # -----------------------
    if model_choice == "UI-TARS-1.5-7B":
        if model_x is None: return "Error: UI-TARS model failed to load on startup.", None
        print("Using UI-TARS Pipeline...")
        
        # 1. Smart Resize
        ip_params = get_image_proc_params(processor_x)
        resized_h, resized_w = smart_resize(
            input_pil_image.height, input_pil_image.width,
            factor=ip_params["patch_size"] * ip_params["merge_size"],
            min_pixels=ip_params["min_pixels"], max_pixels=ip_params["max_pixels"]
        )
        proc_image = input_pil_image.resize((resized_w, resized_h), Image.Resampling.LANCZOS)
        
        # 2. Prompt & Inputs
        messages = get_uitars_prompt(task, proc_image)
        text_prompt = processor_x.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor_x(text=[text_prompt], images=[proc_image], padding=True, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        # 3. Generate
        with torch.no_grad():
            generated_ids = model_x.generate(**inputs, max_new_tokens=128)
            
        generated_ids = [out_ids[len(in_seq):] for in_seq, out_ids in zip(inputs.get("input_ids"), generated_ids)]
        raw_response = processor_x.batch_decode(generated_ids, skip_special_tokens=True)[0]
        
        # 4. Parse & Rescale
        actions = parse_coordinate_response(raw_response)
        
        # Map coordinates from resized space back to original space
        scale_x = orig_w / resized_w
        scale_y = orig_h / resized_h
        for a in actions:
            a['x'] = int(a['x'] * scale_x)
            a['y'] = int(a['y'] * scale_y)

    # -----------------------
    # MODEL: Holo2-8B
    # -----------------------
    elif model_choice == "Holo2-8B":
        if model_h is None: return "Error: Holo2 model failed to load on startup.", None
        print("Using Holo2 Pipeline...")

        # 1. Smart Resize (Holo2 typically uses merge_size=1 or similar logic)
        ip_params = get_image_proc_params(processor_h)
        # Force merge_size to 1 if not detected (as per common practice for this model architecture variant)
        ms = ip_params.get("merge_size", 1) 
        
        resized_h, resized_w = smart_resize(
            input_pil_image.height, input_pil_image.width,
            factor=ip_params["patch_size"] * ms,
            min_pixels=ip_params["min_pixels"], max_pixels=ip_params["max_pixels"]
        )
        proc_image = input_pil_image.resize((resized_w, resized_h), Image.Resampling.LANCZOS)

        # 2. Prompt & Inputs
        messages = get_holo_prompt(proc_image, task)
        text_prompt = apply_chat_template_compat(processor_h, messages)
        
        # Holo2 / Qwen2-VL based inputs
        inputs = processor_h(text=[text_prompt], images=[proc_image], padding=True, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}

        # 3. Generate
        with torch.no_grad():
            generated_ids = model_h.generate(**inputs, max_new_tokens=128)
        
        # Trim input tokens
        generated_ids_trimmed = trim_generated(generated_ids, inputs)
        raw_response = batch_decode_compat(processor_h, generated_ids_trimmed, skip_special_tokens=True)[0]

        # 4. Parse & Rescale
        # Holo2 prompt asks for Click(x,y) similar to UI-TARS
        actions = parse_coordinate_response(raw_response)

        # Map coordinates from resized space back to original space
        scale_x = orig_w / resized_w
        scale_y = orig_h / resized_h
        for a in actions:
            a['x'] = int(a['x'] * scale_x)
            a['y'] = int(a['y'] * scale_y)

    # -----------------------
    # MODEL: Fara-7B
    # -----------------------
    else: 
        if model_v is None: return "Error: Fara model failed to load on startup.", None
        print("Using Fara Pipeline...")
        messages = get_fara_prompt(task, input_pil_image)
        text_prompt = processor_v.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = process_vision_info(messages)
        
        inputs = processor_v(
            text=[text_prompt],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt"
        )
        inputs = inputs.to(device)
        
        with torch.no_grad():
            generated_ids = model_v.generate(**inputs, max_new_tokens=512)
            
        generated_ids = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
        raw_response = processor_v.batch_decode(generated_ids, skip_special_tokens=True)[0]
        
        # Fara usually outputs exact pixels based on original image
        actions = parse_fara_response(raw_response)

    print(f"Raw Output: {raw_response}")
    print(f"Parsed Actions: {actions}")

    # 3. Visualize
    output_image = input_pil_image
    if actions:
        vis = create_localized_image(input_pil_image, actions)
        if vis: output_image = vis
            
    return raw_response, output_image

# -----------------------------------------------------------------------------
# 7. UI SETUP
# -----------------------------------------------------------------------------

with gr.Blocks(theme=steel_blue_theme, css=css) as demo:
    gr.Markdown("# **CUA GUI Agent 🖥️**", elem_id="main-title")
    gr.Markdown("Upload a screenshot, select a model, and provide a task. The model will determine the precise UI coordinates and actions.")

    with gr.Row():
        with gr.Column(scale=2):
            input_image = gr.Image(label="Upload Screenshot", height=500)
            
            with gr.Row():
                model_choice = gr.Radio(
                    choices=["Fara-7B", "UI-TARS-1.5-7B", "Holo2-8B"],
                    label="Select Model",
                    value="Fara-7B",
                    interactive=True
                )
            
            task_input = gr.Textbox(
                label="Task Instruction",
                placeholder="e.g. Input the server address readyforquantum.com...",
                lines=2
            )
            submit_btn = gr.Button("Analyze UI & Generate Action", variant="primary")

        with gr.Column(scale=3):
            output_image = gr.Image(label="Visualized Action Points", elem_id="out_img", height=500)
            output_text = gr.Textbox(label="Raw Model Output", lines=8, show_copy_button=True)

    submit_btn.click(
        fn=process_screenshot,
        inputs=[input_image, task_input, model_choice],
        outputs=[output_text, output_image]
    )
    
    gr.Examples(
        examples=[["./assets/google.png", "Search for 'Hugging Face'", "Fara-7B"]],
        inputs=[input_image, task_input, model_choice],
        label="Quick Examples"
    )

if __name__ == "__main__":
    demo.queue().launch()