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)