app.py

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import re
import json
import html
import time
import tempfile

import spaces
import torch
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoProcessor, AutoModelForImageTextToText

MODEL_ID = "microsoft/Fara1.5-9B"

# Fara's native operating resolution. The model grounds coordinates best when
# the screenshot it sees is 1440x900, so we resize to fit inside this box and
# map predicted coordinates back onto the resized image we display. This is
# the same viewport the Fara harness (github.com/microsoft/fara) drives Playwright at.
FARA_W, FARA_H = 1440, 900

# Fara1.5 (built on Qwen3.5-VL) emits click/drag coordinates in a normalized
# 0-1000 space along each axis, NOT raw screen pixels. To draw the marker on the
# actual screenshot we must scale by (image_dimension / 1000) per axis. The
# scale is inclusive of both endpoints, so we divide by (COORD_SPACE - 1).
COORD_SPACE = 1000

# The verbatim system prompt Fara1.5-9B was trained against (from the model card).
FARA_SYSTEM_PROMPT = (
    "You are Fara, a computer use agent (CUA) specialized for web browsers. "
    "You are developed by Microsoft AI Frontiers. You assist users with "
    "completing and automating tasks that require the use of a web browser.\n\n"
    "The model was trained in the timeframe of January - April 2026. You can "
    "effectively perform tasks even beyond this range by accessing the web "
    "browser and using the latest information on the live web. But your "
    "knowledge cutoff is limited to early 2026, so you may not be aware of "
    "events or developments that occurred after that time, without explicitly "
    "browsing and searching for latest information on the web.\n\n"
    "This edition of the model was trained using SFT on top of Qwen3.5-9B, "
    "using a synthetic data mixture generated and developed by Microsoft AI "
    "Frontiers.\n\n"
    "A critical point is a situation where we must pause and request "
    "information or confirmation from the user before proceeding. There are "
    "three types:\n\n"
    "Case 1: Missing User Information \u2014 The task requires personal "
    "information that the user has not provided (e.g., email, phone number, "
    "address, payment details). Never fabricate or assume personal "
    "information. Fill in only what the user has explicitly provided, then "
    "pause and ask for any missing required fields.\n\n"
    "Case 2: Underspecified Task \u2014 The task description is ambiguous or "
    "missing details needed to make a decision at the current step. Pause and "
    "ask for clarification.\n\n"
    "Case 3: Irreversible Action \u2014 We are about to perform an action that "
    "cannot be undone (e.g., submitting a form, completing a purchase, sending "
    "a message, deleting data). If the user explicitly authorized the action, "
    "proceed. Otherwise, stop and ask for confirmation.\n\n"
    "Only stop at a critical point if (1) required information is missing, (2) "
    "the task is ambiguous, OR (3) an irreversible action lacks explicit user "
    "authorization."
)

# The computer_use tool schema. Passing this as `tools` to the chat template
# makes the model emit the Qwen-style XML function-call block that Fara is
# trained to produce.
COMPUTER_USE_TOOL = {
    "type": "function",
    "function": {
        "name": "computer_use",
        "description": (
            "Use a mouse and keyboard to interact with a web browser, and take "
            "screenshots. This is an interface to a browser GUI. Actions are "
            "grounded in pixel coordinates on a 1440x900 screen with origin at "
            "the top-left corner."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "description": (
                        "The action to perform. One of: left_click, right_click, "
                        "double_click, triple_click, mouse_move, left_click_drag, "
                        "type, key, scroll, hscroll, visit_url, history_back, "
                        "web_search, pause_and_memorize_fact, ask_user_question, "
                        "wait, terminate."
                    ),
                },
                "coordinate": {
                    "type": "array",
                    "description": "[x, y] pixel coordinate for click / move / drag actions.",
                    "items": {"type": "integer"},
                },
                "text": {"type": "string", "description": "Text to type."},
                "key": {"type": "string", "description": "Key or key-combo to press."},
                "amount": {"type": "integer", "description": "Scroll amount."},
                "url": {"type": "string", "description": "URL to visit."},
                "query": {"type": "string", "description": "Web search query."},
                "fact": {"type": "string", "description": "Fact to memorize."},
                "question": {"type": "string", "description": "Question to ask the user."},
                "seconds": {"type": "integer", "description": "Seconds to wait."},
                "answer": {"type": "string", "description": "Final answer on terminate."},
            },
            "required": ["action"],
        },
    },
}

