"""GUI-Decisions: computer-use decisions read from one forward pass (slot reading on Gemma 4 31B). The adapter is loaded on the base model (not the merged model) so the base model can still write the text a type / answer step needs, and run the vanilla comparison, with the adapter switched off. """ import os import gradio as gr import numpy as np import spaces import torch from PIL import Image, ImageDraw from slotread.client import agent_user from slotread.data.aguvis import ACTIONS, KEYS, SDIR from slotread.local import LocalSlotModel from slotread.numbers import unwarp BASE = os.environ.get("BASE", "google/gemma-4-31B-it") ADAPTER = os.environ.get("ADAPTER", "infinitylogesh/GUI-Decisions-31B-LoRA") SLOT, VANILLA = (255, 90, 20), (58, 32, 19) # the Ambient theme: orange the slot read, dark brown the base model writing MAX_PX = 3840 * 2160 WARM = set() # functions that have run once in this process: the first call includes CUDA warm-up def warm_note(name): first = name not in WARM WARM.add(name) return " The first request after a restart includes GPU warm-up; run it again for a warm time." if first else "" m = LocalSlotModel.from_pretrained(BASE, adapter=ADAPTER, processor=ADAPTER, device_map="cpu") m.model.to("cuda") def fit(im): im = im.convert("RGB") if im.width * im.height > MAX_PX: s = (MAX_PX / (im.width * im.height)) ** 0.5 im = im.resize((int(im.width * s), int(im.height * s)), Image.LANCZOS) return im def overlay(im, px, py, color): """The slot distributions as a heatmap: P(x) times P(y), drawn as the colour's opacity.""" heat = np.outer(py, px) heat = (heat / heat.max()) ** 0.5 a = Image.fromarray((heat * 150).astype(np.uint8)).resize(im.size, Image.BILINEAR) layer = Image.new("RGBA", im.size, color + (0,)) layer.putalpha(a) return Image.alpha_composite(im.convert("RGBA"), layer) def mark(d, x, y, color, r=14): d.ellipse([x - r, y - r, x + r, y + r], outline=(255, 255, 255), width=6) d.ellipse([x - r, y - r, x + r, y + r], outline=color, width=3) d.line([x - 4, y, x + 4, y], fill=color, width=2) d.line([x, y - 4, x, y + 4], fill=color, width=2) def probs(v): return torch.softmax(v.float(), -1).numpy() def describe(a, d, text): if a in ("click", "select"): return f"{a}(x={d['x']:.3f}, y={d['y']:.3f})" if a == "swipe": return f"swipe(({d['x']:.3f}, {d['y']:.3f}) → ({d['x2']:.3f}, {d['y2']:.3f}))" if a == "press": return f"press(key={d['key']})" if a == "scroll": return f"scroll(direction={d['sdir']}, amount={d['samt']})" if text is not None: return f'{a}("{text}")' return f"{a}()" @spaces.GPU(size="xlarge", duration=40) def decide(image, instruction, previous, compare): if image is None or not instruction.strip(): raise gr.Error("Add a screenshot and an instruction.") im = fit(image) user_prev = previous.strip() or "None" lp, ms = m.read("agent", im, agent_user(instruction, user_prev)) d = m.tasks["agent"].decode(lp, m.arm) d.update(key=KEYS[int(lp["key"].argmax())], sdir=SDIR[int(lp["sdir"].argmax())], samt=int(lp["samt"].argmax())) a = d["action"] text = m.write_text(im, instruction, a, user_prev) if a in ("type", "answer", "open") else None out = im.convert("RGBA") if a in ("click", "select", "swipe"): out = overlay(out, probs(lp["x"]), probs(lp["y"]), SLOT) dr = ImageDraw.Draw(out) W, H = im.size if a in ("click", "select", "swipe"): mark(dr, d["x"] * W, d["y"] * H, SLOT) if a == "swipe": dr.line([d["x"] * W, d["y"] * H, d["x2"] * W, d["y2"] * H], fill=SLOT, width=4) mark(dr, d["x2"] * W, d["y2"] * H, SLOT, r=8) md = (f"### Slots: `{describe(a, d, text)}`\n\n" f"One forward pass: **{ms:.0f} ms** on this GPU (bf16, transformers" + (", plus the base model writing the text" if text is not None else "") + ")." + warm_note("decide")) if compare: va, vpt, vargs, vtext, raw, vms = m.vanilla_decide(im, instruction, user_prev) if vpt is not None: mark(dr, vpt["x"] * W, vpt["y"] * H, VANILLA) vd = dict(vargs, **(vpt or {"x": float("nan"), "y": float("nan")})) md += (f"\n\n### Base model writing JSON: `{describe(va, vd, vtext)}`\n\n" f"Generating its answer took **{vms:.0f} ms** (Gemma 4's native `\"point\": [y, x]` format):\n\n" f"```\n{raw.strip()[:400]}\n```\n\n" "Orange: the slot read, with its x × y distribution as the heatmap. Dark brown: the base model's point.") label = {k: float(p) for k, p in zip(ACTIONS, probs(lp["action"]))} return out.convert("RGB"), md, label @spaces.GPU(size="xlarge", duration=20) def box(image, description): if image is None or not description.strip(): raise gr.Error("Add an image and describe the element or object.") im = fit(image) lp, ms = m.read("box", im, description) b = m.tasks["box"].decode(lp, m.arm) W, H = im.size out = overlay(im, probs(lp["cx"]), probs(lp["cy"]), SLOT) dr = ImageDraw.Draw(out) xy = [b[0] * W, b[1] * H, b[2] * W, b[3] * H] dr.rectangle(xy, outline=(255, 255, 255), width=6) dr.rectangle(xy, outline=SLOT, width=3) w = unwarp(float(np.argmax(probs(lp["w"])) + 0.5) / m.arm.bins, "log") md = (f"### Box: `[{', '.join(f'{v:.3f}' for v in b)}]`\n\n" f"Centre and log-scaled size, four slots read in one forward pass: **{ms:.0f} ms** on this GPU. " f"The heatmap is the centre's x × y distribution (most likely width {w:.1%} of the image)." + warm_note("box")) return out.convert("RGB"), md HERE = os.path.dirname(os.path.abspath(__file__)) ex = lambda f: os.path.join(HERE, "examples", f) rows = lambda f: [l.split("\t") for l in open(ex(f)).read().splitlines() if l.strip()] # tab-separated examples INTRO = open(os.path.join(HERE, "intro.html")).read() CSS = open(os.path.join(HERE, "ambient.css")).read() HOW = open(os.path.join(HERE, "how.md")).read() with gr.Blocks(title="GUI-Decisions") as demo: gr.HTML(INTRO) gr.HTML('
Timings here are this Space\'s shared GPU running bf16 in plain transformers. ' 'Served on vLLM (NVFP4, one RTX PRO 6000) a step takes ~150 ms, against ~470–550 ms for the same model ' 'writing its native JSON.
') with gr.Tab("Computer use"): with gr.Row(): with gr.Column(): img = gr.Image(type="pil", label="Screenshot") ins = gr.Textbox(label="Instruction", placeholder="open the settings") prev = gr.Textbox(label="Previous actions (optional)", placeholder="1. click(...)", lines=2) cmp = gr.Checkbox(value=True, label="Also run the base model writing its answer as JSON (slower)") go = gr.Button("Decide", variant="primary") with gr.Column(): out_img = gr.Image(label="Decision") gr.HTML('