app.py

import os

# Triton saves its kernel tuning to disk, so a new ZeroGPU worker reuses what the start-up warm-up measured
# instead of re-tuning (about 15 s on its first step). Must be set before torch/triton are imported.
os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")

from deskforge import (DEFAULT_FAMILY, FAMILIES, annotate, parse_action, predict,  # imports spaces first
                       to_pixels)
import web_models  # the web agent's planner (Qwen3.5-9B), resident next to the duel models

import base64
import html
import io
import random
import tempfile
import threading
import time
from pathlib import Path

import gradio as gr
from PIL import Image

import setup_bar
import gpu_bridge
import shot_tool
import web_tab
from desktop_local import THEME_PRESETS, LocalDesktop, describe_theme, random_theme
from duel import CHALLENGES, TOTALS, play
from setup_bar import LOOK_SHUFFLE, RANDOM_TASK
from web_agent import WebAgent

MAX_STEPS = 8
HERE = Path(__file__).parent

desktop = LocalDesktop().start()
desktop.reset("Calculator (Galculator)")
web_agent = WebAgent(web_models.decide)  # a headless Chromium, separate from the Xfce desktop

_shot_lock = threading.Lock()
_last_shot = (0.0, None)


def live_view():
    """Current desktop frame, shared by every viewer (at most ~3 grabs/second)."""
    global _last_shot
    with _shot_lock:
        t, img = _last_shot
        if img is None or time.time() - t > 0.3:
            img = desktop.screenshot()
            _last_shot = (time.time(), img)
        return img


# ---------------------------------------------------------------------------
# Static pieces: top bar (with the view switch) and the welcome message
# ---------------------------------------------------------------------------

LOGO = """<svg viewBox="0 0 24 24" width="15" height="15" aria-hidden="true"><rect x="2.5" y="4" width="13" height="10" rx="2"
fill="none" stroke="#fff" stroke-width="1.8"/><rect x="8.5" y="10" width="13" height="10" rx="2" fill="#fff" stroke="#fff"
stroke-width="1.8"/><path d="M8.5 13h13" stroke="#D2603A" stroke-width="1.5"/></svg>"""

def _icon(body):
    return (f'<svg class="df-view-ico" viewBox="0 0 24 24" width="16" height="16" fill="none" stroke="currentColor" '
            f'stroke-width="1.8" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true">{body}</svg>')


VIEW_ICON = {
    "web": _icon('<circle cx="12" cy="12" r="9"/><path d="M3 12h18M12 3a14 14 0 0 1 0 18M12 3a14 14 0 0 0 0 18"/>'),  # globe
    "live": _icon('<rect x="2.5" y="4" width="19" height="13" rx="2"/><path d="M8 21h8M12 17v4"/>'
                  '<path d="M8.5 9.5l2 1.5-2 1.5M15.5 9.5l-2 1.5 2 1.5"/>'),  # a desktop with two facing cursors
    "shot": _icon('<rect x="3" y="5" width="18" height="14" rx="2"/><circle cx="8.5" cy="10" r="1.6"/>'
                  '<path d="M21 15.5l-4.5-4.5L7 19"/>'),  # a picture
}

VIEW_HINT = {  # a few words on each tab, after a "|"
    "web": "any task, real browser",
    "live": "DeskForge vs. its base",
    "shot": "see where it clicks",
}
VIEW_ABOUT = {  # the full sentence, as each tab's tooltip
    "web": "Give it any task on the web. A large AI model plans each step, and DeskForge finds exactly where to click.",
    "live": "DeskForge races the model it was trained from on a real Linux desktop. See who clicks the right button.",
    "shot": "Upload a screenshot and write an instruction to see where DeskForge would click.",
}

