app.py

"""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' %}"
    "{{- '<role>SYSTEM</role>' + 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' %}{{- '<role>HUMAN</role>' + content + '<|role_end|>' }}"
    "{%- elif message.role == 'assistant' %}{{- '<role>ASSISTANT</role>' + content + '<|role_end|>' }}"
    "{%- endif %}{%- endfor %}"
    "{%- if add_generation_prompt %}{{- '<role>ASSISTANT</role>' }}{%- 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")

# `<role>ASSISTANT</role>` 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 '<image>' 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("<image>", 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<image>")]
    )
    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<sub>{prompt_len} prompt tokens - {elapsed:.1f}s</sub>"
    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 "<think>...</think><action>...</action>"
            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 + "<image>"))

    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"<think>(.*?)</think>", text, flags=re.S)
    if m:
        think = m.group(1).strip()
    m = re.search(r"<action>(.*?)</action>", 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<sub>{prompt_len} prompt tokens - {elapsed:.1f}s</sub>"
    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="<think>...</think>\n<action>Click(box=(500,293))</action>",
                        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)