vllm-omni/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Stream LingBot-World 2.0 pixels from one stepwise request.

One ``AsyncOmni.generate()`` drives the whole rollout: each AR block leaves
``post_decode()`` as one video chunk, decoded through that session's own
temporal VAE cache rather than as an isolated clip. That is what this script
shows end to end -- chunk *N + 1* continues chunk *N*, so the streamed timeline
carries the same frame count as an offline whole-clip decode of the same
latents instead of losing each block's opening frame to a restarted decoder.
Only the session's first latent frame expands to a single raw frame; on the
default checkpoint, whose block is three latent frames at a temporal factor of
four, that is 9 frames for the opening chunk and 12 for every later one.

The decode lives inside the pipeline, keyed by ``request_id``, so nothing here
owns a VAE: this file is a client of the served path, and the same chunks reach
a WebSocket client of ``/v1/realtime/video`` unchanged.

Examples::

    python examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py \
        --image scene.png --prompt "a lit hallway" \
        --action-script actions.json --output-dir /tmp/lingbot-stream

``--action-script`` is a JSON list with one camera action list per AR block,
for example ``[[["w"], ["w"], ["w"]], [["a"], [], []]]`` for a checkpoint that
generates three latent frames per block. The frames per block and temporal
compression are read off the checkpoint rather than assumed, so the script
can be repointed with ``--model``.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import math
import os
import time
from collections.abc import Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    # Only for type hints -- the real imports stay deferred inside run() so
    # --help and the offline validation helpers don't require a CUDA-enabled
    # vLLM install.
    import numpy as np

_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Stream decoded LingBot-World 2.0 video chunks from a single stepwise request.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
    parser.add_argument("--image", required=True, help="Initial RGB image.")
    parser.add_argument("--prompt", required=True, help="Scene prompt for the rollout.")
    parser.add_argument(
        "--action-script",
        required=True,
        help=(
            "JSON file holding one camera action list per AR block. "
            "The per-block frame count is read off the checkpoint, not assumed."
        ),
    )
    parser.add_argument("--output-dir", required=True, help="Directory for decoded frames and metadata.")
    parser.add_argument("--request-id", default="lingbot-world-stream", help="Request id; also the session id.")
    parser.add_argument("--height", type=int, default=480)
    parser.add_argument("--width", type=int, default=832)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--gpu-memory-fraction", type=float, default=0.6)
    parser.add_argument("--tensor-parallel-size", type=int, default=1)
    parser.add_argument("--enforce-eager", action="store_true")
    return parser.parse_args(argv)


def _load_action_script(path: Path) -> list[list[list[str]]]:
    """Parse the per-block camera action script, without assuming its geometry.

    The shape is the one ``sampling_params.extra_args["camera_action_script"]``
    takes: one entry per AR block, each entry one per-frame action list. How
    many frames a block holds is a property of the checkpoint, so it is checked
    in :func:`_validate_against_checkpoint` once the configs have been read.
    """
    try:
        value = json.loads(path.read_text())
    except json.JSONDecodeError as exc:
        raise ValueError(f"--action-script is not valid JSON: {exc.msg}.") from None
    if not isinstance(value, list) or not value:
        raise ValueError("--action-script must be a non-empty JSON list of per-block action lists.")
    script: list[list[list[str]]] = []
    for block_index, block in enumerate(value):
        if not isinstance(block, list) or not block:
            raise ValueError(f"--action-script block {block_index} must be a non-empty list of frame action lists.")
        frames: list[list[str]] = []
        for frame in block:
            if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
                raise ValueError(f"--action-script block {block_index} frames must contain only action strings.")
            frames.append(list(frame))
        script.append(frames)
    return script


def _model_config_json(model: str, relative_path: str) -> dict[str, Any]:
    """Read one config file out of a local checkpoint or a Hub repo."""
    if os.path.isdir(model):
        return json.loads((Path(model) / relative_path).read_text())
    from huggingface_hub import hf_hub_download

    return json.loads(Path(hf_hub_download(repo_id=model, filename=relative_path)).read_text())


