app.py

import os

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

import re
import json
import tempfile

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

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

# Fara's official system prompt (verbatim from the model card) — best grounding.
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.

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.

This edition of the model was trained using SFT on top of Qwen3.5-27B, using a synthetic data mixture generated and developed by Microsoft AI Frontiers.

A critical point is a situation where we must pause and request information or confirmation from the user before proceeding. There are three types:

Case 1: Missing User Information — 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.

Case 2: Underspecified Task — The task description is ambiguous or missing details needed to make a decision at the current step. Pause and ask for clarification.

Case 3: Irreversible Action — 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.

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."""

# Compact but faithful `computer_use` tool schema so the model emits <tool_call> blocks.
COMPUTER_USE_TOOL = {
    "type": "function",
    "function": {
        "name": "computer_use",
        "description": (
            "Interact with a web browser to complete the user's task. Observe the "
            "screenshot and emit a single next action. Coordinates are pixel (x, y) "
            "positions on the current screenshot."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": [
                        "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",
                    ],
                    "description": "The action to perform.",
                },
                "coordinate": {
                    "type": "array",
                    "items": {"type": "integer"},
                    "description": "Target [x, y] pixel coordinate for click/move/drag actions.",
                },
                "text": {
                    "type": "string",
                    "description": "Text to type, key combo to press, URL to visit, "
                    "search query, question, or final answer, depending on the action.",
                },
                "scroll_direction": {
                    "type": "string",
                    "enum": ["up", "down", "left", "right"],
                    "description": "Direction to scroll for scroll/hscroll actions.",
                },
                "scroll_amount": {
                    "type": "integer",
                    "description": "Number of scroll clicks for scroll/hscroll actions.",
                },
            },
            "required": ["action"],
        },
    },
}

# 4-bit quantization keeps the 27B model comfortably within ZeroGPU 'large' (48 GB).
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="cuda",  # bitsandbytes loader is ZeroGPU-aware
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).eval()

# Actions that carry a click/target coordinate we can visualize.
_COORD_ACTIONS = {
    "left_click", "right_click", "double_click", "triple_click",
    "mouse_move", "left_click_drag",
}


_KNOWN_KEYS = ["action", "coordinate", "coord", "text", "url", "scroll_direction", "scroll_amount"]


def _parse_tool_call(text: str):
    """Extract (action_dict, raw_body) from a Fara <tool_call> block.

    Fara's tool-call output is not always well-formed: it may be clean JSON
    (`{"name": ..., "arguments": {...}}`), the chat-template
    <function=...><parameter=...></parameter></function> form, or a mangled
    hybrid (e.g. `<parameter=action": "left_click", "coordinate": [286, 88]}}`).
    We try the clean forms first, then fall back to key-wise regex scraping of
    the body, which is robust to the hybrid mess.
    """
    m = re.search(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", text, re.DOTALL)
    body = m.group(1).strip() if m else text.strip()

    # 1. Clean JSON form.
    try:
        obj = json.loads(body)
        args = obj.get("arguments", obj) if isinstance(obj, dict) else obj
        if isinstance(args, str):
            args = json.loads(args)
        if isinstance(args, dict) and args:
            return args, body
    except Exception:
        pass

    # 2. Clean <parameter=key>value</parameter> form.
    args = {}
    for pm in re.finditer(r"<parameter=([A-Za-z_]+)>\s*(.*?)\s*</parameter>", body, re.DOTALL):
        key, val = pm.group(1).strip(), pm.group(2).strip()
        try:
            args[key] = json.loads(val)
        except Exception:
            args[key] = val
    if args:
        return args, body

    # 3. Robust key-wise regex scrape (handles the mangled hybrid format).
    scraped = {}
    for key in _KNOWN_KEYS:
        # "key": value    or    key": value    or    <parameter=key": value
        km = re.search(rf'{key}["\s]*[:=]\s*', body)
        if not km:
            continue
        rest = body[km.end():].lstrip()
        if rest.startswith("["):
            # array value (coordinate)
            am = re.match(r"\[[^\]]*\]", rest)
            if am:
                try:
                    scraped[key] = json.loads(am.group(0))
                    continue
                except Exception:
                    pass
        if rest.startswith('"'):
            sm = re.match(r'"([^"]*)"', rest)
            if sm:
                scraped[key] = sm.group(1)
                continue
        # bare token (int / word) up to a delimiter
        tm = re.match(r"([^\s,\}\]\"<]+)", rest)
        if tm:
            v = tm.group(1)
            try:
                scraped[key] = json.loads(v)
            except Exception:
                scraped[key] = v
    return (scraped or None), body


def _extract_coordinate(args):
    """Return an (x, y) tuple from parsed args, or None."""
    if not args:
        return None
    coord = args.get("coordinate") or args.get("coord")
    if isinstance(coord, str):
        try:
            coord = json.loads(coord)
        except Exception:
            nums = re.findall(r"-?\d+", coord)
            coord = [int(n) for n in nums[:2]] if len(nums) >= 2 else None
    if isinstance(coord, (list, tuple)) and len(coord) >= 2:
        try:
            return int(coord[0]), int(coord[1])
        except Exception:
            return None
    return None


def _annotate(image: Image.Image, xy, label):
    """Draw a crosshair marker + label at (x, y) on a copy of the image."""
    img = image.convert("RGB").copy()
    draw = ImageDraw.Draw(img)
    x, y = xy
    r = 22
    # Outer glow + inner ring for visibility on any background.
    draw.ellipse([x - r - 3, y - r - 3, x + r + 3, y + r + 3], outline=(255, 255, 255), width=3)
    draw.ellipse([x - r, y - r, x + r, y + r], outline=(220, 30, 60), width=5)
    draw.line([x - r - 12, y, x - r + 6, y], fill=(220, 30, 60), width=4)
    draw.line([x + r - 6, y, x + r + 12, y], fill=(220, 30, 60), width=4)
    draw.line([x, y - r - 12, x, y - r + 6], fill=(220, 30, 60), width=4)
    draw.line([x, y + r - 6, x, y + r + 12], fill=(220, 30, 60), width=4)
    draw.ellipse([x - 3, y - 3, x + 3, y + 3], fill=(220, 30, 60))

    try:
        font = ImageFont.truetype(
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20
        )
    except Exception:
        font = ImageFont.load_default()
    caption = f"{label} ({x}, {y})"
    tb = draw.textbbox((0, 0), caption, font=font)
    tw, th = tb[2] - tb[0], tb[3] - tb[1]
    ty = y + r + 12
    tx = max(4, min(x - tw // 2, img.width - tw - 8))
    if ty + th + 8 > img.height:
        ty = y - r - 12 - th - 8
    draw.rectangle([tx - 4, ty - 4, tx + tw + 4, ty + th + 8], fill=(220, 30, 60))
    draw.text((tx, ty), caption, fill=(255, 255, 255), font=font)
    return img


def _format_action(args):
    """Human-readable one-liner describing the predicted action."""
    if not args:
        return "No structured action parsed."
    action = args.get("action", "unknown")
    parts = [f"**Action:** `{action}`"]
    xy = _extract_coordinate(args)
    if xy:
        parts.append(f"**Coordinate:** `({xy[0]}, {xy[1]})`")
    if args.get("text"):
        parts.append(f"**Text:** {args['text']}")
    if args.get("url"):
        parts.append(f"**URL:** {args['url']}")
    if args.get("scroll_direction"):
        amt = args.get("scroll_amount", "")
        parts.append(f"**Scroll:** {args['scroll_direction']} {amt}".strip())
    return "  \n".join(parts)


def _estimate_duration(screenshot, task, max_new_tokens=512, *args, **kwargs):
    try:
        n = int(max_new_tokens)
    except Exception:
        n = 512
    # ~15s image encode + load overhead, ~0.08s/token generated; capped.
    return int(min(110, 20 + n * 0.08))


@spaces.GPU(duration=_estimate_duration)
def predict_action(screenshot, task, max_new_tokens=512):
    """Predict the next browser action a computer-use agent should take.

    Given a browser screenshot and a natural-language task, Microsoft's
    Fara-1.5-27B computer-use agent returns its chain-of-thought reasoning
    and a structured `computer_use` tool call (click coordinates, typing,
    navigation, etc.). Click/move targets are drawn on the screenshot.

    Args:
        screenshot: A browser screenshot (ideally 1440x900) to act on.
        task: The natural-language goal for the agent to work toward.
        max_new_tokens: Maximum number of tokens to generate.

    Returns:
        annotated_image, raw_model_output, parsed_action_summary
    """
    if screenshot is None:
        raise gr.Error("Please provide a browser screenshot.")
    if not task or not task.strip():
        raise gr.Error("Please describe the task for the agent.")

    image = screenshot.convert("RGB")

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

    inputs = processor.apply_chat_template(
        messages,
        tools=[COMPUTER_USE_TOOL],
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
            temperature=None,
            top_p=None,
            top_k=None,
        )
    out_ids = generated[0][inputs["input_ids"].shape[-1]:]
    text = processor.decode(out_ids, skip_special_tokens=True).strip()

    args, _ = _parse_tool_call(text)
    xy = _extract_coordinate(args)
    action_name = (args or {}).get("action", "action")

    annotated = image
    if xy and (args.get("action") in _COORD_ACTIONS or args.get("coordinate")):
        annotated = _annotate(image, xy, action_name)

    summary = _format_action(args)
    return annotated, text, summary


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

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🖥️ Fara-1.5-27B · Computer-Use Agent
            Microsoft's **[Fara-1.5-27B](https://huggingface.co/microsoft/Fara1.5-27B)** is a
            vision-only web computer-use agent (SFT on Qwen3.5-27B). Give it a **browser
            screenshot** and a **task**; it returns its reasoning plus a structured
            `computer_use` action. Click/move targets are drawn on the screenshot.

            *Single-step demo — predicts the next action for one screenshot. Runs 4-bit
            quantized on ZeroGPU. Best grounding at 1440×900 screenshots.*
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                screenshot = gr.Image(
                    label="Browser screenshot", type="pil", height=420
                )
                task = gr.Textbox(
                    label="Task",
                    placeholder="e.g. Search Wikipedia for 'Alan Turing'",
                    lines=2,
                )
                run = gr.Button("Predict next action", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    max_new_tokens = gr.Slider(
                        64, 1024, value=512, step=32, label="Max new tokens"
                    )
            with gr.Column(scale=1):
                out_image = gr.Image(label="Predicted target", height=420)
                out_action = gr.Markdown(label="Parsed action")
                out_text = gr.Textbox(
                    label="Raw model output (reasoning + tool call)", lines=12
                )

        gr.Examples(
            examples=[
                ["example_wikipedia_home.png", "Search Wikipedia for 'Alan Turing'"],
                ["example_hackernews.png", "Open the comments for the top story"],
                ["example_wikipedia.png", "Click the English language link"],
            ],
            inputs=[screenshot, task],
            outputs=[out_image, out_text, out_action],
            fn=predict_action,
            cache_examples=True,
            cache_mode="lazy",
        )

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

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