"""RynnValue-8B — robotic value model demo. Given a robot manipulation video and the task instruction, predicts how much time is left until the task is finished at every point along the clip, plus a short analysis block (description / instruction match / success). Follows the official reference implementation (https://github.com/alibaba-damo-academy/RynnValue, `rynn_infer/inference.py`): prefix-uniform sampling — for evaluation step *i* the prefix `frames[0:i]` is resampled to `num_frames` frames and the model's **last** prediction slot is read out, so every score only conditions on frames seen so far. """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 (must precede torch / transformers) import re # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 import gradio as gr # noqa: E402 import imageio.v2 as imageio # noqa: E402 import numpy as np # noqa: E402 import torch # noqa: E402 from PIL import Image # noqa: E402 from transformers import AutoConfig, AutoModel, AutoProcessor # noqa: E402 from plot_utils import save_video_with_trend # noqa: E402 MODEL_ID = "Alibaba-DAMO-Academy/RynnValue-8B" # Defaults — each is also the default of its UI component, so clicking an # example and pressing "Analyze" behave identically. DEFAULT_ROBOT = "a Franka single-arm robot" DEFAULT_CAMERA = "the main camera" DEFAULT_NUM_STEPS = 16 DEFAULT_NUM_FRAMES = 32 DEFAULT_MAX_SIDE = 384 DEFAULT_MAX_NEW_TOKENS = 128 # Rendering budget: the input clip is temporally subsampled to at most this many # frames and the playback fps is scaled to match, so wall-clock duration and the # ground-truth "remaining time" reference curve are unchanged. MAX_RENDER_FRAMES = 320 DISPLAY_MAX_SIDE = 640 # Longest ZeroGPU slot this demo will ever request. GPU_BUDGET_S = 240 ROBOT_CHOICES = [ "a Franka single-arm robot", "an SO-101 single-arm robot", "a WidowX single-arm robot", "a Jaco single-arm robot", "a Koch dual-arm robot", "an xArm single-arm robot", "an Trossen dual-arm robot", ] CAMERA_CHOICES = [ "the main camera", "the side camera", "the top-down camera", "the wrist-mounted camera", "the main left camera", "the main right camera", ] # --------------------------------------------------------------------------- # # Model # --------------------------------------------------------------------------- # print(f"Loading {MODEL_ID} ...", flush=True) _config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True) # The exported config predates the attn-impl field, so force the custom # prediction-slot isolation attention (mirrors rynn_infer/inference.py). _config._attn_implementation = "pred_slot_isolated_eager" model = AutoModel.from_pretrained( MODEL_ID, config=_config, trust_remote_code=True, torch_dtype=torch.bfloat16, ) # `torch_dtype=` leaves the value heads in fp32 (they are constructed with an # explicit dtype), so the trailing `dtype=` cast is required — the reference # script does the same `model.to(device=..., dtype=...)`. model = model.to(device="cuda", dtype=torch.bfloat16).eval() processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) tokenizer = processor.tokenizer EOS_TOKEN_ID = tokenizer.convert_tokens_to_ids("<|im_end|>") print("Model ready.", flush=True) # --------------------------------------------------------------------------- # # Video / sampling helpers (ported from rynn_infer/inference.py) # --------------------------------------------------------------------------- # def _resize_long_side(img: Image.Image, max_side: int) -> Image.Image: w, h = img.size if max_side <= 0 or max(w, h) <= max_side: return img scale = max_side / float(max(w, h)) return img.resize((max(1, int(round(w * scale))), max(1, int(round(h * scale)))), Image.BICUBIC) def load_video(video_path: str): """Decode a video into ``(frames, fps)``. Frames are downscaled to ``DISPLAY_MAX_SIDE`` on the fly and strided so that at most ``MAX_RENDER_FRAMES`` are kept; ``fps`` is scaled accordingly so the rendered clip keeps real-time playback speed. """ if not video_path or not os.path.isfile(video_path): raise gr.Error("Please upload a video first.") reader = imageio.get_reader(video_path) try: meta = reader.get_meta_data() src_fps = float(meta.get("fps") or 30.0) duration = float(meta.get("duration") or 0.0) est_total = int(duration * src_fps) if duration > 0 else 0 stride = 1 if est_total > MAX_RENDER_FRAMES: stride = int(np.ceil(est_total / float(MAX_RENDER_FRAMES))) frames = [] for i, frame in enumerate(reader): if i % stride: continue frames.append(_resize_long_side(Image.fromarray(frame).convert("RGB"), DISPLAY_MAX_SIDE)) if len(frames) >= MAX_RENDER_FRAMES + 8: break finally: reader.close() if not frames: raise gr.Error("Could not decode any frame from this video.") return frames, src_fps / float(stride) def sample_frame_indices(total: int, num_frames: int): """Uniformly pick ``num_frames`` indices out of ``total`` frames.""" if num_frames <= 0 or num_frames >= total: return list(range(total)) if num_frames == 1: return [total - 1] step = (total - 1) / (num_frames - 1) return [int(round(j * step)) for j in range(num_frames)] _DESCRIPTION_RE = re.compile(r"-\s*Video Description:\s*(.+)", re.IGNORECASE) _MATCH_RE = re.compile(r"-\s*Match:\s*(Yes|No)", re.IGNORECASE) _SUCCESS_RE = re.compile(r"-\s*Success:\s*(Yes|No)", re.IGNORECASE) def parse_analysis(text: str) -> dict: """Extract description / match / success from the generated Analysis block.""" def _first(pattern): m = pattern.search(text) return m.group(1).strip() if m else None return { "description": _first(_DESCRIPTION_RE), "match": _first(_MATCH_RE), "success": _first(_SUCCESS_RE), } def _estimate_duration( video_path=None, instruction="", robot_description=DEFAULT_ROBOT, camera_description=DEFAULT_CAMERA, num_steps=DEFAULT_NUM_STEPS, num_frames=DEFAULT_NUM_FRAMES, max_image_side=DEFAULT_MAX_SIDE, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, *args, **kwargs, ): """ZeroGPU duration estimate. The cost is driven by the prefix sequence length, which depends on the video's aspect ratio as well as the sliders — so the clip is probed for its resolution and the visual-token count is computed exactly. The linear and quadratic coefficients were fitted on the live Space (measured: 24.6 s for 240x240 @ 32 frames, 31.8 s for 640x360, 41.8 s for 640x480, all 16 steps). """ steps = int(num_steps or DEFAULT_NUM_STEPS) frames = int(num_frames or DEFAULT_NUM_FRAMES) side = int(max_image_side or DEFAULT_MAX_SIDE) width, height, n_src = 640.0, 480.0, float(MAX_RENDER_FRAMES) try: reader = imageio.get_reader(video_path) meta = reader.get_meta_data() reader.close() width, height = (float(v) for v in meta["size"]) n_src = min( float(MAX_RENDER_FRAMES), float(meta.get("duration") or 0.0) * float(meta.get("fps") or 30.0), ) except Exception: # unreadable metadata — fall back to the worst case pass scale = min(1.0, DISPLAY_MAX_SIDE / max(width, height), side / max(width, height)) # 16 px patches merged 2x2 -> one visual token per 32x32 px, plus the 8+8 # / tokens the processor emits per frame. tokens_per_frame = np.ceil(width * scale / 32.0) * np.ceil(height * scale / 32.0) + 16 seq_len = frames * tokens_per_frame + 64 per_step = 4.6e-4 * seq_len + 1.0e-8 * seq_len**2 if _pick_batch_size(frames, side) < 4: per_step *= 1.4 overhead = 8.0 + 0.05 * int(max_new_tokens or 128) + min(9.0, n_src / 40.0) raw = steps * per_step + overhead return int(max(30.0, raw)) if raw <= GPU_BUDGET_S else GPU_BUDGET_S def _pick_batch_size(num_frames: int, max_image_side: int) -> int: """Prefix sub-samples per forward pass. The checkpoint's `pred_slot_isolated_eager` attention materialises the full (batch, heads, L, L) fp32 score matrix, so the batch has to shrink as the per-prefix sequence grows or the GPU runs out of memory. """ cost = (num_frames / 32.0) * (max_image_side / 384.0) ** 2 if cost <= 1.05: return 4 if cost <= 1.6: return 2 return 1 # --------------------------------------------------------------------------- # # Inference # --------------------------------------------------------------------------- # @spaces.GPU(duration=_estimate_duration) def analyze_video( video_path: str, instruction: str, robot_description: str = DEFAULT_ROBOT, camera_description: str = DEFAULT_CAMERA, num_steps: int = DEFAULT_NUM_STEPS, num_frames: int = DEFAULT_NUM_FRAMES, max_image_side: int = DEFAULT_MAX_SIDE, max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, progress=gr.Progress(), ): """Predict how much time is left before a robot finishes a task. Runs RynnValue-8B over a robot manipulation video: for each evaluated step the prefix of the video seen so far is uniformly resampled and the model predicts the remaining time to task completion in seconds. It also generates an Analysis block (what happens in the video, whether the video matches the instruction, whether the task succeeded). Args: video_path: Path to the robot trajectory video (mp4). instruction: Natural-language task the robot is supposed to accomplish. robot_description: Embodiment description, e.g. "a Franka single-arm robot". camera_description: Camera viewpoint description, e.g. "the main camera". num_steps: Number of prefix steps evaluated along the video. num_frames: Frames uniformly resampled from each prefix. max_image_side: Longer side each frame is resized to before the model sees it. max_new_tokens: Token budget for the generated Analysis block. Returns: A video with the synchronized remaining-time curve, and a markdown report. """ if not instruction or not instruction.strip(): raise gr.Error("Please give the task instruction the robot is supposed to follow.") instruction = instruction.strip() robot_description = (robot_description or "").strip() or None camera_description = (camera_description or "").strip() or None if robot_description is None and camera_description is None: raise gr.Error( "This checkpoint was trained with meta information — fill in the " "embodiment and/or viewpoint description." ) num_steps = int(num_steps) num_frames = int(num_frames) max_image_side = int(max_image_side) max_new_tokens = int(max_new_tokens) if _estimate_duration( video_path, instruction, robot_description, camera_description, num_steps, num_frames, max_image_side, max_new_tokens, ) >= GPU_BUDGET_S: raise gr.Error( f"These settings need more than the {GPU_BUDGET_S}s GPU slot this demo " "requests — lower *Prefix steps*, *Frames per step* or *Max image side* " "in Advanced settings (or use a shorter clip)." ) t0 = time.perf_counter() progress(0.05, desc="Decoding video…") frames, fps = load_video(video_path) total = len(frames) model_frames = [_resize_long_side(f, max_image_side) for f in frames] eval_indices = sample_frame_indices(total, num_steps) t_decode = time.perf_counter() - t0 def build_prefix_sample(end_idx): frame_idx = np.linspace(0, end_idx, num_frames, dtype=int) return processor.process_episode( instruction=instruction, images=[model_frames[j] for j in frame_idx], robot_description=robot_description, camera_description=camera_description, ) def run_batch(samples): batch_kwargs = dict( input_ids=torch.cat([s["input_ids"] for s in samples], dim=0).to("cuda").long(), attention_mask=torch.cat([s["attention_mask"] for s in samples], dim=0).to("cuda").long(), pixel_values=torch.cat([s["pixel_values"].flatten(0, 1) for s in samples], dim=0).to("cuda"), image_grid_thw=torch.cat( [s["image_grid_thw"].flatten(0, 1) for s in samples], dim=0 ).to("cuda").long(), ) try: with torch.inference_mode(): outputs = model(**batch_kwargs) except (torch.cuda.OutOfMemoryError, RuntimeError) as exc: torch.cuda.empty_cache() raise gr.Error( "Ran out of GPU memory for these settings — the checkpoint's custom " "attention keeps the full attention matrix in memory, so lower " "*Frames per step* or *Max image side* in Advanced settings. " f"({type(exc).__name__})" ) from exc pred = outputs.value.pred_value if pred.dim() == 2 and pred.shape[0] == 1: pred = pred.reshape(len(samples), -1) if pred.dim() == 3: pred = pred.mean(dim=0) if pred.dim() == 2 and pred.shape[-1] > 1: pred = pred[:, -1] elif pred.dim() == 2: pred = pred[:, 0] return pred.float().reshape(-1).tolist() t1 = time.perf_counter() batch_size = _pick_batch_size(num_frames, max_image_side) pred_value, batch, final_sample = [], [], None for step, end_idx in enumerate(eval_indices): sample = build_prefix_sample(end_idx) if step == len(eval_indices) - 1: final_sample = sample batch.append(sample) if len(batch) >= batch_size or step == len(eval_indices) - 1: pred_value.extend(run_batch(batch)) batch = [] progress( 0.05 + 0.65 * len(pred_value) / max(1, len(eval_indices)), desc=f"Value prediction {len(pred_value)}/{len(eval_indices)}", ) t_value = time.perf_counter() - t1 # ---- Analysis pass on the full-video prefix ---------------------------- # progress(0.72, desc="Generating analysis…") t2 = time.perf_counter() input_ids = final_sample["input_ids"].to("cuda").long() with torch.inference_mode(): gen_out = model.generate( input_ids=input_ids, attention_mask=final_sample["attention_mask"].to("cuda").long(), pixel_values=final_sample["pixel_values"].flatten(0, 1).to("cuda"), image_grid_thw=final_sample["image_grid_thw"].flatten(0, 1).to("cuda").long(), max_new_tokens=max_new_tokens, do_sample=False, num_beams=1, eos_token_id=EOS_TOKEN_ID, pad_token_id=EOS_TOKEN_ID, use_cache=True, ) analysis_text = tokenizer.decode(gen_out[0, input_ids.shape[1]:], skip_special_tokens=True) analysis = parse_analysis(analysis_text) t_gen = time.perf_counter() - t2 # ---- Render ------------------------------------------------------------ # progress(0.85, desc="Rendering trend video…") t3 = time.perf_counter() out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name save_video_with_trend( images=frames, value=pred_value, output_path=out_path, fps=max(1.0, float(fps)), title="Remaining Time (s)", task_title=instruction, sampled_indices=eval_indices, ) t_render = time.perf_counter() - t3 total_time = time.perf_counter() - t0 clip_len = (total - 1) / float(fps) first_pred, last_pred = float(pred_value[0]), float(pred_value[-1]) def _badge(v): if v is None: return "—" return "✅ Yes" if v.lower() == "yes" else "❌ No" report = f"""### Analysis - **Video description:** {analysis['description'] or '—'} - **Matches the instruction:** {_badge(analysis['match'])} - **Task succeeded:** {_badge(analysis['success'])} ### Predicted remaining time | | predicted | actual (clip) | |---|---|---| | at the first frame | **{first_pred:.2f} s** | {clip_len:.2f} s | | at the last frame | **{last_pred:.2f} s** | 0.00 s | {len(eval_indices)} prefix steps · {num_frames} frames/step · {max_image_side} px · decode {t_decode:.1f}s · value {t_value:.1f}s · analysis {t_gen:.1f}s · render {t_render:.1f}s · total {total_time:.1f}s
Raw generation ``` {analysis_text.strip()} ```
""" print( f"[timing] decode={t_decode:.2f}s value={t_value:.2f}s gen={t_gen:.2f}s " f"render={t_render:.2f}s total={total_time:.2f}s steps={len(eval_indices)} " f"frames={num_frames} side={max_image_side}", flush=True, ) return out_path, report # --------------------------------------------------------------------------- # # UI # --------------------------------------------------------------------------- # CSS = """ #col-container { max-width: 1200px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="RynnValue-8B") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """# RynnValue-8B — how much longer will the robot take? A general-purpose value model for robot manipulation from Alibaba DAMO Academy. Give it a trajectory video plus the task instruction: it predicts the **remaining time to completion (seconds)** at every step along the clip, and says whether the video actually matches the instruction and whether the task succeeded. [Model](https://huggingface.co/Alibaba-DAMO-Academy/RynnValue-8B) · [Code](https://github.com/alibaba-damo-academy/RynnValue) · [Project page](https://alibaba-damo-academy.github.io/RynnValue.github.io/) """ ) with gr.Row(): with gr.Column(scale=1): video_in = gr.Video(label="Robot trajectory video", sources=["upload"]) instruction_in = gr.Textbox( label="Task instruction", placeholder="Put the box in the drawer and close it", ) with gr.Row(): robot_in = gr.Dropdown( ROBOT_CHOICES, value=DEFAULT_ROBOT, label="Embodiment", allow_custom_value=True, scale=1, ) camera_in = gr.Dropdown( CAMERA_CHOICES, value=DEFAULT_CAMERA, label="Viewpoint", allow_custom_value=True, scale=1, ) run_btn = gr.Button("Analyze trajectory", variant="primary") with gr.Column(scale=1): video_out = gr.Video(label="Remaining-time curve", autoplay=True) report_out = gr.Markdown() with gr.Accordion("Advanced settings", open=False): with gr.Row(): num_steps_in = gr.Slider( 4, 48, value=DEFAULT_NUM_STEPS, step=1, label="Prefix steps", info="How many points along the video are evaluated", ) num_frames_in = gr.Slider( 8, 48, value=DEFAULT_NUM_FRAMES, step=8, label="Frames per step", info="Frames uniformly resampled from each prefix", ) with gr.Row(): image_side_in = gr.Slider( 224, 512, value=DEFAULT_MAX_SIDE, step=32, label="Max image side (px)", info="Frames are downscaled to this before the model sees them", ) tokens_in = gr.Slider( 32, 256, value=DEFAULT_MAX_NEW_TOKENS, step=16, label="Analysis max new tokens", ) gr.Markdown( "The reference implementation evaluates one prefix step per frame with 64 " "frames at 640 px; the defaults here are trimmed so a run fits comfortably in a " "ZeroGPU slot." ) gr.Examples( examples=[ [ "examples/franka_box_into_drawer.mp4", "Put the box in the drawer and close it", "a Franka single-arm robot", "the main camera", ], [ "examples/soar_put_green_stick_in_brown_bowl.mp4", "Put green stick in brown bowl", "a WidowX single-arm robot", "the main camera", ], [ "examples/berkeley_rpt_stack_cup.mp4", "Pick up the yellow cup and stack it on the other cup", "a Franka single-arm robot", "the wrist-mounted camera", ], [ "examples/jaco_play_pick_up_green_cup.mp4", "Pick up the green cup", "a Jaco single-arm robot", "the main camera", ], [ "examples/so101_lego_into_box.mp4", "Put the pink lego brick into the transparent box", "an SO-101 single-arm robot", "the side camera", ], [ "examples/franka_box_into_drawer.mp4", "Fold the towel and put it in the basket", "a Franka single-arm robot", "the main camera", ], ], inputs=[video_in, instruction_in, robot_in, camera_in], outputs=[video_out, report_out], fn=analyze_video, cache_examples=True, cache_mode="lazy", examples_per_page=6, label="Examples (the last row deliberately mismatches video and instruction)", ) gr.Markdown( """Example clips — `franka_box_into_drawer` is the demo clip bundled with the [RynnValue repo](https://github.com/alibaba-damo-academy/RynnValue) (Apache-2.0); `soar_*`, `berkeley_rpt_*` and `jaco_play_*` are the RoboMeter benchmark clips bundled in the same repository (MIT), originating from [Open X-Embodiment](https://robotics-transformer-x.github.io/) (CC BY 4.0); `so101_lego_into_box` is episode 1 of [lerobot/svla_so101_pickplace](https://huggingface.co/datasets/lerobot/svla_so101_pickplace) (Apache-2.0).""" ) run_btn.click( fn=analyze_video, inputs=[ video_in, instruction_in, robot_in, camera_in, num_steps_in, num_frames_in, image_side_in, tokens_in, ], outputs=[video_out, report_out], api_name="analyze_video", ) if __name__ == "__main__": demo.launch(mcp_server=True)