TOPBAR = f"""
<div class="df-top">
  <div class="df-brand"><div class="df-logo">{LOGO}</div><div class="df-word">DeskForge</div></div>
  <div class="df-sep"></div>
  <div class="df-views">
    <button type="button" class="df-view on" data-view="web" title="{VIEW_ABOUT['web']}">{VIEW_ICON['web']}Web agent<span class="df-view-sep">|</span><span class="df-view-hint">{VIEW_HINT['web']}</span></button>
    <button type="button" class="df-view" data-view="live" title="{VIEW_ABOUT['live']}">{VIEW_ICON['live']}Live duel<span class="df-view-sep">|</span><span class="df-view-hint">{VIEW_HINT['live']}</span></button>
    <button type="button" class="df-view" data-view="shot" title="{VIEW_ABOUT['shot']}">{VIEW_ICON['shot']}Try a screenshot<span class="df-view-sep">|</span><span class="df-view-hint">{VIEW_HINT['shot']}</span></button>
  </div>
  <div class="df-spacer"></div>
  <div class="df-status"><span class="df-led"></span>Running on ZeroGPU</div>
  <a class="df-btn" href="https://arxiv.org/abs/2610.02320" target="_blank">Paper</a>
  <a class="df-btn" href="https://huggingface.co/datasets/docling-project/DeskForge-1M" target="_blank">Dataset</a>
  <a class="df-btn" href="https://saidgurbuz.github.io/deskforge/" target="_blank">Project</a>
</div>"""

TOPBAR_JS = """
element.addEventListener('click', e => {
  const b = e.target.closest('[data-view]');
  if (!b) return;
  element.querySelectorAll('[data-view]').forEach(x => x.classList.toggle('on', x === b));
  const views = { web: ['df-view-web'], live: ['df-view-live', 'df-side'], shot: ['df-view-shot'] };
  for (const [view, ids] of Object.entries(views)) for (const id of ids) {
    const el = document.getElementById(id);
    if (el) el.style.display = view === b.dataset.view ? (id === 'df-view-shot' ? 'block' : 'flex') : 'none';
  }
});
"""

INTRO = """
<div class="df-msg"><b>What is this?</b> DeskForge is a small AI model trained to find and click the right
button on a computer screen. Here it races the model it was trained from, on a real Linux desktop running in
this page.</div>
<div class="df-msg">Pick a model, a look and a task at the top, then press <b>Start</b>. Each step appears here:
the instruction, where both models clicked, and whether they hit the right button (the <b class="df-c-target">green
box</b>, read from the app itself). DeskForge's click is the one that runs.</div>
<div class="df-aside">One desktop is shared by everyone here, so runs take turns. It has no browser, terminal or internet.</div>
"""


# ---------------------------------------------------------------------------
# Rendering: score and the step-by-step chat
# ---------------------------------------------------------------------------

def _pct(hits, n):
    return f"{100 * hits / n:.0f}%" if n else "–"


def render_score(family=DEFAULT_FAMILY, score=None):
    score = score or {"deskforge": [0, 0], "base": [0, 0]}

    def side(which, name, cls):
        hits, n = score[which]
        all_hits, all_n = TOTALS[family][which]
        return (f'<div class="df-vs-side"><div class="df-vs-name"><span class="df-dot {cls}"></span>{html.escape(name)}</div>'
                f'<div class="df-vs-num {cls}-text">{hits}<span> / {n}</span></div>'
                f'<div class="df-vs-all">{_pct(all_hits, all_n)} for all visitors</div></div>')

    return (f'<div class="df-vs">{side("deskforge", f"DeskForge {family}", "df-df")}'
            f'{side("base", FAMILIES[family].base_label, "df-base")}</div>')


def _crop(image, step):
    """A zoomed crop centred on the true target (or DeskForge's point when unscored), as a data URI."""
    if step.box:
        x0, y0, x1, y1 = step.box
        cx, cy, span = (x0 + x1) / 2, (y0 + y1) / 2, max(x1 - x0, (y1 - y0) * 2.2)
    else:
        p = to_pixels(parse_action(step.codes["deskforge"]), image.width, image.height)
        if not p:
            return None
        cx, cy, span = p[0], p[1], 0
    w = min(image.width, max(380, span * 3))
    h = min(image.height, w / 2.2)
    left = int(min(max(cx - w / 2, 0), image.width - w))
    top = int(min(max(cy - h / 2, 0), image.height - h))
    crop = image.crop((left, top, left + int(w), top + int(h)))
    crop.thumbnail((720, 330))
    buf = io.BytesIO()
    crop.convert("RGB").save(buf, format="JPEG", quality=85)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


