Spaces:
Running on Zero
Running on Zero
| import spaces # must be imported before torch | |
| import os | |
| import time | |
| import gradio as gr | |
| import torch | |
| from PIL import Image, ImageDraw, ImageFont | |
| from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration | |
| from gamescaling import ( | |
| MOUSE_AIM_LINE, | |
| OPENING, | |
| OUTPUT_FORMAT, | |
| TAIL, | |
| USER_TURN_TEXT, | |
| parse_action, | |
| split_reasoning, | |
| ) | |
| ROOT = os.path.dirname(os.path.abspath(__file__)) | |
| # --------------------------------------------------------------------------- | |
| # Models | |
| # --------------------------------------------------------------------------- | |
| MODELS = { | |
| "GameScaling-9B-S3 (mixed games)": "THUSI-Lab/GameScaling-9B-S3", | |
| "GameScaling-2B-S3 (mixed games, fast)": "THUSI-Lab/GameScaling-2B-S3", | |
| } | |
| DEFAULT_MODEL = "GameScaling-9B-S3 (mixed games)" | |
| LOADED = {} | |
| for label, repo in MODELS.items(): | |
| proc = AutoProcessor.from_pretrained(repo) | |
| model = Qwen3_5ForConditionalGeneration.from_pretrained(repo, dtype=torch.bfloat16).to("cuda").eval() | |
| LOADED[label] = (model, proc) | |
| THINK_PREFILL = "<think>" | |
| NO_THINK_PREFILL = "<think>\n\n</think>" | |
| MAX_SIDE = 1280 # reference frames are 1280x720 | |
| # --------------------------------------------------------------------------- | |
| # Game presets (system-prompt slots: game name, rules, key bindings) | |
| # --------------------------------------------------------------------------- | |
| _doom = open(os.path.join(ROOT, "prompts", "example_doom_battle1.txt")).read().rstrip("\n") | |
| _DOOM_RULES = _doom.split("**Game Rules**\n", 1)[1].split("\n**Common Keys in", 1)[0] | |
| _DOOM_KEYS = _doom.split("**Common Keys in Doom Battle-1**\n", 1)[1].split("\n**Output Format**", 1)[0] | |
| PRESETS = { | |
| "Doom Battle-1": dict( | |
| name="Doom Battle-1", | |
| rules=_DOOM_RULES, | |
| keys=_DOOM_KEYS, | |
| ), | |
| "Minecraft (voxel sandbox)": dict( | |
| name="Minecraft", | |
| rules=( | |
| "You are playing Minecraft, a first-person voxel sandbox game.\n\n" | |
| "## What you see on screen\n" | |
| "- First-person view; the crosshair at screen centre marks the block you are looking at.\n" | |
| "- The hotbar along the bottom shows the items you carry; the selected slot is highlighted.\n\n" | |
| "## Task\n" | |
| "- Chop down a tree to collect wood logs: walk up to a tree trunk, aim the crosshair at it " | |
| "and hold the left mouse button until the block breaks." | |
| ), | |
| keys="\n".join([ | |
| "- W -> Move forward", | |
| "- A -> Strafe left", | |
| "- S -> Move backward", | |
| "- D -> Strafe right", | |
| "- Space -> Jump", | |
| "- Shift -> Sneak", | |
| "- Ctrl -> Sprint", | |
| "- LB -> Attack / break block", | |
| "- RB -> Use item / place block", | |
| "- E -> Open inventory", | |
| "- one -> Select hotbar slot 1", | |
| "- two -> Select hotbar slot 2", | |
| "- three -> Select hotbar slot 3", | |
| MOUSE_AIM_LINE, | |
| ]), | |
| ), | |
| "SuperTuxKart (racing)": dict( | |
| name="SuperTuxKart", | |
| rules=( | |
| "You are playing SuperTuxKart, a 3D kart racing game, seen from behind your kart.\n\n" | |
| "## Rules\n" | |
| "- Finish the laps as fast as possible and overtake the other karts.\n" | |
| "- Stay on the track: driving off-road slows you down.\n" | |
| "- Item boxes on the track give you an item that you can fire at other karts." | |
| ), | |
| keys="\n".join([ | |
| "- Up -> Accelerate", | |
| "- Down -> Brake / reverse", | |
| "- Left -> Steer left", | |
| "- Right -> Steer right", | |
| "- Space -> Fire item", | |
| "- N -> Use nitro", | |
| "- V -> Drift", | |
| ]), | |
| ), | |
| } | |
| PRESET_NAMES = list(PRESETS) + ["Custom"] | |
| def build_system_prompt(name: str, rules: str, keys: str) -> str: | |
| """Same skeleton as gamescaling.build_system_prompt (official prompt template).""" | |
| parts = [OPENING.format(name=name)] | |
| if rules.strip(): | |
| parts.append("**Game Rules**\n" + rules.strip()) | |
| if keys.strip(): | |
| parts.append(f"**Common Keys in {name}**\n" + keys.strip()) | |
| parts.append("**Output Format**\n" + OUTPUT_FORMAT) | |
| parts.append(TAIL) | |
| return "\n".join(parts) | |
| # sanity check: preset reproduces the authors' Doom prompt byte-for-byte | |
| assert build_system_prompt(**PRESETS["Doom Battle-1"]) == _doom | |
| # --------------------------------------------------------------------------- | |
| # Visualisation | |
| # --------------------------------------------------------------------------- | |
| def _font(size): | |
| for p in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", | |
| "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf"): | |
| if os.path.exists(p): | |
| return ImageFont.truetype(p, size) | |
| try: | |
| return ImageFont.load_default(size=size) | |
| except TypeError: | |
| return ImageFont.load_default() | |
| def render_overlay(frame: Image.Image, action) -> Image.Image: | |
| """Draw the predicted mouse move on the frame and a 6-step key timeline below it.""" | |
| img = frame.convert("RGB").copy() | |
| W, H = img.size | |
| s = max(W, H) / 1280 | |
| d = ImageDraw.Draw(img, "RGBA") | |
| cx, cy = W / 2, H / 2 | |
| x, y, z = action.mouse | |
| tx, ty = cx + x / 1000 * W, cy + y / 1000 * H | |
| r = 10 * s | |
| d.ellipse([cx - r, cy - r, cx + r, cy + r], outline=(255, 255, 255, 230), width=max(2, int(3 * s))) | |
| if abs(tx - cx) + abs(ty - cy) > 2: | |
| d.line([cx, cy, tx, ty], fill=(255, 60, 60, 255), width=max(3, int(5 * s))) | |
| import math | |
| ang = math.atan2(ty - cy, tx - cx) | |
| L = 18 * s | |
| for da in (2.6, -2.6): | |
| d.line([tx, ty, tx + L * math.cos(ang + da), ty + L * math.sin(ang + da)], | |
| fill=(255, 60, 60, 255), width=max(3, int(5 * s))) | |
| f = _font(int(22 * s)) | |
| label = f"mouse X={x:+.0f} Y={y:+.0f} scroll={z:+.0f}" | |
| d.rectangle([10 * s, 10 * s, 10 * s + d.textlength(label, font=f) + 16 * s, 46 * s], fill=(0, 0, 0, 170)) | |
| d.text((18 * s, 14 * s), label, fill=(255, 255, 255), font=f) | |
| # timeline strip | |
| strip_h = int(150 * s) | |
| out = Image.new("RGB", (W, H + strip_h), (24, 24, 27)) | |
| out.paste(img, (0, 0)) | |
| d = ImageDraw.Draw(out) | |
| fh = _font(int(18 * s)) | |
| fk = _font(int(20 * s)) | |
| col_w = W / 6 | |
| for i, keys in enumerate(action.steps): | |
| x0 = i * col_w | |
| d.rectangle([x0 + 4 * s, H + 6 * s, x0 + col_w - 4 * s, H + strip_h - 6 * s], | |
| outline=(80, 80, 90), width=max(1, int(2 * s))) | |
| d.text((x0 + 12 * s, H + 10 * s), f"#{i + 1} {i * 33}-{(i + 1) * 33} ms", fill=(170, 170, 180), font=fh) | |
| ky = H + 40 * s | |
| if not keys: | |
| d.text((x0 + 12 * s, ky), "(no keys)", fill=(110, 110, 120), font=fk) | |
| for k in keys: | |
| tw = d.textlength(k, font=fk) | |
| d.rounded_rectangle([x0 + 12 * s, ky, x0 + 28 * s + tw, ky + 26 * s], radius=6 * s, | |
| fill=(250, 140, 40) if k in ("LB", "RB", "MB") else (70, 160, 90)) | |
| d.text((x0 + 20 * s, ky + 2 * s), k, fill=(255, 255, 255), font=fk) | |
| ky += 30 * s | |
| return out | |
| # --------------------------------------------------------------------------- | |
| # Inference | |
| # --------------------------------------------------------------------------- | |
| def _gpu_duration(frame, preset="", game_name="", game_rules="", key_bindings="", model_name="", | |
| reasoning=True, temperature=1.0, max_new_tokens=512, *args, **kwargs): | |
| # measured on the Space (9B): ~3 s prefill + ~0.036 s/token decode; +15% margin | |
| return min(120, int(4 + 0.04 * int(max_new_tokens)) + 1) | |
| def predict( | |
| frame: Image.Image, | |
| preset: str = "Doom Battle-1", | |
| game_name: str = "", | |
| game_rules: str = "", | |
| key_bindings: str = "", | |
| model_name: str = DEFAULT_MODEL, | |
| reasoning: bool = True, | |
| temperature: float = 1.0, | |
| max_new_tokens: int = 512, | |
| seed: int = 0, | |
| ): | |
| """Predict the next 200 ms of keyboard/mouse actions for a game frame. | |
| Args: | |
| frame: the current game screen. | |
| preset: game preset used for any empty prompt field ("Doom Battle-1", "Minecraft (voxel sandbox)", "SuperTuxKart (racing)", "Custom"). | |
| game_name: game name for the system prompt (empty = from preset). | |
| game_rules: game rules / goal text (empty = from preset). | |
| key_bindings: one "- <key> -> <effect>" line per key (empty = from preset). | |
| model_name: which GameScaling checkpoint to use. | |
| reasoning: prefill "<think>" so the model may think before acting; otherwise act directly. | |
| temperature: sampling temperature (authors use 1.0). | |
| max_new_tokens: generation budget. | |
| seed: random seed. | |
| Returns: | |
| Annotated frame, parsed action summary (markdown), the model's reasoning, raw model reply. | |
| """ | |
| if frame is None: | |
| raise gr.Error("Please upload a game frame.") | |
| p = PRESETS.get(preset, PRESETS["Doom Battle-1"]) | |
| name = game_name.strip() or p["name"] | |
| rules = game_rules if game_rules.strip() else p["rules"] | |
| keys = key_bindings if key_bindings.strip() else p["keys"] | |
| system = build_system_prompt(name, rules, keys) | |
| img = frame.convert("RGB") | |
| if max(img.size) > MAX_SIDE: | |
| img.thumbnail((MAX_SIDE, MAX_SIDE), Image.BILINEAR) | |
| model, proc = LOADED.get(model_name, LOADED[DEFAULT_MODEL]) | |
| messages = [ | |
| {"role": "system", "content": system}, | |
| {"role": "user", "content": [{"type": "text", "text": USER_TURN_TEXT}, {"type": "image"}]}, | |
| ] | |
| prefill = THINK_PREFILL if reasoning else NO_THINK_PREFILL | |
| text = proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + prefill | |
| inputs = proc(text=[text], images=[img], return_tensors="pt").to("cuda") | |
| torch.manual_seed(int(seed)) | |
| t0 = time.perf_counter() | |
| with torch.inference_mode(): | |
| out = model.generate( | |
| **inputs, | |
| max_new_tokens=int(max_new_tokens), | |
| do_sample=temperature > 0, | |
| temperature=max(float(temperature), 1e-5), | |
| top_p=1.0, | |
| top_k=0, | |
| stop_strings=["<|action_end|>"], | |
| tokenizer=proc.tokenizer, | |
| ) | |
| dt = time.perf_counter() - t0 | |
| gen = out[0, inputs["input_ids"].shape[1]:] | |
| reply = proc.tokenizer.decode(gen, skip_special_tokens=False) | |
| reply = reply.replace("<|im_end|>", "").replace("<|endoftext|>", "") | |
| full = prefill + reply | |
| action = parse_action(full) | |
| reasoning_text, _, closed = split_reasoning(full) | |
| vis = render_overlay(img, action) | |
| if action.parsed: | |
| x, y, z = action.mouse | |
| rows = "\n".join( | |
| f"| {i + 1} | {i * 33}-{(i + 1) * 33} ms | {' '.join(f'`{k}`' for k in g) or '—'} |" | |
| for i, g in enumerate(action.steps) | |
| ) | |
| summary = ( | |
| f"**Mouse:** X = {x:+.0f}‰ of width ({'right' if x > 0 else 'left' if x < 0 else 'none'}), " | |
| f"Y = {y:+.0f}‰ of height ({'down' if y > 0 else 'up' if y < 0 else 'none'}), scroll = {z:+.0f}\n\n" | |
| f"| step | time | keys held |\n|---|---|---|\n{rows}\n\n" | |
| f"*{model_name} · {len(gen)} tokens · {dt:.1f}s*" | |
| ) | |
| if action.dropped: | |
| summary += f"\n\nDropped out-of-vocabulary keys: {action.dropped}" | |
| else: | |
| summary = ("**Could not parse an action** (" + ("generation was cut off while thinking — raise max tokens" | |
| if action.truncated else "malformed reply") + "). Executed as a no-op.") | |
| return vis, summary, (reasoning_text if closed else reasoning_text + " …") or "(acted directly, no thinking)", full | |
| def fill_preset(preset): | |
| if preset == "Custom": | |
| return gr.update(), gr.update(), gr.update() | |
| p = PRESETS[preset] | |
| return p["name"], p["rules"], p["keys"] | |
| 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( | |
| "# 🎮 GameScaling: game frame → next 200 ms of keyboard & mouse\n" | |
| "GameScaling models are Qwen3.5 VLMs continued-pretrained on 1,600 h of keyboard-and-mouse gameplay. " | |
| "Give it a game screen and a description of the game; it returns one **200 ms action chunk**: " | |
| "a mouse move plus the keys held in each of six 33 ms steps. " | |
| "[Model](https://huggingface.co/THUSI-Lab/GameScaling) · " | |
| "[Project page](https://gamescaling.github.io) · [Paper](https://arxiv.org/abs/2610.28698)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| frame = gr.Image(type="pil", label="Game frame", height=360) | |
| with gr.Row(): | |
| preset = gr.Dropdown(PRESET_NAMES, value="Doom Battle-1", label="Game preset", scale=2) | |
| model_name = gr.Dropdown(list(MODELS), value=DEFAULT_MODEL, label="Checkpoint", scale=2) | |
| reasoning = gr.Checkbox(value=True, label="Allow thinking (prefill <think>)") | |
| run = gr.Button("Predict action", variant="primary") | |
| with gr.Column(): | |
| vis = gr.Image(type="pil", label="Predicted action (red arrow = mouse move, bottom = 6 key steps)") | |
| summary = gr.Markdown() | |
| with gr.Accordion("Game description (system prompt slots)", open=False): | |
| game_name = gr.Textbox(label="Game name", value=PRESETS["Doom Battle-1"]["name"]) | |
| game_rules = gr.Textbox(label="Game rules", lines=8, value=PRESETS["Doom Battle-1"]["rules"]) | |
| key_bindings = gr.Textbox( | |
| label="Key bindings (one '- <key> -> <effect>' per line; training key names, e.g. W, LB, Space, one, Up)", | |
| lines=8, value=PRESETS["Doom Battle-1"]["keys"], | |
| ) | |
| with gr.Accordion("Model output", open=False): | |
| thinking = gr.Textbox(label="Reasoning", lines=4) | |
| raw = gr.Textbox(label="Raw reply", lines=4) | |
| with gr.Accordion("Advanced settings", open=False): | |
| temperature = gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="Temperature (authors use 1.0)") | |
| max_new_tokens = gr.Slider(64, 2048, value=512, step=64, label="Max new tokens (only used up when the model thinks)") | |
| seed = gr.Number(value=0, precision=0, label="Seed") | |
| preset.change(fill_preset, preset, [game_name, game_rules, key_bindings], api_visibility="private") | |
| outputs = [vis, summary, thinking, raw] | |
| run.click( | |
| predict, | |
| [frame, preset, game_name, game_rules, key_bindings, model_name, reasoning, temperature, max_new_tokens, seed], | |
| outputs, | |
| api_name="predict", | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/doom_battle1.jpg", "Doom Battle-1"], | |
| ["examples/minetest.jpg", "Minecraft (voxel sandbox)"], | |
| ["examples/supertuxkart.jpg", "SuperTuxKart (racing)"], | |
| ], | |
| inputs=[frame, preset], | |
| outputs=outputs, | |
| fn=predict, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| gr.Markdown( | |
| "<small>Single-frame demo: in a real game loop the model also sees the last 5 frames and its own past " | |
| "actions. Doom frame from the authors' repo; Minetest screenshot by Kotolegokot (CC BY-SA 3.0) and " | |
| "SuperTuxKart screenshot by the STK team (CC BY 3.0), via Wikimedia Commons.</small>" | |
| ) | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |