app.py

import os

# Allocator config to avoid fragmentation asserts during pixel-space VAE decode of
# 21-frame latents across three branches.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TORCHINDUCTOR_MAX_WORKERS", "1")

import spaces  # noqa: E402  MUST come before torch / diffusers

import gc  # noqa: E402
import json  # noqa: E402
import tempfile  # noqa: E402
from copy import deepcopy  # noqa: E402
from typing import Optional  # noqa: E402

import numpy as np  # noqa: E402
import torch  # noqa: E402
import gradio as gr  # noqa: E402
from PIL import Image  # noqa: E402
from safetensors.torch import load_file  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402

from diffusers import WanImageToVideoPipeline  # noqa: E402
from diffusers.utils import BaseOutput, export_to_video  # noqa: E402
from diffusers.models.transformers.transformer_wan import WanTimeTextImageEmbedding  # noqa: E402

from core.finetune.models.wan_i2v.module import (  # noqa: E402
    patched_wan_time_text_image_embedding_forward,
)
from core.finetune.models.wan_i2v.module_joint import JointRynnWorld4DTransformer3DModel  # noqa: E402

# The RynnWorld4D transformer overrides the time/text embedding forward so that a
# per-token timestep tensor (0 on the conditioning first frame, t elsewhere) works.
WanTimeTextImageEmbedding.forward = patched_wan_time_text_image_embedding_forward

# ─────────────────────────────────────────────────────────────────────────────
# Config — must match the released Stage-3 checkpoint's training config.
# (See the model card "Critical Config Flags" table and inference-sft.py.)
# ─────────────────────────────────────────────────────────────────────────────
BACKBONE_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
CKPT_ID = "Alibaba-DAMO-Academy/RynnWorld-4D"
DTYPE = torch.bfloat16
FPS = 16

JOINT_KW = dict(
    share_ffn=False,
    joint_start_layer=0,
    joint_end_layer=30,
    joint_every_n_layers=3,
    joint_frame_wise=True,
    joint_use_rope=True,
)

DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")

# ─────────────────────────────────────────────────────────────────────────────
# Sample library — precomputed VAE latents + text embeddings shipped with the
# official repo. These are the authors' validated inference inputs: a single
# reference-image latent + text embedding condition the tri-branch generation.
# ─────────────────────────────────────────────────────────────────────────────
with open(os.path.join(DATA_DIR, "sample.json")) as f:
    _RAW_SAMPLES = json.load(f)

SAMPLES = []
for i, s in enumerate(_RAW_SAMPLES):
    SAMPLES.append(
        {
            "id": i,
            "label": f"Sample {i + 1}",
            "prompt": s["prompt"],
            "rgb_latents": os.path.join(os.path.dirname(os.path.abspath(__file__)), s["rgb_latents"]),
            "flow_depth_latents": os.path.join(os.path.dirname(os.path.abspath(__file__)), s["flow_depth_latents"]),
        }
    )
SAMPLE_CHOICES = [f'{s["label"]}: {s["prompt"][:70]}…' for s in SAMPLES]


def _randn(shape, generator, device, dtype):
    latents = torch.randn(shape, generator=generator, device=device, dtype=dtype)
    return latents


# ─────────────────────────────────────────────────────────────────────────────
# Tri-branch denoising pipeline — a faithful port of inference-sft.py's
# RynnWorld4DInferencePipeline, generating synchronized RGB / depth / flow.
# ─────────────────────────────────────────────────────────────────────────────
class RynnWorld4DInferencePipeline(WanImageToVideoPipeline):
    @torch.no_grad()
    def generate(
        self,
        img_latent: torch.Tensor,
        depth_latent_cond: torch.Tensor,
        flow_latent_cond: torch.Tensor,
        prompt_embeds: torch.Tensor,
        num_latent_frames: int,
        num_inference_steps: int = 50,
        generator: Optional[torch.Generator] = None,
        progress=None,
    ):
        device = self._execution_device
        transformer_dtype = self.transformer.dtype

        _, num_channels, _, latent_h, latent_w = img_latent.shape

        shape = (1, num_channels, num_latent_frames, latent_h, latent_w)
        # All three branches start from the SAME noise, matching training.
        noise = _randn(shape, generator, device, transformer_dtype)
        latents_video = noise.clone()
        latents_depth = noise.clone()
        latents_flow = noise.clone()
        del noise

        first_frame_mask = torch.ones(
            1, 1, num_latent_frames, latent_h, latent_w, device=device, dtype=transformer_dtype
        )
        first_frame_mask[:, :, 0] = 0

        prompt_embeds = prompt_embeds.to(device=device, dtype=transformer_dtype)
        img_latent = img_latent.to(device=device, dtype=transformer_dtype)
        depth_latent_cond = depth_latent_cond.to(device=device, dtype=transformer_dtype)
        flow_latent_cond = flow_latent_cond.to(device=device, dtype=transformer_dtype)

        self.scheduler.set_timesteps(num_inference_steps, device=device)
        depth_scheduler = deepcopy(self.scheduler)
        flow_scheduler = deepcopy(self.scheduler)
        depth_scheduler.set_timesteps(num_inference_steps, device=device)
        flow_scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        ts_mask = first_frame_mask[0][0][:, ::2, ::2].flatten().unsqueeze(0)  # (1, num_tokens)

        total = len(timesteps)
        for i, t in enumerate(timesteps):
            latents_video[:, :, 0:1, :, :] = img_latent
            latents_depth[:, :, 0:1, :, :] = depth_latent_cond
            latents_flow[:, :, 0:1, :, :] = flow_latent_cond

            timestep = ts_mask * t

            video_pred, depth_pred, flow_pred = self.transformer(
                hidden_states=latents_video,
                hidden_states_depth=latents_depth,
                hidden_states_flow=latents_flow,
                timestep=timestep,
                encoder_hidden_states=prompt_embeds,
                encoder_hidden_states_image=None,
                return_dict=False,
            )

            latents_video = self.scheduler.step(video_pred, t, latents_video, return_dict=False)[0]
            latents_depth = depth_scheduler.step(depth_pred, t, latents_depth, return_dict=False)[0]
            latents_flow = flow_scheduler.step(flow_pred, t, latents_flow, return_dict=False)[0]

            if progress is not None:
                progress((i + 1) / total, desc=f"Denoising step {i + 1}/{total}")

        latents_video[:, :, 0:1, :, :] = img_latent
        latents_depth[:, :, 0:1, :, :] = depth_latent_cond
        latents_flow[:, :, 0:1, :, :] = flow_latent_cond

        latents_mean = (
            torch.tensor(self.vae.config.latents_mean)
            .view(1, self.vae.config.z_dim, 1, 1, 1)
            .to(device=device, dtype=self.vae.dtype)
        )
        latents_std = 1.0 / (
            torch.tensor(self.vae.config.latents_std)
            .view(1, self.vae.config.z_dim, 1, 1, 1)
            .to(device=device, dtype=self.vae.dtype)
        )

        def decode(lat):
            lat = lat.to(self.vae.dtype)
            lat = lat / latents_std + latents_mean
            vid = self.vae.decode(lat, return_dict=False)[0]
            return self.video_processor.postprocess_video(vid, output_type="np")

        video_out = decode(latents_video)
        del latents_video
        torch.cuda.empty_cache()
        depth_out = decode(latents_depth)
        del latents_depth
        torch.cuda.empty_cache()
        flow_out = decode(latents_flow)
        del latents_flow
        torch.cuda.empty_cache()

        return (
            BaseOutput(frames=video_out),
            BaseOutput(frames=depth_out),
            BaseOutput(frames=flow_out),
        )


# ─────────────────────────────────────────────────────────────────────────────
# Model loading at module scope (ZeroGPU packs weights, streams to VRAM on
# first @spaces.GPU entry). We skip the T5 text_encoder / tokenizer entirely —
# the sample library ships precomputed text embeddings — which keeps VRAM well
# inside the ZeroGPU budget.
# ─────────────────────────────────────────────────────────────────────────────
print("Downloading RynnWorld-4D checkpoint …", flush=True)
CKPT_DIR = snapshot_download(CKPT_ID, allow_patterns=["pytorch_model/*.pt"])
CKPT_STATES = os.path.join(CKPT_DIR, "pytorch_model", "mp_rank_00_model_states.pt")

print("Building tri-branch transformer from backbone …", flush=True)
transformer = JointRynnWorld4DTransformer3DModel.from_pretrained(
    BACKBONE_ID,
    subfolder="transformer",
    torch_dtype=DTYPE,
    eps=1e-5,
    low_cpu_mem_usage=True,
    **JOINT_KW,
)

print("Assembling pipeline (VAE + scheduler; no text encoder) …", flush=True)
pipe = RynnWorld4DInferencePipeline.from_pretrained(
    BACKBONE_ID,
    transformer=transformer,
    text_encoder=None,
    tokenizer=None,
    torch_dtype=DTYPE,
    low_cpu_mem_usage=True,
)


def _load_sft_checkpoint(path, transformer):
    checkpoint = torch.load(path, map_location="cpu", weights_only=False, mmap=True)
    state_dict = checkpoint["module"]
    keys_to_rename = [k for k in state_dict if "module." in k or ".base_layer." in k]
    for k in keys_to_rename:
        clean_key = k.replace("module.", "").replace(".base_layer.", ".")
        if clean_key != k:
            state_dict[clean_key] = state_dict.pop(k)
    model_sd = transformer.state_dict()
    loaded = 0
    for k in list(state_dict.keys()):
        if k in model_sd:
            model_sd[k].copy_(state_dict[k])
            loaded += 1
    print(f"Loaded {loaded} SFT params.", flush=True)
    del model_sd, checkpoint
    state_dict.clear()
    gc.collect()


print("Loading Stage-3 SFT weights …", flush=True)
_load_sft_checkpoint(CKPT_STATES, pipe.transformer)
gc.collect()

if hasattr(pipe.vae, "enable_slicing"):
    pipe.vae.enable_slicing()

print("Moving pipeline to CUDA …", flush=True)
pipe.to("cuda")
print("Pipeline ready.", flush=True)

VAE_Z_DIM = pipe.vae.config.z_dim
_LATENTS_MEAN = torch.tensor(pipe.vae.config.latents_mean).view(1, VAE_Z_DIM, 1, 1, 1)
_LATENTS_STD = torch.tensor(pipe.vae.config.latents_std).view(1, VAE_Z_DIM, 1, 1, 1)


# ─────────────────────────────────────────────────────────────────────────────
# Decode a single sample's first-frame RGB latent to show the reference image.
# ─────────────────────────────────────────────────────────────────────────────
@spaces.GPU(duration=60)
def preview_first_frame(sample_idx: int):
    """Decode and return the conditioning reference image of a chosen sample.

    Args:
        sample_idx: index into the built-in sample library.
    """
    s = SAMPLES[int(sample_idx)]
    rgb = load_file(s["rgb_latents"])
    with torch.no_grad():
        video_latents = rgb["video_latents"].to(DTYPE)  # [C, T, H, W]
        first = video_latents[:, :1, :, :].unsqueeze(0).to("cuda")  # [1, C, 1, H, W]
        mean = _LATENTS_MEAN.to(device="cuda", dtype=pipe.vae.dtype)
        std_inv = 1.0 / _LATENTS_STD.to(device="cuda", dtype=pipe.vae.dtype)
        lat = first.to(pipe.vae.dtype) / std_inv + mean
        vid = pipe.vae.decode(lat, return_dict=False)[0]
        frames = pipe.video_processor.postprocess_video(vid, output_type="np")[0]
    img = (np.clip(frames[0], 0.0, 1.0) * 255).astype(np.uint8)
    torch.cuda.empty_cache()
    return Image.fromarray(img)


def _estimate_duration(sample_idx, num_inference_steps=50, *args, **kwargs):
    return min(600, 90 + int(num_inference_steps) * 8)