def _checkpoint_geometry(model: str) -> tuple[int, int]:
    """Read ``(frames_per_block, temporal_compression)`` off the checkpoint.

    The pipeline takes both from the model --
    ``transformer.config.num_frames_per_block`` and
    ``vae.config.scale_factor_temporal`` -- so hardcoding 3 and 4 here would
    silently build a wrong action script and an inconsistent ``num_frames`` the
    moment ``--model`` points at a checkpoint with different geometry.
    """
    transformer_config = _model_config_json(model, "transformer/config.json")
    vae_config = _model_config_json(model, "vae/config.json")
    if "num_frames_per_block" not in transformer_config:
        raise ValueError(
            f"{model}'s transformer/config.json declares no num_frames_per_block, "
            "so this checkpoint does not generate AR blocks and cannot be streamed this way."
        )
    frames_per_block = int(transformer_config["num_frames_per_block"])
    temporal_compression = int(vae_config.get("scale_factor_temporal", 4))
    if frames_per_block <= 0 or temporal_compression <= 0:
        raise ValueError(
            "the checkpoint reports a non-positive block geometry "
            f"(num_frames_per_block={frames_per_block}, scale_factor_temporal={temporal_compression})."
        )
    return frames_per_block, temporal_compression


def _validate_against_checkpoint(
    script: list[list[list[str]]], *, frames_per_block: int, temporal_compression: int
) -> int:
    """Check the script against the checkpoint and return the frame count to request."""
    for block_index, block in enumerate(script):
        if len(block) != frames_per_block:
            raise ValueError(
                f"--action-script block {block_index} holds {len(block)} frame action lists, but this "
                f"checkpoint generates {frames_per_block} latent frames per AR block."
            )
    return (len(script) * frames_per_block - 1) * temporal_compression + 1


def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
    image = Path(args.image).expanduser().resolve()
    action_script = Path(args.action_script).expanduser().resolve()
    output_dir = Path(args.output_dir).expanduser().resolve()
    if not image.is_file():
        raise ValueError("--image must point to an existing file.")
    if not action_script.is_file():
        raise ValueError("--action-script must point to an existing JSON file.")
    if not args.prompt.strip():
        raise ValueError("--prompt must contain non-whitespace text.")
    if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
        raise ValueError("--height and --width must be positive multiples of 16.")
    if args.tensor_parallel_size <= 0:
        raise ValueError("--tensor-parallel-size must be positive.")
    if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
        raise ValueError("--gpu-memory-fraction must be in (0, 1].")
    if not args.request_id.strip():
        raise ValueError("--request-id must contain non-whitespace text.")
    return image, action_script, output_dir


def _as_frames(value: Any) -> np.ndarray:
    """Normalize one streamed chunk to ``[T, H, W, 3]`` uint8.

    ``output_type="np"`` reaches the client as a float array in ``[0, 1]``,
    with or without the leading batch axis depending on how the formatter
    unwrapped it; a nested list is the ``"pil"`` shape. Anything else means
    the chunk is not pixels -- most likely the request still asked for
    latents -- and is worth failing on rather than saving.
    """
    import numpy as np
    import PIL.Image
    import torch

    if isinstance(value, list):
        if len(value) == 1 and isinstance(value[0], list):
            value = value[0]
        if value and all(isinstance(frame, PIL.Image.Image) for frame in value):
            return np.stack([np.asarray(frame.convert("RGB")) for frame in value])
    if isinstance(value, torch.Tensor):
        value = value.detach().float().cpu().numpy()
    if not isinstance(value, np.ndarray):
        raise RuntimeError(f"Expected decoded pixels from each AR block, got {type(value).__name__}.")
    if value.ndim == 5 and value.shape[0] == 1:
        value = value[0]
    if value.ndim != 4 or value.shape[-1] != 3:
        raise RuntimeError(f"Expected one [T, H, W, 3] pixel chunk per AR block, got shape {tuple(value.shape)}.")
    if value.dtype != np.uint8:
        value = (np.clip(value, 0.0, 1.0) * 255.0).round().astype(np.uint8)
    return value


def _chunk_metadata(output: Any) -> dict[str, Any]:
    multimodal = getattr(output, "multimodal_output", None) or {}
    metadata = multimodal.get("metadata") if isinstance(multimodal, dict) else None
    ar_diffusion = metadata.get("ar_diffusion") if isinstance(metadata, dict) else None
    if not isinstance(ar_diffusion, dict):
        raise RuntimeError("Streamed chunk is missing its ar_diffusion metadata envelope.")
    return dict(ar_diffusion)


