Spaces:
Running on Zero
Running on Zero
| """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}()" | |
| 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 | |
| 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('<p class="amb-note">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.</p>') | |
| 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('<div class="amb-legend"><span><i style="background:#ff5a14"></i>slot read · P(x) × P(y)</span>' | |
| '<span><i style="background:#3a2013"></i>base model writing JSON</span></div>') | |
| out_md = gr.Markdown() | |
| out_lab = gr.Label(label="action slot", num_top_classes=4) | |
| go.click(decide, [img, ins, prev, cmp], [out_img, out_md, out_lab]) | |
| gr.Examples([[ex(f), i, p.replace("\\n", "\n"), True] for f, i, p in rows("agent.tsv")], [img, ins, prev, cmp]) | |
| with gr.Tab("Boxes"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| bimg = gr.Image(type="pil", label="Image") | |
| bdesc = gr.Textbox(label="Element or object", placeholder="the red car on the left") | |
| bgo = gr.Button("Locate", variant="primary") | |
| with gr.Column(): | |
| bout = gr.Image(label="Box") | |
| bmd = gr.Markdown() | |
| bgo.click(box, [bimg, bdesc], [bout, bmd]) | |
| gr.Examples([[ex(f), t] for f, t in rows("box.tsv")], [bimg, bdesc]) | |
| with gr.Tab("How it works"): | |
| gr.Markdown(HOW) | |
| # the Ambient theme: Inter / Inter Tight / JetBrains Mono, cream canvas, dark-brown ink, orange accent (light only) | |
| ORANGE = gr.themes.Color(c50="#fff4ee", c100="#ffe5d8", c200="#ffc9ad", c300="#ffa77d", c400="#ff7a2e", c500="#ff5a14", | |
| c600="#e64a08", c700="#bf3c06", c800="#8f2e07", c900="#5e1f06", c950="#3a1304") | |
| TAUPE = gr.themes.Color(c50="#fdf6f1", c100="#f8ece4", c200="#efe2d9", c300="#d9c6ba", c400="#9a877e", c500="#6b5850", | |
| c600="#4f3d35", c700="#3a2013", c800="#2a170e", c900="#1c100b", c950="#1a110d") | |
| LIGHT = dict( | |
| body_background_fill="#fdf6f1", body_text_color="#1a110d", body_text_color_subdued="#6b5850", | |
| background_fill_primary="#ffffff", background_fill_secondary="#fbf1ea", | |
| block_background_fill="#ffffff", block_border_color="rgba(58,28,12,0.11)", block_border_width="1px", | |
| block_label_background_fill="#fbf1ea", block_label_text_color="#6b5850", block_title_text_color="#6b5850", | |
| block_shadow="0 1px 2px rgba(58,28,12,0.05), 0 12px 32px -14px rgba(58,28,12,0.18)", | |
| border_color_primary="rgba(58,28,12,0.11)", border_color_accent="#ff5a14", | |
| input_background_fill="#fdf6f1", input_border_color="rgba(58,28,12,0.11)", input_border_color_focus="#ff7a2e", | |
| button_primary_background_fill="#ff5a14", button_primary_background_fill_hover="#ff7a2e", | |
| button_primary_text_color="#ffffff", button_primary_border_color="#ff5a14", | |
| button_secondary_background_fill="#ffffff", button_secondary_background_fill_hover="#fbf1ea", | |
| button_secondary_text_color="#1a110d", button_secondary_border_color="rgba(58,28,12,0.11)", | |
| color_accent="#ff5a14", color_accent_soft="#ffe5d8", link_text_color="#ff5a14", link_text_color_hover="#e64a08", | |
| checkbox_background_color_selected="#ff5a14", checkbox_border_color_selected="#ff5a14", slider_color="#ff5a14", | |
| loader_color="#ff5a14", code_background_fill="#f3e7df", table_even_background_fill="#ffffff", | |
| table_odd_background_fill="#fbf1ea", table_border_color="rgba(58,28,12,0.11)", | |
| ) | |
| THEME = gr.themes.Base(primary_hue=ORANGE, secondary_hue=TAUPE, neutral_hue=TAUPE, radius_size=gr.themes.sizes.radius_lg, | |
| font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"], | |
| font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"]) | |
| THEME.set(block_radius="16px", **LIGHT, **{k + "_dark": v for k, v in LIGHT.items() if hasattr(THEME, k + "_dark")}) | |
| HEAD = ('<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700' | |
| '&family=Inter+Tight:wght@700;800&family=JetBrains+Mono:wght@400;500;600&display=swap">') | |
| LIGHT_ONLY = "() => { document.body.classList.remove('dark'); document.documentElement.classList.remove('dark'); }" | |
| if __name__ == "__main__": | |
| demo.queue().launch(theme=THEME, css=CSS, head=HEAD, js=LIGHT_ONLY) | |