"""LLaDA-UI — GUI grounding & agent demo (block-wise diffusion VLM).
Faithful port of the official reference inference
(`inference/inference_hf.py` in inclusionAI/LLaDA-UI) onto ZeroGPU:
same prompt templates, same image preprocessing tier, same block-diffusion
decode loop (`model.generate`) and the same coordinate parsing.
"""
import os
import math
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import spaces # noqa: E402 (must precede torch / CUDA-touching imports)
import re # noqa: E402
import time # noqa: E402
import gradio as gr # noqa: E402
import torch # noqa: E402
from PIL import Image, ImageDraw # noqa: E402
from transformers import AutoImageProcessor, AutoTokenizer # noqa: E402
from agent_prompts import PLATFORMS # noqa: E402
from configuration_llada2_vl import LLaDA2VLMoEConfig # noqa: E402
from modeling_llada2_vl_moe import LLaDA2MoE_VLForConditionalGeneration # noqa: E402
MODEL_ID = "inclusionAI/LLaDA-UI"
# Block-diffusion special tokens (reference: inference_hf.py)
MASK_ID, EOS_ID = 156895, 156892
IGNORE_INDEX = -100
# Image tier from the model card quick start (IMAGE_MAX_PIXELS=12845056, highres)
IMAGE_MAX_PIXELS = 12845056
IMAGE_MIN_PIXELS = 20 * 28 * 28
# The text backbone was trained with max_position_embeddings=8192; keep the visual
# token budget well inside that so prompt + generation always fit.
MAX_IMAGE_TOKENS = 6144
# Grounding prompt — verbatim from the reference implementation.
POINT_TPL = (
"Output the center point of the position corresponding to the following instruction:\n"
"{instr}\n\n"
"The output should just be the coordinates of a point, in the format [x,y]. "
"Additionally, if the task is infeasible (e.g., the task is not related to the image), "
"the output should be [-1,-1]."
)
SYSTEM_MESSAGE = "You are a helpful assistant."
# The checkpoint's own chat_template injects tokens ("detailed thinking off") that the
# model never saw in training and which produce garbage. The reference overrides it with
# the exact training template; this is that template (plus an add_generation_prompt
# branch, matching inference/sglang_server/llada2_bd_chat_template.jinja).
TRAIN_CHAT_TEMPLATE = (
"{%- if messages and messages[0].role == 'system' %}"
"{{- 'SYSTEM' + messages[0].content + '\n' }}{%- endif %}"
"{%- for message in messages %}"
"{%- if message.content is string %}{%- set content = message.content %}"
"{%- else %}{%- set content = '' %}{%- endif %}"
"{%- if message.role == 'user' %}{{- 'HUMAN' + content + '<|role_end|>' }}"
"{%- elif message.role == 'assistant' %}{{- 'ASSISTANT' + content + '<|role_end|>' }}"
"{%- endif %}{%- endfor %}"
"{%- if add_generation_prompt %}{{- 'ASSISTANT' }}{%- endif %}"
)
COORD_RE = re.compile(
r"\[\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*"
r"(?:,\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*)?\]"
)
BOX_RE = re.compile(
r"\(\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*"
r"(?:,\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*)?\)"
)
# --------------------------------------------------------------------------------------
# Load once, at module scope, straight onto CUDA (ZeroGPU streams the weights in).
# --------------------------------------------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.chat_template = TRAIN_CHAT_TEMPLATE
image_processor = AutoImageProcessor.from_pretrained(
MODEL_ID,
use_fast=False, # the reference runs the slow Qwen2VLImageProcessor
min_pixels=IMAGE_MIN_PIXELS,
max_pixels=IMAGE_MAX_PIXELS,
)
image_processor.min_pixels = IMAGE_MIN_PIXELS
image_processor.max_pixels = IMAGE_MAX_PIXELS
image_processor.size = {"shortest_edge": IMAGE_MIN_PIXELS, "longest_edge": IMAGE_MAX_PIXELS}
config = LLaDA2VLMoEConfig.from_pretrained(MODEL_ID)
# config.json ships no `vision_pad_token_id`; the reference hardcodes 157187, which is
# exactly `image_token_id`.
config.vision_pad_token_id = config.image_token_id
model = LLaDA2MoE_VLForConditionalGeneration.from_pretrained(
MODEL_ID,
config=config,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
)
model.eval().to("cuda")
# `ASSISTANT` header the reference leaves un-masked at the end of the prompt.
ASSISTANT_HEADER_IDS = tokenizer.apply_chat_template([{"role": "assistant", "content": ""}])[:5]
# --------------------------------------------------------------------------------------
# Vision-tower memoisation.
#
# The block-diffusion loop calls the full `forward` once per denoising step, and the
# outer forward re-runs the ViT on the (unchanged) pixel values every single time. Only
# `input_ids` changes between steps, so caching the image embeddings is numerically
# identical and removes the dominant cost (32-64x fewer ViT passes per request).
# --------------------------------------------------------------------------------------
_visual_cache = {"key": None, "val": None, "src": None}
_orig_visual_forward = model.model.visual.forward
def _cached_visual_forward(hidden_states, grid_thw=None):
key = (
hidden_states.data_ptr(),
tuple(hidden_states.shape),
None if grid_thw is None else tuple(grid_thw.flatten().tolist()),
)
if _visual_cache["key"] == key:
return _visual_cache["val"]
out = _orig_visual_forward(hidden_states, grid_thw)
_visual_cache.update(key=key, val=out, src=hidden_states)
return out
model.model.visual.forward = _cached_visual_forward
def _reset_visual_cache():
_visual_cache.update(key=None, val=None, src=None)
# --------------------------------------------------------------------------------------
# Prompt / batch construction — mirrors preprocess_qwen_2_visual + build_batch.
# --------------------------------------------------------------------------------------
def _fit_image(image: Image.Image) -> Image.Image:
image = image.convert("RGB")
budget = MAX_IMAGE_TOKENS * 28 * 28
w, h = image.size
if w * h > budget:
scale = (budget / (w * h)) ** 0.5
image = image.resize((max(28, int(w * scale)), max(28, int(h * scale))), Image.LANCZOS)
return image
def build_batch(image: Image.Image, system_message: str, turns):
"""`turns` is a list of (role, content) with exactly one '' placeholder."""
image = _fit_image(image)
vp = image_processor.preprocess(image, return_tensors="pt")
pixel_values = vp["pixel_values"]
grid_thw = vp["image_grid_thw"][0]
n_image_tokens = int(grid_thw.prod() // image_processor.merge_size**2)
placeholder = "<|vision_start|>" + "<|image_pad|>" * n_image_tokens + "<|vision_end|>"
ids = list(tokenizer.apply_chat_template([{"role": "system", "content": system_message}]))
for role, content in turns:
content = content.replace("", placeholder)
ids += tokenizer.apply_chat_template([{"role": role, "content": content}])
ids += ASSISTANT_HEADER_IDS
input_ids = torch.tensor(ids, dtype=torch.long).unsqueeze(0)
data = {
"input_ids": input_ids.to("cuda"),
"pixel_values": pixel_values.to("cuda", torch.bfloat16),
"image_grid_thw": grid_thw.unsqueeze(0).to("cuda"),
"position_ids": None,
"attention_mask": None,
}
return data, len(ids), image
def _decode(out, prompt_len):
raw = out[0][prompt_len:].tolist()
return tokenizer.decode(
[t for t in raw if t not in (MASK_ID, EOS_ID)], skip_special_tokens=True
).strip()
def _run_diffusion(data, prompt_len, gen_length, steps, block_length, temperature, threshold):
_reset_visual_cache()
# lm_head only needs the active block — exact, and it avoids materialising a
# (1, seq_len, 157188) logits tensor every step.
data = dict(data)
data["logits_to_keep"] = int(block_length)
t0 = time.time()
with torch.inference_mode():
out = model.generate(
data=data,
gen_length=int(gen_length),
steps=int(steps),
block_length=int(block_length),
temperature=float(temperature),
threshold=float(threshold),
eos_id=EOS_ID,
mask_id=MASK_ID,
)
elapsed = time.time() - t0
_reset_visual_cache()
return _decode(out, prompt_len), elapsed
# --------------------------------------------------------------------------------------
# Output parsing / drawing
# --------------------------------------------------------------------------------------
def parse_point(text, pattern=COORD_RE):
"""Reference parse_point: 4 numbers -> bbox centre, 2 numbers -> point, /1000."""
if not text:
return None
matches = list(pattern.finditer(text))
if not matches:
return None
nums = [float(x) for x in matches[-1].groups() if x is not None]
if len(nums) >= 4:
x1, y1, x2, y2 = nums[:4]
if x1 == -1 and y1 == -1:
return [-1.0, -1.0]
return [(x1 + x2) / 2 / 1000.0, (y1 + y2) / 2 / 1000.0]
if len(nums) >= 2:
x, y = nums[:2]
if x == -1 and y == -1:
return [-1.0, -1.0]
return [x / 1000.0, y / 1000.0]
return None
def draw_point(image: Image.Image, point):
canvas = image.convert("RGB").copy()
if point is None or point == [-1.0, -1.0]:
return canvas
w, h = canvas.size
cx, cy = point[0] * w, point[1] * h
draw = ImageDraw.Draw(canvas, "RGBA")
r = max(10, int(min(w, h) * 0.022))
draw.ellipse([cx - r * 2.2, cy - r * 2.2, cx + r * 2.2, cy + r * 2.2], fill=(255, 61, 0, 60))
draw.line([cx - r * 2.6, cy, cx + r * 2.6, cy], fill=(255, 61, 0, 235), width=max(2, r // 4))
draw.line([cx, cy - r * 2.6, cx, cy + r * 2.6], fill=(255, 61, 0, 235), width=max(2, r // 4))
draw.ellipse(
[cx - r, cy - r, cx + r, cy + r],
outline=(255, 255, 255, 255),
width=max(2, r // 3),
)
draw.ellipse(
[cx - r * 0.55, cy - r * 0.55, cx + r * 0.55, cy + r * 0.55], fill=(255, 61, 0, 255)
)
return canvas
# --------------------------------------------------------------------------------------
# ZeroGPU duration estimation
#
# Measured on this Space (RTX PRO 6000, bf16): runtime is essentially linear in the
# number of denoising forward passes and near-flat in prompt length, because the vision
# tower is memoised across steps and lm_head only runs over the active block.
#
# 32 denoising steps (4737 prompt tokens) -> 6.4 s
# 256 denoising steps (5574 prompt tokens) -> 36.1 s
# => ~0.133 s/step + ~2.2 s fixed
#
# A cold GPU lease additionally streams the 33.8 GB of packed weights into VRAM, which
# is charged to the call (first measured call: 18.6 s vs 6.4 s warm) -> ~12 s allowance.
# --------------------------------------------------------------------------------------
GPU_COLD_START_S = 14.0
GPU_SECONDS_PER_STEP = 0.14
GPU_SAFETY = 1.35
GPU_MAX_DURATION = 120
def _total_steps(gen_length: int, steps: int, block_length: int) -> int:
blocks = max(1, math.ceil(float(gen_length) / max(1.0, float(block_length))))
return int(blocks * max(1, int(steps)))
def _estimate_duration(gen_length, steps, block_length) -> int:
secs = GPU_COLD_START_S + GPU_SECONDS_PER_STEP * _total_steps(gen_length, steps, block_length)
return int(min(GPU_MAX_DURATION, math.ceil(secs * GPU_SAFETY)))
def _ground_duration(image, instruction, gen_length=32, steps=32, block_length=32, *a, **k):
return _estimate_duration(gen_length, steps, block_length)
def _act_duration(
image, task, platform="Mobile", previous_action="",
gen_length=256, steps=32, block_length=32, *a, **k,
):
return _estimate_duration(gen_length, steps, block_length)
# --------------------------------------------------------------------------------------
# Mode 1 — GUI grounding
# --------------------------------------------------------------------------------------
@spaces.GPU(duration=_ground_duration)
def ground(
image,
instruction: str,
gen_length: int = 32,
steps: int = 32,
block_length: int = 32,
temperature: float = 1.0,
threshold: float = 0.99,
):
"""Locate the UI element an instruction refers to and return its centre point.
Args:
image: a screenshot (mobile, desktop or web).
instruction: what to point at, e.g. "click the search button".
gen_length: number of tokens to denoise.
steps: denoising steps per block.
block_length: block-diffusion block size.
temperature: 0 = greedy, 1.0 = the reference setting.
threshold: confidence needed to commit a token in a denoising step.
Returns:
The screenshot with the predicted point marked, the point itself, and the raw
model output.
"""
if image is None:
raise gr.Error("Please provide a screenshot.")
if not instruction or not instruction.strip():
raise gr.Error("Please provide an instruction.")
prompt = POINT_TPL.format(instr=instruction.strip())
# Reference order: prompt text first, image placeholder last (~17 pts better).
data, prompt_len, fitted = build_batch(
image, SYSTEM_MESSAGE, [("user", f"{prompt}\n")]
)
text, elapsed = _run_diffusion(
data, prompt_len, gen_length, steps, block_length, temperature, threshold
)
point = parse_point(text)
if point == [-1.0, -1.0]:
summary = "Model reported the instruction as **infeasible** for this screenshot (`[-1,-1]`)."
elif point is None:
summary = "Could not parse a coordinate from the model output."
else:
w, h = fitted.size
summary = (
f"**Point (0-999):** `[{round(point[0] * 1000)}, {round(point[1] * 1000)}]` \n"
f"**Normalised:** `[{point[0]:.3f}, {point[1]:.3f}]` \n"
f"**Pixels ({w}x{h}):** `({round(point[0] * w)}, {round(point[1] * h)})`"
)
summary += f" \n{prompt_len} prompt tokens - {elapsed:.1f}s"
return draw_point(fitted, point), summary, text
# --------------------------------------------------------------------------------------
# Mode 2 — GUI agent (think + action)
# --------------------------------------------------------------------------------------
@spaces.GPU(duration=_act_duration)
def act(
image,
task: str,
platform: str = "Mobile",
previous_action: str = "",
gen_length: int = 256,
steps: int = 32,
block_length: int = 32,
temperature: float = 0.0,
threshold: float = 0.95,
):
"""Predict the next GUI action for a task, given the current screenshot.
Args:
image: the current screenshot.
task: the high-level user task, e.g. "create a new project".
platform: which released agent prompt to use - Mobile, Desktop or Web.
previous_action: the agent's previous raw "......"
turn, if continuing a rollout.
gen_length: maximum tokens to denoise.
steps: denoising steps per block.
block_length: block-diffusion block size.
temperature: 0 = greedy (the released agent setting).
threshold: confidence needed to commit a token in a denoising step.
Returns:
The screenshot with any predicted coordinate marked, the parsed action, the
model's reasoning, and the raw output.
"""
if image is None:
raise gr.Error("Please provide a screenshot.")
if not task or not task.strip():
raise gr.Error("Please describe the task.")
system_tpl, user_tpl = PLATFORMS[platform]
system_message = system_tpl.replace("{task}", task.strip())
user_text = user_tpl.replace("{task}", task.strip())
turns = []
if previous_action and previous_action.strip():
# Released contract: an empty historical user turn carrying the previous
# assistant response verbatim.
turns.append(("user", ""))
turns.append(("assistant", previous_action.strip()))
turns.append(("user", user_text + ""))
data, prompt_len, fitted = build_batch(image, system_message, turns)
text, elapsed = _run_diffusion(
data, prompt_len, gen_length, steps, block_length, temperature, threshold
)
think = ""
m = re.search(r"(.*?)", text, flags=re.S)
if m:
think = m.group(1).strip()
m = re.search(r"(.*?)", text, flags=re.S)
action = m.group(1).strip() if m else text.strip()
point = parse_point(action, BOX_RE)
marked = draw_point(fitted, point)
action_md = f"```\n{action}\n```\n{prompt_len} prompt tokens - {elapsed:.1f}s"
return marked, action_md, think, text
# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1180px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""
GROUND_EXAMPLES = [
["examples/desktop_windows_start.png", "open the Microsoft store"],
["examples/web_gitlab.png", "click the button to create a new project"],
["examples/mobile_ios_home.png", "open facetime app"],
["examples/desktop_solitaire.png", "view solitaire daily challenges"],
["examples/mobile_android_calendar.png", "add new event on calendar"],
]
AGENT_EXAMPLES = [
["examples/web_gitlab.png", "Create a new project called 'llada-ui-demo'.", "Web"],
["examples/mobile_android_calendar.png", "Add a dentist appointment to my calendar.", "Mobile"],
["examples/desktop_solitaire.png", "Start today's Solitaire daily challenge.", "Desktop"],
]
with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="LLaDA-UI") as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(
"# LLaDA-UI - GUI grounding & agent\n"
"[inclusionAI/LLaDA-UI](https://huggingface.co/inclusionAI/LLaDA-UI) is a 16.9B MoE "
"vision-language GUI agent that decodes with **block-wise diffusion** instead of "
"autoregression. Give it a screenshot and it points at UI elements or predicts the "
"next action. Coordinates are normalised to `0-999`."
)
with gr.Tabs():
with gr.Tab("Grounding"):
with gr.Row():
with gr.Column():
g_image = gr.Image(label="Screenshot", type="pil", height=360)
with gr.Row():
g_instr = gr.Textbox(
show_label=False,
placeholder="click the search button",
container=False,
scale=4,
)
g_btn = gr.Button("Locate", variant="primary", scale=1)
with gr.Column():
g_out_img = gr.Image(label="Predicted point", type="pil", height=360)
g_out_md = gr.Markdown(label="Coordinates")
g_out_raw = gr.Textbox(label="Raw model output", lines=2)
with gr.Accordion("Advanced settings", open=False):
with gr.Row():
g_gen = gr.Slider(8, 64, value=32, step=8, label="Generation length")
g_steps = gr.Slider(8, 64, value=32, step=1, label="Denoising steps / block")
g_block = gr.Slider(16, 64, value=32, step=16, label="Block length")
with gr.Row():
g_temp = gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="Temperature")
g_thr = gr.Slider(
0.5, 1.0, value=0.99, step=0.01, label="Confidence threshold"
)
gr.Examples(
examples=GROUND_EXAMPLES,
inputs=[g_image, g_instr],
outputs=[g_out_img, g_out_md, g_out_raw],
fn=ground,
cache_examples=True,
cache_mode="lazy",
)
with gr.Tab("GUI agent"):
with gr.Row():
with gr.Column():
a_image = gr.Image(label="Current screenshot", type="pil", height=360)
a_task = gr.Textbox(
label="User task",
placeholder="Search for a one-way flight from Calgary to New York.",
)
with gr.Row():
a_platform = gr.Radio(
["Mobile", "Desktop", "Web"], value="Mobile", label="Agent prompt"
)
a_btn = gr.Button("Predict action", variant="primary")
with gr.Column():
a_out_img = gr.Image(label="Predicted target", type="pil", height=360)
a_out_action = gr.Markdown(label="Action")
a_out_think = gr.Textbox(label="Reasoning", lines=4)
a_out_raw = gr.Textbox(label="Raw model output", lines=3)
with gr.Accordion("Advanced settings", open=False):
a_prev = gr.Textbox(
label="Previous agent turn (optional)",
value="",
placeholder="...\nClick(box=(500,293))",
lines=2,
)
with gr.Row():
a_gen = gr.Slider(32, 384, value=256, step=32, label="Max generation length")
a_steps = gr.Slider(8, 40, value=32, step=1, label="Denoising steps / block")
a_block = gr.Slider(32, 64, value=32, step=32, label="Block length")
with gr.Row():
a_temp = gr.Slider(0.0, 1.5, value=0.0, step=0.05, label="Temperature")
a_thr = gr.Slider(
0.5, 1.0, value=0.95, step=0.01, label="Confidence threshold"
)
gr.Examples(
examples=AGENT_EXAMPLES,
inputs=[a_image, a_task, a_platform],
outputs=[a_out_img, a_out_action, a_out_think, a_out_raw],
fn=act,
cache_examples=True,
cache_mode="lazy",
)
gr.Markdown(
"Example screenshots come from "
"[OS-Copilot/ScreenSpot-v2](https://huggingface.co/datasets/OS-Copilot/ScreenSpot-v2) "
"(Apache-2.0)."
)
g_inputs = [g_image, g_instr, g_gen, g_steps, g_block, g_temp, g_thr]
g_outputs = [g_out_img, g_out_md, g_out_raw]
g_btn.click(ground, inputs=g_inputs, outputs=g_outputs, api_name="ground")
g_instr.submit(ground, inputs=g_inputs, outputs=g_outputs, api_name=False)
a_inputs = [a_image, a_task, a_platform, a_prev, a_gen, a_steps, a_block, a_temp, a_thr]
a_outputs = [a_out_img, a_out_action, a_out_think, a_out_raw]
a_btn.click(act, inputs=a_inputs, outputs=a_outputs, api_name="act")
if __name__ == "__main__":
demo.queue(max_size=20).launch(mcp_server=True, show_error=True)