app.py

"""LoopVL demo: a recurrent vision-language model with a user-controlled visual compute budget.

Loads the public TierFlow/LoopVL checkpoint (~1.6B parameters) with the authors'
single-file loader, then answers questions about an image. The distinctive control
is the recurrent schedule: the model reuses its 16-layer transformer stack, and the
H (high) and L (low) cycle counts decide how many times the stack is applied per
forward pass — the "inference budget" the paper sweeps across.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before any torch / CUDA-touching import

import time
from PIL import Image, ImageFile

import gradio as gr
import torch

ImageFile.LOAD_TRUNCATED_IMAGES = True

MODEL_ID = "TierFlow/LoopVL"
DEVICE = "cuda"

# Vendored from https://github.com/Tier-Flow/LoopVL (Apache-2.0; Penguin encoder
# keeps its Tencent LICENSE.txt alongside).
from modeling_loopvl import LoopVLForConditionalGeneration, load_tokenizer_and_processor
from hrm_penguin.generation import greedy_generate
from hrm_penguin.inference import build_inference_batch
from hrm_penguin.recurrent.visual_anchor import VisualAnchor

TRIGGER = "\nExplain your reasoning, then give the final answer."


def _download_weights() -> str:
    """Fetch config + tokenizer + single safetensors checkpoint into a local dir."""
    from huggingface_hub import snapshot_download

    return snapshot_download(
        MODEL_ID,
        allow_patterns=[
            "config.json",
            "generation_config.json",
            "preprocessor_config.json",
            "tokenizer.json",
            "tokenizer_config.json",
            "model.safetensors",
        ],
    )


# --- Load everything at module scope, eagerly on "cuda" (string device). ------
_model_dir = _download_weights()
model = LoopVLForConditionalGeneration.from_pretrained(_model_dir, device=DEVICE)
tokenizer, processor = load_tokenizer_and_processor(_model_dir, model.config)

# Base visual-anchor coefficients (H_cycles=2 in the shipped checkpoint). Any
# schedule change rebuilds VisualAnchor, re-copying these values (clamped).
# Kept on CPU: schedule changes run inside the ZeroGPU worker, where a plain
# CPU tensor is trivially inherited across forks and copied cross-device.
recurrent = model.recurrent_model
BASE_ANCHOR_VALUES = recurrent.visual_anchor.raw_alpha.detach().to("cpu").clone()
ANCHOR_MAX_SCALE = recurrent.visual_anchor.max_scale
LAST_SCHEDULE = {"h": int(recurrent.config.H_cycles), "l": int(recurrent.config.L_cycles)}


def _apply_schedule(h_cycles: int, l_cycles: int) -> None:
    """Set the recurrent inference schedule (the authors' configure() pattern).

    The forward pass reads config.H_cycles / config.L_cycles each call, and the
    visual anchor holds one learned coefficient per high cycle, so changing H
    rebuilds it and re-copies the base coefficients (clamped, as in the paper).
    """
    h_cycles, l_cycles = int(h_cycles), int(l_cycles)
    if h_cycles < 1 or l_cycles < 1:
        raise ValueError("H and L cycle counts must be positive")
    if LAST_SCHEDULE["h"] == h_cycles and LAST_SCHEDULE["l"] == l_cycles:
        return  # already configured; avoids rebuilding the anchor needlessly
    recurrent.config.H_cycles = h_cycles
    recurrent.config.L_cycles = l_cycles
    old = recurrent.visual_anchor
    new = VisualAnchor(h_cycles, max_scale=ANCHOR_MAX_SCALE).to(
        device=old.raw_alpha.device, dtype=old.raw_alpha.dtype
    )
    with torch.no_grad():
        for i in range(h_cycles):
            new.raw_alpha[i].copy_(
                BASE_ANCHOR_VALUES[min(i, len(BASE_ANCHOR_VALUES) - 1)]
            )
    recurrent.visual_anchor = new
    LAST_SCHEDULE["h"], LAST_SCHEDULE["l"] = h_cycles, l_cycles


# Coherent schedules from the paper's Table 4 sweep: (label, H, L) -> unrolled
# depth 16 * H * (L+1). H2L3 is the trained schedule and clearly the best; the
# H1/H4 rows in the paper collapse (MMStar 0.5/24), so they are not offered.
SCHEDULES = [
    ("H2L3 · 128 layers (trained default, best)", 2, 3),
    ("H2L2 · 96 layers (lighter, still solid)", 2, 2),
    ("H3L3 · 192 layers (heavier, slightly lower)", 3, 3),
    ("H3L2 · 144 layers", 3, 2),
    ("H2L1 · 64 layers (noticeably degraded)", 2, 1),
    ("H3L1 · 96 layers (noticeably degraded)", 3, 1),
]
_SCHEDULE_LOOKUP = {name: (h, l) for name, h, l in SCHEDULES}


def _estimate_duration(*args, **kwargs) -> int:
    """Callable GPU duration, calibrated from measured calls on this Space.

    Measured on ZeroGPU (cold first call per visitor): 14.1 s at H2L3/32 tok,
    19.3 s at H2L3/96 tok, 39.7 s at H3L3/96 tok. Per-token cost grows
    superlinearly with unrolled depth, so the coefficient is scaled by
    depth/128 with a fixed base that also covers worker init.

    Duration = 25 + budget x 0.2 s/token scaled by depth/128, capped at 120 s.
    Worst offered case (H3L3, 96 tok) resolves to ~54 s vs 39.7 s measured.
    """
    schedule = kwargs.get("schedule") or (args[2] if len(args) > 2 else None)
    budget = kwargs.get("max_tokens") or (args[3] if len(args) > 3 else 32)
    try:
        h, l = _SCHEDULE_LOOKUP.get(str(schedule), (2, 3))
    except Exception:
        h, l = 2, 3
    depth = 16 * h * (l + 1)
    return min(120, 25 + int(float(budget) * 0.2 * depth / 128))


@spaces.GPU(duration=_estimate_duration)
def answer(
    image: Image.Image,
    prompt: str,
    schedule: str = "H2L3 · 128 layers (trained default, best)",
    max_tokens: int = 32,
    use_trigger: bool = False,
    progress=gr.Progress(track_tqdm=True),
) -> str:
    """Answer a question about an image using LoopVL, a recurrent vision-language model.

    Args:
        image: the input image to reason about.
        prompt: the question or instruction about the image.
        schedule: recurrent compute schedule (H and L cycle counts). Higher means
            more effective transformer layers per forward pass.
        max_tokens: maximum number of tokens to generate.
        use_trigger: append the authors' reasoning trigger to the instruction.
        progress: Gradio progress bar.

    Returns:
        The model's answer text.
    """
    if image is None:
        return "Please provide an image."
    prompt = (prompt or "").strip()
    if not prompt:
        return "Please provide a question or instruction about the image."

    h, l = _SCHEDULE_LOOKUP.get(schedule, (2, 3))
    _apply_schedule(h, l)

    instruction = prompt + (TRIGGER if use_trigger else "")
    pil_image = image if isinstance(image, Image.Image) else Image.fromarray(image)
    batch = build_inference_batch(
        tokenizer,
        processor,
        pil_image.convert("RGB"),
        instruction,
        model.config,
        torch.device(DEVICE),
    )
    max_tokens = max(1, min(int(max_tokens), 96))
    start = time.perf_counter()
    with torch.inference_mode():
        generated = greedy_generate(
            model,
            batch,
            max_new_tokens=max_tokens,
            eos_token_id=tokenizer.eos_token_id,
        )
    elapsed = time.perf_counter() - start
    prediction = tokenizer.decode(generated[0], skip_special_tokens=True).strip()
    depth = 16 * h * (l + 1)
    label = f"Answer (schedule {schedule.split(' ')[0]} → {depth} effective layers, {elapsed:.1f}s):"
    return f"{label}\n{prediction}" if prediction else f"{label}\n(no tokens generated)"


CSS = """
#col-container { max-width: 1100px; 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(
            """
            # LoopVL: Recurrent Vision-Language Model

            **LoopVL** (~1.6B params) reuses its 16-layer transformer stack in loops
            instead of stacking 128 layers. You control the compute budget: pick an
            **H/L schedule** to decide how many times the stack runs per forward pass,
            then ask anything about the image.

            [Model](https://huggingface.co/TierFlow/LoopVL) ·
            [Code](https://github.com/Tier-Flow/LoopVL)
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(
                    type="pil", label="Image", sources=["upload", "clipboard"]
                )
                prompt_in = gr.Textbox(
                    label="Question / instruction",
                    placeholder="What is in the image?",
                    lines=2,
                )
                schedule_in = gr.Dropdown(
                    choices=[name for name, _, _ in SCHEDULES],
                    value="H2L3 · 128 layers (trained default, best)",
                    label="Recurrent schedule (compute budget)",
                    info="Unrolled depth = 16 × H × (L+1) effective layers per forward pass. "
                         "H2L3 is the trained schedule; other coherent schedules from the "
                         "paper's Table 4 sweep are offered for comparison.",
                )
                run_btn = gr.Button("Answer", variant="primary")
                with gr.Accordion("Advanced", open=False):
                    max_tokens_in = gr.Slider(
                        minimum=8, maximum=96, step=8, value=32,
                        label="Max new tokens",
                    )
                    trigger_in = gr.Checkbox(
                        value=False,
                        label="Append reasoning trigger",
                        info='Adds "Explain your reasoning, then give the final answer." '
                             "(the authors' math-reasoning prompt)",
                    )
            with gr.Column(scale=1):
                answer_out = gr.Textbox(label="Answer", lines=8)

        gr.Examples(
            examples=[
                ["examples/cat_tabby.jpg", "What animal is this, and what is it doing?"],
                ["examples/rubber_ducks.jpg", "How many objects are in this image? Describe them."],
                ["examples/landscape.jpg", "Describe the scene and the weather."],
                ["examples/dog.jpg", "What breed of dog does this look like?"],
            ],
            inputs=[image_in, prompt_in],
            outputs=answer_out,
            fn=answer,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples (click to load, then press Answer)",
        )

    run_btn.click(
        answer,
        inputs=[image_in, prompt_in, schedule_in, max_tokens_in, trigger_in],
        outputs=answer_out,
        api_name="generate",
        concurrency_limit=1,
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)