def _result(which, step, family):
    name = "DeskForge" if which == "deskforge" else FAMILIES[family].base_label
    hit = step.hits.get(which)
    verdict = ('<span class="df-verdict df-hit">✓ hit</span>' if hit
               else '<span class="df-verdict df-miss">✗ miss</span>' if hit is False else "")
    code = html.escape(step.codes[which].replace("pyautogui.", "").replace("computer.", ""))
    secs = step.times.get(which)
    return f"""
    <div class="df-res">
      <div class="df-res-head"><span class="df-dot {'df-df' if which == 'deskforge' else 'df-base'}"></span>
        <span class="df-res-name">{html.escape(name)}</span>{verdict}<span class="df-spacer"></span>
        <span class="df-res-time">{f'{secs:.1f}s' if secs else ''}</span></div>
      <div class="df-res-code">{code}</div>
    </div>"""


def render_feed(steps=(), crops=None, status=None, running=False, error=None, family=DEFAULT_FAMILY):
    crops = crops or {}
    parts = [] if steps or status else [INTRO]
    if status:
        parts.append(f'<div class="df-msg">{status}</div>')
    for s in steps:
        chip = f'<div class="df-ref">{html.escape(s.note.replace("Theme → ", "New look: "))}</div>' if s.note.startswith("Theme") else ""
        parts.append(f'<div class="df-user-wrap"><div class="df-user">{chip}<div>{html.escape(s.instruction)}</div></div></div>')
        if crops.get(s.n):
            parts.append(f'<img class="df-crop" src="{crops[s.n]}" alt="Where both models clicked">')
        parts.append('<div class="df-res-pair">' + _result("deskforge", s, family) + _result("base", s, family) + "</div>")
        if "referee" in s.note:
            parts.append('<div class="df-aside">DeskForge missed, so the referee clicked the right button to keep going.</div>')
        elif "not scored" in s.note:
            parts.append('<div class="df-aside">That button was not on screen, so this step is not scored.</div>')
    if error:
        parts.append(f'<div class="df-msg df-err">The run stopped: {html.escape(error)}</div>')
    if running:
        parts.append('<div class="df-busy"><div class="df-spin"></div><div class="df-aside">Looking at the screen…</div></div>')
    return '<div class="df-feed">' + "".join(parts) + "</div>"


# ---------------------------------------------------------------------------
# Running duels
# ---------------------------------------------------------------------------

def _apply_look(look):
    return desktop.apply_theme(random_theme() if look == LOOK_SHUFFLE else THEME_PRESETS[look])


def run_challenge(family, name, look, roulette=False, tokens=128):
    if name == RANDOM_TASK:
        name = random.choice(list(CHALLENGES))
    challenge = CHALLENGES[name]
    status = f"Task: <b>{html.escape(name)}</b> · {html.escape(family)}"
    yield render_feed(status=status, running=True, family=family), render_score(family), gr.skip()
    desktop.reset(challenge.scene)
    _apply_look(look)
    yield from _stream(family, status, play(desktop, challenge.steps, family, roulette, tokens))


def run_free(family, scene, steps_text, look, roulette=False, tokens=128):
    lines = [s.strip() for s in (steps_text or "").splitlines() if s.strip()]
    if not lines:
        raise gr.Error("Write at least one instruction, one per line.")
    if len(lines) > MAX_STEPS:
        gr.Warning(f"Only the first {MAX_STEPS} instructions will run.")
        lines = lines[:MAX_STEPS]
    status = f"Your instructions on <b>{html.escape(scene)}</b> · not scored"
    yield render_feed(status=status, running=True, family=family), render_score(family), gr.skip()
    desktop.reset(scene)
    _apply_look(look)
    yield from _stream(family, status, play(desktop, [(l, None) for l in lines], family, roulette, tokens))