async def run(argv: Sequence[str] | None = None) -> Path:
    args = parse_args(argv)
    image, action_script_path, output_dir = _validate_args(args)
    camera_action_script = _load_action_script(action_script_path)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Imports stay below pure input validation so --help and helper tests do not
    # require a CUDA-enabled vLLM installation.
    import numpy as np

    from vllm_omni.entrypoints.async_omni import AsyncOmni
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams

    frames_per_block, temporal_compression = _checkpoint_geometry(args.model)
    num_chunks = len(camera_action_script)
    num_frames = _validate_against_checkpoint(
        camera_action_script,
        frames_per_block=frames_per_block,
        temporal_compression=temporal_compression,
    )

    engine = AsyncOmni(
        model=args.model,
        engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
        enforce_eager=args.enforce_eager,
        tensor_parallel_size=args.tensor_parallel_size,
        max_num_seqs=1,
        # One request, many chunks: the stepwise contract from #6844. Streaming
        # output is what turns each post_decode() into a delivered chunk.
        step_execution=True,
        diffusion_streaming_output=True,
        model_config={
            "ar_diffusion_height": args.height,
            "ar_diffusion_width": args.width,
            "ar_diffusion_kv_config": {
                "gpu_memory_fraction": args.gpu_memory_fraction,
                "warmup_cudagraph": True,
            },
        },
    )
    sampling = OmniDiffusionSamplingParams(
        height=args.height,
        width=args.width,
        num_frames=num_frames,
        num_inference_steps=4,
        max_sequence_length=512,
        seed=args.seed,
        # Pixels, not latents: this is what routes the chunk through the
        # session's streaming VAE decode instead of handing back a tensor.
        output_type="np",
        extra_args={"flow_shift": 5.0, "camera_action_script": camera_action_script},
    )
    prompt = {"prompt": args.prompt, "multi_modal_data": {"image": str(image)}}

    measurements: list[dict[str, Any]] = []
    t_start = time.perf_counter()
    first_frame_at: float | None = None
    chunk_started = t_start
    try:
        async for output in engine.generate(prompt, sampling, request_id=args.request_id):
            images = getattr(output, "images", None)
            if not images:
                continue
            if len(images) != 1:
                raise RuntimeError("Expected exactly one decoded chunk per AR block.")
            frames = _as_frames(images[0])
            now = time.perf_counter()
            if first_frame_at is None:
                first_frame_at = now - t_start
            chunk_index = len(measurements)
            metadata = _chunk_metadata(output)

            np.save(output_dir / f"chunk_{chunk_index:03d}.npy", frames)
            (output_dir / f"chunk_{chunk_index:03d}.json").write_text(
                json.dumps(metadata, indent=2, sort_keys=True) + "\n"
            )
            measurements.append(
                {
                    "chunk_index": chunk_index,
                    "chunk_latency_seconds": now - chunk_started,
                    "frames": int(frames.shape[0]),
                    "frame_shape": list(frames.shape),
                    "metadata": metadata,
                }
            )
            chunk_started = now
            print(json.dumps(measurements[-1], sort_keys=True), flush=True)
    finally:
        engine.shutdown()

    delivered_frames = sum(int(chunk["frames"]) for chunk in measurements)
    summary_path = output_dir / "summary.json"
    summary_path.write_text(
        json.dumps(
            {
                "request_id": args.request_id,
                "chunks_delivered": len(measurements),
                "chunks_expected": num_chunks,
                "frames_delivered": delivered_frames,
                "frames_expected": num_frames,
                # The point of session-owned decode: a streamed rollout carries
                # the same frame timeline as an offline decode of the same
                # latents. Per-chunk decode would deliver fewer, because every
                # block would re-expand its own opening frame.
                "matches_offline_frame_timeline": delivered_frames == num_frames,
                "time_to_first_frame_seconds": first_frame_at,
                "total_seconds": time.perf_counter() - t_start,
                "chunks": measurements,
            },
            indent=2,
            sort_keys=True,
        )
        + "\n"
    )
    return summary_path


def main(argv: Sequence[str] | None = None) -> Path:
    return asyncio.run(run(argv))


if __name__ == "__main__":
    main()