print(f"Loading {MODEL_ID} ...")
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()
print("Model loaded.")


# --------------------------------------------------------------------------- #
# Tool-call parsing
# --------------------------------------------------------------------------- #
def _coerce(v):
    v = v.strip()
    if re.fullmatch(r"-?\d+", v):
        return int(v)
    try:
        return json.loads(v)
    except Exception:
        return v


ACTION_NAMES = (
    "left_click_drag", "left_click", "right_click", "double_click",
    "triple_click", "mouse_move", "type", "key", "hscroll", "scroll",
    "visit_url", "history_back", "web_search", "pause_and_memorize_fact",
    "ask_user_question", "wait", "terminate",
)


# Sentinel tokens that delimit the tool-call block in the model's output.
# We use these instead of literal XML to avoid confusing the heredoc.
TOOL_OPEN = "\u003ctool\u003e"   # <tool>
TOOL_CLOSE = "\u003c/tool\u003e"  # </tool>

# The actual delimiters the model emits (Qwen-hermes style):
#   <function=computer_use> ... </function>  wrapped in <tool> ... </tool>
# But Fara also emits a JSON-object form:
#   {"name": "computer_use", "arguments": {...}}
# We handle both.


def _extract_coord_anywhere(text):
    """Find the first [x, y] pair near a 'coordinate' mention, else any pair."""
    m = re.search(
        r"coordinate[^\[\]]{0,20}\[\s*(-?\d+)\s*,\s*(-?\d+)\s*\]",
        text, re.IGNORECASE | re.DOTALL,
    )
    if not m:
        m = re.search(r"\[\s*(-?\d+)\s*,\s*(-?\d+)\s*\]", text)
    if m:
        return int(m.group(1)), int(m.group(2))
    return None


def parse_tool_call(text):
    """Parse Fara's tool-call output into (name, args_dict) or None.

    Fara emits one of two formats:
      1. Qwen-hermes style:
         <tool><function=computer_use>
           <parameter=action>left_click</parameter>
           <parameter=coordinate>[720, 450]</parameter>
         </function></tool>
      2. JSON object:
         {"name": "computer_use", "arguments": {"action": "left_click", "coordinate": [720, 450]}}

    In practice the formatting is sometimes malformed, so this parser is
    deliberately forgiving.
    """
    # Try JSON form first
    jm = re.search(r'\{[^{}]*"name"\s*:\s*"computer_use"[^{}]*\}', text, re.DOTALL)
    if jm:
        try:
            obj = json.loads(jm.group(0))
            jargs = obj.get("arguments", obj.get("parameters", {})) or {}
            if jargs.get("action"):
                return "computer_use", jargs
        except Exception:
            pass

    # Try a broader JSON search
    jm2 = re.search(r'\{.*?"action".*?\}', text, re.DOTALL)
    if jm2:
        try:
            obj = json.loads(jm2.group(0))
            jargs = obj.get("arguments", obj.get("parameters", obj)) or {}
            if isinstance(jargs, dict) and jargs.get("action"):
                return "computer_use", jargs
        except Exception:
            pass

    # Qwen-hermes XML form: find the block between <tool> and </tool>
    # or between <function= and </function>
    block = text
    m_tool = re.search(r"<tool>(.*?)(?:</tool>|$)", text, re.DOTALL)
    if m_tool:
        block = m_tool.group(1)

    args = {}
    # Well-formed <parameter=name>value</parameter> pairs
    for pm in re.finditer(
        r"<parameter=([^>\s]+)\s*>(.*?)</parameter>", block, re.DOTALL
    ):
        args[pm.group(1).strip()] = _coerce(pm.group(2))

    # Function / action name
    name = "computer_use"
    fn = re.search(r"<function=([^>\s]+)\s*>", block)
    if fn:
        name = fn.group(1).strip()
    else:
        fn2 = re.search(r"<(computer_use)\s*>", block)
        if fn2:
            name = fn2.group(1)

    # Recover the action verb if it wasn't a clean parameter
    action = args.get("action")
    if isinstance(action, str):
        tok = re.match(r"\s*([a-z_]+)", action)
        if tok:
            args["action"] = tok.group(1)
    if not args.get("action") or args.get("action") not in ACTION_NAMES:
        for a in ACTION_NAMES:
            if re.search(r"\b" + re.escape(a) + r"\b", block):
                args["action"] = a
                break

    # Recover a coordinate if it wasn't captured as a clean parameter
    coord = None
    if "coordinate" in args and isinstance(args["coordinate"], (list, tuple)):
        pass
    else:
        coord = _extract_coord_anywhere(block)
        if coord is not None:
            args["coordinate"] = [coord[0], coord[1]]

    if not args.get("action") and "coordinate" not in args:
        return None
    return name, args