def _stream(family, status, duel):
    """Relay duel steps to the chat; on failure (e.g. GPU quota) keep what ran and say why."""
    steps, score, crops = [], None, {}
    try:
        for image, steps, score in duel:
            if steps and steps[-1].n not in crops and image is not None:
                crops[steps[-1].n] = _crop(image, steps[-1])
            yield (render_feed(steps, crops, status, running=True, family=family), render_score(family, score),
                   describe_theme(desktop.theme))
        yield render_feed(steps, crops, status, family=family), render_score(family, score), describe_theme(desktop.theme)
    except Exception as e:
        message = gpu_bridge.friendly_error(str(e) or type(e).__name__)
        gr.Warning(f"The run stopped: {message}")
        yield (render_feed(steps, crops, status, error=message, family=family), render_score(family, score),
               describe_theme(desktop.theme))


# ---------------------------------------------------------------------------
# Single screenshot grounding (the original demo)
# ---------------------------------------------------------------------------

EXAMPLES = [
    ("homebank_file_menu.png", "Select the Export as QIF... option from the File menu."),
    ("qalculate_factorial.png", "Add the factorial function to the current expression."),
    ("calculator_financial_mode.png", "Select Financial Mode in the calculator's mode dropdown."),
    ("mousepad_encoding_dropdown.png", "Open the file encoding dropdown to view available encoding options."),
    ("image_viewer_menu.png", "Open the Image Viewer menu."),
    ("applications_menu.png", "Open the Applications menu to view application categories."),
    ("nautilus_show_sidebar.png", "Enable the sidebar in the file manager."),
    ("transmission_sort_by_size.png", "Select the Sort by Size option in the Transmission View menu."),
]
gr.set_static_paths(paths=[HERE / "examples"])


def ground(instruction: str, screenshot, max_new_tokens: int = 128, family: str = DEFAULT_FAMILY):
    """Predict the next computer-use action for a desktop screenshot.

    Args:
        instruction: the task to perform on the screen, e.g. "Open the File menu".
        screenshot: a PIL image of the desktop to act on.
        max_new_tokens: generation length cap for the emitted action.
        family: which DeskForge model to use, "Gemma4-E4B" or "Qwen3.5-4B".

    Returns:
        Tuple of (annotated screenshot with the predicted click point,
        predicted action code, fractional coordinates, inference seconds).
    """
    if screenshot is None:
        raise gr.Error("Please provide a desktop screenshot first.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter an instruction, e.g. 'Open the File menu'.")
    if family not in FAMILIES:
        raise gr.Error(f"Unknown model family {family!r}; pick one of {', '.join(FAMILIES)}.")
    screenshot = screenshot.convert("RGB")
    code, elapsed = predict(instruction, screenshot, int(max_new_tokens), family)
    action = parse_action(code)
    coords = f"x={action.x:.4f}, y={action.y:.4f}" if action.has_point else "— (no point action)"
    return annotate(screenshot, action), code, coords, f"{elapsed:.1f}s"


def ground_screenshot(args: list) -> dict:
    """Server function behind the "Try a screenshot" view; the script's arguments arrive as one list."""
    path, instruction, family = (list(args) + ["", "", DEFAULT_FAMILY])[:3]
    if path.startswith("example:"):
        name = path.split(":", 1)[1]
        if name not in dict(EXAMPLES):
            raise gr.Error("Unknown example.")
        file = HERE / "examples" / name
    else:
        file = Path(path).resolve()
        if not file.is_relative_to(Path(tempfile.gettempdir()).resolve()):  # only files uploaded to this app
            raise gr.Error("Please upload the screenshot again.")
    image, code, coords, secs = ground(instruction, Image.open(file), 128, family)
    image.thumbnail((1400, 1400))
    buf = io.BytesIO()
    image.convert("RGB").save(buf, format="JPEG", quality=88)
    return {"image": "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode(),
            "code": code, "coords": coords, "time": secs}


# ---------------------------------------------------------------------------
# Theme and styles (palette, type and components from the Docling Agent design)
# ---------------------------------------------------------------------------

