Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/observation.images.top.png filter=lfs diff=lfs merge=lfs -text
+examples/observation.images.wrist.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,43 @@
 ---
 title: Next Forcing World Model
-emoji: 💻
-colorFrom: indigo
-colorTo: blue
+emoji: 🎬
+colorFrom: purple
+colorTo: red
 sdk: gradio
 sdk_version: 6.25.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Causal world model for robot video generation
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Next Forcing: Causal World Modeling with Multi-Chunk Prediction
+
+This Space demonstrates **Next Forcing**, a causal autoregressive world model that generates robot manipulation video
+from initial camera observations and a text instruction.
+
+## How it works
+
+1. Provide top-down and wrist camera images (the initial scene).
+2. Describe the manipulation task in natural language.
+3. The model autoregressively generates future video chunks (4 frames per chunk).
+
+The model is a 5.1B parameter causal video-action transformer (based on the Wan2.1 architecture) that jointly
+denoises video latents and action latents in a causal, chunk-by-chunk fashion.
+
+## Model
+
+- **Checkpoint**: [`gangweix/next-forcing-base`](https://huggingface.co/gangweix/next-forcing-base) (5.1B, BF16)
+- **Paper**: [arXiv:2606.11187](https://arxiv.org/abs/2606.11187)
+- **Code**: [github.com/gangweix/next-forcing](https://github.com/gangweix/next-forcing)
+
+## Reference
+
+```
+@article{xu2026next,
+  title={Next Forcing: Causal World Modeling with Multi-Chunk Prediction},
+  author={Xu, Gangwei and Zhang, Qihang and Zhou, Jiaming and Zhu, Xing and Shen, Yujun and Yang, Xin and Xu, Yinghao},
+  journal={arXiv preprint arXiv:2606.11187},
+  year={2026}
+}
+```

app.py

@@ -0,0 +1,651 @@
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+import spaces
+import sys
+import time
+import tempfile
+import numpy as np
+import torch
+import torch.nn.functional as F
+import gradio as gr
+from PIL import Image
+from einops import rearrange
+from tqdm import tqdm
+
+# ------------------------------------------------------------------ #
+#  Model + config setup (module scope, eagerly on GPU)
+# ------------------------------------------------------------------ #
+
+MODEL_ID = "gangweix/next-forcing-base"
+
+from huggingface_hub import snapshot_download
+
+_model_path = snapshot_download(
+    MODEL_ID,
+    repo_type="model",
+    allow_patterns=[
+        "transformer/*",
+        "vae/*",
+        "text_encoder/*",
+        "tokenizer/*",
+    ],
+)
+
+from wan_va.modules.utils import (
+    WanVAEStreamingWrapper,
+    load_text_encoder,
+    load_tokenizer,
+    load_transformer,
+    load_vae,
+)
+from wan_va.utils.scheduler import FlowMatchScheduler
+from wan_va.utils.utils import get_mesh_id, data_seq_to_patch
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+
+# ---- Demo config (matches va_demo_cfg.py) ----
+CONFIG = dict(
+    attn_window=30,
+    frame_chunk_size=4,
+    env_type="none",
+    height=256,
+    width=256,
+    action_dim=30,
+    action_per_frame=8,
+    obs_cam_keys=["observation.images.top", "observation.images.wrist"],
+    guidance_scale=5,
+    action_guidance_scale=1,
+    num_inference_steps=5,
+    video_exec_step=-1,
+    action_num_inference_steps=10,
+    snr_shift=5.0,
+    action_snr_shift=1.0,
+    patch_size=(1, 2, 2),
+    used_action_channel_ids=list(range(0, 5)) + list(range(28, 29)),
+    action_norm_method="quantiles",
+    norm_stat={
+        "q01": [
+            -90.60303497314453,
+            -98.73043060302734,
+            -79.9008560180664,
+            48.95470428466797,
+            -32.794578552246094,
+        ] + [0.0] * 23 + [0.8250824809074402, 0],
+        "q99": [
+            71.735107421875,
+            65.89081573486328,
+            92.87967681884766,
+            100.0,
+            22.784151077270508,
+        ] + [0.0] * 23 + [100.0, 0],
+    },
+)
+
+# Inverse action channel mapping
+inverse_used_action_channel_ids = [len(CONFIG["used_action_channel_ids"])] * CONFIG[
+    "action_dim"
+]
+for i, j in enumerate(CONFIG["used_action_channel_ids"]):
+    inverse_used_action_channel_ids[j] = i
+CONFIG["inverse_used_action_channel_ids"] = inverse_used_action_channel_ids
+
+# ---- Load model components ----
+vae = load_vae(os.path.join(_model_path, "vae"), torch_dtype=DTYPE, torch_device=DEVICE)
+streaming_vae = WanVAEStreamingWrapper(vae)
+tokenizer = load_tokenizer(os.path.join(_model_path, "tokenizer"))
+text_encoder = load_text_encoder(
+    os.path.join(_model_path, "text_encoder"), torch_dtype=DTYPE, torch_device=DEVICE
+)
+transformer = load_transformer(
+    os.path.join(_model_path, "transformer"),
+    torch_dtype=DTYPE,
+    torch_device=DEVICE,
+    attn_mode="torch",
+    disable_mcp=True,
+)
+transformer.eval().requires_grad_(False)
+
+scheduler = FlowMatchScheduler(shift=CONFIG["snr_shift"], sigma_min=0.0, extra_one_step=True)
+action_scheduler = FlowMatchScheduler(
+    shift=CONFIG["action_snr_shift"], sigma_min=0.0, extra_one_step=True
+)
+scheduler.set_timesteps(1000, training=True)
+action_scheduler.set_timesteps(1000, training=True)
+
+action_mask = torch.zeros([CONFIG["action_dim"]]).bool()
+action_mask[CONFIG["used_action_channel_ids"]] = True
+
+actions_q01 = torch.tensor(CONFIG["norm_stat"]["q01"], dtype=torch.float32).reshape(-1, 1, 1)
+actions_q99 = torch.tensor(CONFIG["norm_stat"]["q99"], dtype=torch.float32).reshape(-1, 1, 1)
+
+from diffusers.video_processor import VideoProcessor
+
+video_processor = VideoProcessor(vae_scale_factor=1)
+
+
+# ------------------------------------------------------------------ #
+#  Inference helpers
+# ------------------------------------------------------------------ #
+
+def _get_t5_prompt_embeds(prompt, max_sequence_length=512):
+    from diffusers.pipelines.wan.pipeline_wan import prompt_clean
+
+    prompt_list = [prompt] if isinstance(prompt, str) else prompt
+    prompt_list = [prompt_clean(u) for u in prompt_list]
+    batch_size = len(prompt_list)
+
+    text_inputs = tokenizer(
+        prompt_list,
+        padding="max_length",
+        max_length=max_sequence_length,
+        truncation=True,
+        add_special_tokens=True,
+        return_attention_mask=True,
+        return_tensors="pt",
+    )
+    text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+    seq_lens = mask.gt(0).sum(dim=1).long()
+
+    text_encoder_device = next(text_encoder.parameters()).device
+    prompt_embeds = text_encoder(
+        text_input_ids.to(text_encoder_device), mask.to(text_encoder_device)
+    ).last_hidden_state
+    prompt_embeds = prompt_embeds.to(dtype=DTYPE, device=DEVICE)
+    prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+    prompt_embeds = torch.stack(
+        [
+            torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))])
+            for u in prompt_embeds
+        ],
+        dim=0,
+    )
+    _, seq_len, _ = prompt_embeds.shape
+    prompt_embeds = prompt_embeds.repeat(1, 1, 1)
+    prompt_embeds = prompt_embeds.view(batch_size, seq_len, -1)
+    return prompt_embeds.to(DEVICE)
+
+
+def encode_prompt(prompt):
+    prompt_embeds = _get_t5_prompt_embeds(prompt)
+    neg_prompt_embeds = _get_t5_prompt_embeds("")
+    return prompt_embeds, neg_prompt_embeds
+
+
+def normalize_latents(latents, latents_mean, latents_std):
+    latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(device=latents.device)
+    latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device)
+    latents = ((latents.float() - latents_mean) * latents_std).to(latents)
+    return latents
+
+
+def encode_obs(images_dict_list):
+    """Encode observation images into latent space.
+
+    Args:
+        images_dict_list: list of dicts, each mapping cam_key -> np.ndarray(H,W,3) uint8
+    """
+    images = images_dict_list
+    if not isinstance(images, list):
+        images = [images]
+    if len(images) < 1:
+        return None
+    videos = []
+    for k_i, k in enumerate(CONFIG["obs_cam_keys"]):
+        height_i, width_i = CONFIG["height"], CONFIG["width"]
+        history_video_k = (
+            torch.from_numpy(np.stack([each[k] for each in images]))
+            .float()
+            .permute(3, 0, 1, 2)
+        )
+        history_video_k = F.interpolate(
+            history_video_k,
+            size=(height_i, width_i),
+            mode="bilinear",
+            align_corners=False,
+        ).unsqueeze(0)
+        videos.append(history_video_k)
+
+    videos = torch.cat(videos, dim=0) / 255.0 * 2.0 - 1.0
+    vae_device = next(streaming_vae.vae.parameters()).device
+    videos_chunk = videos.to(vae_device).to(DTYPE)
+    enc_out = streaming_vae.encode_chunk(videos_chunk)
+
+    mu, logvar = torch.chunk(enc_out, 2, dim=1)
+    latents_mean = torch.tensor(vae.config.latents_mean).to(mu.device)
+    latents_std = torch.tensor(vae.config.latents_std).to(mu.device)
+    mu_norm = normalize_latents(mu, latents_mean, 1.0 / latents_std)
+    video_latent = torch.cat(mu_norm.split(1, dim=0), dim=-1)
+    return video_latent.to(DEVICE)
+
+
+def _repeat_input_for_cfg(input_dict, use_cfg, prompt_embeds, negative_prompt_embeds):
+    if use_cfg:
+        input_dict["noisy_latents"] = input_dict["noisy_latents"].repeat(2, 1, 1, 1, 1)
+        input_dict["text_emb"] = torch.cat(
+            [prompt_embeds.to(DTYPE).clone(), negative_prompt_embeds.to(DTYPE).clone()],
+            dim=0,
+        )
+        input_dict["grid_id"] = input_dict["grid_id"][None].repeat(2, 1, 1)
+        input_dict["timesteps"] = input_dict["timesteps"][None].repeat(2, 1)
+    else:
+        input_dict["grid_id"] = input_dict["grid_id"][None]
+        input_dict["timesteps"] = input_dict["timesteps"][None]
+    return input_dict
+
+
+def _prepare_latent_input(
+    latent_model_input,
+    action_model_input,
+    latent_t=0,
+    action_t=0,
+    latent_cond=None,
+    action_cond=None,
+    frame_st_id=0,
+    patch_size=(1, 2, 2),
+    prompt_embeds=None,
+    use_cfg=False,
+    negative_prompt_embeds=None,
+):
+    input_dict = dict()
+    if latent_model_input is not None:
+        input_dict["latent_res_lst"] = {
+            "noisy_latents": latent_model_input,
+            "timesteps": torch.ones(
+                [latent_model_input.shape[2]], dtype=torch.float32, device=DEVICE
+            ) * latent_t,
+            "grid_id": get_mesh_id(
+                latent_model_input.shape[-3] // patch_size[0],
+                latent_model_input.shape[-2] // patch_size[1],
+                latent_model_input.shape[-1] // patch_size[2],
+                0,
+                1,
+                frame_st_id,
+            ).to(DEVICE),
+            "text_emb": prompt_embeds.to(DTYPE).clone(),
+        }
+        if latent_cond is not None:
+            input_dict["latent_res_lst"]["noisy_latents"][:, :, 0:1] = latent_cond[:, :, 0:1]
+            input_dict["latent_res_lst"]["timesteps"][0:1] *= 0
+
+    if action_model_input is not None:
+        input_dict["action_res_lst"] = {
+            "noisy_latents": action_model_input,
+            "timesteps": torch.ones(
+                [action_model_input.shape[2]], dtype=torch.float32, device=DEVICE
+            ) * action_t,
+            "grid_id": get_mesh_id(
+                action_model_input.shape[-3],
+                action_model_input.shape[-2],
+                action_model_input.shape[-1],
+                1,
+                1,
+                frame_st_id,
+                action=True,
+            ).to(DEVICE),
+            "text_emb": prompt_embeds.to(DTYPE).clone(),
+        }
+        if action_cond is not None:
+            input_dict["action_res_lst"]["noisy_latents"][:, :, 0:1] = action_cond[:, :, 0:1]
+            input_dict["action_res_lst"]["timesteps"][0:1] *= 0
+        input_dict["action_res_lst"]["noisy_latents"][:, ~action_mask] *= 0
+    return input_dict
+
+
+def infer_chunk(
+    init_latent,
+    frame_st_id,
+    prompt_embeds,
+    negative_prompt_embeds,
+    use_cfg,
+    guidance_scale,
+    action_guidance_scale,
+    num_chunks_to_infer,
+):
+    """Generate one video chunk (video latents + action latents)."""
+    frame_chunk_size = CONFIG["frame_chunk_size"]
+    latent_height = CONFIG["height"] // 16
+    latent_width = (CONFIG["width"] // 16) * len(CONFIG["obs_cam_keys"])
+
+    latents = torch.randn(
+        1, 48, frame_chunk_size, latent_height, latent_width, device=DEVICE, dtype=DTYPE
+    )
+    actions = torch.randn(
+        1,
+        CONFIG["action_dim"],
+        frame_chunk_size,
+        CONFIG["action_per_frame"],
+        1,
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+
+    video_inference_step = CONFIG["num_inference_steps"]
+    action_inference_step = CONFIG["action_num_inference_steps"]
+    video_step = CONFIG["video_exec_step"]
+
+    scheduler.set_timesteps(video_inference_step)
+    action_scheduler.set_timesteps(action_inference_step)
+    timesteps = scheduler.timesteps
+    action_timesteps = action_scheduler.timesteps
+
+    timesteps = F.pad(timesteps, (0, 1), mode="constant", value=0)
+    if video_step != -1:
+        timesteps = timesteps[:video_step]
+    action_timesteps = F.pad(action_timesteps, (0, 1), mode="constant", value=0)
+
+    with torch.no_grad():
+        # 1. Video generation loop
+        for i, t in enumerate(timesteps):
+            last_step = i == len(timesteps) - 1
+            latent_cond = init_latent[:, :, 0:1].to(DTYPE) if frame_st_id == 0 else None
+            input_dict = _prepare_latent_input(
+                latents,
+                None,
+                t,
+                t,
+                latent_cond,
+                None,
+                frame_st_id=frame_st_id,
+                patch_size=CONFIG["patch_size"],
+                prompt_embeds=prompt_embeds,
+                use_cfg=use_cfg,
+                negative_prompt_embeds=negative_prompt_embeds,
+            )
+
+            video_noise_pred = transformer(
+                _repeat_input_for_cfg(
+                    input_dict["latent_res_lst"],
+                    use_cfg,
+                    prompt_embeds,
+                    negative_prompt_embeds,
+                ),
+                update_cache=1 if last_step else 0,
+                cache_name="pos",
+                action_mode=False,
+            )
+
+            if not last_step or video_step != -1:
+                video_noise_pred = data_seq_to_patch(
+                    CONFIG["patch_size"],
+                    video_noise_pred,
+                    frame_chunk_size,
+                    latent_height,
+                    latent_width,
+                    batch_size=2 if use_cfg else 1,
+                )
+                if guidance_scale > 1:
+                    video_noise_pred = video_noise_pred[1:] + guidance_scale * (
+                        video_noise_pred[:1] - video_noise_pred[1:]
+                    )
+                else:
+                    video_noise_pred = video_noise_pred[:1]
+                latents = scheduler.step(video_noise_pred, t, latents, return_dict=False)
+
+            latents[:, :, 0:1] = (
+                latent_cond if frame_st_id == 0 else latents[:, :, 0:1]
+            )
+
+        # 2. Action generation loop
+        for i, t in enumerate(action_timesteps):
+            last_step = i == len(action_timesteps) - 1
+            action_cond = (
+                torch.zeros(
+                    [1, CONFIG["action_dim"], 1, CONFIG["action_per_frame"], 1],
+                    device=DEVICE,
+                    dtype=DTYPE,
+                )
+                if frame_st_id == 0
+                else None
+            )
+            input_dict = _prepare_latent_input(
+                None,
+                actions,
+                t,
+                t,
+                None,
+                action_cond,
+                frame_st_id=frame_st_id,
+                patch_size=CONFIG["patch_size"],
+                prompt_embeds=prompt_embeds,
+                use_cfg=use_cfg,
+                negative_prompt_embeds=negative_prompt_embeds,
+            )
+            action_noise_pred = transformer(
+                _repeat_input_for_cfg(
+                    input_dict["action_res_lst"],
+                    use_cfg,
+                    prompt_embeds,
+                    negative_prompt_embeds,
+                ),
+                update_cache=1 if last_step else 0,
+                cache_name="pos",
+                action_mode=True,
+            )
+
+            if not last_step:
+                action_noise_pred = rearrange(
+                    action_noise_pred, "b (f n) c -> b c f n 1", f=frame_chunk_size
+                )
+                if action_guidance_scale > 1:
+                    action_noise_pred = action_noise_pred[1:] + action_guidance_scale * (
+                        action_noise_pred[:1] - action_noise_pred[1:]
+                    )
+                else:
+                    action_noise_pred = action_noise_pred[:1]
+                actions = action_scheduler.step(
+                    action_noise_pred, t, actions, return_dict=False
+                )
+
+            actions[:, :, 0:1] = (
+                action_cond if frame_st_id == 0 else actions[:, :, 0:1]
+            )
+
+    actions[:, ~action_mask] *= 0
+    return actions, latents
+
+
+def decode_video(pred_latent):
+    """Decode latent tensor to video frames."""
+    latents = pred_latent.to(vae.dtype)
+    latents_mean = (
+        torch.tensor(vae.config.latents_mean)
+        .view(1, vae.config.z_dim, 1, 1, 1)
+        .to(latents.device, latents.dtype)
+    )
+    latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(
+        1, vae.config.z_dim, 1, 1, 1
+    ).to(latents.device, latents.dtype)
+    latents = latents / latents_std + latents_mean
+    video = vae.decode(latents, return_dict=False)[0]
+    video = video_processor.postprocess_video(video, output_type="np")[0]
+    return video
+
+
+# ------------------------------------------------------------------ #
+#  Gradio inference function
+# ------------------------------------------------------------------ #
+
+
+@spaces.GPU(duration=120)
+def generate(
+    top_img: "np.ndarray",
+    wrist_img: "np.ndarray",
+    prompt: str,
+    num_chunks: int = 5,
+    seed: int = 0,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate a robot manipulation video from initial observations and a text prompt.
+
+    Args:
+        top_img: Top-down camera observation image.
+        wrist_img: Wrist camera observation image.
+        prompt: Natural language instruction for the robot task.
+        num_chunks: Number of video chunks to generate autoregressively (each chunk = 4 frames).
+        seed: Random seed for reproducibility.
+    """
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+    use_cfg = CONFIG["guidance_scale"] > 1 or CONFIG["action_guidance_scale"] > 1
+
+    # Prepare observations
+    obs = [
+        {
+            CONFIG["obs_cam_keys"][0]: top_img,
+            CONFIG["obs_cam_keys"][1]: wrist_img,
+        }
+    ]
+
+    # Reset KV cache
+    transformer.clear_cache("pos")
+    streaming_vae.clear_cache()
+
+    # Encode initial observation
+    init_latent = encode_obs(obs)
+
+    # Encode prompt
+    prompt_embeds, negative_prompt_embeds = encode_prompt(prompt)
+
+    # Latent dimensions
+    latent_height = CONFIG["height"] // 16
+    latent_width = (CONFIG["width"] // 16) * len(CONFIG["obs_cam_keys"])
+    patch_size = CONFIG["patch_size"]
+    latent_token_per_chunk = (
+        CONFIG["frame_chunk_size"] * latent_height * latent_width
+    ) // (patch_size[0] * patch_size[1] * patch_size[2])
+    action_token_per_chunk = CONFIG["frame_chunk_size"] * CONFIG["action_per_frame"]
+
+    # Create KV cache
+    transformer.create_empty_cache(
+        "pos",
+        CONFIG["attn_window"],
+        latent_token_per_chunk,
+        action_token_per_chunk,
+        device=DEVICE,
+        dtype=DTYPE,
+        batch_size=2 if use_cfg else 1,
+    )
+
+    # Autoregressive chunk generation
+    pred_latent_lst = []
+    for chunk_id in range(num_chunks):
+        frame_st_id = chunk_id * CONFIG["frame_chunk_size"]
+        actions, latents = infer_chunk(
+            init_latent,
+            frame_st_id,
+            prompt_embeds,
+            negative_prompt_embeds,
+            use_cfg,
+            CONFIG["guidance_scale"],
+            CONFIG["action_guidance_scale"],
+            num_chunks,
+        )
+        pred_latent_lst.append(latents)
+
+    pred_latent = torch.cat(pred_latent_lst, dim=2)
+
+    # Clean up caches
+    transformer.clear_cache("pos")
+    streaming_vae.clear_cache()
+    torch.cuda.empty_cache()
+
+    # Decode video
+    video = decode_video(pred_latent)
+
+    # Save to temp file
+    from diffusers.utils import export_to_video
+
+    tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    tmp_file.close()
+    export_to_video(video, tmp_file.name, fps=10)
+
+    return tmp_file.name
+
+
+# ------------------------------------------------------------------ #
+#  Gradio UI
+# ------------------------------------------------------------------ #
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.Markdown(
+        """
+        # Next Forcing: Causal World Modeling with Multi-Chunk Prediction
+
+        Generate robot manipulation video from initial observations and a text instruction.
+        Upload top-down and wrist camera images, describe the task, and the model autoregressively
+        predicts future video frames.
+
+        [Paper](https://arxiv.org/abs/2606.11187) | [Code](https://github.com/gangweix/next-forcing) | [Model](https://huggingface.co/gangweix/next-forcing-base)
+        """
+    )
+
+    with gr.Row():
+        with gr.Column():
+            top_img = gr.Image(
+                label="Top Camera",
+                type="numpy",
+                height=256,
+            )
+            wrist_img = gr.Image(
+                label="Wrist Camera",
+                type="numpy",
+                height=256,
+            )
+            prompt = gr.Textbox(
+                label="Task Instruction",
+                placeholder="e.g. Pick the green cube and place it inside the blue box",
+                lines=2,
+            )
+            with gr.Accordion("Advanced Settings", open=False):
+                num_chunks = gr.Slider(
+                    label="Number of chunks (4 frames each)",
+                    minimum=1,
+                    maximum=10,
+                    value=5,
+                    step=1,
+                )
+                seed = gr.Number(label="Seed", value=0, precision=0)
+            run_btn = gr.Button("Generate Video", variant="primary")
+
+        with gr.Column():
+            video_out = gr.Video(label="Generated Video")
+
+    gr.Examples(
+        examples=[
+            [
+                "examples/observation.images.top.png",
+                "examples/observation.images.wrist.png",
+                "Pick the green cube and place it inside the blue box",
+                5,
+                0,
+            ],
+            [
+                "examples/observation.images.top.png",
+                "examples/observation.images.wrist.png",
+                "Move the red block to the left side of the table",
+                5,
+                42,
+            ],
+        ],
+        inputs=[top_img, wrist_img, prompt, num_chunks, seed],
+        outputs=video_out,
+        fn=generate,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    run_btn.click(
+        fn=generate,
+        inputs=[top_img, wrist_img, prompt, num_chunks, seed],
+        outputs=video_out,
+        api_name="generate",
+    )
+
+demo.launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,15 @@
+diffusers>=0.36.0
+transformers>=4.55.0
+accelerate
+einops
+easydict
+numpy
+pillow
+tqdm
+imageio[ffmpeg]
+ftfy
+safetensors
+torchvision
+websockets
+msgpack
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl

wan_va/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from . import configs, distributed, modules

wan_va/build_dataset_index.py

@@ -0,0 +1,41 @@
+"""Build reusable latent sample indexes without loading the training model."""
+
+import argparse
+
+from .configs import VA_CONFIGS
+from .dataset import MultiLatentLeRobotDataset
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--config-name', default='robotwin_train')
+    parser.add_argument(
+        '--rebuild',
+        action='store_true',
+        help='Validate all latent paths again and replace the matching caches.',
+    )
+    parser.add_argument('--init-worker', type=int)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    config = VA_CONFIGS[args.config_name]
+    config.rank = 0
+    config.world_size = 1
+    config.rebuild_dataset_index_cache = args.rebuild
+    if args.init_worker is not None:
+        config.init_worker = args.init_worker
+
+    dataset = MultiLatentLeRobotDataset(config=config)
+    print(
+        f'Dataset index ready: {len(dataset)} samples from '
+        f'{len(dataset._datasets)} datasets '
+        f'({dataset.index_cache_hits} cache hits, '
+        f'{dataset.hf_cache_hits} direct Arrow loads, '
+        f'{dataset.index_cache_misses} rebuilt)'
+    )
+
+
+if __name__ == '__main__':
+    main()

wan_va/configs/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from .va_franka_cfg import va_franka_cfg
+from .va_robotwin_cfg import va_robotwin_cfg
+from .va_franka_i2va import va_franka_i2va_cfg
+from .va_robotwin_i2va import va_robotwin_i2va_cfg
+from .va_robotwin_train_cfg import va_robotwin_train_cfg
+from .va_demo_train_cfg import va_demo_train_cfg
+from .va_demo_cfg import va_demo_cfg
+from .va_demo_i2va import va_demo_i2va_cfg
+
+VA_CONFIGS = {
+    'robotwin': va_robotwin_cfg,
+    'franka': va_franka_cfg,
+    'robotwin_i2va': va_robotwin_i2va_cfg,
+    'franka_i2va': va_franka_i2va_cfg,
+    'robotwin_train': va_robotwin_train_cfg,
+    'demo': va_demo_cfg,
+    'demo_train': va_demo_train_cfg,
+    'demo_i2va': va_demo_i2va_cfg,
+}

wan_va/configs/mcp_train_config.py

@@ -0,0 +1,13 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from easydict import EasyDict
+
+
+mcp_train_cfg = EasyDict()
+mcp_train_cfg.enable_mcp = True
+mcp_train_cfg.num_mcp_depths = 3
+mcp_train_cfg.mcp_blocks_per_depth = 3
+# Zero-based indices for Transformer layers 4, 12, 20, and 30 in the paper.
+mcp_train_cfg.mcp_hidden_collect_layers = [3, 11, 19, 29]
+mcp_train_cfg.mcp_snr_shift = 10.0
+mcp_train_cfg.mcp_loss_weights = [0.5, 0.2, 0.1]
+mcp_train_cfg.mcp_init_from_backbone = True

wan_va/configs/shared_config.py

@@ -0,0 +1,21 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import torch
+from easydict import EasyDict
+
+va_shared_cfg = EasyDict()
+
+va_shared_cfg.host = '0.0.0.0'
+va_shared_cfg.port = 29536
+va_shared_cfg.infer_mode = 'server'
+
+va_shared_cfg.param_dtype = torch.bfloat16
+va_shared_cfg.save_root = './train_out'
+
+va_shared_cfg.patch_size = (1, 2, 2)
+
+va_shared_cfg.enable_offload = False
+
+# Cache the validated latent sample index so repeated training runs do not scan
+# every latent path on network storage again.
+va_shared_cfg.enable_dataset_index_cache = True
+va_shared_cfg.rebuild_dataset_index_cache = False

wan_va/configs/va_demo_cfg.py

@@ -0,0 +1,56 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import torch
+from easydict import EasyDict
+
+from .shared_config import va_shared_cfg
+
+va_demo_cfg = EasyDict(__name__='Config: VA demo')
+va_demo_cfg.update(va_shared_cfg)
+
+va_demo_cfg.wan22_pretrained_model_name_or_path = "/path/to/pretrained/model"
+
+va_demo_cfg.attn_window = 30
+va_demo_cfg.frame_chunk_size = 4
+va_demo_cfg.env_type = 'none'
+
+va_demo_cfg.height = 256
+va_demo_cfg.width = 256
+va_demo_cfg.action_dim = 30
+va_demo_cfg.action_per_frame = 8
+va_demo_cfg.obs_cam_keys = [
+    'observation.images.top', 'observation.images.wrist'
+]
+va_demo_cfg.guidance_scale = 5
+va_demo_cfg.action_guidance_scale = 1
+
+va_demo_cfg.num_inference_steps = 5
+va_demo_cfg.video_exec_step = -1
+va_demo_cfg.action_num_inference_steps = 10
+
+va_demo_cfg.snr_shift = 5.0
+va_demo_cfg.action_snr_shift = 1.0
+
+va_demo_cfg.used_action_channel_ids = list(range(0, 5)) + list(range(28, 29))
+inverse_used_action_channel_ids = [len(va_demo_cfg.used_action_channel_ids)
+                                   ] * va_demo_cfg.action_dim
+for i, j in enumerate(va_demo_cfg.used_action_channel_ids):
+    inverse_used_action_channel_ids[j] = i
+va_demo_cfg.inverse_used_action_channel_ids = inverse_used_action_channel_ids
+
+va_demo_cfg.action_norm_method = 'quantiles'
+va_demo_cfg.norm_stat = {
+    "q01": [
+        -90.60303497314453,
+        -98.73043060302734,
+        -79.9008560180664,
+        48.95470428466797,
+        -32.794578552246094,
+    ] + [0.] * 23 + [0.8250824809074402, 0],
+    "q99": [
+        71.735107421875,
+        65.89081573486328,
+        92.87967681884766,
+        100.0,
+        22.784151077270508,
+    ] + [0.] * 23 + [100.0, 0],
+}

wan_va/configs/va_demo_i2va.py

@@ -0,0 +1,11 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from easydict import EasyDict
+from .va_demo_cfg import va_demo_cfg
+
+va_demo_i2va_cfg = EasyDict(__name__='Config: VA demo i2va')
+va_demo_i2va_cfg.update(va_demo_cfg)
+
+va_demo_i2va_cfg.input_img_path = 'example/demo'
+va_demo_i2va_cfg.num_chunks_to_infer = 10
+va_demo_i2va_cfg.prompt = 'Pick the green cube and place it inside the blue box'
+va_demo_i2va_cfg.infer_mode = 'i2va'

wan_va/configs/va_demo_train_cfg.py

@@ -0,0 +1,27 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from easydict import EasyDict
+from .mcp_train_config import mcp_train_cfg
+from .va_demo_cfg import va_demo_cfg
+import os
+
+va_demo_train_cfg = EasyDict(__name__='Config: VA demo train')
+va_demo_train_cfg.update(va_demo_cfg)
+va_demo_train_cfg.update(mcp_train_cfg)
+
+va_demo_train_cfg.dataset_path = '/path/to/your/dataset'
+va_demo_train_cfg.empty_emb_path = os.path.join(va_demo_train_cfg.dataset_path, 'empty_emb.pt')
+va_demo_train_cfg.enable_wandb = True
+va_demo_train_cfg.load_worker = 16
+va_demo_train_cfg.save_interval = 50
+va_demo_train_cfg.gc_interval = 50
+va_demo_train_cfg.cfg_prob = 0.1
+
+# Training parameters
+va_demo_train_cfg.learning_rate = 1e-4
+va_demo_train_cfg.beta1 = 0.9
+va_demo_train_cfg.beta2 = 0.95
+va_demo_train_cfg.weight_decay = 1e-1
+va_demo_train_cfg.warmup_steps = 10
+va_demo_train_cfg.batch_size = 1 
+va_demo_train_cfg.gradient_accumulation_steps = 8
+va_demo_train_cfg.num_steps = 2000 

wan_va/configs/va_franka_cfg.py

@@ -0,0 +1,58 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import torch
+from easydict import EasyDict
+
+from .shared_config import va_shared_cfg
+
+va_franka_cfg = EasyDict(__name__='Config: VA franka')
+va_franka_cfg.update(va_shared_cfg)
+
+va_franka_cfg.wan22_pretrained_model_name_or_path = "/path/to/pretrained/model"
+
+va_franka_cfg.attn_window = 30
+va_franka_cfg.frame_chunk_size = 4
+va_franka_cfg.env_type = 'none'
+
+va_franka_cfg.height = 224
+va_franka_cfg.width = 320
+va_franka_cfg.action_dim = 30
+va_franka_cfg.action_per_frame = 20
+va_franka_cfg.obs_cam_keys = [
+    'observation.images.cam_high', 'observation.images.cam_left_wrist',
+    'observation.images.cam_right_wrist'
+]
+va_franka_cfg.guidance_scale = 5
+va_franka_cfg.action_guidance_scale = 1
+
+va_franka_cfg.num_inference_steps = 5
+va_franka_cfg.video_exec_step = -1
+va_franka_cfg.action_num_inference_steps = 10
+
+va_franka_cfg.snr_shift = 5.0
+va_franka_cfg.action_snr_shift = 1.0
+
+va_franka_cfg.used_action_channel_ids = list(range(0, 7)) + list(range(
+    28, 29)) + list(range(7, 14)) + list(range(29, 30))
+inverse_used_action_channel_ids = [len(va_franka_cfg.used_action_channel_ids)
+                                   ] * va_franka_cfg.action_dim
+for i, j in enumerate(va_franka_cfg.used_action_channel_ids):
+    inverse_used_action_channel_ids[j] = i
+va_franka_cfg.inverse_used_action_channel_ids = inverse_used_action_channel_ids
+
+va_franka_cfg.action_norm_method = 'quantiles'
+va_franka_cfg.norm_stat = {
+    "q01": [
+        0.3051295876502991, -0.22647984325885773, 0.19957000017166138,
+        -0.022680532187223434, -0.05553057789802551, -0.2693849802017212,
+        -0.29341773986816405, 0.2935442328453064, -0.4431332051753998,
+        0.21256473660469055, -0.7962440848350525, -0.40816226601600647,
+        -0.28359392285346985, -0.44507765769958496
+    ] + [0.] * 16,
+    "q99": [
+        0.7572150230407715, 0.47736290097236633, 0.6428080797195435,
+        0.9835678935050964, 0.9927203059196472, 0.28041139245033264,
+        0.47529348731040877, 0.7564866304397571, 0.04082797020673729,
+        0.5355993628501885, 0.9976375699043274, 0.8973174452781656,
+        0.6016915678977965, 0.5027598619461056
+    ] + [0.] * 14 + [1.0, 1.0],
+}

wan_va/configs/va_franka_i2va.py

@@ -0,0 +1,11 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from easydict import EasyDict
+from .va_franka_cfg import va_franka_cfg
+
+va_franka_i2va_cfg = EasyDict(__name__='Config: VA franka i2va')
+va_franka_i2va_cfg.update(va_franka_cfg)
+
+va_franka_i2va_cfg.input_img_path = 'example/franka'
+va_franka_i2va_cfg.num_chunks_to_infer = 10
+va_franka_i2va_cfg.prompt = 'pick bunk'
+va_franka_i2va_cfg.infer_mode = 'i2va'

wan_va/configs/va_robotwin_cfg.py

@@ -0,0 +1,59 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import os
+
+from easydict import EasyDict
+
+from .shared_config import va_shared_cfg
+
+va_robotwin_cfg = EasyDict(__name__='Config: VA robotwin')
+va_robotwin_cfg.update(va_shared_cfg)
+
+va_robotwin_cfg.wan22_pretrained_model_name_or_path = os.environ.get(
+    "NEXT_FORCING_MODEL_PATH",
+    "/path/to/next-forcing-posttrain-robotwin",
+)
+
+va_robotwin_cfg.attn_window = 72
+va_robotwin_cfg.frame_chunk_size = 2
+va_robotwin_cfg.env_type = 'robotwin_tshape'
+
+va_robotwin_cfg.height = 256
+va_robotwin_cfg.width = 320
+va_robotwin_cfg.action_dim = 30
+va_robotwin_cfg.action_per_frame = 16
+va_robotwin_cfg.obs_cam_keys = [
+    'observation.images.cam_high', 'observation.images.cam_left_wrist',
+    'observation.images.cam_right_wrist'
+]
+va_robotwin_cfg.guidance_scale = 5
+va_robotwin_cfg.action_guidance_scale = 1
+
+va_robotwin_cfg.num_inference_steps = 25
+va_robotwin_cfg.video_exec_step = -1
+va_robotwin_cfg.action_num_inference_steps = 50
+
+va_robotwin_cfg.snr_shift = 5.0
+va_robotwin_cfg.action_snr_shift = 1.0
+
+va_robotwin_cfg.used_action_channel_ids = list(range(0, 7)) + list(
+    range(28, 29)) + list(range(7, 14)) + list(range(29, 30))
+inverse_used_action_channel_ids = [
+    len(va_robotwin_cfg.used_action_channel_ids)
+] * va_robotwin_cfg.action_dim
+for i, j in enumerate(va_robotwin_cfg.used_action_channel_ids):
+    inverse_used_action_channel_ids[j] = i
+va_robotwin_cfg.inverse_used_action_channel_ids = inverse_used_action_channel_ids
+
+va_robotwin_cfg.action_norm_method = 'quantiles'
+va_robotwin_cfg.norm_stat = {
+    "q01": [
+        -0.06172713458538055, -3.6716461181640625e-05, -0.08783501386642456,
+        -1, -1, -1, -1, -0.3547105032205582, -1.3113021850585938e-06,
+        -0.11975435614585876, -1, -1, -1, -1
+    ] + [0.] * 16,
+    "q99": [
+        0.3462600058317184, 0.39966784834861746, 0.14745532035827624, 1, 1, 1,
+        1, 0.034201726913452024, 0.39142737388610793, 0.1792279863357542, 1, 1,
+        1, 1
+    ] + [0.] * 14 + [1.0, 1.0],
+}

wan_va/configs/va_robotwin_i2va.py

@@ -0,0 +1,11 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from easydict import EasyDict
+from .va_robotwin_cfg import va_robotwin_cfg
+
+va_robotwin_i2va_cfg = EasyDict(__name__='Config: VA robotwin i2va')
+va_robotwin_i2va_cfg.update(va_robotwin_cfg)
+
+va_robotwin_i2va_cfg.input_img_path = 'example/robotwin'
+va_robotwin_i2va_cfg.num_chunks_to_infer = 10
+va_robotwin_i2va_cfg.prompt = 'Grab the medium-sized white mug, rotate it, place it on the table, and hook it onto the smooth dark gray rack.'
+va_robotwin_i2va_cfg.infer_mode = 'i2va'

wan_va/configs/va_robotwin_train_cfg.py

@@ -0,0 +1,40 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from easydict import EasyDict
+from .mcp_train_config import mcp_train_cfg
+from .va_robotwin_cfg import va_robotwin_cfg
+import os
+
+va_robotwin_train_cfg = EasyDict(__name__='Config: VA robotwin train')
+va_robotwin_train_cfg.update(va_robotwin_cfg)
+va_robotwin_train_cfg.update(mcp_train_cfg)
+
+va_robotwin_train_cfg.wan22_pretrained_model_name_or_path = os.environ.get(
+    'NEXT_FORCING_PRETRAINED_MODEL_PATH',
+    '/path/to/pretrained/model',
+)
+va_robotwin_train_cfg.dataset_path = os.environ.get(
+    'NEXT_FORCING_DATASET_PATH',
+    '/path/to/your/dataset',
+)
+va_robotwin_train_cfg.empty_emb_path = os.path.join(
+    va_robotwin_train_cfg.dataset_path, 'empty_emb.pt')
+va_robotwin_train_cfg.save_root = os.environ.get(
+    'NEXT_FORCING_SAVE_ROOT',
+    '/path/to/your/output',
+)
+va_robotwin_train_cfg.enable_wandb = False
+va_robotwin_train_cfg.init_worker = 1
+va_robotwin_train_cfg.load_worker = 16
+va_robotwin_train_cfg.save_interval = 1000
+va_robotwin_train_cfg.gc_interval = 50
+va_robotwin_train_cfg.cfg_prob = 0.1
+
+# Training parameters
+va_robotwin_train_cfg.learning_rate = 2e-5
+va_robotwin_train_cfg.beta1 = 0.9
+va_robotwin_train_cfg.beta2 = 0.95
+va_robotwin_train_cfg.weight_decay = 0.1
+va_robotwin_train_cfg.warmup_steps = 100
+va_robotwin_train_cfg.batch_size = 1 
+va_robotwin_train_cfg.gradient_accumulation_steps = 1
+va_robotwin_train_cfg.num_steps = 50000 

wan_va/dataset/__init__.py

@@ -0,0 +1,10 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from .lerobot_latent_dataset import (
+    MultiLatentLeRobotDataset,
+    dataset_indexes_ready,
+)
+
+__all__ = [
+    'MultiLatentLeRobotDataset',
+    'dataset_indexes_ready',
+]

wan_va/dataset/lerobot_latent_dataset.py

@@ -0,0 +1,548 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.utils import get_episode_data_index
+from lerobot.datasets.compute_stats import aggregate_stats, compute_episode_stats
+import datasets
+import hashlib
+import json
+import numpy as np
+from pathlib import Path
+from collections.abc import Callable
+import os
+from tqdm import tqdm
+from multiprocessing import Pool
+from functools import partial
+import torch
+from einops import rearrange
+from torch.utils.data import DataLoader
+from scipy.spatial.transform import Rotation as R
+from lerobot.constants import HF_LEROBOT_HOME
+
+
+DATASET_INDEX_CACHE_VERSION = 1
+
+
+def dataset_index_fingerprint(dataset_root, latent_root, video_keys):
+    episodes_path = Path(dataset_root) / 'meta' / 'episodes.jsonl'
+    episodes_stat = episodes_path.stat()
+    return {
+        'version': DATASET_INDEX_CACHE_VERSION,
+        'episodes_size': episodes_stat.st_size,
+        'episodes_mtime_ns': episodes_stat.st_mtime_ns,
+        'latent_root': str(Path(latent_root).resolve()),
+        'video_keys': list(video_keys),
+    }
+
+
+def dataset_index_cache_path(dataset_root, fingerprint):
+    cache_key = json.dumps(fingerprint, sort_keys=True).encode('utf-8')
+    cache_key = hashlib.sha256(cache_key).hexdigest()[:16]
+    return (
+        Path(dataset_root)
+        / '.cache'
+        / 'next_forcing'
+        / f'valid_metas_{cache_key}.json'
+    )
+
+
+def dataset_indexes_ready(config):
+    if not getattr(config, 'enable_dataset_index_cache', True):
+        return False
+    if getattr(config, 'rebuild_dataset_index_cache', False):
+        return False
+
+    repo_list = recursive_find_file(config.dataset_path, 'info.json')
+    repo_list = [Path(path.split('/meta/info.json')[0]) for path in repo_list]
+    if not repo_list:
+        return False
+
+    for dataset_root in repo_list:
+        fingerprint = dataset_index_fingerprint(
+            dataset_root,
+            dataset_root / 'latents',
+            config.obs_cam_keys,
+        )
+        cache_path = dataset_index_cache_path(dataset_root, fingerprint)
+        try:
+            with cache_path.open('r', encoding='utf-8') as handle:
+                payload = json.load(handle)
+        except (FileNotFoundError, json.JSONDecodeError, OSError):
+            return False
+
+        cache_files = payload.get('hf_cache_files')
+        if payload.get('fingerprint') != fingerprint:
+            return False
+        if not isinstance(payload.get('valid_metas'), list):
+            return False
+        if not cache_files or not all(Path(path).is_file() for path in cache_files):
+            return False
+    return True
+
+def recursive_find_file(directory, filename='info.json'):
+    result = []
+    try:
+        for root, dirs, files in os.walk(directory):
+            if filename in files:
+                full_path = os.path.join(root, filename)
+                result.append(full_path)
+    except PermissionError:
+        print(f"Error: can not access {directory}")
+    except Exception as e:
+        print(f"Error: {e}")
+    return result
+
+def construct_lerobot(
+    repo_id,
+    config,
+):
+    return LatentLeRobotDataset(
+        repo_id=repo_id,
+        config=config,
+    )
+
+def construct_lerobot_multi_processor(config, 
+                                      num_init_worker=8,
+                                      ):
+    datasets_out_lst = []
+    construct_func = partial(
+        construct_lerobot,
+        config=config,
+    )
+    repo_list = recursive_find_file(config.dataset_path, 'info.json')
+    repo_list = [v.split('/meta/info.json')[0] for v in repo_list]
+    if not repo_list:
+        raise FileNotFoundError(
+            f"No LeRobot datasets found under {config.dataset_path}")
+    num_init_worker = min(max(int(num_init_worker), 1), len(repo_list))
+    if num_init_worker == 1:
+        dataset_iterator = map(construct_func, repo_list)
+        if getattr(config, 'rank', 0) == 0:
+            dataset_iterator = tqdm(
+                dataset_iterator,
+                total=len(repo_list),
+                desc='Loading dataset indexes',
+            )
+        datasets_out_lst = list(dataset_iterator)
+    else:
+        with Pool(num_init_worker) as pool:
+            dataset_iterator = pool.imap(construct_func, repo_list)
+            if getattr(config, 'rank', 0) == 0:
+                dataset_iterator = tqdm(
+                    dataset_iterator,
+                    total=len(repo_list),
+                    desc='Loading dataset indexes',
+                )
+            datasets_out_lst = list(dataset_iterator)
+                
+    return datasets_out_lst
+
+def get_relative_pose(pose):
+    if torch.is_tensor(pose):
+        pose = pose.detach().cpu().numpy()
+    
+    rot = R.from_quat(pose[:, 3:7])
+    first_rot = R.from_quat(np.tile(pose[:1, 3:7], (pose.shape[0], 1)))
+    trans = pose[:, :3]
+    relative_trans = trans - trans[0:1]
+
+    relative_rot = first_rot.inv() * rot
+    relative_quat = relative_rot.as_quat()
+
+    relative_pose = np.concatenate([relative_trans, relative_quat], axis=1)
+    return torch.from_numpy(relative_pose)
+
+class MultiLatentLeRobotDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        config,
+        num_init_worker=None,
+    ):
+        if num_init_worker is None:
+            num_init_worker = getattr(config, 'init_worker', 8)
+        self._datasets = construct_lerobot_multi_processor(config, 
+                                                           num_init_worker, 
+                                                           )
+        self.index_cache_hits = sum(
+            dataset.index_cache_hit for dataset in self._datasets)
+        self.index_cache_misses = len(self._datasets) - self.index_cache_hits
+        self.hf_cache_hits = sum(
+            dataset.hf_cache_hit for dataset in self._datasets)
+        self.item_id_to_dataset_id, self.acc_dset_num = (
+            self._get_item_id_to_dataset_id()
+        )
+
+    def __len__(
+        self,
+    ):
+        return sum(len(v) for v in self._datasets)
+
+    def _get_item_id_to_dataset_id(self):
+        item_id_to_dataset_id = {}
+        acc_dset_num = {}
+        acc_nums = [0]
+        id = 0
+        for dset_id, dset in enumerate(self._datasets):
+            acc_nums.append(acc_nums[-1] + len(dset))
+            for _ in range(len(dset)):
+                item_id_to_dataset_id[id] = dset_id
+                id += 1
+        for did in range(len(self._datasets)):
+            acc_dset_num[did] = acc_nums[did]
+        return item_id_to_dataset_id, acc_dset_num
+
+    def __getitem__(self, idx) -> dict:
+        assert idx < len(self)
+        cur_dset = self._datasets[self.item_id_to_dataset_id[idx]]
+        local_idx = idx - self.acc_dset_num[self.item_id_to_dataset_id[idx]]
+        return cur_dset[local_idx]
+
+class LatentLeRobotDataset(LeRobotDataset):
+    def __init__(
+        self,
+        repo_id,
+        config=None,
+    ):
+        self.repo_id = repo_id
+        self.root = HF_LEROBOT_HOME / repo_id
+        self.image_transforms = None
+        self.delta_timestamps = None
+        self.episodes = None
+        self.tolerance_s = 1e-4
+        self.revision = "v2.1"
+        self.video_backend = 'pyav'
+        self.delta_indices = None
+        self.batch_encoding_size = 1
+        self.episodes_since_last_encoding = 0
+        self.image_writer = None
+        self.episode_buffer = None
+        self.root.mkdir(exist_ok=True, parents=True)
+        self.meta = LeRobotDatasetMetadata(
+            self.repo_id, self.root, self.revision, force_cache_sync=False
+        )
+        if self.episodes is not None and self.meta._version >= packaging.version.parse("v2.1"):
+            episodes_stats = [self.meta.episodes_stats[ep_idx] for ep_idx in self.episodes]
+            self.stats = aggregate_stats(episodes_stats)
+        
+        self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes)
+        
+        self.latent_path = Path(repo_id) / 'latents'
+        self.empty_emb = torch.load(config.empty_emb_path, weights_only=False)
+        self.config = config
+        self.cfg_prob = config.cfg_prob
+        self.used_video_keys = config.obs_cam_keys
+        self.q01 = np.array(config.norm_stat['q01'], dtype='float')[None]
+        self.q99 = np.array(config.norm_stat['q99'], dtype='float')[None]
+        fingerprint = self._dataset_index_fingerprint()
+        cache_payload = self._read_dataset_index_cache(fingerprint)
+        self.hf_dataset = self._load_cached_hf_dataset(cache_payload)
+        self.hf_cache_hit = self.hf_dataset is not None
+        if self.hf_dataset is None:
+            try:
+                assert all((self.root / fpath).is_file() for fpath in self.get_episodes_file_paths())
+                self.hf_dataset = self.load_hf_dataset()
+            except (AssertionError, FileNotFoundError, NotADirectoryError):
+                self.revision = get_safe_version(self.repo_id, self.revision)
+                self.download_episodes(download_videos)
+                self.hf_dataset = self.load_hf_dataset()
+        self._hf_torch_view = self.hf_dataset.with_format(
+                type='torch',
+                columns=['action'],
+                output_all_columns=False
+            )
+        self.index_cache_hit = self.parse_meta(
+            fingerprint=fingerprint,
+            cache_payload=cache_payload,
+        )
+        if self.index_cache_hit and not self.hf_cache_hit:
+            self._save_dataset_index_cache(fingerprint, self.new_metas)
+
+    def _dataset_index_fingerprint(self):
+        return dataset_index_fingerprint(
+            self.root,
+            self.latent_path,
+            self.used_video_keys,
+        )
+
+    def _dataset_index_cache_path(self, fingerprint):
+        return dataset_index_cache_path(self.root, fingerprint)
+
+    def _read_dataset_index_cache(self, fingerprint):
+        if not getattr(self.config, 'enable_dataset_index_cache', True):
+            return None
+        if getattr(self.config, 'rebuild_dataset_index_cache', False):
+            return None
+
+        cache_path = self._dataset_index_cache_path(fingerprint)
+        try:
+            with cache_path.open('r', encoding='utf-8') as handle:
+                payload = json.load(handle)
+        except (FileNotFoundError, json.JSONDecodeError, OSError):
+            return None
+
+        if payload.get('fingerprint') != fingerprint:
+            return None
+        return payload
+
+    def _load_dataset_index_cache(self, fingerprint, cache_payload=None):
+        payload = cache_payload
+        if payload is None:
+            payload = self._read_dataset_index_cache(fingerprint)
+        if payload is None:
+            return None
+        valid_metas = payload.get('valid_metas')
+        if not isinstance(valid_metas, list):
+            return None
+        return valid_metas
+
+    def _load_cached_hf_dataset(self, cache_payload):
+        if cache_payload is None:
+            return None
+        cache_files = cache_payload.get('hf_cache_files')
+        if not cache_files or not all(Path(path).is_file() for path in cache_files):
+            return None
+
+        cached_datasets = [
+            datasets.Dataset.from_file(path) for path in cache_files
+        ]
+        if len(cached_datasets) == 1:
+            return cached_datasets[0]
+        return datasets.concatenate_datasets(cached_datasets)
+
+    def _hf_cache_files(self):
+        hf_dataset = getattr(self, 'hf_dataset', None)
+        if hf_dataset is None:
+            return []
+        return [
+            cache_file['filename']
+            for cache_file in hf_dataset.cache_files
+            if Path(cache_file['filename']).is_file()
+        ]
+
+    def _save_dataset_index_cache(self, fingerprint, valid_metas):
+        if not getattr(self.config, 'enable_dataset_index_cache', True):
+            return
+
+        cache_path = self._dataset_index_cache_path(fingerprint)
+        temporary_path = cache_path.with_name(
+            f'.{cache_path.name}.tmp-{os.getpid()}'
+        )
+        try:
+            cache_path.parent.mkdir(parents=True, exist_ok=True)
+            with temporary_path.open('w', encoding='utf-8') as handle:
+                json.dump(
+                    {
+                        'fingerprint': fingerprint,
+                        'valid_metas': valid_metas,
+                        'hf_cache_files': self._hf_cache_files(),
+                    },
+                    handle,
+                    separators=(',', ':'),
+                )
+            os.replace(temporary_path, cache_path)
+        except OSError as exc:
+            try:
+                temporary_path.unlink()
+            except FileNotFoundError:
+                pass
+            raise RuntimeError(
+                f'Failed to write dataset index cache: {cache_path}'
+            ) from exc
+
+    def parse_meta(self, fingerprint=None, cache_payload=None):
+        if fingerprint is None:
+            fingerprint = self._dataset_index_fingerprint()
+        cached_metas = self._load_dataset_index_cache(
+            fingerprint,
+            cache_payload=cache_payload,
+        )
+        if cached_metas is not None:
+            self.new_metas = cached_metas
+            return True
+
+        out = []
+        for key, value in self.meta.episodes.items():
+            episode_index = value["episode_index"]
+            tasks = value["tasks"]
+            action_config = value["action_config"]
+            for acfg in action_config:
+                cur_meta = {
+                    "episode_index": episode_index,
+                    "tasks": tasks,
+                }
+                cur_meta.update(acfg)
+
+                check_statu = self._check_meta(
+                    cur_meta["start_frame"],
+                    cur_meta["end_frame"],
+                    cur_meta["episode_index"],
+                )
+
+                if check_statu:
+                    out.append(cur_meta)
+        self.new_metas = out
+        self._save_dataset_index_cache(fingerprint, out)
+        return False
+
+    def _check_meta(self, start_frame, end_frame, episode_index):
+        episode_chunk = self.meta.get_episode_chunk(episode_index)
+        latent_path = Path(self.latent_path) / f"chunk-{episode_chunk:03d}"
+        for key in self.used_video_keys:
+            cur_path = latent_path / key
+            latent_file = (
+                cur_path / f"episode_{episode_index:06d}_{start_frame}_{end_frame}.pth"
+            )
+            if not os.path.exists(latent_file):
+                return False
+        return True
+
+    def _get_global_idx(self, episode_index: int, local_index: int):
+        ep_start = self.episode_data_index["from"][episode_index]
+        return local_index + ep_start
+
+    def _get_range_hf_data(self, start_frame, end_frame):
+        batch = self._hf_torch_view[start_frame:end_frame]
+        return batch
+
+    def _flatten_latent_dict(self, latent_dict):
+        out = {}
+        for key, value in latent_dict.items():
+            for inner_key, inner_value in value.items():
+                new_key = f"{key}.{inner_key}"
+                out[new_key] = inner_value
+        return out
+
+    def _get_range_latent_data(self, start_frame, end_frame, episode_index):
+        episode_chunk = self.meta.get_episode_chunk(episode_index)
+        latent_path = Path(self.latent_path) / f"chunk-{episode_chunk:03d}"
+        out = {}
+        for key in self.used_video_keys:
+            cur_path = latent_path / key
+            latent_file = (
+                cur_path / f"episode_{episode_index:06d}_{start_frame}_{end_frame}.pth"
+            )
+            assert os.path.exists(latent_file)
+            latent_data = torch.load(latent_file, weights_only=False)
+            out[key] = latent_data
+        
+        return self._flatten_latent_dict(out)
+    
+        
+    def _cat_video_latents(self,
+                           data_dict
+                           ):
+        latent_lst = []
+        for key in self.used_video_keys:
+            latent= data_dict[f"{key}.latent"]
+            latent_num_frames = data_dict[f"{key}.latent_num_frames"]
+            latent_height = data_dict[f"{key}.latent_height"]
+            latent_width = data_dict[f"{key}.latent_width"]
+            latent = rearrange(latent, 
+                                 '(f h w) c -> f h w c', 
+                                 f=latent_num_frames, 
+                                 h=latent_height, 
+                                 w=latent_width)
+            latent_lst.append(latent)
+        if self.config.env_type == 'robotwin_tshape':
+            wrist_latent = torch.cat(latent_lst[1:], dim=2)
+            cat_latent = torch.cat([wrist_latent, latent_lst[0]], dim=1)
+        else:
+            cat_latent = torch.cat(latent_lst, dim=2)
+
+        text_emb = data_dict[f"{self.used_video_keys[0]}.text_emb"]
+        if torch.rand(1).item() < self.cfg_prob:
+            text_emb = self.empty_emb
+
+        out_dict = dict(
+            latents = cat_latent,
+            text_emb = text_emb,
+        )
+        return out_dict
+    
+    def _action_post_process(self, local_start_frame, local_end_frame, latent_frame_ids, action):
+        act_shift = int(latent_frame_ids[0] - local_start_frame)
+        frame_stride = latent_frame_ids[1] - latent_frame_ids[0]
+        action = action[act_shift:]
+        if self.config.env_type == 'robotwin_tshape': ## TODO support get_relative_pose for other dataset, currently only support robotwin 
+            left_action = get_relative_pose(action[:, :7])
+            right_action = get_relative_pose(action[:, 8:15])
+            action = np.concatenate([left_action, action[:, 7:8], right_action, action[:, 15:16]], axis=1)
+        action = np.pad(action, pad_width=((frame_stride * 4, 0), (0, 0)), mode='constant', constant_values=0)
+
+        latent_frame_num = (len(latent_frame_ids) - 1) // 4 + 1
+        required_action_num = latent_frame_num * frame_stride * 4
+
+        action = action[:required_action_num]
+        action_mask = np.ones_like(action, dtype='bool')
+        assert action.shape[0] == required_action_num
+
+
+        action_paded = np.pad(action, ((0, 0), (0, 1)), mode='constant', constant_values=0)
+        action_mask_padded = np.pad(action_mask, ((0, 0), (0, 1)), mode='constant', constant_values=0)
+
+        action_aligned = action_paded[:, self.config.inverse_used_action_channel_ids]
+        action_mask_aligned = action_mask_padded[:, self.config.inverse_used_action_channel_ids]
+        action_aligned = (action_aligned - self.q01) / (
+                self.q99 - self.q01 + 1e-6) * 2. - 1.
+        action_aligned = np.clip(action_aligned, -1.5, 1.5)
+        action_aligned = rearrange(action_aligned, "(f n) c -> c f n 1", f=latent_frame_num)
+        action_mask_aligned = rearrange(action_mask_aligned, "(f n) c -> c f n 1", f=latent_frame_num)
+        action_aligned *= action_mask_aligned
+        return torch.from_numpy(action_aligned).float(), torch.from_numpy(action_mask_aligned).bool()
+
+    def __getitem__(self, idx) -> dict:
+        idx = idx % len(self.new_metas)
+        cur_meta = self.new_metas[idx]
+        episode_index = cur_meta["episode_index"]
+        start_frame = cur_meta["start_frame"]
+        end_frame = cur_meta["end_frame"]
+        local_start_frame = start_frame
+        local_end_frame = end_frame
+
+        ori_data_dict = self._get_range_latent_data(start_frame, end_frame, episode_index)
+
+        latent_frame_ids = ori_data_dict[f"{self.used_video_keys[0]}.frame_ids"]
+        start_frame = self._get_global_idx(episode_index, start_frame)
+        end_frame = self._get_global_idx(episode_index, end_frame)
+
+        hf_data_frames = self._get_range_hf_data(start_frame, end_frame)
+        ori_data_dict.update(hf_data_frames)
+        out_dict = self._cat_video_latents(ori_data_dict)
+
+        out_dict['actions'], out_dict['actions_mask'] = self._action_post_process(local_start_frame, local_end_frame, latent_frame_ids, ori_data_dict['action'])
+
+        out_dict['latents'] = out_dict['latents'].permute(3, 0, 1, 2)
+        return out_dict
+
+    def __len__(self):
+        return len(self.new_metas)
+
+if __name__ == '__main__':
+    from wan_va.configs import VA_CONFIGS
+    from tqdm import tqdm
+    dset = MultiLatentLeRobotDataset(
+        VA_CONFIGS['demo_train']
+    )
+    for key, value in dset[0].items():
+        if isinstance(value, torch.Tensor):
+            print(f'{key}: {value.shape} tensor')
+        elif isinstance(value, np.ndarray):
+            print(f'{key}: {value.shape} np')
+        else:
+            print(f'{key}: {value}')
+    print(len(dset))
+    dloader = DataLoader(
+            dset,
+            batch_size=1,
+            shuffle=True,
+            num_workers=32,
+        )
+    max_l = 0
+    action_list = []
+    for data in tqdm(dloader):
+        _, _, F, H, W = data['latents'].shape
+        max_l = max(max_l, F*H*W)
+        action_list.append(data['actions'].flatten(2).permute(0, 2, 1).flatten(0, 1))
+    action_all = torch.cat(action_list, dim=0)
+    print(max_l)
+    print(action_all.shape, action_all.mean(dim=0), action_all.min(dim=0)[0], action_all.max(dim=0)[0])
+    

wan_va/distributed/__init__.py

@@ -0,0 +1 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.

wan_va/distributed/fsdp.py

@@ -0,0 +1,55 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import gc
+
+import torch
+from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
+
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    checkpoint_wrapper as ptd_checkpoint_wrapper,
+)
+
+def apply_ac(model):
+    """Apply activation checkpointing to the model."""
+    for layer_id, transformer_block in enumerate(model.blocks):
+        transformer_block = ptd_checkpoint_wrapper(transformer_block, preserve_rng_state=False)
+        model.blocks[layer_id] = transformer_block
+    if getattr(model, 'enable_mcp', False):
+        for group_id, mcp_group in enumerate(model.mcp_blocks):
+            for block_id, transformer_block in enumerate(mcp_group):
+                transformer_block = ptd_checkpoint_wrapper(
+                    transformer_block, preserve_rng_state=False)
+                model.mcp_blocks[group_id][block_id] = transformer_block
+
+
+def shard_model(model,
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.float32):
+    mp_policy = MixedPrecisionPolicy(
+        param_dtype=param_dtype,
+        reduce_dtype=reduce_dtype,
+        cast_forward_inputs=False,
+    )
+    fsdp_config = {"mp_policy": mp_policy, "reshard_after_forward": True}
+
+    for block in model.blocks:
+        fully_shard(block.attn1, **fsdp_config)
+        fully_shard(block.attn2, **fsdp_config)
+        fully_shard(block.ffn, **fsdp_config)
+        fully_shard(block, **fsdp_config)
+
+    if getattr(model, 'enable_mcp', False):
+        for mcp_group in model.mcp_blocks:
+            for block in mcp_group:
+                fully_shard(block.attn1, **fsdp_config)
+                fully_shard(block.attn2, **fsdp_config)
+                fully_shard(block.ffn, **fsdp_config)
+                fully_shard(block, **fsdp_config)
+
+    fully_shard(model, **fsdp_config)
+    return model
+
+
+def free_model(model):
+    del model
+    gc.collect()
+    torch.cuda.empty_cache()

wan_va/distributed/util.py

@@ -0,0 +1,40 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+import torch.distributed as dist
+
+
+def _configure_model(model, shard_fn, param_dtype, device, eval_mode=True):
+    """
+    TODO
+    """
+    if eval_mode:
+        model.eval().requires_grad_(False)
+    if dist.is_initialized():
+        dist.barrier()
+
+    if dist.is_initialized():
+        model = shard_fn(model)
+    else:
+        model.to(param_dtype)
+        model.to(device)
+
+    return model
+
+
+def init_distributed(world_size, local_rank, rank):
+    # if world_size > 1:
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend="nccl",
+                            init_method="env://",
+                            rank=rank,
+                            world_size=world_size)
+
+def dist_mean(local_tensor):
+    if dist.is_initialized():
+        dist.all_reduce(local_tensor, op=dist.ReduceOp.AVG)
+    return local_tensor
+
+def dist_max(local_tensor):
+    if dist.is_initialized():
+        dist.all_reduce(local_tensor, op=dist.ReduceOp.MAX)
+    return local_tensor

wan_va/mcp.py

@@ -0,0 +1,59 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import torch
+
+
+def validate_mcp_settings(
+    num_mcp_depths,
+    mcp_blocks_per_depth,
+    mcp_hidden_collect_layers,
+    mcp_loss_weights=None,
+    num_layers=None,
+):
+    if num_mcp_depths <= 0:
+        raise ValueError("num_mcp_depths must be positive")
+    if mcp_blocks_per_depth <= 0:
+        raise ValueError("mcp_blocks_per_depth must be positive")
+    if len(mcp_hidden_collect_layers) == 0:
+        raise ValueError("mcp_hidden_collect_layers cannot be empty")
+    if len(set(mcp_hidden_collect_layers)) != len(mcp_hidden_collect_layers):
+        raise ValueError("mcp_hidden_collect_layers must be unique")
+    if any(layer < 0 for layer in mcp_hidden_collect_layers):
+        raise ValueError("mcp_hidden_collect_layers must be non-negative")
+    if num_layers is not None and max(mcp_hidden_collect_layers) >= num_layers:
+        raise ValueError(
+            "mcp_hidden_collect_layers must be smaller than the number of model layers"
+        )
+    if num_layers is not None and mcp_blocks_per_depth > num_layers:
+        raise ValueError(
+            "mcp_blocks_per_depth cannot exceed the number of model layers")
+    if mcp_loss_weights is not None:
+        if len(mcp_loss_weights) != num_mcp_depths:
+            raise ValueError("mcp_loss_weights must match num_mcp_depths")
+        if any(weight < 0 for weight in mcp_loss_weights):
+            raise ValueError("mcp_loss_weights must be non-negative")
+
+
+def shift_latents_for_mcp(latents, frame_shift):
+    """Shift video latents into the future and mark targets backed by real frames."""
+    if latents.ndim != 5:
+        raise ValueError("latents must have shape [B, C, F, H, W]")
+    if frame_shift <= 0:
+        raise ValueError("frame_shift must be positive")
+
+    batch_size, _, num_frames, _, _ = latents.shape
+    if num_frames == 0:
+        raise ValueError("latents must contain at least one frame")
+
+    pad_frames = min(frame_shift, num_frames)
+    shifted = latents[:, :, pad_frames:]
+    last_frame = latents[:, :, -1:]
+    shifted = torch.cat(
+        [shifted, last_frame.expand(-1, -1, pad_frames, -1, -1)], dim=2
+    )
+
+    valid_frames = torch.arange(num_frames, device=latents.device) + frame_shift
+    valid_frames = valid_frames < num_frames
+    valid_mask = valid_frames.view(1, 1, num_frames, 1, 1).expand(
+        batch_size, -1, -1, -1, -1
+    )
+    return shifted, valid_mask

wan_va/modules/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from .utils import (
+    WanVAEStreamingWrapper,
+    load_text_encoder,
+    load_tokenizer,
+    load_transformer,
+    load_vae,
+)
+
+__all__ = [
+    'load_transformer', 'load_text_encoder', 'load_tokenizer', 'load_vae',
+    'WanVAEStreamingWrapper'
+]

wan_va/modules/model.py

@@ -0,0 +1,1196 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import math
+from copy import deepcopy
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention import FeedForward
+from diffusers.models.embeddings import (
+    PixArtAlphaTextProjection,
+    TimestepEmbedding,
+    Timesteps,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import FP32LayerNorm
+from einops import rearrange
+from typing import Callable, ClassVar
+from torch.nn.attention.flex_attention import (
+    _mask_mod_signature,
+    BlockMask,
+    create_block_mask,
+    flex_attention,
+    and_masks,
+    or_masks
+)
+from functools import partial
+
+try:
+    from flash_attn_interface import flash_attn_func
+except:
+    from flash_attn import flash_attn_func
+
+__all__ = ['WanTransformer3DModel']
+
+
+def custom_sdpa(q, k, v):
+    out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2),
+                                         v.transpose(1, 2))
+    return out.transpose(1, 2)
+
+class FlexAttnFunc(nn.Module):
+    flex_attn: ClassVar[Callable] = torch.compile(
+        flex_attention, dynamic=True, 
+    )
+    compiled_create_block_mask: ClassVar[Callable] = torch.compile(create_block_mask)
+    attention_mask: ClassVar[BlockMask] = None
+    cross_attention_mask: ClassVar[BlockMask] = None
+
+    def __init__(
+        self, 
+        is_cross=False,
+    ) -> None:
+        super().__init__()
+        self.is_cross = is_cross
+    
+    def forward(
+        self, 
+        query: torch.Tensor, 
+        key: torch.Tensor, 
+        value: torch.Tensor,
+        dtype=torch.bfloat16,
+    ) -> torch.Tensor:
+        q_varlen = rearrange(query[0], "s n d -> 1 n s d")
+        k_varlen = rearrange(key[0], "s n d -> 1 n s d")
+        v_varlen = rearrange(value[0], "s n d -> 1 n s d")
+
+        half_dtypes = (torch.float16, torch.bfloat16)
+        assert dtype in half_dtypes
+        def half(x):
+            return x if x.dtype in half_dtypes else x.to(dtype)
+        
+        q_varlen = half(q_varlen)
+        k_varlen = half(k_varlen)
+        v_varlen = half(v_varlen)
+        q_varlen = q_varlen.to(v_varlen.dtype)
+        k_varlen = k_varlen.to(v_varlen.dtype)
+
+        block_mask = FlexAttnFunc.cross_attention_mask if self.is_cross else FlexAttnFunc.attention_mask
+
+        x_out = FlexAttnFunc.flex_attn(q_varlen, k_varlen, v_varlen, block_mask=block_mask, kernel_options = {
+                                                    "BLOCK_M": 64,
+                                                    "BLOCK_N": 64,
+                                                    "BLOCK_M1": 32,
+                                                    "BLOCK_N1": 64,
+                                                    "BLOCK_M2": 64,
+                                                    "BLOCK_N2": 32,
+                                                })
+
+        x_out = rearrange(x_out, "b n s d -> b s n d")
+        return x_out
+
+    @staticmethod
+    @torch.no_grad()
+    def init_mask(
+        latent_shape, 
+        action_shape, 
+        padded_length, 
+        chunk_size,
+        window_size,
+        patch_size,
+        device,
+    ):
+        torch._inductor.config.realize_opcount_threshold = 100
+        B, _, L_F, L_H, L_W = latent_shape
+        _, _, A_F, A_H, A_W = action_shape
+
+        latent_seq_id = torch.arange(B)[:, None, None, None].\
+            expand(-1, L_F // patch_size[0], L_H // patch_size[1], L_W // patch_size[2]).flatten()
+        action_seq_id = torch.arange(B)[:, None, None, None].expand(-1, A_F, A_H, A_W).flatten()
+        seq_ids = torch.cat([latent_seq_id] * 2 + [action_seq_id] * 2)
+
+        latent_frame_id = torch.arange(L_F)[None, :, None, None].expand(B, -1, L_H // patch_size[1], L_W // patch_size[2])[None].flatten()
+        action_frame_id = torch.arange(A_F)[None, :, None, None].expand(B, -1, A_H, A_W)[None].flatten()
+        frame_ids = torch.cat([latent_frame_id // chunk_size * 2] * 2 + [action_frame_id // chunk_size * 2 + 1] * 2)
+
+        noise_ids = torch.cat(
+            [
+                torch.zeros_like(latent_frame_id),
+                torch.ones_like(latent_frame_id),
+                torch.zeros_like(action_frame_id),
+                torch.ones_like(action_frame_id),
+            ]
+        )
+
+        seq_ids = F.pad(seq_ids, (0, padded_length), value=-1)
+        frame_ids = F.pad(frame_ids, (0, padded_length), value=-1)
+        noise_ids = F.pad(noise_ids, (0, padded_length), value=-1)
+
+        mask_mod = FlexAttnFunc._get_mask_mod(seq_ids.long().to(device), frame_ids.long().to(device), noise_ids.long().to(device), window_size)
+        block_mask = FlexAttnFunc.compiled_create_block_mask(
+                mask_mod, 1, 1, len(seq_ids), len(seq_ids), device=device, _compile=True
+            )
+        FlexAttnFunc.attention_mask = block_mask
+
+        text_seq_ids = torch.arange(B)[:, None].expand(-1, 512).flatten()
+        mask_mod_cross = FlexAttnFunc._get_cross_mask_mod(seq_ids.long().to(device), text_seq_ids.long().to(device))
+        block_mask_cross = FlexAttnFunc.compiled_create_block_mask(
+                mask_mod_cross, 1, 1, len(seq_ids), len(text_seq_ids), device=device, _compile=True
+            )
+        FlexAttnFunc.cross_attention_mask = block_mask_cross
+    
+    @staticmethod
+    @torch.no_grad()
+    def _get_cross_mask_mod(seq_ids, text_seq_ids):
+        def seq_mask(
+            b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ):
+            return (seq_ids[q_idx] == text_seq_ids[kv_idx]) & (seq_ids[q_idx] >=0 ) & (text_seq_ids[kv_idx] >= 0)
+        return seq_mask
+    
+    @staticmethod
+    @torch.no_grad()
+    def _get_mask_mod(seq_ids, frame_ids, noise_ids, window_size):
+        def seq_mask(
+            b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ):
+            return (seq_ids[q_idx] == seq_ids[kv_idx]) & (seq_ids[q_idx] >=0 ) & (seq_ids[kv_idx] >= 0)
+        
+        def block_causal_mask(
+            b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ):
+            return (frame_ids[kv_idx] <= frame_ids[q_idx])
+        
+        def block_causal_mask_exclude_self(
+            b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ):
+            return (frame_ids[kv_idx] < frame_ids[q_idx])
+        
+        def block_self_mask(
+            b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ):
+            return (frame_ids[kv_idx] == frame_ids[q_idx])
+        
+        def clean2clean_mask(
+                b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ):
+            return (noise_ids[q_idx] == 1) & (noise_ids[kv_idx] == 1)
+        
+        def noise2clean_mask(
+            b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ):
+            return (noise_ids[q_idx] == 0) & (noise_ids[kv_idx] == 1)
+        def noise2noise_mask(
+            b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ):
+            return (noise_ids[q_idx] == 0) & (noise_ids[kv_idx] == 0)
+        
+        def block_window_mask(
+            b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor, window_size: int
+        ):
+            return ((frame_ids[q_idx] - frame_ids[kv_idx]).abs() <= window_size)
+
+        mask_list = []
+        mask_list.append(and_masks(clean2clean_mask, block_causal_mask))
+        mask_list.append(and_masks(noise2clean_mask, block_causal_mask_exclude_self))
+        mask_list.append(and_masks(noise2noise_mask, block_self_mask))
+        mask = or_masks(*mask_list)
+        mask = and_masks(mask, seq_mask)
+        mask = and_masks(mask, partial(block_window_mask, window_size=window_size))
+        return mask
+       
+class WanTimeTextImageEmbedding(nn.Module):
+
+    def __init__(
+        self,
+        dim,
+        time_freq_dim,
+        time_proj_dim,
+        text_embed_dim,
+        pos_embed_seq_len,
+    ):
+        super().__init__()
+
+        self.timesteps_proj = Timesteps(num_channels=time_freq_dim,
+                                        flip_sin_to_cos=True,
+                                        downscale_freq_shift=0)
+        self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim,
+                                               time_embed_dim=dim)
+        self.act_fn = nn.SiLU()
+        self.time_proj = nn.Linear(dim, time_proj_dim)
+        self.text_embedder = PixArtAlphaTextProjection(text_embed_dim,
+                                                       dim,
+                                                       act_fn="gelu_tanh")
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        dtype=None,
+    ):
+        B, L = timestep.shape
+        timestep = timestep.reshape(-1)
+        timestep = self.timesteps_proj(timestep)
+        # time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype
+        time_embedder_dtype = self.time_embedder.linear_1.weight.dtype
+        if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
+            timestep = timestep.to(time_embedder_dtype)
+        temb = self.time_embedder(timestep).to(dtype=dtype)
+        timestep_proj = self.time_proj(self.act_fn(temb))
+        return temb.reshape(B, L, -1), timestep_proj.reshape(B, L, -1)
+
+
+class WanRotaryPosEmbed(nn.Module):
+    def __init__(
+        self,
+        attention_head_dim: int,
+        patch_size,
+        max_seq_len: int,
+        theta: float = 10000.0,
+    ):
+        super().__init__()
+
+        self.attention_head_dim = attention_head_dim
+        self.patch_size = patch_size
+        self.max_seq_len = max_seq_len
+        self.theta = theta
+
+        self.f_dim = self.attention_head_dim - 2 * (self.attention_head_dim // 3)
+        self.h_dim = self.attention_head_dim // 3
+        self.w_dim = self.attention_head_dim // 3
+
+        # Precompute and register buffers
+        f_freqs_base, h_freqs_base, w_freqs_base = self._precompute_freqs_base()
+        self.f_freqs_base = f_freqs_base
+        self.h_freqs_base = h_freqs_base
+        self.w_freqs_base = w_freqs_base
+
+    def _precompute_freqs_base(self):
+        # freqs_base = 1.0 / (theta ** (2k / dim))
+        f_freqs_base = 1.0 / (self.theta**(torch.arange(
+            0, self.f_dim, 2)[:(self.f_dim // 2)].double() / self.f_dim))
+        h_freqs_base = 1.0 / (self.theta**(torch.arange(
+            0, self.h_dim, 2)[:(self.h_dim // 2)].double() / self.h_dim))
+        w_freqs_base = 1.0 / (self.theta**(torch.arange(
+            0, self.w_dim, 2)[:(self.w_dim // 2)].double() / self.w_dim))
+        return f_freqs_base, h_freqs_base, w_freqs_base
+
+    def forward(self, grid_ids):
+        with torch.no_grad():
+            f_freqs = grid_ids[:, 0, :].unsqueeze(-1) * self.f_freqs_base.to(grid_ids.device)
+            h_freqs = grid_ids[:, 1, :].unsqueeze(-1) * self.h_freqs_base.to(grid_ids.device)
+            w_freqs = grid_ids[:, 2, :].unsqueeze(-1) * self.w_freqs_base.to(grid_ids.device)
+            freqs = torch.cat([f_freqs, h_freqs, w_freqs], dim=-1).float()
+            freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+
+        return freqs_cis
+
+
+class WanAttention(torch.nn.Module):
+
+    def __init__(
+        self,
+        dim,
+        heads=8,
+        dim_head=64,
+        eps=1e-5,
+        dropout=0.0,
+        cross_attention_dim_head=None,
+        attn_mode='torch',
+    ):
+        super().__init__()
+        if attn_mode == 'torch':
+            self.attn_op = custom_sdpa
+        elif attn_mode == 'flashattn':
+            self.attn_op = flash_attn_func
+        elif attn_mode == 'flex':
+            self.attn_op = FlexAttnFunc(cross_attention_dim_head is not None)
+        else:
+            raise ValueError(
+                f"Unsupported attention mode: {attn_mode}, only support torch and flashattn"
+            )
+
+        self.inner_dim = dim_head * heads
+        self.heads = heads
+        self.cross_attention_dim_head = cross_attention_dim_head
+        self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads
+
+        self.to_q = torch.nn.Linear(dim, self.inner_dim, bias=True)
+        self.to_k = torch.nn.Linear(dim, self.kv_inner_dim, bias=True)
+        self.to_v = torch.nn.Linear(dim, self.kv_inner_dim, bias=True)
+        self.to_out = torch.nn.ModuleList([
+            torch.nn.Linear(self.inner_dim, dim, bias=True),
+            torch.nn.Dropout(dropout),
+        ])
+        self.norm_q = torch.nn.RMSNorm(dim_head * heads,
+                                       eps=eps,
+                                       elementwise_affine=True)
+        self.norm_k = torch.nn.RMSNorm(dim_head * heads,
+                                       eps=eps,
+                                       elementwise_affine=True)
+        self.attn_caches = {} if cross_attention_dim_head is None else None
+
+    def clear_pred_cache(self, cache_name):
+        if self.attn_caches is None:
+            return
+        cache = self.attn_caches[cache_name]
+        is_pred = cache['is_pred']
+        cache['mask'][is_pred] = False
+
+    def clear_cache(self, cache_name):
+        if self.attn_caches is None:
+            return
+        self.attn_caches[cache_name] = None
+
+    def init_kv_cache(self, cache_name, total_tolen, num_head, head_dim,
+                      device, dtype, batch_size):
+        if self.attn_caches is None:
+            return
+        self.attn_caches[cache_name] = {
+            'k':
+            torch.empty([batch_size, total_tolen, num_head, head_dim],
+                        device=device,
+                        dtype=dtype),
+            'v':
+            torch.empty([batch_size, total_tolen, num_head, head_dim],
+                        device=device,
+                        dtype=dtype),
+            'id':
+            torch.full((total_tolen, ), -1, device=device),
+            "mask":
+            torch.zeros((total_tolen, ), dtype=torch.bool, device=device),
+            "is_pred":
+            torch.zeros((total_tolen, ), dtype=torch.bool, device=device),
+        }
+
+    def allocate_slots(self, cache_name, key_size):
+        cache = self.attn_caches[cache_name]
+        mask = cache["mask"]
+        ids = cache["id"]
+        free = (~mask).nonzero(as_tuple=False).squeeze(-1)
+
+        if free.numel() < key_size:
+            used = mask.nonzero(as_tuple=False).squeeze(-1)
+
+            used_ids = ids[used]
+            order = torch.argsort(used_ids)
+            need = key_size - free.numel()
+            to_free = used[order[:need]]
+
+            mask[to_free] = False
+            ids[to_free] = -1
+            free = (~mask).nonzero(as_tuple=False).squeeze(-1)
+
+        assert free.numel() >= key_size
+        return free[:key_size]
+
+    def _next_cache_id(self, cache_name):
+        ids = self.attn_caches[cache_name]['id']
+        mask = self.attn_caches[cache_name]['mask']
+
+        if mask.any():
+            return ids[mask].max() + 1
+        else:
+            return torch.tensor(0, device=ids.device, dtype=ids.dtype)
+
+    def update_cache(self, cache_name, key, value, is_pred):
+        cache = self.attn_caches[cache_name]
+
+        key_size = key.shape[1]
+        slots = self.allocate_slots(cache_name, key_size)
+
+        new_id = self._next_cache_id(cache_name)
+
+        cache['k'][:, slots] = key
+        cache['v'][:, slots] = value
+        cache['mask'][slots] = True
+        cache['id'][slots] = new_id
+        cache['is_pred'][slots] = is_pred
+        return slots
+
+    def restore_cache(self, cache_name, slots):
+        self.attn_caches[cache_name]['mask'][slots] = False
+
+    def forward(
+        self,
+        q,
+        k,
+        v,
+        rotary_emb,
+        update_cache=0,
+        cache_name='pos',
+    ):
+        kv_cache = self.attn_caches[
+            cache_name] if (self.attn_caches is not None) and (cache_name in self.attn_caches) else None
+
+        query, key, value = self.to_q(q), self.to_k(k), self.to_v(v)
+        query = self.norm_q(query)
+        query = query.unflatten(2, (self.heads, -1))
+        key = self.norm_k(key)
+        key = key.unflatten(2, (self.heads, -1))
+        value = value.unflatten(2, (self.heads, -1))
+        if rotary_emb is not None:
+
+            def apply_rotary_emb(x, freqs):
+                x_out = torch.view_as_complex(
+                    x.to(torch.float64).reshape(x.shape[0], x.shape[1],
+                                                x.shape[2], -1, 2))
+                x_out = torch.view_as_real(x_out * freqs).flatten(3)
+                return x_out.to(x.dtype)
+            query = apply_rotary_emb(query, rotary_emb)
+            key = apply_rotary_emb(key, rotary_emb)
+        slots = None
+        if kv_cache is not None and kv_cache['k'] is not None:
+            slots = self.update_cache(cache_name,
+                                      key,
+                                      value,
+                                      is_pred=(update_cache == 1))
+            key_pool = self.attn_caches[cache_name]['k']
+            value_pool = self.attn_caches[cache_name]['v']
+            mask = self.attn_caches[cache_name]['mask']
+            valid = mask.nonzero(as_tuple=False).squeeze(-1)
+            key = key_pool[:, valid]
+            value = value_pool[:, valid]
+
+        hidden_states = self.attn_op(query, key, value)
+
+        if update_cache == 0:
+            if kv_cache is not None and kv_cache['k'] is not None:
+                self.restore_cache(cache_name, slots)
+
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.type_as(query)
+        hidden_states = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)
+        return hidden_states
+
+
+class WanTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        dim,
+        ffn_dim,
+        num_heads,
+        cross_attn_norm=False,
+        eps=1e-6,
+        attn_mode: str = "flashattn",
+    ):
+        super().__init__()
+        self.attn_mode = attn_mode
+
+        # 1. Self-attention
+        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.attn1 = WanAttention(
+            dim=dim,
+            heads=num_heads,
+            dim_head=dim // num_heads,
+            eps=eps,
+            cross_attention_dim_head=None,
+            attn_mode=attn_mode,
+        )
+
+        # 2. Cross-attention
+        self.attn2 = WanAttention(
+            dim=dim,
+            heads=num_heads,
+            dim_head=dim // num_heads,
+            eps=eps,
+            cross_attention_dim_head=dim // num_heads,
+            attn_mode=attn_mode,
+        )
+        self.norm2 = FP32LayerNorm(
+            dim, eps,
+            elementwise_affine=True) if cross_attn_norm else nn.Identity()
+
+        # 3. Feed-forward
+        self.ffn = FeedForward(dim,
+                               inner_dim=ffn_dim,
+                               activation_fn="gelu-approximate")
+        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states,
+        temb,
+        rotary_emb,
+        update_cache=0,
+        cache_name='pos',
+    ) -> torch.Tensor:
+        temb_scale_shift_table = self.scale_shift_table[None] + temb.float()
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = \
+            rearrange(temb_scale_shift_table, 'b l n c -> b n l c').chunk(6, dim=1)
+        shift_msa = shift_msa.squeeze(1)
+        scale_msa = scale_msa.squeeze(1)
+        gate_msa = gate_msa.squeeze(1)
+        c_shift_msa = c_shift_msa.squeeze(1)
+        c_scale_msa = c_scale_msa.squeeze(1)
+        c_gate_msa = c_gate_msa.squeeze(1)
+        # 1. Self-attention
+        norm_hidden_states = (self.norm1(hidden_states.float()) *
+                              (1. + scale_msa) +
+                              shift_msa).type_as(hidden_states)
+        attn_output = self.attn1(norm_hidden_states,
+                                 norm_hidden_states,
+                                 norm_hidden_states,
+                                 rotary_emb,
+                                 update_cache=update_cache,
+                                 cache_name=cache_name)
+        hidden_states = (hidden_states.float() +
+                         attn_output * gate_msa).type_as(hidden_states)
+
+        # 2. Cross-attention
+        norm_hidden_states = self.norm2(
+            hidden_states.float()).type_as(hidden_states)
+        attn_output = self.attn2(norm_hidden_states,
+                                 encoder_hidden_states,
+                                 encoder_hidden_states,
+                                 None,
+                                 update_cache=0,
+                                 cache_name=cache_name)
+        hidden_states = hidden_states + attn_output
+
+        # 3. Feed-forward
+        norm_hidden_states = (self.norm3(hidden_states.float()) *
+                              (1. + c_scale_msa) +
+                              c_shift_msa).type_as(hidden_states)
+
+        ff_output = self.ffn(norm_hidden_states)
+
+        hidden_states = (hidden_states.float() +
+                         ff_output.float() * c_gate_msa).type_as(hidden_states)
+        return hidden_states
+
+
+class WanTransformer3DModel(ModelMixin, ConfigMixin):
+    r"""
+    TODO
+    """
+    _supports_gradient_checkpointing = True
+    _skip_layerwise_casting_patterns = [
+                                        # "patch_embedding", 
+                                        "patch_embedding_mlp",
+                                        "condition_embedder", 
+                                        'condition_embedder_action',
+                                        "norm"]
+    _no_split_modules = ["WanTransformerBlock"]
+    _keep_in_fp32_modules = ["time_embedder", 
+                             "scale_shift_table", 
+                             "scale_shift_table_action",
+                             "norm1", 
+                             'action_norm1',
+                             'text_norm1',
+                             "norm2", 
+                             'action_norm2',
+                             'text_norm2',
+                             "norm3",
+                             'action_norm3',
+                             'text_norm3'
+                             ]
+    _keys_to_ignore_on_load_unexpected = ["norm_added_q"]
+    _repeated_blocks = ["WanTransformerBlock"]
+
+    @register_to_config
+    def __init__(self,
+                 patch_size=[1, 2, 2],
+                 num_attention_heads=24,
+                 attention_head_dim=128,
+                 in_channels=48,
+                 out_channels=48,
+                 action_dim=30,
+                 text_dim=4096,
+                 freq_dim=256,
+                 ffn_dim=14336,
+                 num_layers=30,
+                 cross_attn_norm=True,
+                 eps=1e-06,
+                 rope_max_seq_len=1024,
+                 pos_embed_seq_len=None,
+                 attn_mode="torch",
+                 enable_mcp=False,
+                 num_mcp_depths=3,
+                 mcp_blocks_per_depth=3,
+                 mcp_hidden_collect_layers=(3, 11, 19, 29)):
+        r"""
+        TODO
+        """
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+        inner_dim = num_attention_heads * attention_head_dim
+        self.inner_dim = inner_dim
+        self.ffn_dim = ffn_dim
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        self.attn_mode = attn_mode
+        self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size,
+                                      rope_max_seq_len)
+        self.patch_embedding_mlp = nn.Linear(
+            in_channels * patch_size[0] * patch_size[1] * patch_size[2],
+            inner_dim)
+        self.action_embedder = nn.Linear(action_dim, inner_dim)
+        self.condition_embedder = WanTimeTextImageEmbedding(
+            dim=inner_dim,
+            time_freq_dim=freq_dim,
+            time_proj_dim=inner_dim * 6,
+            text_embed_dim=text_dim,
+            pos_embed_seq_len=pos_embed_seq_len,
+        )
+        self.condition_embedder_action = deepcopy(self.condition_embedder)
+
+        self.blocks = nn.ModuleList([
+            WanTransformerBlock(inner_dim,
+                                ffn_dim,
+                                num_attention_heads,
+                                cross_attn_norm,
+                                eps,
+                                attn_mode=attn_mode) for _ in range(num_layers)
+        ])
+
+        self.enable_mcp = enable_mcp
+        self.num_mcp_depths = num_mcp_depths
+        self.mcp_blocks_per_depth = mcp_blocks_per_depth
+        self.mcp_hidden_collect_layers = list(mcp_hidden_collect_layers)
+        if self.enable_mcp:
+            self._build_mcp_modules()
+
+        self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False)
+        self.proj_out = nn.Linear(inner_dim,
+                                  out_channels * math.prod(patch_size))
+        self.action_proj_out = nn.Linear(inner_dim, action_dim)
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+
+    def _validate_mcp_architecture(self):
+        if self.num_mcp_depths <= 0:
+            raise ValueError("num_mcp_depths must be positive")
+        if self.mcp_blocks_per_depth <= 0:
+            raise ValueError("mcp_blocks_per_depth must be positive")
+        if not self.mcp_hidden_collect_layers:
+            raise ValueError("mcp_hidden_collect_layers cannot be empty")
+        if len(set(self.mcp_hidden_collect_layers)) != len(
+                self.mcp_hidden_collect_layers):
+            raise ValueError("mcp_hidden_collect_layers must be unique")
+        if min(self.mcp_hidden_collect_layers) < 0 or max(
+                self.mcp_hidden_collect_layers) >= len(self.blocks):
+            raise ValueError(
+                "mcp_hidden_collect_layers must reference existing model layers")
+        if self.mcp_blocks_per_depth > len(self.blocks):
+            raise ValueError(
+                "mcp_blocks_per_depth cannot exceed the number of model layers")
+
+    def _build_mcp_modules(self):
+        self._validate_mcp_architecture()
+        self.mcp_hidden_fuser = nn.Sequential(
+            nn.Linear(
+                self.inner_dim * len(self.mcp_hidden_collect_layers),
+                self.inner_dim,
+            ),
+            nn.SiLU(),
+            nn.Linear(self.inner_dim, self.inner_dim),
+        )
+        self.mcp_input_projections = nn.ModuleList([
+            nn.Linear(self.inner_dim * 2, self.inner_dim)
+            for _ in range(self.num_mcp_depths)
+        ])
+        self.mcp_blocks = nn.ModuleList([
+            nn.ModuleList([
+                WanTransformerBlock(
+                    self.inner_dim,
+                    self.ffn_dim,
+                    self.num_attention_heads,
+                    self.cross_attn_norm,
+                    self.eps,
+                    attn_mode=self.attn_mode,
+                ) for _ in range(self.mcp_blocks_per_depth)
+            ]) for _ in range(self.num_mcp_depths)
+        ])
+
+    def enable_mcp_training(self,
+                            num_mcp_depths,
+                            mcp_blocks_per_depth,
+                            mcp_hidden_collect_layers,
+                            init_from_backbone=True):
+        requested_layers = list(mcp_hidden_collect_layers)
+        if self.enable_mcp:
+            current = (
+                self.num_mcp_depths,
+                self.mcp_blocks_per_depth,
+                self.mcp_hidden_collect_layers,
+            )
+            requested = (
+                num_mcp_depths,
+                mcp_blocks_per_depth,
+                requested_layers,
+            )
+            if current != requested:
+                raise ValueError(
+                    f"MCP checkpoint architecture {current} does not match {requested}")
+            return False
+
+        self.enable_mcp = True
+        self.num_mcp_depths = num_mcp_depths
+        self.mcp_blocks_per_depth = mcp_blocks_per_depth
+        self.mcp_hidden_collect_layers = requested_layers
+        self._build_mcp_modules()
+
+        reference = next(self.blocks[0].parameters())
+        self.mcp_hidden_fuser.to(device=reference.device,
+                                 dtype=reference.dtype)
+        self.mcp_input_projections.to(device=reference.device,
+                                      dtype=reference.dtype)
+        self.mcp_blocks.to(device=reference.device, dtype=reference.dtype)
+        self._initialize_mcp_projection_weights()
+        if init_from_backbone:
+            self.initialize_mcp_blocks_from_backbone()
+
+        used_default_values = set(
+            self.config.get("_use_default_values", []))
+        used_default_values.difference_update({
+            "enable_mcp",
+            "num_mcp_depths",
+            "mcp_blocks_per_depth",
+            "mcp_hidden_collect_layers",
+        })
+        self.register_to_config(
+            enable_mcp=True,
+            num_mcp_depths=num_mcp_depths,
+            mcp_blocks_per_depth=mcp_blocks_per_depth,
+            mcp_hidden_collect_layers=requested_layers,
+            _use_default_values=sorted(used_default_values),
+        )
+        return True
+
+    def _initialize_mcp_projection_weights(self):
+        for module in [self.mcp_hidden_fuser, self.mcp_input_projections]:
+            for layer in module.modules():
+                if isinstance(layer, nn.Linear):
+                    nn.init.normal_(layer.weight, std=0.02)
+                    if layer.bias is not None:
+                        nn.init.zeros_(layer.bias)
+
+    def initialize_mcp_blocks_from_backbone(self):
+        if not self.enable_mcp:
+            return
+        source_blocks = self.blocks[-self.mcp_blocks_per_depth:]
+        for group in self.mcp_blocks:
+            for target_block, source_block in zip(group, source_blocks):
+                target_block.load_state_dict(source_block.state_dict())
+
+    def disable_mcp_modules(self):
+        if not self.enable_mcp:
+            return
+        del self.mcp_hidden_fuser
+        del self.mcp_input_projections
+        del self.mcp_blocks
+        self.enable_mcp = False
+        self.register_to_config(enable_mcp=False)
+
+    def clear_cache(self, cache_name):
+        for block in self.blocks:
+            block.attn1.clear_cache(cache_name)
+
+    def clear_pred_cache(self, cache_name):
+        for block in self.blocks:
+            block.attn1.clear_pred_cache(cache_name)
+
+    def create_empty_cache(self, cache_name, attn_window,
+                           latent_token_per_chunk, action_token_per_chunk,
+                           device, dtype, batch_size):
+        total_tolen = (attn_window // 2) * latent_token_per_chunk + (
+            attn_window // 2) * action_token_per_chunk
+        for block in self.blocks:
+            block.attn1.init_kv_cache(cache_name, total_tolen,
+                                      self.num_attention_heads,
+                                      self.attention_head_dim, device, dtype, batch_size)
+    
+    def _input_embed(self, latents, input_type='latent'):
+        if input_type == 'latent':
+            hidden_states = rearrange(
+                latents,
+                'b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)',
+                p1=self.patch_size[0],
+                p2=self.patch_size[1],
+                p3=self.patch_size[2])
+            hidden_states = self.patch_embedding_mlp(hidden_states)
+        elif input_type == 'action':
+            hidden_states = rearrange(latents, 'b c f h w -> b (f h w) c')
+            hidden_states = self.action_embedder(hidden_states)
+        elif input_type == 'text':
+            hidden_states = self.condition_embedder.text_embedder(latents)
+        else:
+            raise ValueError(f"Unsupported input type: {input_type}")
+        return hidden_states
+
+    def _time_embed(self, timesteps, H, W, dtype, action_mode=False):
+        pach_scale_h, pach_scale_w = (1, 1) if action_mode else (
+            self.patch_size[1], self.patch_size[2])
+        latent_time_steps = torch.repeat_interleave(
+            timesteps,
+            (H // pach_scale_h) *
+            (W // pach_scale_w), dim=1)  # L
+        current_condition_embedder = self.condition_embedder_action if action_mode else self.condition_embedder
+        temb, timestep_proj = current_condition_embedder(
+            latent_time_steps, dtype=dtype)
+        timestep_proj = timestep_proj.unflatten(2, (6, -1))  # B L 6 C
+        return temb, timestep_proj
+
+    def _forward_mcp(
+        self,
+        input_dict,
+        hidden_states,
+        collected_hidden_states,
+        text_hidden_states,
+        latent_grid_id,
+        action_grid_id,
+        latent_timestep_proj,
+        action_timestep_proj,
+        split_list,
+        batch_size,
+    ):
+        mcp_latent_dicts = input_dict.get('mcp_latent_dicts')
+        if not self.enable_mcp or not mcp_latent_dicts:
+            return []
+        if len(mcp_latent_dicts) != self.num_mcp_depths:
+            raise ValueError(
+                "MCP input depth count must match num_mcp_depths")
+        if len(collected_hidden_states) != len(
+                self.mcp_hidden_collect_layers):
+            raise RuntimeError(
+                "MCP did not collect all configured backbone hidden states")
+
+        latent_length, clean_latent_length, action_length, clean_action_length, padded_length = split_list
+        if latent_length != clean_latent_length:
+            raise ValueError("MCP requires noisy and clean video token lengths to match")
+
+        video_length = latent_length + clean_latent_length
+        video_hidden_states = [
+            states[:, :video_length] for states in collected_hidden_states
+        ]
+        fused_hidden_states = self.mcp_hidden_fuser(
+            torch.cat(video_hidden_states, dim=-1))
+        previous_hidden_states = fused_hidden_states[:, :latent_length]
+        clean_hidden_states = fused_hidden_states[:, latent_length:]
+
+        action_start = video_length
+        action_end = action_start + action_length + clean_action_length
+        action_hidden_states = hidden_states[:, action_start:action_end]
+        clean_latent_timestep_proj = latent_timestep_proj[
+            :, latent_length:video_length]
+
+        outputs = []
+        for depth in range(self.num_mcp_depths):
+            mcp_latent_dict = mcp_latent_dicts[depth]
+            noisy_latents = mcp_latent_dict['noisy_latents'].to(torch.bfloat16)
+            noisy_hidden_states = self._input_embed(
+                noisy_latents, input_type='latent').flatten(0, 1)[None]
+            if noisy_hidden_states.shape[1] != latent_length:
+                raise ValueError(
+                    "MCP future video token length must match the main video token length")
+
+            future_hidden_states = self.mcp_input_projections[depth](torch.cat(
+                [previous_hidden_states, noisy_hidden_states], dim=-1))
+            mcp_hidden_states = torch.cat([
+                future_hidden_states,
+                clean_hidden_states,
+                action_hidden_states,
+            ], dim=1)
+            mcp_hidden_states = F.pad(
+                mcp_hidden_states, (0, 0, 0, padded_length))
+
+            future_grid_id = mcp_latent_dict['grid_id'].permute(
+                1, 0, 2).flatten(1)[None]
+            full_grid_id = torch.cat([
+                future_grid_id,
+                latent_grid_id,
+                action_grid_id,
+                action_grid_id,
+            ], dim=2)
+            mcp_rotary_emb = self.rope(full_grid_id)[:, :, None]
+            mcp_rotary_emb = F.pad(
+                mcp_rotary_emb, (0, 0, 0, 0, 0, padded_length))
+
+            future_time_steps = mcp_latent_dict['timesteps'].flatten()[None]
+            future_temb, future_timestep_proj = self._time_embed(
+                future_time_steps,
+                noisy_latents.shape[-2],
+                noisy_latents.shape[-1],
+                dtype=mcp_hidden_states.dtype,
+                action_mode=False,
+            )
+            mcp_timestep_proj = torch.cat([
+                future_timestep_proj,
+                clean_latent_timestep_proj,
+                action_timestep_proj,
+            ], dim=1)
+            mcp_timestep_proj = F.pad(
+                mcp_timestep_proj, (0, 0, 0, 0, 0, padded_length))
+
+            for block in self.mcp_blocks[depth]:
+                mcp_hidden_states = block(
+                    mcp_hidden_states,
+                    text_hidden_states,
+                    mcp_timestep_proj,
+                    mcp_rotary_emb,
+                    update_cache=False,
+                )
+
+            previous_hidden_states = mcp_hidden_states[:, :latent_length]
+            mcp_output = previous_hidden_states
+            temb_scale_shift_table = (
+                self.scale_shift_table[None] + future_temb[:, :, None, ...])
+            shift, scale = rearrange(
+                temb_scale_shift_table,
+                'b l n c -> b n l c',
+            ).chunk(2, dim=1)
+            shift = shift.to(mcp_output.device).squeeze(1)
+            scale = scale.to(mcp_output.device).squeeze(1)
+            mcp_output = (
+                self.norm_out(mcp_output.float()) * (1. + scale) + shift
+            ).type_as(mcp_output)
+            mcp_output = self.proj_out(mcp_output)
+            mcp_output = rearrange(
+                mcp_output,
+                '1 (b l) (n c) -> b (l n) c',
+                n=math.prod(self.patch_size),
+                b=batch_size,
+            )
+            outputs.append(mcp_output)
+
+        return outputs
+
+    def forward_train(self, input_dict):
+        input_dict['latent_dict']['noisy_latents'] = input_dict['latent_dict']['noisy_latents'].to(torch.bfloat16)
+        input_dict['latent_dict']['latent'] = input_dict['latent_dict']['latent'].to(torch.bfloat16)
+        input_dict['action_dict']['noisy_latents'] = input_dict['action_dict']['noisy_latents'].to(torch.bfloat16)
+        input_dict['action_dict']['latent'] = input_dict['action_dict']['latent'].to(torch.bfloat16)
+        if self.enable_mcp:
+            for mcp_latent_dict in input_dict.get('mcp_latent_dicts', []):
+                mcp_latent_dict['noisy_latents'] = mcp_latent_dict[
+                    'noisy_latents'].to(torch.bfloat16)
+
+        latent_dict = input_dict['latent_dict']
+        action_dict = input_dict['action_dict']
+        batch_size = latent_dict['noisy_latents'].shape[0]
+
+        latent_hidden_states = self._input_embed(latent_dict['noisy_latents'], input_type='latent').flatten(0, 1)[None]
+        action_hidden_states = self._input_embed(action_dict['noisy_latents'], input_type='action').flatten(0, 1)[None]
+        text_hidden_states = self._input_embed(latent_dict["text_emb"], input_type='text')
+
+        text_hidden_states = text_hidden_states.flatten(0, 1)[None]
+
+        condition_latent_hidden_states = self._input_embed(latent_dict['latent'], input_type='latent').flatten(0, 1)[None]
+        condition_action_hidden_states = self._input_embed(action_dict['latent'], input_type='action').flatten(0, 1)[None]
+
+        hidden_states = torch.cat([latent_hidden_states, 
+                                   condition_latent_hidden_states,
+                                   action_hidden_states, 
+                                   condition_action_hidden_states], dim=1)
+
+
+        latent_grid_id = latent_dict['grid_id'].permute(1, 0, 2).flatten(1)[None]
+        action_grid_id = action_dict['grid_id'].permute(1, 0, 2).flatten(1)[None]
+        full_grid_id = torch.cat([latent_grid_id] * 2 + [action_grid_id] * 2, dim=2)
+
+        rotary_emb = self.rope(full_grid_id)[:, :, None] 
+
+        latent_time_steps = torch.cat(
+            [latent_dict['timesteps'].flatten(0, 1), latent_dict['cond_timesteps'].flatten(0, 1)]
+        )[None]
+        action_time_steps = torch.cat(
+            [action_dict['timesteps'].flatten(0, 1), action_dict['cond_timesteps'].flatten(0, 1)]
+        )[None]
+        latent_temb, latent_timestep_proj =self._time_embed(latent_time_steps, 
+                        latent_dict['noisy_latents'].shape[-2], 
+                        latent_dict['noisy_latents'].shape[-1], 
+                        dtype=hidden_states.dtype, 
+                        action_mode=False)
+        action_temb, action_timestep_proj = self._time_embed(action_time_steps,
+                        action_dict['noisy_latents'].shape[-2], 
+                        action_dict['noisy_latents'].shape[-1], 
+                        dtype=hidden_states.dtype, 
+                        action_mode=True)
+        temb = torch.cat([latent_temb, action_temb], dim=1)
+        timestep_proj = torch.cat([latent_timestep_proj, action_timestep_proj], dim=1)
+
+        total_length = hidden_states.shape[1]
+        padded_length = (128 - total_length % 128) % 128
+        hidden_states = F.pad(hidden_states, (0, 0, 0, padded_length))
+        rotary_emb = F.pad(rotary_emb, (0, 0, 0, 0, 0, padded_length))
+        temb = F.pad(temb, (0, 0, 0, padded_length))
+        timestep_proj = F.pad(timestep_proj, (0, 0, 0, 0, 0, padded_length))
+
+        split_list = [latent_hidden_states.shape[1], 
+                      condition_latent_hidden_states.shape[1], 
+                      action_hidden_states.shape[1], 
+                      condition_action_hidden_states.shape[1],
+                      padded_length]
+
+        FlexAttnFunc.init_mask(latent_dict['noisy_latents'].shape, 
+                               action_dict['noisy_latents'].shape, 
+                               padded_length, 
+                               input_dict["chunk_size"],
+                               window_size=input_dict['window_size'],
+                               patch_size=self.patch_size,
+                               device=hidden_states.device
+                               )
+
+        collected_hidden_states = {}
+        for layer_id, block in enumerate(self.blocks):
+            hidden_states = block(hidden_states,
+                                         text_hidden_states,
+                                         timestep_proj,
+                                         rotary_emb,
+                                         update_cache=False)
+            if self.enable_mcp and layer_id in self.mcp_hidden_collect_layers:
+                collected_hidden_states[layer_id] = hidden_states
+
+        collected_hidden_states = [
+            collected_hidden_states[layer_id]
+            for layer_id in self.mcp_hidden_collect_layers
+            if layer_id in collected_hidden_states
+        ]
+
+        mcp_outputs = self._forward_mcp(
+            input_dict=input_dict,
+            hidden_states=hidden_states,
+            collected_hidden_states=collected_hidden_states,
+            text_hidden_states=text_hidden_states,
+            latent_grid_id=latent_grid_id,
+            action_grid_id=action_grid_id,
+            latent_timestep_proj=latent_timestep_proj,
+            action_timestep_proj=action_timestep_proj,
+            split_list=split_list,
+            batch_size=batch_size,
+        )
+        temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...]
+        shift, scale = rearrange(temb_scale_shift_table,
+                                 'b l n c -> b n l c').chunk(2, dim=1)
+        shift = shift.to(hidden_states.device).squeeze(1)
+        scale = scale.to(hidden_states.device).squeeze(1)
+        hidden_states = (self.norm_out(hidden_states.float()) *
+                                (1. + scale) +
+                                shift).type_as(hidden_states)
+        latent_hidden_states, _, action_hidden_states, _, _ = torch.split(hidden_states, split_list, dim=1)
+        latent_hidden_states = self.proj_out(latent_hidden_states)
+        latent_hidden_states = rearrange(latent_hidden_states,
+                                             '1 (b l) (n c) -> b (l n) c',
+                                             n=math.prod(self.patch_size), b=batch_size)  #
+        action_hidden_states = self.action_proj_out(action_hidden_states)
+        action_hidden_states = rearrange(action_hidden_states,
+                                             '1 (b l) c -> b l c',
+                                             b=batch_size)  #
+
+        if self.enable_mcp:
+            return latent_hidden_states, action_hidden_states, mcp_outputs
+        return latent_hidden_states, action_hidden_states
+
+    def forward(
+        self,
+        input_dict,
+        update_cache=0,
+        cache_name="pos",
+        action_mode=False,
+        train_mode=False,
+    ):
+        r"""
+        Forward pass through the diffusion model
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+        if train_mode:
+            return self.forward_train(input_dict)
+        if action_mode:  # action input emb
+            latent_hidden_states = rearrange(input_dict['noisy_latents'],
+                                             'b c f h w -> b (f h w) c')
+            latent_hidden_states = self.action_embedder(
+                latent_hidden_states)  # B L1 C
+        else:  # latent input emb
+            latent_hidden_states = rearrange(
+                input_dict['noisy_latents'],
+                'b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)',
+                p1=self.patch_size[0],
+                p2=self.patch_size[1],
+                p3=self.patch_size[2])
+            latent_hidden_states = self.patch_embedding_mlp(
+                latent_hidden_states)
+        text_hidden_states = self.condition_embedder.text_embedder(
+            input_dict["text_emb"])  # B L2 C
+
+        latent_grid_id = input_dict['grid_id']
+        rotary_emb = self.rope(latent_grid_id)[:, :, None]  # 1 L 1 C
+        pach_scale_h, pach_scale_w = (1, 1) if action_mode else (
+            self.patch_size[1], self.patch_size[2])
+
+        latent_time_steps = torch.repeat_interleave(
+            input_dict['timesteps'],
+            (input_dict['noisy_latents'].shape[-2] // pach_scale_h) *
+            (input_dict['noisy_latents'].shape[-1] // pach_scale_w), dim=1)  # L
+        current_condition_embedder = self.condition_embedder_action if action_mode else self.condition_embedder
+        temb, timestep_proj = current_condition_embedder(
+            latent_time_steps, dtype=latent_hidden_states.dtype)
+        timestep_proj = timestep_proj.unflatten(2, (6, -1))  # B L 6 C
+
+        for block in self.blocks:
+            latent_hidden_states = block(latent_hidden_states,
+                                         text_hidden_states,
+                                         timestep_proj,
+                                         rotary_emb,
+                                         update_cache=update_cache,
+                                         cache_name=cache_name)
+        temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...]
+        shift, scale = rearrange(temb_scale_shift_table,
+                                 'b l n c -> b n l c').chunk(2, dim=1)
+        shift = shift.to(latent_hidden_states.device).squeeze(1)
+        scale = scale.to(latent_hidden_states.device).squeeze(1)
+        latent_hidden_states = (self.norm_out(latent_hidden_states.float()) *
+                                (1. + scale) +
+                                shift).type_as(latent_hidden_states)
+
+        if action_mode:
+            latent_hidden_states = self.action_proj_out(latent_hidden_states)
+        else:
+            latent_hidden_states = self.proj_out(latent_hidden_states)
+            latent_hidden_states = rearrange(latent_hidden_states,
+                                             'b l (n c) -> b (l n) c',
+                                             n=math.prod(self.patch_size))  #
+
+        return latent_hidden_states
+
+
+if __name__ == '__main__':
+    model = WanTransformer3DModel(patch_size=[1, 2, 2],
+                                  num_attention_heads=24,
+                                  attention_head_dim=128,
+                                  in_channels=48,
+                                  out_channels=48,
+                                  action_dim=30,
+                                  text_dim=4096,
+                                  freq_dim=256,
+                                  ffn_dim=14336,
+                                  num_layers=30,
+                                  cross_attn_norm=True,
+                                  eps=1e-6,
+                                  rope_max_seq_len=1024,
+                                  pos_embed_seq_len=None,
+                                  attn_mode="torch")
+    print(model)

wan_va/modules/utils.py

@@ -0,0 +1,100 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import torch
+from diffusers import AutoencoderKLWan
+from transformers import (
+    T5TokenizerFast,
+    UMT5EncoderModel,
+)
+
+from .model import WanTransformer3DModel
+
+
+def load_vae(
+    vae_path,
+    torch_dtype,
+    torch_device,
+):
+    vae = AutoencoderKLWan.from_pretrained(
+        vae_path,
+        torch_dtype=torch_dtype,
+    )
+    return vae.to(torch_device)
+
+
+def load_text_encoder(
+    text_encoder_path,
+    torch_dtype,
+    torch_device,
+):
+    text_encoder = UMT5EncoderModel.from_pretrained(
+        text_encoder_path,
+        torch_dtype=torch_dtype,
+    )
+    return text_encoder.to(torch_device)
+
+
+def load_tokenizer(tokenizer_path, ):
+    tokenizer = T5TokenizerFast.from_pretrained(tokenizer_path, )
+    return tokenizer
+
+
+def load_transformer(
+    transformer_path,
+    torch_dtype,
+    torch_device,
+    disable_mcp=False,
+    **kwargs
+):
+    model = WanTransformer3DModel.from_pretrained(
+        transformer_path,
+        torch_dtype=torch_dtype,
+        **kwargs
+    )
+    if disable_mcp and getattr(model, 'enable_mcp', False):
+        model.disable_mcp_modules()
+    return model.to(torch_device)
+
+
+def patchify(x, patch_size):
+    if patch_size is None or patch_size == 1:
+        return x
+    batch_size, channels, frames, height, width = x.shape
+    x = x.view(batch_size, channels, frames, height // patch_size, patch_size,
+               width // patch_size, patch_size)
+    x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous()
+    x = x.view(batch_size, channels * patch_size * patch_size, frames,
+               height // patch_size, width // patch_size)
+    return x
+
+
+class WanVAEStreamingWrapper:
+
+    def __init__(self, vae_model):
+        self.vae = vae_model
+        self.encoder = vae_model.encoder
+        self.quant_conv = vae_model.quant_conv
+
+        if hasattr(self.vae, "_cached_conv_counts"):
+            self.enc_conv_num = self.vae._cached_conv_counts["encoder"]
+        else:
+            count = 0
+            for m in self.encoder.modules():
+                if m.__class__.__name__ == "WanCausalConv3d":
+                    count += 1
+            self.enc_conv_num = count
+
+        self.clear_cache()
+
+    def clear_cache(self):
+        self.feat_cache = [None] * self.enc_conv_num
+
+    def encode_chunk(self, x_chunk):
+        if hasattr(self.vae.config,
+                   "patch_size") and self.vae.config.patch_size is not None:
+            x_chunk = patchify(x_chunk, self.vae.config.patch_size)
+        feat_idx = [0]
+        out = self.encoder(x_chunk,
+                           feat_cache=self.feat_cache,
+                           feat_idx=feat_idx)
+        enc = self.quant_conv(out)
+        return enc

wan_va/train.py

@@ -0,0 +1,822 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import argparse
+import os
+from pathlib import Path
+import wandb
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, DistributedSampler
+from tqdm import tqdm
+from torch.distributed.checkpoint.state_dict import (
+    get_model_state_dict,
+    get_optimizer_state_dict,
+    set_optimizer_state_dict,
+    StateDictOptions,
+)
+from safetensors.torch import save_file, load_file
+import json
+
+from .configs import VA_CONFIGS
+from .distributed.fsdp import shard_model, apply_ac
+from .distributed.util import (
+    _configure_model, 
+    init_distributed, 
+    dist_mean, 
+    dist_max
+)
+from einops import rearrange
+from .modules.utils import (
+    load_transformer,
+)
+from .utils import (
+    init_logger, 
+    logger, 
+    get_mesh_id, 
+    sample_timestep_id,
+    data_seq_to_patch,
+    warmup_constant_lambda,
+    FlowMatchScheduler
+)
+
+from .dataset import MultiLatentLeRobotDataset, dataset_indexes_ready
+from .mcp import shift_latents_for_mcp, validate_mcp_settings
+import gc
+
+
+class Trainer:
+    def __init__(self, config):
+        if config.enable_wandb and config.rank == 0:
+            wandb.login(host=os.environ['WANDB_BASE_URL'], key=os.environ['WANDB_API_KEY'])
+            self.wandb = wandb
+            self.wandb.init(
+                entity=os.environ["WANDB_TEAM_NAME"],
+                project=os.getenv("WANDB_PROJECT", "va_robotwin"),
+                # dir=log_dir,
+                config=config,
+                mode="online",
+                name='test_lln'
+                # name=os.path.basename(os.path.normpath(job_config.job.dump_folder))
+            )
+            logger.info("WandB logging enabled")
+        self.step = 0
+        self.config = config
+        self.device = torch.device(f"cuda:{config.local_rank}")
+        self.dtype = config.param_dtype
+        self.patch_size = config.patch_size
+        self.enable_mcp = getattr(config, 'enable_mcp', True)
+
+        if self.enable_mcp:
+            validate_mcp_settings(
+                num_mcp_depths=config.num_mcp_depths,
+                mcp_blocks_per_depth=config.mcp_blocks_per_depth,
+                mcp_hidden_collect_layers=config.mcp_hidden_collect_layers,
+                mcp_loss_weights=config.mcp_loss_weights,
+            )
+
+        # Load models
+        logger.info("Loading models...")
+
+        # Load and shard transformer with FSDP
+        logger.info("Loading transformer...")
+
+        if hasattr(config, 'resume_from') and config.resume_from:
+            transformer_path = os.path.join(config.resume_from, 'transformer')
+            if config.rank == 0:
+                logger.info(f"Resuming from checkpoint: {transformer_path}")
+        else:
+            transformer_path = os.path.join(config.wan22_pretrained_model_name_or_path, 'transformer')
+
+        self.transformer = load_transformer(
+            transformer_path,
+            torch_dtype=torch.float32,
+            torch_device='cpu',
+            attn_mode="flex",
+            disable_mcp=not self.enable_mcp,
+        )
+
+        if self.enable_mcp:
+            validate_mcp_settings(
+                num_mcp_depths=config.num_mcp_depths,
+                mcp_blocks_per_depth=config.mcp_blocks_per_depth,
+                mcp_hidden_collect_layers=config.mcp_hidden_collect_layers,
+                mcp_loss_weights=config.mcp_loss_weights,
+                num_layers=len(self.transformer.blocks),
+            )
+            initialized = self.transformer.enable_mcp_training(
+                num_mcp_depths=config.num_mcp_depths,
+                mcp_blocks_per_depth=config.mcp_blocks_per_depth,
+                mcp_hidden_collect_layers=config.mcp_hidden_collect_layers,
+                init_from_backbone=config.mcp_init_from_backbone,
+            )
+            if initialized:
+                total_mcp_blocks = (
+                    config.num_mcp_depths * config.mcp_blocks_per_depth)
+                if config.mcp_init_from_backbone:
+                    num_backbone_blocks = len(self.transformer.blocks)
+                    source_start = (
+                        num_backbone_blocks - config.mcp_blocks_per_depth)
+                    logger.info(
+                        f"Initializing {config.num_mcp_depths} MCP depths "
+                        f"({config.mcp_blocks_per_depth} blocks per depth) "
+                        f"from backbone blocks[{source_start}:"
+                        f"{num_backbone_blocks}]"
+                    )
+                    for depth in range(config.num_mcp_depths):
+                        logger.info(
+                            f"  MCP depth {depth + 1}: initialized from "
+                            f"backbone blocks[{source_start}:"
+                            f"{num_backbone_blocks}]"
+                        )
+                    logger.info(
+                        f"Initialized {config.num_mcp_depths} x "
+                        f"{config.mcp_blocks_per_depth} = "
+                        f"{total_mcp_blocks} MCP blocks"
+                    )
+                else:
+                    logger.info(
+                        f"Initialized {total_mcp_blocks} MCP blocks from "
+                        "scratch"
+                    )
+                logger.info(
+                    "MCP hidden-state collection uses backbone indices "
+                    f"{list(config.mcp_hidden_collect_layers)}"
+                )
+                logger.info(
+                    "Initialized MCP hidden fuser and input projections "
+                    "from scratch"
+                )
+            else:
+                logger.info(
+                    f"Loaded MCP modules from checkpoint: "
+                    f"{config.num_mcp_depths} depths x "
+                    f"{config.mcp_blocks_per_depth} blocks per depth"
+                )
+
+        logger.info("Setting up activation checkpointing ...")
+        apply_ac(self.transformer)
+
+        logger.info("Setting up FSDP...")
+        shard_fn = shard_model
+        self.transformer = _configure_model(
+            model=self.transformer,
+            shard_fn=shard_fn,
+            param_dtype=self.dtype,
+            device=self.device,
+            eval_mode=False,
+        )
+        self.transformer.train()
+        self.transformer.requires_grad_(True)
+
+        # Optimizer
+        self.optimizer = torch.optim.AdamW(
+            [p for p in self.transformer.parameters() if p.requires_grad],
+            lr=config.learning_rate,
+            betas=(config.beta1, config.beta2),
+            eps=1e-8,
+            weight_decay=config.weight_decay,
+            fused=True,
+            foreach=False,
+        )
+
+        self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, 
+            lr_lambda=lambda step: warmup_constant_lambda(step, warmup_steps=config.warmup_steps))
+
+        # Setup dataloaders
+        logger.info("Setting up datasets...")
+        cache_ready = False
+        if (
+            config.world_size > 1
+            and getattr(config, 'enable_dataset_index_cache', True)
+        ):
+            cache_ready_flag = torch.zeros(1, device=self.device, dtype=torch.int)
+            if config.rank == 0:
+                cache_ready_flag.fill_(int(dataset_indexes_ready(config)))
+            dist.broadcast(cache_ready_flag, src=0)
+            cache_ready = bool(cache_ready_flag.item())
+
+        use_rank_zero_indexing = (
+            config.world_size > 1
+            and getattr(config, 'enable_dataset_index_cache', True)
+            and not cache_ready
+        )
+        if config.rank == 0 and cache_ready:
+            logger.info(
+                "Dataset index and Arrow caches are complete; "
+                "loading all ranks concurrently"
+            )
+        if use_rank_zero_indexing and config.rank != 0:
+            dist.barrier()
+            # Rank 0 has completed any requested rebuild at this point.
+            config.rebuild_dataset_index_cache = False
+
+        train_dataset = MultiLatentLeRobotDataset(config=config)
+
+        if use_rank_zero_indexing:
+            if config.rank == 0:
+                dist.barrier()
+            dist.barrier()
+
+        if config.rank == 0:
+            logger.info(
+                "Dataset ready: %d samples from %d datasets "
+                "(%d index cache hits, %d direct Arrow loads, %d rebuilt)",
+                len(train_dataset),
+                len(train_dataset._datasets),
+                train_dataset.index_cache_hits,
+                train_dataset.hf_cache_hits,
+                train_dataset.index_cache_misses,
+            )
+        train_sampler = DistributedSampler(
+            train_dataset,
+            num_replicas=config.world_size,
+            rank=config.rank,
+            shuffle=True,
+            seed=42
+        ) if config.world_size > 1 else None
+        self.train_loader = DataLoader(
+            train_dataset,
+            batch_size=config.batch_size,
+            shuffle=(train_sampler is None), 
+            num_workers=config.load_worker,
+            sampler=train_sampler,
+        )
+
+        self.train_scheduler_latent = FlowMatchScheduler(shift=self.config.snr_shift, sigma_min=0.0, extra_one_step=True)
+        self.train_scheduler_latent.set_timesteps(1000, training=True)
+        self.train_scheduler_action = FlowMatchScheduler(shift=self.config.action_snr_shift, sigma_min=0.0, extra_one_step=True)
+        self.train_scheduler_action.set_timesteps(1000, training=True)
+        self.train_scheduler_mcp = None
+        if self.enable_mcp:
+            self.train_scheduler_mcp = FlowMatchScheduler(
+                shift=self.config.mcp_snr_shift,
+                sigma_min=0.0,
+                extra_one_step=True,
+            )
+            self.train_scheduler_mcp.set_timesteps(1000, training=True)
+
+        self.save_dir = Path(config.save_root) / "checkpoints"
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+
+        self.gradient_accumulation_steps = getattr(config, 'gradient_accumulation_steps', 1)
+        self.train_loader_iter = None
+        # if hasattr(config, 'resume_from') and config.resume_from:
+        #     self._load_training_state(config.resume_from)
+    
+    def _get_next_batch(self):
+        """Get next batch from iterator, reset if epoch is finished."""
+        if self.train_loader_iter is None:
+            self.train_loader_iter = iter(self.train_loader)
+        
+        try:
+            batch = next(self.train_loader_iter)
+        except StopIteration:
+            # Reset sampler and iterator when epoch finishes
+            if hasattr(self.train_loader.sampler, 'set_epoch'):
+                self.train_loader.sampler.set_epoch(self.train_loader.sampler.epoch + 1)
+            self.train_loader_iter = iter(self.train_loader)
+            batch = next(self.train_loader_iter)
+        
+        return batch
+
+    @torch.no_grad()
+    def _add_noise(self, latent, train_scheduler, action_mask=False,
+                   action_mode=False, noisy_cond_prob=0., frame_shift=0):
+        B, C, F, H, W = latent.shape
+
+        timestep_ids = sample_timestep_id(batch_size=F, num_train_timesteps=train_scheduler.num_train_timesteps)
+        noise = torch.zeros_like(latent).normal_()
+        timesteps = train_scheduler.timesteps[timestep_ids].to(device=self.device)
+        noisy_latents =train_scheduler.add_noise(latent, noise, timesteps, t_dim=2)
+        targets =train_scheduler.training_target(latent, noise, timesteps)
+
+        patch_f, patch_h, patch_w = self.patch_size
+        if action_mode:
+            patch_f = patch_h = patch_w = 1
+        
+        latent_grid_id = get_mesh_id(
+            latent.shape[-3] // patch_f,  # F
+            latent.shape[-2] // patch_h,  # H
+            latent.shape[-1] // patch_w,  # W
+            t=1 if action_mode else 0,  # 1 for action mode (0 for latent), not used
+            f_w=1,
+            f_shift=frame_shift,
+            action=action_mode
+        ).to(self.device)  # shape: [4, seq_len]
+        latent_grid_id = latent_grid_id[None].repeat(B, 1, 1)
+
+        if torch.rand(1).item() < noisy_cond_prob:
+            cond_timestep_ids = sample_timestep_id(
+                    batch_size=F,
+                    min_timestep_bd=0.5, 
+                    max_timestep_bd=1.0, 
+                    num_train_timesteps=train_scheduler.num_train_timesteps,
+                )
+            noise = torch.zeros_like(latent).normal_()
+            cond_timesteps = train_scheduler.timesteps[cond_timestep_ids].to(device=self.device)
+            latent = train_scheduler.add_noise(latent, noise, cond_timesteps, t_dim=2)
+        else:
+            cond_timesteps = torch.zeros_like(timesteps)
+
+        if action_mask is not None:
+            noisy_latents *= action_mask.float()
+            targets *= action_mask.float()
+            latent *= action_mask.float()
+
+        return dict(
+            timesteps=timesteps[None].repeat(B, 1),
+            noisy_latents=noisy_latents,
+            targets=targets,
+            latent=latent,
+            cond_timesteps=cond_timesteps[None].repeat(B, 1),
+            grid_id=latent_grid_id,
+        )
+
+    @torch.no_grad()
+    def _prepare_input_dict(self, batch_dict):
+        """Prepare input dict following infer code pattern from wan_va_server.py."""
+        chunk_size = torch.randint(1, 5, (1,)).item()
+        # Generate grid_id following infer code (no batch dimension yet)
+        # For action mode: get_mesh_id(shape[-3], shape[-2], shape[-1], t=1, f_w=1, f_shift, action=True)
+        latent_dict = self._add_noise(
+            latent=batch_dict['latents'], 
+            train_scheduler=self.train_scheduler_latent, 
+            action_mask=None, 
+            action_mode=False,
+            noisy_cond_prob=0.5)
+        
+        action_dict = self._add_noise(
+            latent=batch_dict['actions'], 
+            train_scheduler=self.train_scheduler_action, 
+            action_mask=batch_dict['actions_mask'], 
+            action_mode=True,
+            noisy_cond_prob=0.0)
+
+        latent_dict['text_emb'] = batch_dict['text_emb']
+        action_dict['text_emb'] = batch_dict['text_emb']
+        action_dict['actions_mask'] = batch_dict['actions_mask']
+
+        input_dict = {
+            'latent_dict': latent_dict,
+            'action_dict': action_dict,
+            'chunk_size': chunk_size,
+            'window_size': torch.randint(4, 65, (1,)).item(),
+        }
+        if self.enable_mcp:
+            mcp_latent_dicts = []
+            for depth in range(self.config.num_mcp_depths):
+                frame_shift = (depth + 1) * chunk_size
+                shifted_latents, valid_mask = shift_latents_for_mcp(
+                    batch_dict['latents'], frame_shift)
+                mcp_latent_dict = self._add_noise(
+                    latent=shifted_latents,
+                    train_scheduler=self.train_scheduler_mcp,
+                    action_mask=None,
+                    action_mode=False,
+                    noisy_cond_prob=0.0,
+                    frame_shift=frame_shift,
+                )
+                mcp_latent_dict.pop('latent')
+                mcp_latent_dict.pop('cond_timesteps')
+                mcp_latent_dict['valid_mask'] = valid_mask
+                mcp_latent_dicts.append(mcp_latent_dict)
+            input_dict['mcp_latent_dicts'] = mcp_latent_dicts
+        return input_dict
+
+    def convert_input_format(self, input_dict):
+        """Convert input dict to match transformer input format if needed."""
+        for key, value in input_dict.items():
+            input_dict[key] = value.to(self.device)#.to(self.dtype)
+        return input_dict
+
+    def compute_loss(self,
+        input_dict,
+        pred
+    ):
+        if self.enable_mcp:
+            latent_pred, action_pred, mcp_pred_list = pred
+        else:
+            latent_pred, action_pred = pred
+            mcp_pred_list = []
+        if self.enable_mcp and len(mcp_pred_list) != self.config.num_mcp_depths:
+            raise RuntimeError(
+                "MCP output depth count must match num_mcp_depths")
+        action_pred = rearrange(action_pred, 'b (f n) c -> b c f n 1', f=input_dict['action_dict']['targets'].shape[-3])
+        latent_pred = data_seq_to_patch(
+                        self.patch_size, latent_pred,
+                        input_dict['latent_dict']['targets'].shape[-3], input_dict['latent_dict']['targets'].shape[-2],
+                        input_dict['latent_dict']['targets'].shape[-1], batch_size=latent_pred.shape[0])
+        Bn, Fn = input_dict['latent_dict']['timesteps'].shape
+        latent_loss_weight = self.train_scheduler_latent.training_weight(input_dict['latent_dict']['timesteps'].flatten()).reshape(Bn, Fn)
+        action_loss_weight = self.train_scheduler_action.training_weight(input_dict['action_dict']['timesteps'].flatten()).reshape(Bn, Fn)
+
+        # Frame-wise video loss calculation
+        latent_loss = F.mse_loss(latent_pred.float(), input_dict['latent_dict']['targets'].float().detach(), reduction='none')
+        latent_loss = latent_loss * latent_loss_weight[:, None, :, None, None]
+        # Permute to (B, F, H, W, C) and flatten to (B*F, H*W*C)
+        latent_loss = latent_loss.permute(0, 2, 3, 4, 1)  # (B, C, F, H, W) -> (B, F, H, W, C)
+        latent_loss = latent_loss.flatten(0, 1).flatten(1)  # (B, F, H, W, C) -> (B*F, H*W*C)
+        # Sum per frame and compute mask per frame
+        latent_loss_per_frame = latent_loss.sum(dim=1)  # (B*F,)
+        latent_mask_per_frame = torch.ones_like(latent_loss).sum(dim=1)  # (B*F,)
+        latent_loss = (latent_loss_per_frame / (latent_mask_per_frame + 1e-6)).mean()
+
+        # Frame-wise action loss calculation
+        action_loss = F.mse_loss(action_pred.float(), input_dict['action_dict']['targets'].float().detach(), reduction='none')
+        action_loss = action_loss * action_loss_weight[:, None, :, None, None]
+        action_loss = action_loss * input_dict['action_dict']['actions_mask'].float()
+        # Permute to (B, F, H, W, C) and flatten to (B*F, H*W*C)
+        action_loss = action_loss.permute(0, 2, 3, 4, 1)  # (B, C, F, H, W) -> (B, F, H, W, C)
+        action_mask = input_dict['action_dict']['actions_mask'].float().permute(0, 2, 3, 4, 1)  # (B, C, F, H, W) -> (B, F, H, W, C)
+        action_loss = action_loss.flatten(0, 1).flatten(1)  # (B, F, H, W, C) -> (B*F, H*W*C)
+        action_mask = action_mask.flatten(0, 1).flatten(1)  # (B, F, H, W, C) -> (B*F, H*W*C)
+        # Sum per frame and normalize by mask per frame
+        action_loss_per_frame = action_loss.sum(dim=1)  # (B*F,)
+        action_mask_per_frame = action_mask.sum(dim=1)  # (B*F,)
+        action_loss = (action_loss_per_frame / (action_mask_per_frame + 1e-6)).mean()
+
+        mcp_losses = []
+        for mcp_pred, mcp_latent_dict in zip(
+                mcp_pred_list, input_dict.get('mcp_latent_dicts', [])):
+            mcp_pred = data_seq_to_patch(
+                self.patch_size,
+                mcp_pred,
+                mcp_latent_dict['targets'].shape[-3],
+                mcp_latent_dict['targets'].shape[-2],
+                mcp_latent_dict['targets'].shape[-1],
+                batch_size=mcp_pred.shape[0],
+            )
+            mcp_batch_size, mcp_num_frames = mcp_latent_dict[
+                'timesteps'].shape
+            mcp_loss_weight = self.train_scheduler_mcp.training_weight(
+                mcp_latent_dict['timesteps'].flatten()).reshape(
+                    mcp_batch_size, mcp_num_frames)
+            mcp_loss = F.mse_loss(
+                mcp_pred.float(),
+                mcp_latent_dict['targets'].float().detach(),
+                reduction='none',
+            )
+            mcp_loss = mcp_loss * mcp_loss_weight[:, None, :, None, None]
+            valid_mask = mcp_latent_dict['valid_mask'].to(
+                device=mcp_loss.device, dtype=mcp_loss.dtype)
+            valid_count = valid_mask.expand_as(mcp_loss).sum()
+            mcp_loss = (mcp_loss * valid_mask).sum() / valid_count.clamp_min(1.)
+            mcp_losses.append(mcp_loss / self.gradient_accumulation_steps)
+
+        return (
+            latent_loss / self.gradient_accumulation_steps,
+            action_loss / self.gradient_accumulation_steps,
+            mcp_losses,
+        )
+
+    def _train_step(self, batch, batch_idx):
+        """Train a single batch, returns losses for logging."""
+        batch = self.convert_input_format(batch)
+        input_dict = self._prepare_input_dict(batch)
+        
+        should_sync = (batch_idx + 1) % self.gradient_accumulation_steps == 0
+        
+        if not should_sync:
+            self.transformer.set_requires_gradient_sync(False)
+        else:
+            self.transformer.set_requires_gradient_sync(True)
+
+        output = self.transformer(input_dict, train_mode=True)
+        latent_loss, action_loss, mcp_losses = self.compute_loss(
+            input_dict, output)
+        mcp_loss = sum(
+            weight * depth_loss
+            for weight, depth_loss in zip(
+                self.config.mcp_loss_weights, mcp_losses)
+        ) if mcp_losses else latent_loss.new_zeros(())
+        loss = latent_loss + action_loss + mcp_loss
+
+        loss.backward()
+
+        losses = {
+            'latent_loss': latent_loss.detach(),
+            'action_loss': action_loss.detach(),
+            'mcp_losses': [depth_loss.detach() for depth_loss in mcp_losses],
+            'mcp_loss': mcp_loss.detach(),
+        }
+        
+        # Only update weights after accumulating gradients
+        if should_sync:
+            total_norm = torch.nn.utils.clip_grad_norm_(self.transformer.parameters(), 2.0)
+            self.optimizer.step()
+            self.lr_scheduler.step()
+            self.optimizer.zero_grad()
+            
+            losses['total_norm'] = total_norm
+            losses['should_log'] = True
+        else:
+            losses['should_log'] = False
+
+        return losses
+
+    def save_checkpoint(self,):
+        """Save model checkpoint in the same format as pretrained model."""
+        try:
+            state_dict = get_model_state_dict(
+                self.transformer,
+                options=StateDictOptions(full_state_dict=True, cpu_offload=True),
+            )
+            state_dict_bf16 = {k: v.to(torch.bfloat16) for k, v in state_dict.items()}
+            # optim_state = get_optimizer_state_dict(
+            #         self.transformer, self.optimizer,
+            #         options=StateDictOptions(full_state_dict=True, cpu_offload=True),
+            #     )
+
+            # Only rank 0 saves the checkpoint
+            if self.config.rank == 0:
+                checkpoint_dir = self.save_dir / f"checkpoint_step_{self.step}"
+                checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+                # Save transformer in the same format as pretrained model
+                transformer_dir = checkpoint_dir / "transformer"
+                transformer_dir.mkdir(parents=True, exist_ok=True)
+
+                logger.info(f"Saving transformer to {transformer_dir}")
+
+                # Manually save in diffusers format (outside FSDP context to avoid deadlock)
+                # Save model weights
+                model_file = transformer_dir / "diffusion_pytorch_model.safetensors"
+                save_file(state_dict_bf16, model_file)
+
+                # Save config (copy from original transformer config and update _name_or_path)
+                config_file = transformer_dir / "config.json"
+                config_dict = dict(self.transformer.config)
+                config_dict.pop('_name_or_path', None)
+                with open(config_file, 'w') as f:
+                    json.dump(config_dict, f, indent=2)
+
+                # # Save optimizer state and training metadata in PyTorch format
+                # training_state_path = checkpoint_dir / "training_state.pt"
+                # logger.info(f"Saving training state to {training_state_path}")
+                # torch.save({
+                #     'step': self.step,
+                #     'optimizer_state_dict': optim_state,
+                #     'config': vars(self.config),
+                # }, training_state_path)
+
+                logger.info(f"Checkpoint saved successfully at step {self.step}")
+
+            # Synchronize all processes after saving
+            if dist.is_initialized():
+                dist.barrier()
+
+        except Exception as e:
+            if self.config.rank == 0:
+                logger.error(f"Failed to save checkpoint: {e}")
+                import traceback
+                logger.error(traceback.format_exc())
+            # Ensure all processes stay synchronized even on error
+            if dist.is_initialized():
+                dist.barrier()
+
+    def _load_training_state(self, checkpoint_path):
+        """Load training state (optimizer + step) after FSDP and optimizer creation."""
+        checkpoint_dir = Path(checkpoint_path)
+        training_state_path = checkpoint_dir / "training_state.pt"
+
+        if not training_state_path.exists():
+            if self.config.rank == 0:
+                logger.warning(f"Training state not found: {training_state_path}, starting from step 0")
+            return
+
+        if self.config.rank == 0:
+            logger.info(f"Loading training state from {training_state_path}")
+
+        # All ranks load the training state directly
+        training_state = torch.load(training_state_path, map_location='cpu', weights_only=False)
+
+        # All ranks load optimizer state (required for FSDP)
+        set_optimizer_state_dict(
+            self.transformer, self.optimizer,
+            optim_state_dict=training_state['optimizer_state_dict'],
+            options=StateDictOptions(full_state_dict=True, strict=False)
+        )
+        self.step = training_state.get('step', 0)
+
+        if self.config.rank == 0:
+            logger.info(f"Training state loaded, resuming from step {self.step}")
+
+        # Synchronize all ranks
+        if dist.is_initialized():
+            dist.barrier()
+
+    def train(self):
+        """Main training loop - train by steps instead of epochs."""
+        logger.info(f"Starting training for {self.config.num_steps} steps...")
+        self.transformer.train()
+
+        progress_bar = tqdm(
+            total=self.config.num_steps,
+            desc="Training",
+            disable=(self.config.rank != 0),
+            leave=True,
+            dynamic_ncols=True,
+            initial=self.step
+        )
+
+        self.optimizer.zero_grad()
+        accumulated_latent_losses = []
+        accumulated_action_losses = []
+        accumulated_mcp_losses = [
+            [] for _ in range(self.config.num_mcp_depths)
+        ] if self.enable_mcp else []
+        accumulated_mcp_total_losses = []
+        step_in_accumulation = 0
+
+        while self.step < self.config.num_steps:
+            # Get next batch (handles epoch reset automatically)
+            batch = self._get_next_batch()
+            
+            losses = self._train_step(batch, step_in_accumulation)
+            
+            # Accumulate losses for logging
+            accumulated_latent_losses.append(losses['latent_loss'])
+            accumulated_action_losses.append(losses['action_loss'])
+            for depth, mcp_loss in enumerate(losses['mcp_losses']):
+                accumulated_mcp_losses[depth].append(mcp_loss)
+            accumulated_mcp_total_losses.append(losses['mcp_loss'])
+            step_in_accumulation += 1
+
+            # Log and checkpoint when optimizer steps
+            if losses['should_log']:
+                lr = self.lr_scheduler.get_last_lr()[0]
+
+                # Average accumulated losses
+                latent_loss_show = dist_mean(torch.stack(accumulated_latent_losses).sum()).detach().cpu().item()
+                action_loss_show = dist_mean(torch.stack(accumulated_action_losses).sum()).detach().cpu().item()
+                max_latent_loss_show = dist_max(torch.stack(accumulated_latent_losses).sum()).detach().cpu().item()
+                max_action_loss_show = dist_max(torch.stack(accumulated_action_losses).sum()).detach().cpu().item()
+                mcp_loss_shows = [
+                    dist_mean(torch.stack(depth_losses).sum()).detach().cpu().item()
+                    for depth_losses in accumulated_mcp_losses
+                ]
+                mcp_total_loss_show = dist_mean(
+                    torch.stack(accumulated_mcp_total_losses).sum()
+                ).detach().cpu().item()
+
+                # Clear accumulated losses
+                accumulated_latent_losses = []
+                accumulated_action_losses = []
+                accumulated_mcp_losses = [
+                    [] for _ in range(self.config.num_mcp_depths)
+                ] if self.enable_mcp else []
+                accumulated_mcp_total_losses = []
+                step_in_accumulation = 0
+
+                torch.cuda.synchronize()
+                if self.step % self.config.gc_interval == 0:
+                    torch.cuda.empty_cache()
+                    gc.collect()
+
+                if self.config.rank == 0:
+                    total_norm = losses['total_norm']
+                    progress_bar.n += 1
+                    postfix = {
+                        'latent_loss': f'{latent_loss_show:.4f}',
+                        'action_loss': f'{action_loss_show:.4f}',
+                        'step': self.step,
+                        'grad_norm': f'{total_norm.item():.2f}',
+                        'lr': f'{lr:.2e}'
+                    }
+                    if self.enable_mcp:
+                        postfix['mcp_loss'] = f'{mcp_total_loss_show:.4f}'
+                    progress_bar.set_postfix(postfix)
+                    if self.config.enable_wandb:
+                        log_values = {
+                            'loss_metrics/global_avg_video_loss': latent_loss_show,
+                            'loss_metrics/global_avg_action_loss': action_loss_show,
+                            'loss_metrics/global_max_video_loss': max_latent_loss_show,
+                            'loss_metrics/global_max_action_loss': max_action_loss_show,
+                            'grad_norm': total_norm.item(),
+                            'lr': lr,
+                        }
+                        if self.enable_mcp:
+                            log_values['loss_metrics/mcp_weighted_total'] = (
+                                mcp_total_loss_show)
+                            for depth, mcp_loss_show in enumerate(
+                                    mcp_loss_shows):
+                                log_values[
+                                    f'loss_metrics/mcp_depth_{depth + 1}'] = (
+                                        mcp_loss_show)
+                        self.wandb.log(log_values, step=self.step)
+                
+                self.step += 1
+                
+                if self.step % self.config.save_interval == 0:
+                    if self.config.rank == 0:
+                        logger.info(f"Starting save model at step {self.step}")
+                    self.save_checkpoint()
+
+            if dist.is_initialized():
+                dist.barrier()
+
+        progress_bar.close()
+        logger.info("Training completed!")
+
+
+def run(args):
+    """Main entry point."""
+    config = VA_CONFIGS[args.config_name]
+
+    overrides = {
+        'wan22_pretrained_model_name_or_path': args.pretrained_model_path,
+        'dataset_path': args.dataset_path,
+        'empty_emb_path': args.empty_emb_path,
+        'learning_rate': args.learning_rate,
+        'cfg_prob': args.cfg_prob,
+        'init_worker': args.init_worker,
+        'load_worker': args.load_worker,
+        'batch_size': args.batch_size,
+        'gradient_accumulation_steps': args.gradient_accumulation_steps,
+        'num_steps': args.num_steps,
+        'save_interval': args.save_interval,
+        'save_root': args.save_root,
+    }
+    for key, value in overrides.items():
+        if value is not None:
+            config[key] = value
+
+    if args.dataset_path is not None and args.empty_emb_path is None:
+        config.empty_emb_path = os.path.join(args.dataset_path, 'empty_emb.pt')
+    if args.disable_wandb:
+        config.enable_wandb = False
+
+    rank = int(os.getenv("RANK", 0))
+    local_rank = int(os.environ.get('LOCAL_RANK', 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+
+    init_distributed(world_size, local_rank, rank)
+
+    config.rank = rank
+    config.local_rank = local_rank
+    config.world_size = world_size
+
+    if rank == 0:
+        logger.info(f"Using config: {args.config_name}")
+        logger.info(f"World size: {world_size}, Local rank: {local_rank}")
+
+    trainer = Trainer(config)
+    trainer.train()
+
+
+def main():
+    """Parse arguments and run training."""
+    parser = argparse.ArgumentParser(description="Train WAN model for robotics")
+    parser.add_argument(
+        "--config-name",
+        type=str,
+        default='robotwin_train',
+        help="Config name",
+    )
+    parser.add_argument(
+        "--save-root",
+        type=str,
+        default=None,
+        help="Root directory for saving checkpoints",
+    )
+    parser.add_argument(
+        "--pretrained-model-path",
+        type=str,
+        default=None,
+        help="Pretrained model root containing the transformer directory",
+    )
+    parser.add_argument(
+        "--dataset-path",
+        type=str,
+        default=None,
+        help="Root directory containing open-format LeRobot datasets",
+    )
+    parser.add_argument(
+        "--empty-emb-path",
+        type=str,
+        default=None,
+        help="Path to empty_emb.pt (defaults to DATASET_PATH/empty_emb.pt)",
+    )
+    parser.add_argument(
+        "--disable-wandb",
+        action="store_true",
+        help="Disable Weights & Biases logging",
+    )
+    parser.add_argument("--learning-rate", type=float, default=None)
+    parser.add_argument("--cfg-prob", type=float, default=None)
+    parser.add_argument("--init-worker", type=int, default=None)
+    parser.add_argument("--load-worker", type=int, default=None)
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument(
+        "--gradient-accumulation-steps", type=int, default=None)
+    parser.add_argument("--num-steps", type=int, default=None)
+    parser.add_argument("--save-interval", type=int, default=None)
+
+    args = parser.parse_args()
+    run(args)
+
+
+if __name__ == "__main__":
+    init_logger()
+    main()

wan_va/utils/Simple_Remote_Infer/LEGAL.md

@@ -0,0 +1,7 @@
+Legal Disclaimer
+
+Within this source code, the comments in Chinese shall be the original, governing version. Any comment in other languages are for reference only. In the event of any conflict between the Chinese language version comments and other language version comments, the Chinese language version shall prevail.
+
+法律免责声明
+
+关于代码注释部分，中文注释为官方版本，其它语言注释仅做参考。中文注释可能与其它语言注释存在不一致，当中文注释与其它语言注释存在不一致时，请以中文注释为准。

wan_va/utils/Simple_Remote_Infer/deploy/msgpack_numpy.py

@@ -0,0 +1,61 @@
+"""Adds NumPy array support to msgpack.
+
+msgpack is good for (de)serializing data over a network for multiple reasons:
+- msgpack is secure (as opposed to pickle/dill/etc which allow for arbitrary code execution)
+- msgpack is widely used and has good cross-language support
+- msgpack does not require a schema (as opposed to protobuf/flatbuffers/etc) which is convenient in dynamically typed
+    languages like Python and JavaScript
+- msgpack is fast and efficient (as opposed to readable formats like JSON/YAML/etc); I found that msgpack was ~4x faster
+    than pickle for serializing large arrays using the below strategy
+
+The code below is adapted from https://github.com/lebedov/msgpack-numpy. The reason not to use that library directly is
+that it falls back to pickle for object arrays.
+"""
+
+import functools
+
+import msgpack
+import numpy as np
+
+
+def pack_array(obj):
+    if (isinstance(
+            obj,
+        (np.ndarray, np.generic))) and obj.dtype.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {obj.dtype}")
+
+    if isinstance(obj, np.ndarray):
+        return {
+            b"__ndarray__": True,
+            b"data": obj.tobytes(),
+            b"dtype": obj.dtype.str,
+            b"shape": obj.shape,
+        }
+
+    if isinstance(obj, np.generic):
+        return {
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
+        }
+
+    return obj
+
+
+def unpack_array(obj):
+    if b"__ndarray__" in obj:
+        return np.ndarray(buffer=obj[b"data"],
+                          dtype=np.dtype(obj[b"dtype"]),
+                          shape=obj[b"shape"])
+
+    if b"__npgeneric__" in obj:
+        return np.dtype(obj[b"dtype"]).type(obj[b"data"])
+
+    return obj
+
+
+Packer = functools.partial(msgpack.Packer, default=pack_array)
+packb = functools.partial(msgpack.packb, default=pack_array)
+
+Unpacker = functools.partial(msgpack.Unpacker, object_hook=unpack_array)
+unpackb = functools.partial(msgpack.unpackb, object_hook=unpack_array)

wan_va/utils/Simple_Remote_Infer/deploy/websocket_policy_server.py

@@ -0,0 +1,96 @@
+import asyncio
+import http
+import logging
+import time
+import traceback
+
+import websockets.asyncio.server as _server
+import websockets.frames
+
+from .msgpack_numpy import Packer, unpackb
+
+logger = logging.getLogger(__name__)
+
+
+class WebsocketPolicyServer:
+    """Serves a policy using the websocket protocol. See websocket_client_policy.py for a client implementation.
+
+    Currently only implements the `load` and `infer` methods.
+    """
+
+    def __init__(
+        self,
+        policy,
+        host: str = "0.0.0.0",
+        port: int | None = None,
+        metadata: dict | None = None,
+    ) -> None:
+        self._policy = policy
+        self._host = host
+        self._port = port
+        self._metadata = metadata or {}
+        logging.getLogger("websockets.server").setLevel(logging.INFO)
+
+    def serve_forever(self) -> None:
+        asyncio.run(self.run())
+
+    async def run(self):
+        async with _server.serve(
+                self._handler,
+                self._host,
+                self._port,
+                compression=None,
+                max_size=None,
+                process_request=_health_check,
+                ping_interval=None,
+                ping_timeout=None,
+        ) as server:
+            await server.serve_forever()
+
+    async def _handler(self, websocket: _server.ServerConnection):
+        logger.info(f"Connection from {websocket.remote_address} opened")
+        packer = Packer()
+
+        await websocket.send(packer.pack(self._metadata))
+
+        prev_total_time = None
+        while True:
+            try:
+                start_time = time.monotonic()
+                obs = unpackb(await websocket.recv())
+
+                infer_time = time.monotonic()
+                action = self._policy.infer(obs)
+                infer_time = time.monotonic() - infer_time
+
+                action["server_timing"] = {
+                    "infer_ms": infer_time * 1000,
+                }
+                if prev_total_time is not None:
+                    # We can only record the last total time since we also want to include the send time.
+                    action["server_timing"][
+                        "prev_total_ms"] = prev_total_time * 1000
+
+                await websocket.send(packer.pack(action))
+                prev_total_time = time.monotonic() - start_time
+
+            except websockets.ConnectionClosed:
+                logger.info(
+                    f"Connection from {websocket.remote_address} closed")
+                break
+            except Exception:
+                await websocket.send(traceback.format_exc())
+                await websocket.close(
+                    code=websockets.frames.CloseCode.INTERNAL_ERROR,
+                    reason=
+                    "Internal server error. Traceback included in previous frame.",
+                )
+                raise
+
+
+def _health_check(connection: _server.ServerConnection,
+                  request: _server.Request) -> _server.Response | None:
+    if request.path == "/healthz":
+        return connection.respond(http.HTTPStatus.OK, "OK\n")
+    # Continue with the normal request handling.
+    return None

wan_va/utils/__init__.py

@@ -0,0 +1,10 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+from .logging import init_logger, logger
+from .scheduler import FlowMatchScheduler
+from .server_utils import run_async_server_mode
+from .utils import data_seq_to_patch, get_mesh_id, save_async, sample_timestep_id, warmup_constant_lambda
+
+__all__ = [
+    'logger', 'init_logger', 'get_mesh_id', 'save_async', 'data_seq_to_patch',
+    'FlowMatchScheduler', 'run_async_server_mode', 'sample_timestep_id', 'warmup_constant_lambda'
+]

wan_va/utils/logging.py

@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+
+logger = logging.getLogger()
+
+
+def init_logger():
+    logger.setLevel(logging.INFO)
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.INFO)
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    ch.setFormatter(formatter)
+    logger.addHandler(ch)
+
+    # suppress verbose torch.profiler logging
+    os.environ["KINETO_LOG_LEVEL"] = "5"

wan_va/utils/scheduler.py

@@ -0,0 +1,135 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import math
+import torch
+
+class FlowMatchScheduler():
+
+    def __init__(
+        self,
+        num_inference_steps=100,
+        num_train_timesteps=1000,
+        shift=3.0,
+        sigma_max=1.0,
+        sigma_min=0.003 / 1.002,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        exponential_shift=False,
+        exponential_shift_mu=None,
+        shift_terminal=None,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.exponential_shift = exponential_shift
+        self.exponential_shift_mu = exponential_shift_mu
+        self.shift_terminal = shift_terminal
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self,
+                      num_inference_steps=100,
+                      denoising_strength=1.0,
+                      training=False,
+                      shift=None,
+                      dynamic_shift_len=None):
+        if shift is not None:
+            self.shift = shift
+        sigma_start = self.sigma_min + (self.sigma_max -
+                                        self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
+                                         num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
+                                         num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        if self.exponential_shift:
+            mu = self.calculate_shift(
+                dynamic_shift_len
+            ) if dynamic_shift_len is not None else self.exponential_shift_mu
+            self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1))
+        else:
+            self.sigmas = self.shift * self.sigmas / (
+                1 + (self.shift - 1) * self.sigmas)
+        if self.shift_terminal is not None:
+            one_minus_z = 1 - self.sigmas
+            scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+            self.sigmas = 1 - (one_minus_z / scale_factor)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(
+                -2 * ((x - num_inference_steps / 2) / num_inference_steps)**2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * (num_inference_steps /
+                                           y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+            self.training = True
+        else:
+            self.training = False
+
+    def step(self, model_output, timestep, sample, to_final=False, **kwargs):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        if to_final or timestep_id + 1 >= len(self.timesteps):
+            sigma_ = 1 if (self.inverse_timesteps
+                           or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1]
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def return_to_timestep(self, timestep, sample, sample_stablized):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        model_output = (sample - sample_stablized) / sigma
+        return model_output
+
+    def add_noise(self, original_samples, noise, timestep, t_dim=2):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep = timestep[None]
+        timestep_id = torch.argmin((self.timesteps[:, None] - timestep).abs(),
+                                   dim=0)
+        shape = [1] * noise.ndim
+        shape[t_dim] = timestep_id.shape[0]
+        sigma = self.sigmas[timestep_id].to(original_samples).view(shape)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin(
+            (self.timesteps[:, None].to(timestep.device) -
+             timestep[None]).abs(),
+            dim=0)
+        weights = self.linear_timesteps_weights.to(
+            timestep.device)[timestep_id].to(timestep.device)
+        return weights
+
+    def calculate_shift(
+        self,
+        image_seq_len,
+        base_seq_len: int = 256,
+        max_seq_len: int = 8192,
+        base_shift: float = 0.5,
+        max_shift: float = 0.9,
+    ):
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        mu = image_seq_len * m + b
+        return mu

wan_va/utils/server_utils.py

@@ -0,0 +1,83 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import torch
+import torch.distributed as dist
+
+from .logging import logger
+from .Simple_Remote_Infer.deploy.websocket_policy_server import WebsocketPolicyServer
+
+
+class DistributedModelWrapper:
+    """
+    TODO
+    """
+
+    def __init__(self, model, local_rank):
+        self.model = model
+        self.local_rank = local_rank
+
+    def infer(self, obs):
+        return distributed_infer(self.model, obs, self.local_rank)
+
+
+def distributed_infer(model, obs, local_rank):
+    """
+    TODO
+    """
+    rank = dist.get_rank()
+    assert rank == local_rank, "distributed_infer can only run at（rank 0)"
+
+    cmd = torch.tensor(1,
+                       dtype=torch.int64,
+                       device='cuda' if torch.cuda.is_available() else 'cpu')
+    dist.broadcast(cmd, src=0)
+
+    obj_list = [obs]
+    dist.broadcast_object_list(obj_list, src=0)
+
+    result = model.infer(obs)
+
+    return result
+
+
+def worker_loop(model, local_rank):
+    """
+    TODO
+    """
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    rank = dist.get_rank()
+
+    while True:
+        cmd = torch.zeros(1, dtype=torch.int64, device=device)
+        dist.broadcast(cmd, src=0)
+        cmd_val = cmd.item()
+
+        if cmd_val == -1:
+            break
+        elif cmd_val == 1:
+            obj_list = [None]
+            dist.broadcast_object_list(obj_list, src=0)
+            obs = obj_list[0]
+            _ = model.infer(obs)
+        else:
+            pass
+
+    logger.info(f"[worker_loop] Rank {rank} exiting.")
+
+
+def run_async_server_mode(model, local_rank, host, port):
+    logger.info("Running in ASYNC SERVER mode")
+    if local_rank == 0:
+        dist_model = DistributedModelWrapper(model, local_rank=local_rank)
+        model_server = WebsocketPolicyServer(dist_model, host=host, port=port)
+        model_server.serve_forever()
+
+        cmd = torch.tensor(
+            -1,
+            dtype=torch.int64,
+            device='cuda' if torch.cuda.is_available() else 'cpu')
+        dist.broadcast(cmd, src=0)
+    else:
+        try:
+            worker_loop(model, local_rank)
+        except KeyboardInterrupt:
+            logger.info(f"Rank {local_rank}: Shutting down")

wan_va/utils/utils.py

@@ -0,0 +1,92 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import concurrent.futures
+
+import numpy as np
+import torch
+
+executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+
+__all__ = ['get_mesh_id', 'save_async', 'data_seq_to_patch']
+
+
+def data_seq_to_patch(
+    patch_size,
+    data_seq,
+    latent_num_frames,
+    latent_height,
+    latent_width,
+    batch_size=1,
+):
+    p_t, p_h, p_w = patch_size
+    post_patch_num_frames = latent_num_frames // p_t
+    post_patch_height = latent_height // p_h
+    post_patch_width = latent_width // p_w
+
+    data_patch = data_seq.reshape(batch_size, post_patch_num_frames,
+                                  post_patch_height, post_patch_width, p_t,
+                                  p_h, p_w, -1)
+    data_patch = data_patch.permute(0, 7, 1, 4, 2, 5, 3, 6)
+    data_patch = data_patch.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+    return data_patch
+
+
+def get_mesh_id(f, h, w, t, f_w=1, f_shift=0, action=False):
+    f_idx = torch.arange(f_shift, f + f_shift) * f_w
+    h_idx = torch.arange(h)
+    w_idx = torch.arange(w)
+    ff, hh, ww = torch.meshgrid(f_idx, h_idx, w_idx, indexing='ij')
+    if action:
+        ff_offset = (torch.ones([h]).cumsum(0) / (h + 1)).view(1, -1, 1)
+        ff = ff + ff_offset
+        hh = torch.ones_like(hh) * -1
+        ww = torch.ones_like(ww) * -1
+
+    grid_id = torch.cat(
+        [
+            ff.unsqueeze(0),
+            hh.unsqueeze(0),
+            ww.unsqueeze(0),
+        ],
+        dim=0,
+    ).flatten(1)
+    grid_id = torch.cat([grid_id, torch.full_like(grid_id[:1], t)], dim=0)
+    return grid_id
+
+
+def save_async(obj, file_path):
+    """
+    todo
+    """
+    if torch.is_tensor(obj) or (isinstance(obj, dict) and any(
+            torch.is_tensor(v) for v in obj.values())):
+        if torch.is_tensor(obj):
+            if obj.is_cuda:
+                obj = obj.cpu()
+        elif isinstance(obj, dict):
+            obj = {
+                k: v.cpu() if torch.is_tensor(v) else v
+                for k, v in obj.items()
+            }
+        executor.submit(torch.save, obj, file_path)
+    elif isinstance(obj, np.ndarray):
+        obj_copy = obj.copy()
+        executor.submit(np.save, file_path, obj_copy)
+    else:
+        executor.submit(torch.save, obj, file_path)
+
+def sample_timestep_id(
+    batch_size: int = 1,
+    min_timestep_bd: float = 0.0,
+    max_timestep_bd: float = 1.0,
+    num_train_timesteps: int = 1000,
+):
+    u = torch.rand(size=[batch_size])
+    u = u * (max_timestep_bd - min_timestep_bd) + min_timestep_bd
+    timestep_id = (u * num_train_timesteps).clamp(min=0, max=num_train_timesteps - 1).to(torch.int64)
+    return timestep_id
+
+
+def warmup_constant_lambda(current_step, warmup_steps=1000):
+    if current_step < warmup_steps:
+        return float(current_step) / float(max(1, warmup_steps))
+    return 1.0

wan_va/wan_va_server.py

@@ -0,0 +1,729 @@
+# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
+import argparse
+import os
+import time
+from functools import partial
+from PIL import Image
+from diffusers.video_processor import VideoProcessor
+from diffusers.utils import export_to_video
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers.pipelines.wan.pipeline_wan import prompt_clean
+from einops import rearrange
+from tqdm import tqdm
+
+from .configs import VA_CONFIGS
+from .distributed.fsdp import shard_model
+from .distributed.util import _configure_model, init_distributed
+from .modules.utils import (
+    WanVAEStreamingWrapper,
+    load_text_encoder,
+    load_tokenizer,
+    load_transformer,
+    load_vae,
+)
+from .utils import (
+    FlowMatchScheduler,
+    data_seq_to_patch,
+    get_mesh_id,
+    init_logger,
+    logger,
+    run_async_server_mode,
+    save_async,
+)
+
+
+class VA_Server:
+
+    def __init__(self, job_config):
+        self.cache_name = 'pos'
+        self.job_config = job_config
+        self.save_root = job_config.save_root
+        self.dtype = job_config.param_dtype
+        self.device = torch.device(f"cuda:{job_config.local_rank}")
+        self.enable_offload = getattr(job_config, 'enable_offload', True)  # offload vae & text_encoder to save vram
+
+        self.scheduler = FlowMatchScheduler(shift=self.job_config.snr_shift,
+                                            sigma_min=0.0,
+                                            extra_one_step=True)
+        self.action_scheduler = FlowMatchScheduler(
+            shift=self.job_config.action_snr_shift,
+            sigma_min=0.0,
+            extra_one_step=True)
+        self.scheduler.set_timesteps(1000, training=True)
+        self.action_scheduler.set_timesteps(1000, training=True)
+
+        self.vae = load_vae(
+            os.path.join(job_config.wan22_pretrained_model_name_or_path,
+                         'vae'),
+            torch_dtype=self.dtype,
+            torch_device='cpu' if self.enable_offload else self.device,
+        )
+        self.streaming_vae = WanVAEStreamingWrapper(self.vae)
+
+        self.tokenizer = load_tokenizer(
+            os.path.join(job_config.wan22_pretrained_model_name_or_path,
+                         'tokenizer'), )
+
+        self.text_encoder = load_text_encoder(
+            os.path.join(job_config.wan22_pretrained_model_name_or_path,
+                         'text_encoder'),
+            torch_dtype=self.dtype,
+            torch_device='cpu' if self.enable_offload else self.device,
+        )
+
+        self.transformer = load_transformer(
+            os.path.join(job_config.wan22_pretrained_model_name_or_path,
+                         'transformer'),
+            torch_dtype=self.dtype,
+            torch_device=self.device,
+            attn_mode="torch",
+            disable_mcp=True,
+        )
+        shard_fn = shard_model
+        self.transformer = _configure_model(model=self.transformer,
+                                            shard_fn=shard_fn,
+                                            param_dtype=self.dtype,
+                                            device=self.device,
+                                            eval_mode=True,
+                                            )
+
+        self.env_type = job_config.env_type
+        self.streaming_vae_half = None
+        if self.env_type == 'robotwin_tshape':
+            vae_half = load_vae(
+                os.path.join(job_config.wan22_pretrained_model_name_or_path,
+                             'vae'),
+                torch_dtype=self.dtype,
+                torch_device='cpu' if self.enable_offload else self.device,
+            )
+            self.streaming_vae_half = WanVAEStreamingWrapper(vae_half)
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt=None,
+        num_videos_per_prompt=1,
+        max_sequence_length=512,
+        device=None,
+        dtype=None,
+    ):
+        device = device or self.device
+        dtype = dtype or self.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt = [prompt_clean(u) for u in prompt]
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+        seq_lens = mask.gt(0).sum(dim=1).long()
+
+        text_encoder_device = next(self.text_encoder.parameters()).device
+        prompt_embeds = self.text_encoder(text_input_ids.to(text_encoder_device),
+                                          mask.to(text_encoder_device)).last_hidden_state
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+        prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+        prompt_embeds = torch.stack([
+            torch.cat(
+                [u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))])
+            for u in prompt_embeds
+        ],
+                                    dim=0)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt,
+                                           seq_len, -1)
+
+        return prompt_embeds.to(device)
+
+    def encode_prompt(
+        self,
+        prompt,
+        negative_prompt=None,
+        do_classifier_free_guidance=True,
+        num_videos_per_prompt=1,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        max_sequence_length=226,
+        device=None,
+        dtype=None,
+    ):
+        r"""
+        TODO
+        """
+        device = device or self.device
+        dtype = dtype or self.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = batch_size * [negative_prompt] if isinstance(
+                negative_prompt, str) else negative_prompt
+
+            if prompt is not None and type(prompt) is not type(
+                    negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}.")
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`.")
+
+            negative_prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=negative_prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+        return prompt_embeds, negative_prompt_embeds
+
+    def normalize_latents(
+        self,
+        latents: torch.Tensor,
+        latents_mean: torch.Tensor,
+        latents_std: torch.Tensor,
+    ) -> torch.Tensor:
+        latents_mean = latents_mean.view(1, -1, 1, 1,
+                                         1).to(device=latents.device)
+        latents_std = latents_std.view(1, -1, 1, 1,
+                                       1).to(device=latents.device)
+        latents = ((latents.float() - latents_mean) * latents_std).to(latents)
+        return latents
+
+    def preprocess_action(self, action):
+        action_model_input = torch.from_numpy(action)
+        CA, FA, HA = action_model_input.shape  # C, F, H
+        action_model_input_paded = F.pad(action_model_input,
+                                         [0, 0, 0, 0, 0, 1],
+                                         mode='constant',
+                                         value=0)
+
+        action_model_input = action_model_input_paded[
+            self.job_config.inverse_used_action_channel_ids]
+
+        if self.action_norm_method == 'quantiles':
+            action_model_input = (action_model_input - self.actions_q01) / (
+                self.actions_q99 - self.actions_q01 + 1e-6) * 2. - 1.
+        else:
+            raise NotImplementedError
+        return action_model_input.unsqueeze(0).unsqueeze(-1)  # B, C, F, H, W
+
+    def postprocess_action(self, action):
+        action = action.cpu()  # B, C, F, H, W
+
+        action = action[0, ..., 0]  #C, F, H
+        if self.action_norm_method == 'quantiles':
+            action = (action + 1) / 2 * (self.actions_q99 - self.actions_q01 +
+                                         1e-6) + self.actions_q01
+        else:
+            raise NotImplementedError
+        action = action.squeeze(0).detach().cpu().numpy()
+        return action[self.job_config.used_action_channel_ids]
+    
+    def _repeat_input_for_cfg(self, input_dict):
+        if self.use_cfg:
+            input_dict['noisy_latents'] = input_dict['noisy_latents'].repeat(2, 1, 1, 1, 1)
+            input_dict['text_emb'] = torch.cat([self.prompt_embeds.to(self.dtype).clone(), self.negative_prompt_embeds.to(self.dtype).clone()], dim=0)
+            input_dict['grid_id'] = input_dict['grid_id'][None].repeat(2, 1, 1)
+            input_dict['timesteps'] = input_dict['timesteps'][None].repeat(2, 1)
+        else:
+            input_dict['grid_id'] = input_dict['grid_id'][None]
+            input_dict['timesteps'] = input_dict['timesteps'][None]
+        return input_dict
+
+    def _prepare_latent_input(self,
+                              latent_model_input,
+                              action_model_input,
+                              latent_t=0,
+                              action_t=0,
+                              latent_cond=None,
+                              action_cond=None,
+                              frame_st_id=0,
+                              patch_size=(1, 2, 2)):
+        logger.info(f"FRAME START ID: {frame_st_id}")
+        input_dict = dict()
+        if latent_model_input is not None:
+            input_dict['latent_res_lst'] = {
+                'noisy_latents':
+                latent_model_input,
+                'timesteps':
+                torch.ones([latent_model_input.shape[2]],
+                           dtype=torch.float32,
+                           device=self.device) * latent_t,
+                'grid_id':
+                get_mesh_id(latent_model_input.shape[-3] // patch_size[0],
+                            latent_model_input.shape[-2] // patch_size[1],
+                            latent_model_input.shape[-1] // patch_size[2], 0,
+                            1, frame_st_id).to(self.device),
+                'text_emb':
+                self.prompt_embeds.to(self.dtype).clone(),
+            }
+            if latent_cond is not None:
+                input_dict['latent_res_lst'][
+                    'noisy_latents'][:, :, 0:1] = latent_cond[:, :, 0:1]
+                input_dict['latent_res_lst']['timesteps'][0:1] *= 0
+
+        if action_model_input is not None:
+            input_dict['action_res_lst'] = {
+                'noisy_latents':
+                action_model_input,
+                'timesteps':
+                torch.ones([action_model_input.shape[2]],
+                           dtype=torch.float32,
+                           device=self.device) * action_t,
+                'grid_id':
+                get_mesh_id(action_model_input.shape[-3],
+                            action_model_input.shape[-2],
+                            action_model_input.shape[-1],
+                            1,
+                            1,
+                            frame_st_id,
+                            action=True).to(self.device),
+                'text_emb':
+                self.prompt_embeds.to(self.dtype).clone(),
+            }
+
+            if action_cond is not None:
+                input_dict['action_res_lst'][
+                    'noisy_latents'][:, :, 0:1] = action_cond[:, :, 0:1]
+                input_dict['action_res_lst']['timesteps'][0:1] *= 0
+            input_dict['action_res_lst']['noisy_latents'][:, ~self.
+                                                          action_mask] *= 0
+        return input_dict
+
+    def _encode_obs(self, obs):
+        images = obs['obs']
+        if not isinstance(images, list):
+            images = [images]
+        if len(images) < 1:
+            return None
+        videos = []
+        for k_i, k in enumerate(self.job_config.obs_cam_keys):
+            if self.env_type == 'robotwin_tshape':
+                if k_i == 0:  # camera high
+                    height_i, width_i = self.height, self.width
+                else:
+                    height_i, width_i = self.height // 2, self.width // 2
+            else:
+                height_i, width_i = self.height, self.width
+
+            history_video_k = torch.from_numpy(
+                np.stack([each[k]
+                          for each in images])).float().permute(3, 0, 1, 2)
+            history_video_k = F.interpolate(history_video_k,
+                                            size=(height_i, width_i),
+                                            mode='bilinear',
+                                            align_corners=False).unsqueeze(0)
+            videos.append(history_video_k)
+
+        if self.env_type == 'robotwin_tshape':
+            videos_high = videos[0] / 255.0 * 2.0 - 1.0
+            videos_left_and_right = torch.cat(videos[1:],
+                                              dim=0) / 255.0 * 2.0 - 1.0
+            vae_device = next(self.streaming_vae.vae.parameters()).device
+            enc_out_high = self.streaming_vae.encode_chunk(
+                videos_high.to(vae_device).to(self.dtype))
+            enc_out_left_and_right = self.streaming_vae_half.encode_chunk(
+                videos_left_and_right.to(vae_device).to(self.dtype))
+            enc_out = torch.cat([
+                torch.cat(enc_out_left_and_right.split(1, dim=0), dim=-1),
+                enc_out_high
+            ],
+                                dim=-2)
+        else:
+            videos = torch.cat(videos, dim=0) / 255.0 * 2.0 - 1.0
+            vae_device = next(self.streaming_vae.vae.parameters()).device
+            videos_chunk = videos.to(vae_device).to(self.dtype)
+            enc_out = self.streaming_vae.encode_chunk(videos_chunk)
+
+        mu, logvar = torch.chunk(enc_out, 2, dim=1)
+        latents_mean = torch.tensor(self.vae.config.latents_mean).to(mu.device)
+        latents_std = torch.tensor(self.vae.config.latents_std).to(mu.device)
+        mu_norm = self.normalize_latents(mu, latents_mean, 1.0 / latents_std)
+        video_latent = torch.cat(mu_norm.split(1, dim=0), dim=-1)
+        return video_latent.to(self.device)
+
+    def _reset(self, prompt=None):
+        logger.info('Reset.')
+        self.use_cfg = (self.job_config.guidance_scale > 1) or (self.job_config.action_guidance_scale > 1)
+        #### Reset all parameters
+        self.frame_st_id = 0
+        self.init_latent = None
+        #### clean vae and transformer cache
+        self.transformer.clear_cache(self.cache_name)
+        self.streaming_vae.clear_cache()
+
+        self.action_per_frame = self.job_config.action_per_frame
+        self.height, self.width = self.job_config.height, self.job_config.width
+
+        if self.env_type == 'robotwin_tshape':
+            self.latent_height, self.latent_width = (
+                (self.height // 16) * 3) // 2, self.width // 16
+            self.streaming_vae_half.clear_cache()
+        else:
+            self.latent_height, self.latent_width = self.height // 16, self.width // 16 * len(
+                self.job_config.obs_cam_keys)
+
+        patch_size = self.job_config.patch_size
+        latent_token_per_chunk = (self.job_config.frame_chunk_size *
+                                  self.latent_height * self.latent_width) // (
+                                      patch_size[0] * patch_size[1] *
+                                      patch_size[2])
+        action_token_per_chunk = self.job_config.frame_chunk_size * self.action_per_frame
+        self.transformer.create_empty_cache(self.cache_name,
+                                            self.job_config.attn_window,
+                                            latent_token_per_chunk,
+                                            action_token_per_chunk,
+                                            dtype=self.dtype,
+                                            device=self.device,
+                                            batch_size = 2 if self.use_cfg else 1
+                                            )
+
+        self.action_mask = torch.zeros([self.job_config.action_dim]).bool()
+        self.action_mask[self.job_config.used_action_channel_ids] = True
+
+        self.actions_q01 = torch.tensor(self.job_config.norm_stat['q01'],
+                                        dtype=torch.float32).reshape(-1, 1, 1)
+        self.actions_q99 = torch.tensor(self.job_config.norm_stat['q99'],
+                                        dtype=torch.float32).reshape(-1, 1, 1)
+        self.action_norm_method = self.job_config.action_norm_method
+
+        ##### get prompt
+        if prompt is None:
+            self.prompt_embeds = self.negative_prompt_embeds = None
+        else:
+            self.prompt_embeds, self.negative_prompt_embeds = self.encode_prompt(
+                prompt=prompt,
+                negative_prompt=None,
+                do_classifier_free_guidance=self.job_config.guidance_scale > 1,
+                num_videos_per_prompt=1,
+                prompt_embeds=None,
+                negative_prompt_embeds=None,
+                max_sequence_length=512,
+                device=self.device,
+                dtype=self.dtype,
+            )
+
+        self.exp_name = f"{prompt}_{time.strftime('%Y%m%d_%H%M%S')}" if prompt else "default"
+        self.exp_save_root = os.path.join(self.save_root, 'real', self.exp_name)
+        os.makedirs(self.exp_save_root, exist_ok=True)
+        torch.cuda.empty_cache()
+
+    def _infer(self, obs, frame_st_id=0):
+        frame_chunk_size = self.job_config.frame_chunk_size
+        if frame_st_id == 0:
+            init_latent = self._encode_obs(obs)
+            self.init_latent = init_latent
+
+        latents = torch.randn(1,
+                              48,
+                              frame_chunk_size,
+                              self.latent_height,
+                              self.latent_width,
+                              device=self.device,
+                              dtype=self.dtype)
+        actions = torch.randn(1,
+                              self.job_config.action_dim,
+                              frame_chunk_size,
+                              self.action_per_frame,
+                              1,
+                              device=self.device,
+                              dtype=self.dtype)
+
+        video_inference_step = self.job_config.num_inference_steps
+        action_inference_step = self.job_config.action_num_inference_steps
+        video_step = self.job_config.video_exec_step
+
+        self.scheduler.set_timesteps(video_inference_step)
+        self.action_scheduler.set_timesteps(action_inference_step)
+        timesteps = self.scheduler.timesteps
+        action_timesteps = self.action_scheduler.timesteps
+
+        timesteps = F.pad(timesteps, (0, 1), mode='constant', value=0)
+
+        if video_step != -1:
+            timesteps = timesteps[:video_step]
+
+        action_timesteps = F.pad(
+            action_timesteps,
+            (0,
+             1),  # pad 1 element at the end (right side) of the last dimension
+            mode='constant',
+            value=0)
+
+        with (
+                torch.no_grad(),
+        ):
+            # 1. Video Generation Loop
+            for i, t in enumerate(tqdm(timesteps)):
+                last_step = i == len(timesteps) - 1
+                latent_cond = init_latent[:, :, 0:1].to(
+                    self.dtype) if frame_st_id == 0 else None
+                input_dict = self._prepare_latent_input(
+                    latents,
+                    None,
+                    t,
+                    t,
+                    latent_cond,
+                    None,
+                    frame_st_id=frame_st_id)
+
+                video_noise_pred = self.transformer(
+                    self._repeat_input_for_cfg(input_dict['latent_res_lst']),
+                    update_cache=1 if last_step else 0,
+                    cache_name=self.cache_name,
+                    action_mode=False)
+
+                if not last_step or video_step != -1:
+                    video_noise_pred = data_seq_to_patch(
+                        self.job_config.patch_size, video_noise_pred,
+                        frame_chunk_size, self.latent_height,
+                        self.latent_width, batch_size=2 if self.use_cfg else 1)
+                    if self.job_config.guidance_scale > 1:
+                        video_noise_pred = video_noise_pred[1:] + self.job_config.guidance_scale * (video_noise_pred[:1] - video_noise_pred[1:])
+                    else:
+                        video_noise_pred = video_noise_pred[:1]
+                    latents = self.scheduler.step(video_noise_pred,
+                                                  t,
+                                                  latents,
+                                                  return_dict=False)
+
+                latents[:, :, 0:1] = latent_cond if frame_st_id == 0 else latents[:, :, 0:1]
+
+            for i, t in enumerate(tqdm(action_timesteps)):
+                last_step = i == len(action_timesteps) - 1
+                action_cond = torch.zeros(
+                    [
+                        1, self.job_config.action_dim, 1,
+                        self.action_per_frame, 1
+                    ],
+                    device=self.device,
+                    dtype=self.dtype) if frame_st_id == 0 else None
+
+                input_dict = self._prepare_latent_input(
+                    None,
+                    actions,
+                    t,
+                    t,
+                    None,
+                    action_cond,
+                    frame_st_id=frame_st_id)
+                action_noise_pred = self.transformer(
+                    self._repeat_input_for_cfg(input_dict['action_res_lst']),
+                    update_cache=1 if last_step else 0,
+                    cache_name=self.cache_name,
+                    action_mode=True)
+
+                if not last_step:
+                    action_noise_pred = rearrange(action_noise_pred,
+                                                  'b (f n) c -> b c f n 1',
+                                                  f=frame_chunk_size)
+                    if self.job_config.action_guidance_scale > 1:
+                        action_noise_pred = action_noise_pred[1:] + self.job_config.action_guidance_scale * (action_noise_pred[:1] - action_noise_pred[1:])
+                    else:
+                        action_noise_pred = action_noise_pred[:1]
+                    actions = self.action_scheduler.step(action_noise_pred,
+                                                         t,
+                                                         actions,
+                                                         return_dict=False)
+
+                actions[:, :, 0:1] = action_cond if frame_st_id == 0 else actions[:, :, 0:1]
+
+        actions[:, ~self.action_mask] *= 0
+
+        save_async(latents, os.path.join(self.exp_save_root, f'latents_{frame_st_id}.pt'))
+        save_async(actions, os.path.join(self.exp_save_root, f'actions_{frame_st_id}.pt'))
+
+        actions = self.postprocess_action(actions)
+        torch.cuda.empty_cache()
+        return actions, latents
+
+    def _compute_kv_cache(self, obs):
+        ### optional async save obs for debug
+        self.transformer.clear_pred_cache(self.cache_name)
+        save_async(obs['obs'], os.path.join(self.exp_save_root, f'obs_data_{self.frame_st_id}.pt'))
+        latent_model_input = self._encode_obs(obs)
+        if self.frame_st_id == 0:
+            latent_model_input = torch.cat(
+                [self.init_latent, latent_model_input],
+                dim=2) if latent_model_input is not None else self.init_latent
+
+        action_model_input = self.preprocess_action(obs['state'])
+        action_model_input = action_model_input.to(latent_model_input)
+        logger.info(
+            f"get KV cache obs: {latent_model_input.shape} {action_model_input.shape}"
+        )
+        input_dict = self._prepare_latent_input(latent_model_input,
+                                                action_model_input,
+                                                frame_st_id=self.frame_st_id)
+
+        with (
+                torch.no_grad(),
+        ):
+            self.transformer(self._repeat_input_for_cfg(input_dict['latent_res_lst']),
+                             update_cache=2,
+                             cache_name=self.cache_name,
+                             action_mode=False)
+
+            self.transformer(self._repeat_input_for_cfg(input_dict['action_res_lst']),
+                             update_cache=2,
+                             cache_name=self.cache_name,
+                             action_mode=True)
+        torch.cuda.empty_cache()
+        self.frame_st_id += latent_model_input.shape[2]
+
+    @torch.no_grad()
+    def infer(self, obs):
+        reset = obs.get('reset', False)
+        prompt = obs.get('prompt', None)
+        compute_kv_cache = obs.get('compute_kv_cache', False)
+
+        if reset:
+            logger.info(f"******************* Reset server ******************")
+            self._reset(prompt=prompt)
+            return dict()
+        elif compute_kv_cache:
+            logger.info(
+                f"################# Compute KV Cache #################")
+            self._compute_kv_cache(obs)
+            return dict()
+        else:
+            logger.info(f"################# Infer One Chunk #################")
+            action, _ = self._infer(obs, frame_st_id=self.frame_st_id)
+            return dict(action=action)
+    
+    def decode_one_video(self, latents, output_type):
+        latents = latents.to(self.vae.dtype)
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+            latents.device, latents.dtype
+        )
+        latents = latents / latents_std + latents_mean
+        video = self.vae.decode(latents, return_dict=False)[0]
+        video = self.video_processor.postprocess_video(video, output_type=output_type)
+        return video
+    
+    def load_init_obs(self):
+        imf_dict = {v: np.array(Image.open(os.path.join(self.job_config.input_img_path, f"{v}.png")).convert("RGB")) for v in self.job_config.obs_cam_keys}
+        init_obs = {}
+        init_obs['obs'] = [imf_dict]
+        return init_obs
+    
+    @torch.no_grad()
+    def generate(self):
+        self.video_processor = VideoProcessor(vae_scale_factor=1)
+        self._reset(self.job_config.prompt)
+        init_obs = self.load_init_obs()
+        pred_latent_lst = []
+        pred_action_lst = []
+        for chunk_id in range(self.job_config.num_chunks_to_infer):
+            actions, latents = self._infer(init_obs, frame_st_id=(chunk_id * self.job_config.frame_chunk_size))
+            actions = torch.from_numpy(actions)
+            pred_latent_lst.append(latents)
+            pred_action_lst.append(actions)
+        pred_latent = torch.cat(pred_latent_lst, dim=2)
+        pred_action = torch.cat(pred_action_lst, dim=1).flatten(1)
+        self.transformer.clear_cache(self.cache_name)
+        self.streaming_vae.clear_cache()
+        if self.streaming_vae_half:
+            self.streaming_vae_half.clear_cache()
+        del self.transformer
+        del self.streaming_vae_half
+        del self.text_encoder
+        torch.cuda.empty_cache()
+        
+        # Move VAE to GPU for decoding
+        if self.enable_offload:
+            self.vae = self.vae.to(self.device).to(self.dtype)
+        
+        decoded_video = self.decode_one_video(pred_latent, 'np')[0]
+        export_to_video(decoded_video, os.path.join(self.save_root, "demo.mp4"), fps=10)
+
+def run(args):    
+    
+    config = VA_CONFIGS[args.config_name]
+    port = config.port if args.port is None else args.port
+    if args.save_root is not None:
+        config.save_root = args.save_root
+    rank = int(os.getenv("RANK", 0))
+    local_rank = int(os.environ.get('LOCAL_RANK', 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    init_distributed(world_size, local_rank, rank)
+    config.rank = rank
+    config.local_rank = local_rank
+    config.world_size = world_size
+    model = VA_Server(config)
+    if config.infer_mode == 'i2va':
+        logger.info("******************************USE I2VA mode******************************")
+        model.generate()
+    elif config.infer_mode == 'server':
+        logger.info(f"******************************USE Server mode******************************")
+        run_async_server_mode(model, local_rank, config.host, port)
+    else:
+        raise ValueError(f"Unknown infer mode: {config.infer_mode}")
+
+def main():
+    """
+    TODO
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config-name",
+        type=str,
+        required=False,
+        default='robotwin',
+        help="config name.",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=None,
+        help='(start) port'
+    )
+    parser.add_argument(
+        "--save_root",
+        type=str,
+        default=None,
+        help='save root'
+    )
+    args = parser.parse_args()
+    run(args)
+    logger.info("Finish all process!!!!!!!!!!!!")
+
+
+if __name__ == "__main__":
+    init_logger()
+    main()