def _get_coord(args):
    """Extract an (x, y) integer coordinate from parsed args, if present."""
    for key in ("coordinate", "coord", "position"):
        if key in args:
            c = args[key]
            if isinstance(c, (list, tuple)) and len(c) >= 2:
                try:
                    return int(c[0]), int(c[1])
                except Exception:
                    return None
            if isinstance(c, str):
                cc = _extract_coord_anywhere(c)
                if cc:
                    return cc
    if "x" in args and "y" in args:
        try:
            return int(args["x"]), int(args["y"])
        except Exception:
            return None
    return None


def _norm_to_pixel(coord, width, height):
    """Map a Fara normalized (0-1000) coordinate onto real image pixels.

    Fara1.5 grounds click/drag targets in a 0-1000 normalized space per axis
    rather than raw pixels, so a value of 500 means "halfway across". We scale
    each axis by (dimension - 1) / (COORD_SPACE - 1) so that 0 -> 0 and 1000 ->
    the last pixel, then clamp to the image bounds.
    """
    if coord is None:
        return None
    nx, ny = coord
    px = round(nx * (width - 1) / (COORD_SPACE - 1))
    py = round(ny * (height - 1) / (COORD_SPACE - 1))
    px = max(0, min(width - 1, px))
    py = max(0, min(height - 1, py))
    return px, py


def _rescale_coord_args(args, width, height):
    """Rewrite any normalized coordinate in args to pixel space in place."""
    coord = _get_coord(args)
    if coord is not None:
        args["coordinate"] = list(_norm_to_pixel(coord, width, height))
        args.pop("coord", None)
        args.pop("position", None)
        args.pop("x", None)
        args.pop("y", None)
    return args


# --------------------------------------------------------------------------- #
# Visual overlay
# --------------------------------------------------------------------------- #
def _font(size):
    for p in (
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "DejaVuSans-Bold.ttf",
    ):
        try:
            return ImageFont.truetype(p, size)
        except Exception:
            continue
    return ImageFont.load_default()


ACCENT = (220, 38, 38)      # red marker
ACCENT2 = (37, 99, 235)     # blue for drag origin


