"""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)