@spaces.GPU(duration=_estimate_duration)
def generate(sample_idx: int, num_inference_steps: int = 50, seed: int = 42,
             progress=gr.Progress()):
    """Generate synchronized RGB, depth and optical-flow videos for a sample.

    Args:
        sample_idx: index into the built-in sample library (reference image + prompt).
        num_inference_steps: number of denoising steps (higher = better, slower).
        seed: RNG seed for reproducibility.
    """
    s = SAMPLES[int(sample_idx)]
    rgb_data = load_file(s["rgb_latents"])
    fd_data = load_file(s["flow_depth_latents"])

    video_latents_gt = rgb_data["video_latents"].to(DTYPE)  # [C, T, H, W]
    depth_latents = fd_data["depth_latents"].to(DTYPE)
    flow_latents = fd_data["flow_latents"].to(DTYPE)
    text_embeds = rgb_data["text_embeds"].to(DTYPE)  # [1, 226, 4096]

    img_latent = video_latents_gt[:, :1, :, :].unsqueeze(0)
    depth_latent_cond = depth_latents[:, :1, :, :].unsqueeze(0)
    flow_latent_cond = flow_latents[:, :1, :, :].unsqueeze(0)
    num_latent_frames = video_latents_gt.shape[1]

    generator = torch.Generator(device="cuda").manual_seed(int(seed))

    rgb_out, depth_out, flow_out = pipe.generate(
        img_latent=img_latent,
        depth_latent_cond=depth_latent_cond,
        flow_latent_cond=flow_latent_cond,
        prompt_embeds=text_embeds,
        num_latent_frames=num_latent_frames,
        num_inference_steps=int(num_inference_steps),
        generator=generator,
        progress=progress,
    )

    paths = []
    for tag, out in [("rgb", rgb_out), ("depth", depth_out), ("flow", flow_out)]:
        f = tempfile.NamedTemporaryFile(suffix=f"_{tag}.mp4", delete=False)
        export_to_video(out.frames[0], f.name, fps=FPS)
        paths.append(f.name)

    gc.collect()
    torch.cuda.empty_cache()
    return paths[0], paths[1], paths[2]


# ─────────────────────────────────────────────────────────────────────────────
# 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) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # RynnWorld-4D — 4D Embodied World Model
            Generate **synchronized RGB, depth, and optical-flow** videos from a reference
            image + text prompt with a tri-branch diffusion transformer built on
            [Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers).
            Pick a built-in sample (a reference-image latent + prompt validated by the
            authors), preview its reference frame, then generate the three aligned streams.

            [Model card](https://huggingface.co/Alibaba-DAMO-Academy/RynnWorld-4D) ·
            [Project page](https://alibaba-damo-academy.github.io/RynnWorld-4D.github.io/) ·
            [GitHub](https://github.com/Alibaba-DAMO-Academy/RynnWorld-4D)
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                sample = gr.Dropdown(
                    choices=[(c, i) for i, c in enumerate(SAMPLE_CHOICES)],
                    value=0,
                    label="Sample (reference image + prompt)",
                )
                prompt_box = gr.Textbox(
                    label="Prompt (from selected sample)",
                    value=SAMPLES[0]["prompt"],
                    interactive=False,
                    lines=2,
                )
                ref_image = gr.Image(label="Reference frame", height=240)
                run = gr.Button("Generate 4D video", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    steps = gr.Slider(10, 50, value=50, step=1, label="Denoising steps")
                    seed = gr.Number(value=42, precision=0, label="Seed")
            with gr.Column(scale=2):
                with gr.Row():
                    rgb_video = gr.Video(label="RGB", autoplay=True, loop=True)
                    depth_video = gr.Video(label="Depth", autoplay=True, loop=True)
                    flow_video = gr.Video(label="Optical flow", autoplay=True, loop=True)

        def _on_sample_change(idx):
            return SAMPLES[int(idx)]["prompt"]

        sample.change(_on_sample_change, inputs=sample, outputs=prompt_box)
        sample.change(preview_first_frame, inputs=sample, outputs=ref_image)

        run.click(
            generate,
            inputs=[sample, steps, seed],
            outputs=[rgb_video, depth_video, flow_video],
            api_name="generate",
        )

        gr.Examples(
            examples=[[i] for i in range(len(SAMPLES))],
            inputs=[sample],
            outputs=[rgb_video, depth_video, flow_video],
            fn=generate,
            cache_examples=False,
            run_on_click=True,
        )

if __name__ == "__main__":
    demo.launch(mcp_server=True)