def draw_overlay(img, action, args, step_no=None):
    """Draw the predicted action onto a copy of the (already Fara-sized) image.

    When step_no is given, the click marker is numbered so a sequence of
    actions leaves a readable trail.
    """
    canvas = img.convert("RGB").copy()
    d = ImageDraw.Draw(canvas, "RGBA")
    coord = _get_coord(args)
    label = action

    if coord is not None:
        x, y = coord
        r = 26
        # halo + crosshair + center dot
        d.ellipse([x - r, y - r, x + r, y + r], fill=(220, 38, 38, 60),
                  outline=ACCENT, width=3)
        d.line([x - r - 10, y, x + r + 10, y], fill=ACCENT, width=2)
        d.line([x, y - r - 10, x, y + r + 10], fill=ACCENT, width=2)
        d.ellipse([x - 5, y - 5, x + 5, y + 5], fill=ACCENT)

        if step_no is not None:
            # numbered badge above the marker
            bf = _font(20)
            btxt = str(step_no)
            bb = d.textbbox((0, 0), btxt, font=bf)
            bw, bh = bb[2] - bb[0], bb[3] - bb[1]
            bx, by = x + r - 4, y - r - bh - 8
            d.ellipse([bx - 4, by - 4, bx + bw + 8, by + bh + 8],
                      fill=ACCENT, outline=(255, 255, 255), width=2)
            d.text((bx + 1, by - 1), btxt, font=bf, fill=(255, 255, 255))

        if action in ("left_click_drag", "mouse_move"):
            # draw an arrow from an implied origin (top-left offset) to target
            ox, oy = max(0, x - 160), max(0, y - 110)
            d.line([ox, oy, x, y], fill=ACCENT2, width=4)
            d.ellipse([ox - 6, oy - 6, ox + 6, oy + 6], fill=ACCENT2)

        # coordinate label pill
        f = _font(22)
        prefix = f"{step_no}. " if step_no is not None else ""
        txt = f"{prefix}{label}  ({x}, {y})"
        tb = d.textbbox((0, 0), txt, font=f)
        tw, th = tb[2] - tb[0], tb[3] - tb[1]
        ly = y + r + 14
        if ly + th + 12 > canvas.height:
            ly = y - r - th - 26
        lx = min(max(4, x - tw // 2 - 10), canvas.width - tw - 24)
        d.rounded_rectangle([lx, ly, lx + tw + 20, ly + th + 14], radius=8,
                            fill=(17, 24, 39, 235))
        d.text((lx + 10, ly + 5), txt, font=f, fill=(255, 255, 255))
    else:
        # No coordinate (type/key/scroll/terminate/...). Show a top banner.
        f = _font(24)
        detail = ""
        for k in ("text", "key", "url", "query", "answer", "question",
                  "fact", "amount", "seconds"):
            if k in args and args[k] not in (None, ""):
                detail = f"{k}: {args[k]}"
                break
        prefix = f"{step_no}. " if step_no is not None else ""
        txt = f"{prefix}{label}" + (f"   \u2013   {detail}" if detail else "")
        txt = txt if len(txt) < 90 else txt[:87] + "..."
        tb = d.textbbox((0, 0), txt, font=f)
        tw = tb[2] - tb[0]
        d.rectangle([0, 0, canvas.width, 46], fill=(17, 24, 39, 235))
        d.text((max(14, (canvas.width - tw) // 2), 10), txt, font=f,
               fill=(255, 255, 255))
    return canvas


def fit_to_fara(img):
    """Resize the screenshot to fit inside 1440x900, preserving aspect ratio."""
    img = img.convert("RGB")
    img = img.copy()
    img.thumbnail((FARA_W, FARA_H), Image.LANCZOS)
    return img


def format_action_md(name, args):
    lines = [f"**Action:** `{name}` \u2192 `{args.get('action', name)}`", ""]
    for k, v in args.items():
        if k == "action":
            continue
        lines.append(f"- **{k}**: `{v}`")
    if len(lines) == 2:
        lines.append("_(no arguments)_")
    return "\n".join(lines)


# --------------------------------------------------------------------------- #
# Inference
# --------------------------------------------------------------------------- #
@spaces.GPU(duration=60)
def predict(screenshot, task, max_new_tokens=512):
    """Predict the next grounded action for a screenshot + task goal.

    Upload a browser screenshot and describe what the agent should do.
    Fara1.5-9B will analyze the screenshot and predict the next action
    (click, type, scroll, etc.) with pixel-level coordinate grounding.
    The predicted action is visualized on the image with a marker.

    Args:
        screenshot: A browser screenshot (PNG/JPG). Will be resized to 1440x900.
        task: The natural-language goal the agent should work toward.
        max_new_tokens: Generation budget for the model's response.

    Returns:
        (annotated_image, action_markdown, raw_output, action_json)
    """
    if screenshot is None:
        raise gr.Error("Please upload a screenshot first.")
    if not task or not task.strip():
        raise gr.Error("Please describe the task / goal for the agent.")

    # Resize to Fara's native 1440x900 viewport
    view = fit_to_fara(screenshot)

    # Build the chat messages
    messages = [
        {"role": "system", "content": FARA_SYSTEM_PROMPT},
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": task.strip()},
            ],
        },
    ]

    prompt = processor.apply_chat_template(
        messages,
        tools=[COMPUTER_USE_TOOL],
        add_generation_prompt=True,
        tokenize=False,
    )
    inputs = processor(
        text=[prompt], images=[view], return_tensors="pt"
    ).to(model.device)

    t0 = time.perf_counter()
    with torch.inference_mode():
        gen = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
            temperature=None,
            top_p=None,
            top_k=None,
        )
    elapsed = time.perf_counter() - t0
    out = gen[0][inputs["input_ids"].shape[-1]:]
    decoded = processor.decode(out, skip_special_tokens=True).strip()
    print(f"Fara1.5-9B generated in {elapsed:.1f}s, {len(out)} tokens")

    # Parse the tool call
    parsed = parse_tool_call(decoded)

    if parsed is None:
        # No structured tool call parsed; return raw output
        annotated = view
        action_md = (
            "### No structured tool call parsed\n\n"
            f"**Raw model output:**\n\n```\n{html.escape(decoded[:2000])}\n```"
        )
        return annotated, action_md, decoded, {"parsed": False, "raw": decoded}

    name, args = parsed
    # Convert normalized 0-1000 coordinates to pixel coordinates
    args = _rescale_coord_args(args, view.width, view.height)
    action = args.get("action", name)

    # Draw the overlay on the resized image
    annotated = draw_overlay(view, action, args, step_no=1)

    # Build the markdown output
    reasoning = decoded.rsplit("<tool>", 1)[0].strip() if "<tool>" in decoded else ""
    if not reasoning:
        # Try splitting on the JSON form
        reasoning = re.split(r'\{[^{}]*"name"\s*:\s*"computer_use"', decoded)[0].strip()

    reasoning_escaped = html.escape(reasoning[:1000]) if reasoning else "_(none)_"

    action_md = (
        f"### Predicted Action\n\n"
        f"{format_action_md(name, args)}\n\n"
        f"**Inference time:** {elapsed:.1f}s\n\n"
        f"---\n\n"
        f"### Model Reasoning\n\n"
        f"{reasoning_escaped}"
    )

    action_json = {
        "parsed": True,
        "tool": name,
        "action": action,
        "args": {k: v for k, v in args.items() if k != "action"},
        "coordinate": list(_get_coord(args)) if _get_coord(args) else None,
        "reasoning": reasoning,
        "inference_seconds": round(elapsed, 2),
    }

    return annotated, action_md, decoded, action_json


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Fara1.5-9B Computer Use Agent") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🖱️ Fara1.5-9B — Computer Use Agent (Visual Grounding)

            [Fara1.5-9B](https://huggingface.co/microsoft/Fara1.5-9B) by **Microsoft
            Research AI Frontiers** is a 9B vision-language **computer use agent** that
            predicts pixel-level actions from browser screenshots.

            **How it works:** Upload a screenshot of a web page and describe the task.
            Fara analyzes the screenshot and predicts the next action — a click at a
            specific pixel coordinate, text to type, a scroll, etc. — visualized with
            a marker on the image.

            The model grounds click/drag targets in a normalized 0–1000 coordinate
            space, mapped onto a 1440×900 viewport. Coordinates are rescaled to match
            the actual screenshot dimensions for the overlay.
            """
        )
        with gr.Row(equal_height=False):
            # LEFT: inputs
            with gr.Column(scale=5):
                screenshot = gr.Image(
                    label="📷 Browser screenshot",
                    type="pil",
                    height=400,
                )
                task = gr.Textbox(
                    label="🎯 Task / goal",
                    placeholder="e.g. Click the 'Add to cart' button",
                    lines=2,
                    value="Click the search box at the top of the page, then type 'wireless headphones' and press Enter to search.",
                )
                run = gr.Button("▶ Predict action", variant="primary", size="lg")
                with gr.Accordion("Advanced settings", open=False):
                    max_new_tokens = gr.Slider(
                        64, 2048, value=512, step=32,
                        label="Max new tokens",
                        info="Generation budget for the model's response",
                    )

            # RIGHT: outputs
            with gr.Column(scale=6):
                out_image = gr.Image(
                    label="🖼️ Annotated screenshot with predicted action",
                    height=400,
                    type="pil",
                )
                out_action = gr.Markdown(
                    value="_Upload a screenshot and describe a task to see the predicted action._",
                    label="Predicted action + reasoning",
                )
                out_raw = gr.Textbox(
                    label="Raw model output",
                    lines=6,
                    interactive=False,
                    visible=False,
                )
                out_json = gr.JSON(
                    label="Structured action (API / MCP)",
                    visible=False,
                )

        gr.Markdown("### Try an example — real screenshots with grounded actions")
        gr.Examples(
            examples=[
                ["examples/checkout_form.png",
                 "Click the 'Place Order' button at the bottom of the checkout form to submit the order."],
                ["examples/search_page.png",
                 "Click the search box at the top of the page, then type 'wireless headphones' and press Enter to search."],
            ],
            inputs=[screenshot, task],
            outputs=[out_image, out_action, out_raw, out_json],
            fn=predict,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        predict,
        inputs=[screenshot, task, max_new_tokens],
        outputs=[out_image, out_action, out_raw, out_json],
        api_name="predict",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)