app.py

"""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
    # <value>/<relative_value> 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 |

<sub>{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</sub>

<details><summary>Raw generation</summary>

```
{analysis_text.strip()}
```
</details>"""

    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(
                "<sub>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.</sub>"
            )

        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(
            """<sub>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).</sub>"""
        )

    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)