HEAD = """
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=Hanken+Grotesk:wght@400;500;600;700&family=JetBrains+Mono:wght@500;600&family=Newsreader:ital,opsz,wght@0,6..72,400;0,6..72,600;1,6..72,400&display=swap" rel="stylesheet">
<script>
// The design is light-only: pin Gradio to its light theme.
(() => { const u = new URL(location.href);
  if (u.searchParams.get("__theme") !== "light") { u.searchParams.set("__theme", "light"); location.replace(u); } })();
</script>
"""

THEME = gr.themes.Base(
    font=[gr.themes.GoogleFont("Hanken Grotesk"), gr.themes.Font("system-ui"), gr.themes.Font("sans-serif")],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), gr.themes.Font("ui-monospace"), gr.themes.Font("monospace")],
).set(
    body_background_fill="#FFFDF8",
    body_text_color="#2B2722",
    body_text_size="14px",
    background_fill_primary="#FFFFFF",
    background_fill_secondary="#FBF6EE",
    border_color_primary="#E6E3DF",
    block_background_fill="transparent",
    block_border_width="0px",
    block_shadow="none",
    block_radius="0px",
    color_accent="#D2603A",
    link_text_color="#C0552F",
)

CSS = setup_bar.CSS + shot_tool.CSS + web_tab.CSS + gpu_bridge.CSS + """
body, .gradio-container, .gradio-container button, .gradio-container input, .gradio-container textarea {
  font-family: 'Hanken Grotesk', system-ui, -apple-system, sans-serif; }
.gradio-container { max-width: 100% !important; padding: 0 !important; background: #FFFDF8; }
.gradio-container > .main { padding: 0 !important; }
.gradio-container .main, .gradio-container .wrap { gap: 0; }
footer { border-top: 1px solid #E6E3DF; background: #FBF6EE; }

/* top bar */
.df-top { height: 52px; display: flex; align-items: center; gap: 14px; padding: 0 18px; background: #FBF6EE; border-bottom: 1px solid #E6E3DF; }
.df-brand { display: flex; align-items: center; gap: 9px; }
.df-logo { width: 26px; height: 26px; border-radius: 7px; background: #D2603A; display: flex; align-items: center; justify-content: center; }
.df-word { font-weight: 700; font-size: 15px; letter-spacing: -.01em; color: #2B2722; }
.df-sep { width: 1px; height: 20px; background: #E6E3DF; }
.df-views { display: flex; gap: 4px; }
.df-top .df-view { display: inline-flex; align-items: center; }
.df-top .df-view-ico { flex: none; margin-right: 7px; color: #a39d94; }
.df-top .df-view.on .df-view-ico { color: #D2603A; }
.df-top .df-view-sep { margin: 0 7px; color: #CFCAC3; font-weight: 400; }
.df-top .df-view-hint { font-size: 13px; font-weight: 400; color: #9a9389; }
.df-top .df-view.on .df-view-hint { color: #8a8378; }
@media (max-width: 1100px) { .df-top .df-view-sep, .df-top .df-view-hint { display: none; } }
.df-top .df-view { font: 500 14px 'Hanken Grotesk', system-ui, sans-serif; color: #6f6a62; background: transparent; border: 1px solid transparent;
           border-radius: 8px; padding: 6px 11px; cursor: pointer; }
.df-top .df-view:hover { color: #2B2722; }
.df-top .df-view.on { color: #2B2722; background: #fff; border-color: #E6E3DF; box-shadow: 0 1px 2px rgba(40,30,20,.05); }
.df-spacer { flex: 1; }
.df-status { display: flex; align-items: center; gap: 7px; font-size: 13px; color: #8a8378; }
.df-led { width: 7px; height: 7px; border-radius: 50%; background: #62C554; }
.df-top .df-btn { font-size: 13px; color: #3a332b !important; border: 1px solid #E6E3DF; background: #fff; padding: 6px 12px; border-radius: 8px;
          text-decoration: none !important; }
.df-top .df-btn:hover { border-color: #D6D1CA; }
@media (max-width: 760px) { .df-status, .df-btn { display: none; } }

/* shell: the web agent is the first view; the duel and the screenshot tool are a click away */
#df-view-web { display: flex; }
#df-view-live, #df-side { display: none; }
#df-shell { gap: 0 !important; align-items: stretch; flex-wrap: nowrap; min-height: calc(100vh - 52px); }
#df-main { background: #F3EDE2; border-right: 1px solid #E6E3DF; padding: 0 !important; gap: 0 !important; min-width: 0; }
#df-view-live { gap: 0 !important; }
.df-canvas { padding: 20px 24px !important; gap: 8px !important; }
.df-sheet, .df-sheet img { border-radius: 6px !important; }
.df-sheet { border: 1px solid #E0DCD6 !important; box-shadow: 0 6px 22px rgba(60,40,20,.12) !important; overflow: hidden; background: #fff; }
.df-sheet .icon-button-wrapper, .df-sheet .icon-buttons { display: none !important; }
.df-sheet .image-frame > button, .df-sheet .image-container > button { cursor: default; padding: 0; border: none; background: none; }
.df-caption { font: 500 11.5px 'JetBrains Mono', monospace; color: #9a9389; padding-top: 4px; }
#df-side { background: #FBF6EE; padding: 0 !important; gap: 0 !important; width: 400px; max-width: 400px; min-width: 400px !important; flex: none !important; }
@media (max-width: 1000px) {
  #df-shell { flex-wrap: wrap; }
  #df-side { width: 100%; max-width: 100%; min-width: 0 !important; }
}

/* sidebar: score and chat */
.df-side-head { height: 46px; display: flex; align-items: center; padding: 0 16px; border-bottom: 1px solid #E6E3DF;
                font-size: 14px; font-weight: 700; color: #2B2722; }
#df-score-wrap { padding: 12px 14px; border-bottom: 1px solid #E6E3DF; }
.df-vs { display: grid; grid-template-columns: 1fr 1fr; gap: 8px; }
.df-vs-side { background: #fff; border: 1px solid #E6E3DF; border-radius: 10px; padding: 10px 12px; box-shadow: 0 1px 2px rgba(40,30,20,.04); }
.df-vs-name { display: flex; align-items: center; gap: 7px; font-size: 13px; font-weight: 600; color: #3a332b;
              white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
.df-vs-num { font-size: 20px; font-weight: 600; margin-top: 2px; font-variant-numeric: tabular-nums; }
.df-vs-num span { font-size: 14px; font-weight: 500; color: #a39d94; }
.df-df-text { color: #C0552F; } .df-base-text { color: #3B6EBE; }
.df-vs-all { font-size: 12px; color: #9a9389; margin-top: 1px; }
.df-dot { width: 8px; height: 8px; border-radius: 50%; display: inline-block; flex: none; }
.df-df { background: #D2603A; } .df-base { background: #3B6EBE; }
.df-c-target { color: #3f8a35; }

#df-feed-wrap { padding: 0 !important; }
.df-feed { padding: 16px; display: flex; flex-direction: column; gap: 12px; max-height: calc(100vh - 210px); min-height: 260px; overflow-y: auto; }
.df-feed > * { flex-shrink: 0; }
.df-msg { font-size: 14px; line-height: 1.6; color: #3a332b; }
.df-aside { font-family: 'Newsreader', serif; font-style: italic; font-size: 14px; line-height: 1.5; color: #8a8378; }
.df-err { color: #A8452A; }
.df-user-wrap { display: flex; justify-content: flex-end; animation: df-in .15s ease-out; }
.df-user { max-width: 88%; background: #D2603A; color: #fff; font-size: 14px; line-height: 1.5; padding: 9px 13px;
           border-radius: 13px 13px 4px 13px; }
.df-ref { font: 500 11px 'JetBrains Mono', monospace; background: rgba(255,255,255,.2); padding: 2px 7px; border-radius: 5px;
          margin-bottom: 5px; display: inline-block; }
.df-crop { width: 100%; border-radius: 10px; border: 1px solid #E6E3DF; display: block; }
.df-res-pair { display: flex; flex-direction: column; background: #fff; border: 1px solid #E6E3DF; border-radius: 10px; overflow: hidden;
               box-shadow: 0 1px 2px rgba(40,30,20,.04); animation: df-in .15s ease-out; }
.df-res + .df-res { border-top: 1px solid #EEECE8; }
.df-res { padding: 9px 12px; }
.df-res-head { display: flex; align-items: center; gap: 7px; }
.df-res-name { font-size: 13.5px; font-weight: 600; color: #2B2722; }
.df-res-time { font: 500 11.5px 'JetBrains Mono', monospace; color: #a39d94; }
.df-res-code { font: 500 12px/1.5 'JetBrains Mono', monospace; color: #6f6a62; margin-top: 3px; overflow-wrap: anywhere; }
.df-verdict { font: 600 11.5px 'Hanken Grotesk', system-ui, sans-serif; padding: 1px 7px; border-radius: 5px; }
.df-hit { color: #3f7a34; background: #EAF3E6; } .df-miss { color: #A8452A; background: #F8E9E2; }
.df-busy { display: flex; align-items: center; gap: 9px; }
.df-spin { width: 14px; height: 14px; border-radius: 50%; border: 2px solid #E6E3DF; border-top-color: #D2603A; animation: df-spin .7s linear infinite; }
@keyframes df-spin { to { transform: rotate(360deg); } }
@keyframes df-in { from { opacity: 0; transform: translateY(5px); } to { opacity: 1; transform: none; } }
"""


