# app.py — Fénix Grounding (ShowUI-2B en ZeroGPU)
# Recibe una captura + una descripción de qué clicar, y devuelve las
# coordenadas normalizadas 0-1 del elemento. El agente lo llama SOLO
# cuando necesita acertar un clic fino sobre un icono/botón.
import spaces
import ast
import torch
import gradio as gr
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

MODEL_ID = "showlab/ShowUI-2B"

_SYSTEM = (
    "Based on the screenshot of the page, I give a text description and you give its "
    "corresponding location. The coordinate represents a clickable location [x, y] "
    "for an element, which is a relative coordinate on the screenshot, scaled from 0 to 1."
)

min_pixels = 256 * 28 * 28
max_pixels = 1344 * 28 * 28

# Se carga en CPU al arrancar; ZeroGPU asigna la GPU solo dentro de @spaces.GPU
model = Qwen2VLForConditionalGeneration.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16
)
processor = AutoProcessor.from_pretrained(
    "Qwen/Qwen2-VL-2B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels
)


@spaces.GPU(duration=60)
def localizar(imagen, consulta):
    """imagen: ruta de la captura. consulta: qué elemento clicar (mejor en inglés)."""
    if imagen is None or not (consulta or "").strip():
        return "ERROR: falta imagen o consulta"
    model.to("cuda")
    messages = [{
        "role": "user",
        "content": [
            {"type": "text", "text": _SYSTEM},
            {"type": "image", "image": imagen,
             "min_pixels": min_pixels, "max_pixels": max_pixels},
            {"type": "text", "text": consulta},
        ],
    }]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(text=[text], images=image_inputs, videos=video_inputs,
                       padding=True, return_tensors="pt").to("cuda")
    generated = model.generate(**inputs, max_new_tokens=128)
    trimmed = [out[len(inp):] for inp, out in zip(inputs.input_ids, generated)]
    out_text = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    try:
        coords = ast.literal_eval(out_text)
        return f"{float(coords[0]):.4f},{float(coords[1]):.4f}"
    except Exception:
        return f"ERROR: {out_text}"


with gr.Blocks(title="Fénix Grounding") as demo:
    gr.Markdown("## 🎯 Fénix Grounding — ShowUI-2B\n"
                "Sube una captura y describe el elemento; devuelve x,y (0-1).")
    with gr.Row():
        img = gr.Image(label="Captura", type="filepath")
        with gr.Column():
            q = gr.Textbox(label="Qué clicar (mejor en inglés)",
                           placeholder="Notepad icon in the taskbar")
            btn = gr.Button("Localizar", variant="primary")
            out = gr.Textbox(label="Coordenadas x,y (0-1)")
    btn.click(localizar, [img, q], out, api_name="localizar")

demo.launch()