# ---------------------------------------------------------------------------
# Layout: custom HTML everywhere except the live image
# ---------------------------------------------------------------------------

TASKS = {name: (c.scene.split("(")[1].rstrip(")") if "(" in c.scene else c.scene, len(c.steps))
         for name, c in CHALLENGES.items()}
MODELS = [DEFAULT_FAMILY, *[f for f in FAMILIES if f != DEFAULT_FAMILY]]
SETUP_DEFAULT = setup_bar.default_state(MODELS, list(CHALLENGES))
SETUP_PROPS = setup_bar.props(MODELS, list(THEME_PRESETS), TASKS)


def _settings(evt):
    """Setup-bar state sent with an event, merged over the defaults."""
    data = evt._data if isinstance(getattr(evt, "_data", None), dict) else {}
    return {**SETUP_DEFAULT, **data}


def _caption(text):
    return f'<div class="df-caption">Desktop look: {html.escape(text)}</div>'


with gr.Blocks(title="DeskForge · Live Desktop Duel", fill_width=True) as demo:
    gr.HTML(TOPBAR, js_on_load=TOPBAR_JS, apply_default_css=False)
    with gr.Row(elem_id="df-shell"):
        with gr.Column(elem_id="df-main", scale=1):
            with gr.Column(elem_id="df-view-web"):
                gr.HTML(None, html_template=web_tab.TEMPLATE, js_on_load=gpu_bridge.CLIENT + web_tab.SCRIPT,
                        server_functions=web_agent.server_functions(), suggestions=web_tab.SUGGESTIONS,
                        logo=web_tab.LOGO, apply_default_css=False)
            with gr.Column(elem_id="df-view-live"):
                setup = gr.HTML(SETUP_DEFAULT, html_template=setup_bar.TEMPLATE, js_on_load=setup_bar.SCRIPT,
                                elem_id="df-setup", apply_default_css=False, **SETUP_PROPS)
                with gr.Column(elem_classes="df-canvas"):
                    live = gr.Image(show_label=False, format="jpeg", interactive=False, container=False,
                                    elem_classes="df-sheet", buttons=[])
                    look_now = gr.HTML(_caption(describe_theme(desktop.theme)), apply_default_css=False)
            with gr.Column(elem_id="df-view-shot"):
                gr.HTML(None, html_template=shot_tool.TEMPLATE, js_on_load=gpu_bridge.CLIENT + shot_tool.SCRIPT,
                        server_functions=[ground_screenshot], models=MODELS, apply_default_css=False,
                        examples=[{"name": n, "text": t, "url": f"gradio_api/file={HERE / 'examples' / n}"} for n, t in EXAMPLES])
        with gr.Column(elem_id="df-side"):
            gr.HTML('<div class="df-side-head">Duel, step by step</div>', apply_default_css=False)
            score = gr.HTML(render_score(), elem_id="df-score-wrap", apply_default_css=False)
            feed = gr.HTML(render_feed(), elem_id="df-feed-wrap", apply_default_css=False)

    # The original grounding endpoint (API and MCP), kept without any visible UI.
    with gr.Column(visible=False):
        api_image, api_text = gr.Image(type="pil"), gr.Textbox()
        api_tokens, api_family = gr.Number(128), gr.Textbox(DEFAULT_FAMILY)
        api_out = [gr.Image(type="pil"), gr.Textbox(), gr.Textbox(), gr.Textbox()]
        api_btn = gr.Button()
    api_btn.click(ground, [api_text, api_image, api_tokens, api_family], api_out, api_name="predict")

    settings = gr.State(SETUP_DEFAULT)

    def _relay(gen):
        for feed_html, score_html, look_text in gen:
            yield feed_html, score_html, _caption(look_text) if isinstance(look_text, str) else look_text

    def on_change(evt: gr.EventData):
        st = _settings(evt)
        return st, render_score(st["model"])

    def on_look(evt: gr.EventData):
        return _caption(describe_theme(_apply_look(_settings(evt)["look"])))

    def on_start(evt: gr.EventData):
        st = _settings(evt)
        if (st.get("own") or "").strip():
            task = random.choice(list(CHALLENGES)) if st["random"] else st["task"]
            gen = run_free(st["model"], CHALLENGES[task].scene, st["own"], st["look"], st["roulette"], int(st["tokens"]))
        else:
            task = RANDOM_TASK if st["random"] else st["task"]
            gen = run_challenge(st["model"], task, st["look"], st["roulette"], int(st["tokens"]))
        yield from _relay(gen)

    def duel_api(family: str, task: str, look: str = LOOK_SHUFFLE) -> tuple[str, str]:
        """Run a scored challenge on the shared desktop; returns the final chat HTML and score HTML."""
        feed_html = score_html = ""
        for feed_html, score_html, _ in run_challenge(family, task, look):
            pass
        return feed_html, score_html

    gr.Timer(1.0).tick(live_view, outputs=live, show_progress="hidden", concurrency_limit=None,
                       api_visibility="private")
    demo.load(live_view, outputs=live, api_visibility="private")
    demo.load(lambda: render_score(DEFAULT_FAMILY), outputs=score, api_visibility="private")

    # GPU work requested by the custom views runs here, as a real Gradio event (see gpu_bridge.py).
    web_fns = {fn.__name__: fn for fn in web_agent.server_functions()}
    gpu_ops = {
        "web_plan": lambda d: web_fns["web_plan"](d.get("run", "")),
        "ground_screenshot": lambda d: ground_screenshot([d.get("src", ""), d.get("text", ""), d.get("model", DEFAULT_FAMILY)]),
    }
    bridge = gr.HTML(None, html_template=gpu_bridge.TEMPLATE, js_on_load=gpu_bridge.SCRIPT, elem_id="df-gpu-bridge",
                     apply_default_css=False)
    bridge.submit(gpu_bridge.handler(gpu_ops), None, bridge, concurrency_limit=6, concurrency_id="gpu-bridge",
                  api_visibility="private", show_progress="hidden", queue=True)

    setup.change(on_change, None, [settings, score], queue=False, api_visibility="private")
    setup.input(on_look, None, look_now, concurrency_id="desktop", api_visibility="private")
    start_event = setup.submit(on_start, None, [feed, score, look_now], concurrency_limit=1,
                               concurrency_id="desktop", api_visibility="private", show_progress="hidden")
    setup.stop(None, cancels=[start_event])
    gr.api(duel_api, api_name="duel", concurrency_limit=1, concurrency_id="desktop")

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=THEME, css=CSS, head=HEAD)