Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,9 @@ 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/episode_000000.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/episode_000008.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/episode_000087.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/gt/episode_000000.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/gt/episode_000008.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/gt/episode_000087.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,137 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+tags:
+  - robotics
+  - world-model
+  - video-prediction
+  - action-conditioned
+  - nvidia
+  - cosmos
+base_model: nvidia/Cosmos-Predict2.5-2B
+---
+
+# Cosmos-Predict2.5-2B GR1 Action-Conditioned World Model
+
+Action-conditioned video prediction model fine-tuned from [nvidia/Cosmos-Predict2.5-2B](https://huggingface.co/nvidia/Cosmos-Predict2.5-2B) on GR1 dual-arm robot teleoperation data.
+
+Given an initial observation frame and a sequence of robot actions, this model predicts future video frames that depict the robot executing those actions.
+
+## Model Details
+
+- **Base model:** Cosmos-Predict2.5-2B (video2world pre-trained)
+- **Fine-tuning data:** GR1 dual-arm robot teleoperation (LeRobot format)
+- **Action dimension:** 29 (left_arm:7 + left_hand:6 + right_arm:7 + right_hand:6 + waist:3)
+- **Video frames:** 13 frames per chunk (1 conditional + 12 predicted)
+- **Temporal compression:** 4x (13 image frames = 4 latent frames)
+- **Resolution:** 480x832
+- **Checkpoint format:** EMA weights in bf16 (`model_ema_bf16.pt`, ~4GB)
+
+## Training Configuration
+
+| Parameter | Value |
+|-----------|-------|
+| GPUs | 4x H200 (140GB) |
+| Batch size | 4 per GPU (global=16) |
+| Learning rate | 8e-5 |
+| Max iterations | 4,000 |
+| Save interval | 2,000 |
+| Optimizer | AdamW (cosine schedule) |
+| Precision | bf16 |
+| Episodes per task | 3 (uniform sampling) |
+| VAE | Wan2.1 |
+| Text encoder | Cosmos-Reason1-7B |
+| Tokenizer | Qwen2.5-VL-7B-Instruct |
+
+## Inference
+
+### Prerequisites
+
+This model requires the [Cosmos-Predict2.5](https://github.com/NVIDIA/Cosmos-Predict2.5) codebase and the following dependencies:
+
+- Wan2.1 VAE: `Wan2.1_VAE.pth`
+- Text encoder: Cosmos-Reason1-7B
+- Tokenizer: Qwen2.5-VL-7B-Instruct
+
+### Quick Start
+
+```bash
+python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+  --experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes_release_oss \
+  --ckpt_path=/path/to/model_ema_bf16.pt \
+  --input_video_root=/path/to/eval_data \
+  --save_root=/path/to/output \
+  --resolution 480,832 \
+  --guidance 0 \
+  --chunk_size 12 \
+  --fps_downsample_ratio 2 \
+  --save_fps 10 \
+  --vae_path /path/to/Wan2.1_VAE.pth \
+  --text_encoder_path /path/to/Cosmos-Reason1-7B \
+  --qwen_path /path/to/Qwen2.5-VL-7B-Instruct
+```
+
+### Input Data Format
+
+The inference script expects a directory containing paired files:
+- `episode_XXXXXX.mp4` - Input video (first frame used as conditional)
+- `episode_XXXXXX_actions.npy` - Action array of shape `(T, 29)` in numpy format
+
+### Eval Data Preparation
+
+Use `scripts/prepare_gr1_eval_data.py` to extract eval episodes from a LeRobot dataset:
+
+```bash
+python scripts/prepare_gr1_eval_data.py \
+  --dataset-path /path/to/GR1_robot \
+  --output-dir /path/to/eval_data \
+  --num-episodes 100
+```
+
+## Examples
+
+The `examples/` directory contains sample predictions alongside ground truth videos:
+
+| Episode | Prediction | Ground Truth |
+|---------|-----------|--------------|
+| 000000 | [episode_000000.mp4](examples/episode_000000.mp4) | [episode_000000.mp4](examples/gt/episode_000000.mp4) |
+| 000008 | [episode_000008.mp4](examples/episode_000008.mp4) | [episode_000008.mp4](examples/gt/episode_000008.mp4) |
+| 000087 | [episode_000087.mp4](examples/episode_000087.mp4) | [episode_000087.mp4](examples/gt/episode_000087.mp4) |
+
+## File Structure
+
+```
+.
+├── model_ema_bf16.pt           # Model weights (EMA, bf16)
+├── inference/
+│   ├── inference_gr00t.py      # Main inference script
+│   └── inference_pipeline.py   # Inference pipeline (ActionVideo2WorldInference)
+├── scripts/
+│   ├── eval_gr1_robot.sh       # One-command eval script
+│   ├── prepare_gr1_eval_data.py # Eval data preparation
+│   └── convert_distcp_to_pt.py # DCP -> PT checkpoint converter
+└── examples/
+    ├── episode_000000.mp4      # Predicted videos
+    ├── episode_000008.mp4
+    ├── episode_000087.mp4
+    └── gt/                     # Ground truth videos
+        ├── episode_000000.mp4
+        ├── episode_000008.mp4
+        └── episode_000087.mp4
+```
+
+## Limitations
+
+- Trained only on GR1 dual-arm robot data; may not generalize to other embodiments
+- Prediction quality degrades for long-horizon generation (many chunks)
+- Episodes with fewer than 12 actions cannot be processed (chunk_size=12)
+- No classifier-free guidance (guidance=0) works best for this model
+
+## License
+
+Apache-2.0 (same as base Cosmos-Predict2.5)
+
+## Acknowledgements
+
+- Base model: [NVIDIA Cosmos-Predict2.5](https://github.com/NVIDIA/Cosmos-Predict2.5)
+- Training data: GR1 teleoperation dataset (LeRobot format)
+- VAE: [Wan2.1](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B)

examples/episode_000000.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5f14c2d6e41c7eeae80f354d5b094ab7eccbf32620bd45225959f7ab05ff277
+size 199034

examples/episode_000008.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c1b85ee2929a2616e7f2cd45bb3667b70dc4e4382270c4164d3c1074f121439
+size 394188

examples/episode_000087.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1da2555b494710a119c08760aebbf744e4b281874e01a215a66d19590e44e2d1
+size 623808

examples/gt/episode_000000.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd87b160860d6fdd47846eb9839af86d43684c59e640be88d1df1e120a200066
+size 530414

examples/gt/episode_000008.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ac15a966e56f983a29037b57efc5c8f5ad5ec2a026e5e4d07f7b3471aad407c
+size 540443

examples/gt/episode_000087.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec40150f03290095fca08842f89ae4cd18c6b7651053bdfa24f6a83517285aef
+size 1202359

inference/inference_gr00t.py

@@ -0,0 +1,389 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+
+
+# ---------------------------------- benchmark ----------------------------------
+
+
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame/checkpoints/iter_000014000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame-14k\
+  --resolution 480,832 --guidance 0 --chunk_size 12
+
+CUDA_VISIBLE_DEVICES=7 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame/checkpoints/iter_000020000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame-20k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 80 --end 100
+
+CUDA_VISIBLE_DEVICES=1 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame/checkpoints/iter_000028000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame-28k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 0 --end 100
+
+CUDA_VISIBLE_DEVICES=7 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes/checkpoints/iter_000004000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes-4k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 80 --end 100
+
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_49frame_full \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_49frame_full/checkpoints/iter_000004000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_49frame_full-6k\
+  --resolution 480,832 --guidance 0 --chunk_size 48 --start 80 --end 100
+
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_conditioned_posttrained_rl_merged_action_gr00t_gr1_customized_13frame_full \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_conditioned_posttrained_rl_merged_action_gr00t_gr1_customized_13frame_full/checkpoints/iter_000004000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_conditioned_posttrained_rl_merged_action_gr00t_gr1_customized_13frame_full-4k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 0 --end 100
+
+CUDA_VISIBLE_DEVICES=1 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_conditioned_posttrained_rl_merged_action_gr00t_gr1_customized_13frame_full \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_conditioned_posttrained_rl_merged_action_gr00t_gr1_customized_13frame_full/checkpoints/iter_000008000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_conditioned_posttrained_rl_merged_action_gr00t_gr1_customized_13frame_full-8k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 0 --end 100
+
+CUDA_VISIBLE_DEVICES=6 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_49frame_full \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_49frame_full/checkpoints/iter_000010000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_49frame_full-10k\
+  --resolution 480,832 --guidance 0 --chunk_size 48 --start 90 --end 100
+
+CUDA_VISIBLE_DEVICES=6 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_73frame_full \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_73frame_full/checkpoints/iter_000006000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_73frame_full-6k\
+  --resolution 480,832 --guidance 0 --chunk_size 72 --start 70 --end 80
+
+CUDA_VISIBLE_DEVICES=5 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes/checkpoints/iter_000008000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes-8k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 75 --end 100
+
+
+CUDA_VISIBLE_DEVICES=6 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes/checkpoints/iter_000010000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes-10k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 90 --end 100
+
+CUDA_VISIBLE_DEVICES=7 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes/checkpoints/iter_000014000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes-14k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 90 --end 100
+
+
+CUDA_VISIBLE_DEVICES=7 PYTHONPATH=. python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+--experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes \
+  --ckpt_path s3://bucket/cosmos_predict2_action_conditioned/action_conditional/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes/checkpoints/iter_000016000 \
+  --input_video_root results/gr00t_gr1/gt \
+  --save_root results/gr00t_gr1/cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes-16k\
+  --resolution 480,832 --guidance 0 --chunk_size 12 --start 0 --end 100
+"""
+
+import argparse
+import os
+from glob import glob
+
+import mediapy
+import numpy as np
+import torch
+from loguru import logger
+
+from cosmos_predict2._src.imaginaire.utils import distributed
+from cosmos_predict2._src.predict2.action.inference.inference_pipeline import (
+    _DEFAULT_NEGATIVE_PROMPT,
+    ActionVideo2WorldInference,
+)
+
+_IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", "webp"]
+_VIDEO_EXTENSIONS = [".mp4"]
+
+_ACTION_SCALER = 20.0
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parses command-line arguments for the Video2World inference script."""
+    parser = argparse.ArgumentParser(description="Image2World/Video2World inference script")
+    parser.add_argument("--experiment", type=str, required=True, help="Experiment config")
+    parser.add_argument("--chunk_size", type=int, default=12, help="Chunk size for action conditioning")
+    parser.add_argument(
+        "--num_chunks", type=int, default=12, help="Number of chunks to generate (-1 for all available chunks)"
+    )
+    parser.add_argument("--guidance", type=int, default=7, help="Guidance value")
+    parser.add_argument("--seed", type=int, default=1, help="Guidance value")
+    parser.add_argument(
+        "--ckpt_path",
+        type=str,
+        default="",
+        help="Path to the checkpoint. If not provided, will use the one specify in the config",
+    )
+    parser.add_argument("--s3_cred", type=str, default="credentials/s3_checkpoint.secret")
+    parser.add_argument(
+        "--resolution",
+        type=str,
+        default="none",
+        help="Resolution of the video (H,W). Be default it will use model trained resolution. 9:16",
+    )
+    parser.add_argument("--input_video_root", type=str, default="bridge/annotation/test_100", help="Action root")
+    parser.add_argument("--save_root", type=str, default="results/image2world", help="Save root")
+
+    # for pi dataset
+    parser.add_argument("--camera_id", type=str, default="base", help="Camera id")
+    parser.add_argument("--start", type=int, default=0)
+    parser.add_argument("--end", type=int, default=100)
+    parser.add_argument("--fps_downsample_ratio", type=int, default=1)
+    parser.add_argument("--gripper_scale", type=float, default=1.0)
+    parser.add_argument("--gripper_key", type=str, default="continuous_gripper_state", help="Gripper key")
+    parser.add_argument("--state_key", type=str, default="state", help="State key")
+
+    parser.add_argument("--reverse", action="store_true", help="Reverse the video")
+    parser.add_argument("--single_chunk", action="store_true", help="Single chunk")
+    parser.add_argument("--start_frame_idx", type=int, default=0, help="Start frame index")
+    parser.add_argument("--save_fps", type=int, default=10, help="Save fps")
+
+    parser.add_argument(
+        "--negative_prompt",
+        type=str,
+        default=_DEFAULT_NEGATIVE_PROMPT,
+        help="Custom negative prompt for classifier-free guidance. If not specified, uses default embeddings from S3.",
+    )
+    parser.add_argument(
+        "--num_latent_conditional_frames",
+        type=int,
+        default=1,
+        help="Number of latent conditional frames (0, 1 or 2). For images, both values work by duplicating frames. For videos, uses the first N frames.",
+    )
+    # Context parallel arguments
+    parser.add_argument(
+        "--context_parallel_size",
+        type=int,
+        default=1,
+        help="Context parallel size (number of GPUs to split context over). Set to 8 for 8 GPUs",
+    )
+    # Local model path overrides (avoid downloading from HuggingFace/S3)
+    parser.add_argument("--vae_path", type=str, default="", help="Local path to Wan2.1 VAE .pth file")
+    parser.add_argument("--text_encoder_path", type=str, default="", help="Local path to Cosmos-Reason1-7B directory")
+    parser.add_argument("--qwen_path", type=str, default="", help="Local path to Qwen2.5-VL-7B-Instruct directory")
+    parser.add_argument("--experiment_opts", type=str, nargs="*", default=[], help="Additional Hydra override strings")
+    return parser.parse_args()
+
+
+def get_action_sequence_from_states(
+    data,
+    fps_downsample_ratio=1,
+    use_quat=False,
+    state_key="state",
+    gripper_scale=1.0,
+    gripper_key="continuous_gripper_state",
+):
+    """
+    Get the action sequence from the states.
+    """
+
+    actions = np.array(data["action"])[::fps_downsample_ratio][:-1]
+    return actions
+
+
+def get_video_id(img_path: str):
+    """Extract video ID from image path by removing directory and extension."""
+    return img_path.split("/")[-1].split(".")[0]
+
+
+def main():
+    torch.enable_grad(False)  # Disable gradient calculations for inference
+    args = parse_arguments()
+
+    # Validate num_latent_conditional_frames at the very beginning
+    if args.num_latent_conditional_frames not in [0, 1, 2]:
+        raise ValueError(
+            f"num_latent_conditional_frames must be 0, 1 or 2, but got {args.num_latent_conditional_frames}"
+        )
+
+    # Determine supported extensions based on num_latent_conditional_frames
+    if args.num_latent_conditional_frames > 1:
+        supported_extensions = _VIDEO_EXTENSIONS
+        # Check if input folder contains any videos
+        has_videos = False
+        for file_name in os.listdir(args.input_root):
+            file_ext = os.path.splitext(file_name)[1].lower()
+            if file_ext in _VIDEO_EXTENSIONS:
+                has_videos = True
+                break
+
+        if not has_videos:
+            raise ValueError(
+                f"num_latent_conditional_frames={args.num_latent_conditional_frames} > 1 requires video inputs, "
+                f"but no videos found in {args.input_root}. Found extensions: "
+                f"{set(os.path.splitext(f)[1].lower() for f in os.listdir(args.input_root) if os.path.splitext(f)[1])}"
+            )
+
+        logger.info(f"Using video-only mode with {args.num_latent_conditional_frames} conditional frames")
+    elif args.num_latent_conditional_frames == 1:
+        supported_extensions = _IMAGE_EXTENSIONS + _VIDEO_EXTENSIONS
+        logger.info(f"Using image+video mode with {args.num_latent_conditional_frames} conditional frame")
+
+    # Initialize the inference handler with context parallel support
+    # Build experiment_opts from local path overrides
+    experiment_opts = list(args.experiment_opts)
+    if args.vae_path:
+        experiment_opts.append(f"++model.config.tokenizer.vae_pth={args.vae_path}")
+        logger.info(f"Using local VAE: {args.vae_path}")
+    if args.text_encoder_path:
+        experiment_opts.append(f"++model.config.text_encoder_config.ckpt_path={args.text_encoder_path}")
+        logger.info(f"Using local text encoder: {args.text_encoder_path}")
+    if args.qwen_path:
+        experiment_opts.append(f"++model.config.text_encoder_config.model_config.tokenizer.tokenizer_type={args.qwen_path}")
+        logger.info(f"Using local Qwen tokenizer: {args.qwen_path}")
+
+    video2world_cli = ActionVideo2WorldInference(
+        args.experiment,
+        args.ckpt_path,
+        args.s3_cred,
+        context_parallel_size=args.context_parallel_size,
+        experiment_opts=experiment_opts if experiment_opts else None,
+    )
+
+    mem_bytes = torch.cuda.memory_allocated(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    logger.info(f"GPU memory usage after model dcp.load: {mem_bytes / (1024**3):.2f} GB")
+
+    # get input video and annotation path
+    input_video_path = os.path.join(args.input_video_root)
+
+    # Only process files on rank 0 if using distributed processing
+    rank0 = True
+    if args.context_parallel_size > 1:
+        rank0 = distributed.get_rank() == 0
+
+    # pdb.set_trace()
+    video_list = glob(os.path.join(input_video_path, "*.mp4"))
+    input_json_list = [video_path.replace(".mp4", "_actions.npy") for video_path in video_list]
+
+    # Ensure save directory exists
+    os.makedirs(args.save_root, exist_ok=True)
+
+    # Process each file in the input directory
+    for annotation_path, video_path in zip(input_json_list[args.start : args.end], video_list[args.start : args.end]):
+        actions = np.load(annotation_path)
+
+        # Convert camera_id to integer if it's a string and can be converted to an integer
+
+        actions = actions[: len(actions)]
+        video_array = mediapy.read_video(video_path)
+
+        # Resize video_array with arg.resolution if specified
+        if args.resolution != "none":
+            try:
+                h, w = map(int, args.resolution.split(","))
+                video_array = np.stack([mediapy.resize_image(frame, (h, w)) for frame in video_array], axis=0)
+            except Exception as e:
+                logger.warning(f"Failed to resize video to {args.resolution}: {e}")
+
+        img_array = video_array[args.start_frame_idx]
+        # img_name = annotation_path.split("/")[-1].split(".")[0]
+        img_name = video_path.split("/")[-1].split(".")[0]
+
+        frames = [img_array]
+        chunk_video = []
+        video_array = video_array[:: args.fps_downsample_ratio]
+
+        video_name = f"{args.save_root}/{img_name.replace('.jpg', '.mp4')}"
+        chunk_video_name = f"{args.save_root}/{img_name + '.mp4'}"
+        logger.info(f"Saving video to {video_name}")
+        if os.path.exists(chunk_video_name):
+            logger.info(f"Video already exists: {chunk_video_name}")
+            continue
+
+        # Calculate the maximum number of chunks to generate
+        max_chunks = len(actions) // args.chunk_size
+        if args.num_chunks > 0:
+            max_chunks = min(max_chunks, args.num_chunks)
+
+        logger.info(f"Generating {max_chunks} chunks (chunk_size={args.chunk_size}, total_actions={len(actions)})")
+
+        chunk_count = 0
+        for i in range(args.start_frame_idx, len(actions), args.chunk_size):
+            if actions[i : i + args.chunk_size].shape[0] != args.chunk_size:
+                break
+
+            # Check if we've reached the desired number of chunks
+            if args.num_chunks > 0 and chunk_count >= args.num_chunks:
+                logger.info(f"Reached target number of chunks ({args.num_chunks}), stopping generation")
+                break
+
+            logger.info(f"Generating chunk {chunk_count + 1}/{max_chunks}")
+            next_img_array, video_clamped = video2world_cli.step_inference(
+                img_array=img_array,
+                action=actions[i : i + args.chunk_size],
+                guidance=args.guidance,
+                seed=i,
+            )
+            frames.append(next_img_array)
+            img_array = next_img_array
+            chunk_video.append(video_clamped)
+            chunk_count += 1
+
+            if args.single_chunk:
+                break
+
+        if not chunk_video:
+            logger.warning(f"Skipping {img_name}: not enough actions for a full chunk ({len(actions)} < {args.chunk_size})")
+            continue
+
+        chunk_list = [chunk_video[0]] + [chunk_video[i][: args.chunk_size] for i in range(1, len(chunk_video))]
+        chunk_video = np.concatenate(chunk_list, axis=0)
+        if args.single_chunk:
+            chunk_video_name = f"{args.save_root}/{img_name + '_single_chunk.mp4'}"
+        else:
+            # chunk_video_name = f"{args.save_root}/{img_name + '_chunk.mp4'}"
+            chunk_video_name = f"{args.save_root}/{img_name + '.mp4'}"
+        mediapy.write_video(chunk_video_name, chunk_video, fps=args.save_fps)
+
+        # concat_video = np.concatenate([chunk_video, video_array[: chunk_video.shape[0]]], axis=2)
+        # concat_video_name = f"{args.save_root}/{img_name + '_concat.mp4'}"
+        # mediapy.write_video(concat_video_name, concat_video, fps=args.save_fps)
+
+        logger.info(f"Saved video to {chunk_video_name}")
+
+    # Synchronize all processes before cleanup
+    if args.context_parallel_size > 1:
+        torch.distributed.barrier()
+
+    # Clean up distributed resources
+    video2world_cli.cleanup()
+
+
+if __name__ == "__main__":
+    main()

inference/inference_pipeline.py

@@ -0,0 +1,436 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torchvision
+from loguru import logger
+from megatron.core import parallel_state
+
+from cosmos_predict2._src.imaginaire.utils import distributed
+from cosmos_predict2._src.interactive.utils.model_loader import (
+    load_model_from_checkpoint as load_distilled_model_from_checkpoint,
+)
+from cosmos_predict2._src.predict2.inference.get_t5_emb import get_text_embedding
+from cosmos_predict2._src.predict2.utils.model_loader import load_model_from_checkpoint
+
+_DEFAULT_NEGATIVE_PROMPT = "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."
+
+
+class ActionVideo2WorldInference:
+    """
+    Handles the Video2World inference process, including model loading, data preparation,
+    and video generation from an image/video and text prompt. Now supports context parallelism
+    and distilled model inference.
+    """
+
+    def __init__(
+        self,
+        experiment_name: str,
+        ckpt_path: str,
+        s3_credential_path: str,
+        context_parallel_size: int = 1,
+        distilled: bool = False,
+        num_steps: int = 4,
+        experiment_opts: list[str] | None = None,
+    ):
+        """
+        Initializes the Video2WorldInference class.
+
+        Loads the diffusion model and its configuration based on the provided
+        experiment name and checkpoint path. Sets up distributed processing if needed.
+
+        Args:
+            experiment_name (str): Name of the experiment configuration.
+            ckpt_path (str): Path to the model checkpoint (local or S3).
+            s3_credential_path (str): Path to S3 credentials file (if loading from S3).
+            context_parallel_size (int): Number of GPUs for context parallelism.
+            distilled (bool): Whether to load a distilled model (DMD2).
+            num_steps (int): Number of diffusion steps for inference (default 4 for distilled models).
+            experiment_opts (list[str] | None): Hydra-style override strings (e.g. "++model.config.tokenizer.vae_pth=/path").
+        """
+        self.experiment_name = experiment_name
+        self.ckpt_path = ckpt_path
+        self.s3_credential_path = s3_credential_path
+        self.context_parallel_size = context_parallel_size
+        self.process_group = None
+        self.distilled = distilled
+        self.num_steps = num_steps
+        self._experiment_opts = experiment_opts
+
+        # Initialize distributed processing if context parallel size > 1
+        if self.context_parallel_size > 1:
+            self._init_distributed()
+
+        # Choose the appropriate config file and loader based on whether we're loading a distilled model
+        if self.distilled:
+            config_file = "cosmos_predict2/_src/interactive/configs/registry_predict2p5.py"
+            logger.info(f"Loading distilled model with config: {config_file}")
+            # Use the cosmos3 loader for distilled models (DMD2)
+            model, config = load_distilled_model_from_checkpoint(
+                experiment_name=self.experiment_name,
+                s3_checkpoint_dir=self.ckpt_path,
+                config_file=config_file,
+                load_ema_to_reg=True,
+            )
+        else:
+            config_file = "cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py"
+            # Load the model and config using predict2 loader
+            model, config = load_model_from_checkpoint(
+                experiment_name=self.experiment_name,
+                s3_checkpoint_dir=self.ckpt_path,
+                config_file=config_file,
+                load_ema_to_reg=True,
+                experiment_opts=getattr(self, "_experiment_opts", None),
+            )
+
+        # For distilled models, set net_fake_score to None (not needed for inference)
+        if self.distilled and hasattr(model, "net_fake_score"):
+            logger.info("Setting net_fake_score to None for distilled model inference")
+            model.net_fake_score = None
+
+        # Enable context parallel on the model if using context parallelism
+        if self.context_parallel_size > 1:
+            model.net.enable_context_parallel(self.process_group)
+
+        self.model = model
+        self.config = config
+        self.batch_size = 1
+        self.neg_t5_embeddings = None
+
+    def _init_distributed(self):
+        """Initialize distributed processing for context parallelism."""
+
+        # Initialize distributed environment
+        distributed.init()
+
+        # Initialize model parallel states
+        parallel_state.initialize_model_parallel(
+            context_parallel_size=self.context_parallel_size,
+        )
+
+        # Get the process group for context parallel
+        self.process_group = parallel_state.get_context_parallel_group()
+
+        logger.info(f"Initialized context parallel with size {self.context_parallel_size}")
+        logger.info(f"Current rank: {distributed.get_rank()}, World size: {distributed.get_world_size()}")
+
+    def _get_data_batch_input(
+        self,
+        video: torch.Tensor,
+        prompt: str,
+        num_conditional_frames: int = 1,
+        negative_prompt: str = _DEFAULT_NEGATIVE_PROMPT,
+        use_neg_prompt: bool = True,
+    ):
+        """
+        Prepares the input data batch for the diffusion model.
+
+        Constructs a dictionary containing the video tensor, text embeddings,
+        and other necessary metadata required by the model's forward pass.
+        Optionally includes negative text embeddings.
+
+        Args:
+            video (torch.Tensor): The input video tensor (B, C, T, H, W).
+            prompt (str): The text prompt for conditioning.
+            num_conditional_frames (int): Number of conditional frames to use.
+            negative_prompt (str, optional): Custom negative prompt.
+            use_neg_prompt (bool, optional): Whether to include negative prompt embeddings. Defaults to True.
+
+        Returns:
+            dict: A dictionary containing the prepared data batch, moved to the correct device and dtype.
+        """
+        B, C, T, H, W = video.shape
+
+        data_batch = {
+            "dataset_name": "video_data",
+            "video": video,
+            "fps": torch.randint(16, 32, (self.batch_size,)).float(),  # Random FPS (might be used by model)
+            "padding_mask": torch.zeros(self.batch_size, 1, H, W),  # Padding mask (assumed no padding here)
+            "num_conditional_frames": num_conditional_frames,  # Specify number of conditional frames
+        }
+
+        if use_neg_prompt:
+            assert negative_prompt is not None, "Negative prompt is required when use_neg_prompt is True"
+
+        # Compute text embeddings
+        if self.model.text_encoder is not None:
+            data_batch["ai_caption"] = [prompt]
+            data_batch["t5_text_embeddings"] = self.model.text_encoder.compute_text_embeddings_online(
+                data_batch={"ai_caption": [prompt], "images": None},
+                input_caption_key="ai_caption",
+            )
+            if use_neg_prompt:
+                data_batch["neg_t5_text_embeddings"] = self.model.text_encoder.compute_text_embeddings_online(
+                    data_batch={"ai_caption": [negative_prompt], "images": None},
+                    input_caption_key="ai_caption",
+                )
+        else:
+            data_batch["t5_text_embeddings"] = get_text_embedding(prompt)
+            if use_neg_prompt:
+                data_batch["neg_t5_text_embeddings"] = get_text_embedding(negative_prompt)
+
+        # Move tensors to GPU and convert to bfloat16 if they are floating point
+        for k, v in data_batch.items():
+            if isinstance(v, torch.Tensor) and torch.is_floating_point(data_batch[k]):
+                data_batch[k] = v.cuda().to(dtype=torch.bfloat16)
+
+        return data_batch
+
+    def step_inference_with_latents(
+        self,
+        img_array: np.ndarray,
+        action: np.ndarray = None,
+        guidance: int = 3,
+        seed: int = 1,
+        num_latent_conditional_frames: int = 1,
+        query_steps: list[int] = None,
+    ):
+        """
+        Runs a single inference step to generate the next video frame and the full video given an input image and action.
+        Returns intermediate latents for analysis.
+
+        Note: For distilled models, this method falls back to standard inference without latent collection
+        as distilled models use a different sampling process.
+        """
+
+        num_video_frames = action.shape[0] + 1
+
+        img_tensor = torchvision.transforms.functional.to_tensor(img_array).unsqueeze(0)
+        vid_input = torch.cat([img_tensor, torch.zeros_like(img_tensor).repeat(num_video_frames - 1, 1, 1, 1)], dim=0)
+        vid_input = (vid_input * 255.0).to(torch.uint8)  # Convert to uint8 range if needed (might depend on model)
+        vid_input = vid_input.unsqueeze(0).permute(0, 2, 1, 3, 4)  # Add batch dim B=1 and permute
+
+        # Prepare the data batch with text embeddings
+        data_batch = self._get_data_batch_input(
+            vid_input,
+            prompt="",
+            num_conditional_frames=num_latent_conditional_frames,
+            negative_prompt="",
+            use_neg_prompt=False,
+        )
+
+        data_batch["action"] = torch.from_numpy(action).cuda().to(dtype=torch.bfloat16)[None, ...]
+
+        mem_bytes = torch.cuda.memory_allocated(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        logger.info(f"GPU memory usage after getting data_batch: {mem_bytes / (1024**3):.2f} GB")
+
+        # Generate latent samples using the diffusion model
+        if self.distilled:
+            # Distilled model inference: use generate_samples_from_batch (no latent collection)
+            logger.info(f"Running distilled inference with {self.num_steps} steps (latent collection not supported)")
+            sample = self.model.generate_samples_from_batch(
+                data_batch,
+                n_sample=1,
+                seed=seed,
+                num_steps=self.num_steps,
+            )
+            latents_to_save = None  # Distilled models don't support latent collection
+        else:
+            # Teacher model inference with latent collection
+            sample, latents_to_save = self.model.generate_samples_with_latents_from_batch(
+                data_batch,
+                n_sample=1,  # Generate one sample
+                guidance=guidance,
+                seed=seed,  # Fixed seed for reproducibility
+                is_negative_prompt=True,  # Use classifier-free guidance
+                query_steps=query_steps,
+            )
+
+        # Decode the latent sample into a video tensor
+        video = self.model.decode(sample)
+
+        video_normalized = (video - (-1)) / (1 - (-1))
+        video_clamped = (torch.clamp(video_normalized[0], 0, 1) * 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy()
+        next_frame = torch.clamp(video_normalized[0, :, -1, :, :], 0, 1)
+        next_frame = (next_frame * 255).to(torch.uint8).permute(1, 2, 0).cpu().numpy()
+        return next_frame, video_clamped, latents_to_save
+
+    def step_inference(
+        self,
+        img_array: np.ndarray,
+        action: np.ndarray = None,
+        guidance: int = 3,
+        seed: int = 1,
+        num_latent_conditional_frames: int = 1,
+    ):
+        """
+        Runs a single inference step to generate the next video frame and the full video given an input image and action.
+
+        Args:
+            img_array (np.ndarray): Input image as a numpy array (H, W, C), typically the first frame.
+            action (np.ndarray, optional): Action vector to condition the model. Should be shape (action_dim,) or (chunk_size, action_dim).
+            guidance (int, optional): Guidance scale for classifier-free guidance. Default is 3.
+            seed (int, optional): Random seed for reproducibility. Default is 1.
+            num_latent_conditional_frames (int, optional): Number of conditional frames to use for the model. Default is 1.
+
+        Returns:
+            next_frame (np.ndarray): The next predicted frame as a numpy array (H, W, C), uint8.
+            video_clamped (np.ndarray): The generated video as a numpy array (T, H, W, C), uint8.
+        """
+        num_video_frames = action.shape[0] + 1
+
+        img_tensor = torchvision.transforms.functional.to_tensor(img_array).unsqueeze(0)  # (1, H, W, C)
+        vid_input = torch.cat([img_tensor, torch.zeros_like(img_tensor).repeat(num_video_frames - 1, 1, 1, 1)], dim=0)
+        vid_input = (vid_input * 255.0).to(torch.uint8)  # Convert to uint8 range if needed (might depend on model)
+        vid_input = vid_input.unsqueeze(0).permute(0, 2, 1, 3, 4)  # Add batch dim B=1 and permute
+
+        # Prepare the data batch with text embeddings
+        data_batch = self._get_data_batch_input(
+            vid_input,
+            prompt="",
+            num_conditional_frames=num_latent_conditional_frames,
+            negative_prompt="",
+            use_neg_prompt=False,
+        )
+
+        data_batch["action"] = torch.from_numpy(action).cuda().to(dtype=torch.bfloat16)[None, ...]
+
+        mem_bytes = torch.cuda.memory_allocated(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        logger.info(f"GPU memory usage after getting data_batch: {mem_bytes / (1024**3):.2f} GB")
+
+        # Generate latent samples using the diffusion model
+        # For distilled models, use num_steps; for teacher models, use guidance with more steps
+        if self.distilled:
+            # Distilled model inference: use fewer steps, no guidance needed
+            logger.info(f"Running distilled inference with {self.num_steps} steps")
+            sample = self.model.generate_samples_from_batch(
+                data_batch,
+                n_sample=1,
+                seed=seed,
+                num_steps=self.num_steps,
+            )
+        else:
+            # Teacher model inference: use guidance with standard sampling
+            sample = self.model.generate_samples_from_batch(
+                data_batch,
+                n_sample=1,  # Generate one sample
+                guidance=guidance,
+                seed=seed,  # Fixed seed for reproducibility
+                is_negative_prompt=True,  # Use classifier-free guidance
+            )
+
+        # Decode the latent sample into a video tensor
+        video = self.model.decode(sample)
+
+        video_normalized = (video - (-1)) / (1 - (-1))
+        video_clamped = (torch.clamp(video_normalized[0], 0, 1) * 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy()
+        next_frame = torch.clamp(video_normalized[0, :, -1, :, :], 0, 1)
+        next_frame = (next_frame * 255).to(torch.uint8).permute(1, 2, 0).cpu().numpy()
+        return next_frame, video_clamped
+
+    def step_inference_multi_frame(
+        self,
+        video_array: np.ndarray,
+        action: np.ndarray = None,
+        guidance: int = 3,
+        seed: int = 1,
+        num_latent_conditional_frames: int = 2,
+    ):
+        """
+        Runs a single inference step to generate the next video frame and the full video given an input image and action.
+
+        Args:
+            video_array (np.ndarray): Input video as a numpy array (T, H, W, C).
+            action (np.ndarray, optional): Action vector to condition the model. Should be shape (action_dim,) or (chunk_size, action_dim).
+
+            guidance (int, optional): Guidance scale for classifier-free guidance. Default is 3.
+            seed (int, optional): Random seed for reproducibility. Default is 1.
+            num_latent_conditional_frames (int, optional): Number of conditional frames to use for the model. Default is 1.
+
+        Returns:
+            next_frame (np.ndarray): The next predicted frame as a numpy array (H, W, C), uint8.
+            video_clamped (np.ndarray): The generated video as a numpy array (T, H, W, C), uint8.
+        """
+        num_video_frames = action.shape[0] + 1 + (num_latent_conditional_frames - 1) * 4
+        num_cond_image_frames = (num_latent_conditional_frames - 1) * 4 + 1
+
+        assert num_cond_image_frames == video_array.shape[0], (
+            "Number of conditional frames is not equal to the number of frames in the video"
+        )
+        assert action.shape[0] == num_video_frames - num_cond_image_frames, (
+            "Number of action frames is not equal to the number of frames in the video"
+        )
+
+        video_tensor = torch.stack(
+            [torchvision.transforms.functional.to_tensor(v) for v in video_array]
+        )  # (T, C, H, W)
+        vid_input = torch.cat(
+            [
+                video_tensor,
+                torch.zeros_like(video_tensor[0][None, ...]).repeat(num_video_frames - num_cond_image_frames, 1, 1, 1),
+            ],
+            dim=0,
+        )
+        vid_input = (vid_input * 255.0).to(torch.uint8)  # Convert to uint8 range if needed (might depend on model)
+        vid_input = vid_input.unsqueeze(0).permute(0, 2, 1, 3, 4)  # Add batch dim B=1 and permute
+
+        # Prepare the data batch with text embeddings
+        data_batch = self._get_data_batch_input(
+            vid_input,
+            prompt="",
+            num_conditional_frames=num_latent_conditional_frames,
+            negative_prompt="",
+            use_neg_prompt=False,
+        )
+
+        zero_action = np.zeros(
+            (4 * (num_latent_conditional_frames - 1), action.shape[1])
+        )  # (4 * (num_latent_conditional_frames-1), action_dim)
+        action_padded = np.concatenate([zero_action, action], axis=0)
+
+        data_batch["action"] = torch.from_numpy(action_padded).cuda().to(dtype=torch.bfloat16)[None, ...]
+
+        mem_bytes = torch.cuda.memory_allocated(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        logger.info(f"GPU memory usage after getting data_batch: {mem_bytes / (1024**3):.2f} GB")
+
+        # Generate latent samples using the diffusion model
+        # For distilled models, use num_steps; for teacher models, use guidance with more steps
+        if self.distilled:
+            # Distilled model inference: use fewer steps, no guidance needed
+            logger.info(f"Running distilled inference with {self.num_steps} steps")
+            sample = self.model.generate_samples_from_batch(
+                data_batch,
+                n_sample=1,
+                seed=seed,
+                num_steps=self.num_steps,
+            )
+        else:
+            # Teacher model inference: use guidance with standard sampling
+            sample = self.model.generate_samples_from_batch(
+                data_batch,
+                n_sample=1,  # Generate one sample
+                guidance=guidance,
+                seed=seed,  # Fixed seed for reproducibility
+                is_negative_prompt=True,  # Use classifier-free guidance
+            )
+
+        # Decode the latent sample into a video tensor
+        video = self.model.decode(sample)
+
+        video_normalized = (video - (-1)) / (1 - (-1))
+        video_clamped = (torch.clamp(video_normalized[0], 0, 1) * 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy()
+        next_frame = torch.clamp(video_normalized[0, :, -1, :, :], 0, 1)
+        next_frame = (next_frame * 255).to(torch.uint8).permute(1, 2, 0).cpu().numpy()
+        return next_frame, video_clamped
+
+    def cleanup(self):
+        """Clean up distributed resources."""
+        if self.context_parallel_size > 1:
+            if parallel_state.is_initialized():
+                parallel_state.destroy_model_parallel()
+            dist.destroy_process_group()

model_ema_bf16.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94d4a1491a58ae392e4ef7059589f1295503063c9f5407091e2e66544f598e0a
+size 4256645510

scripts/convert_distcp_to_pt.py

@@ -0,0 +1,134 @@
+#!/usr/bin/env -S uv run --script
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "numpy",
+#   "s5cmd",
+#   "torch",
+#   "tyro",
+# ]
+# [tool.uv.sources]
+# torch = [{ index = "pytorch" }]
+# [[tool.uv.index]]
+# name = "pytorch"
+# url = "https://download.pytorch.org/whl/cpu"
+# explicit = true
+# ///
+
+"""Download distributed checkpoint from S3 and convert to pytorch checkpoint.
+
+Usage:
+
+```python
+./scripts/convert_distcp_to_pt.py "s3://bucket/cosmos_predict2_multiview/cosmos2_mv/buttercup_predict2p5_2b_mv_7views_res720p_fps30_t8_from16kfps10mv_jointalpamayov2mads720pmulticaps29frames-0/checkpoints/iter_000028000" "checkpoints/buttercup_predict2p5_2b_mv_7views_res720p_fps30_t8_from16kfps10mv_jointalpamayov2mads720pmulticaps29frames-0_iter_000028000"
+```
+"""
+
+import shlex
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import tyro
+from torch.distributed.checkpoint.format_utils import dcp_to_torch_save
+
+
+@dataclass(frozen=True)
+class Args:
+    input_dir: tyro.conf.Positional[str]
+    """S3 URI of the checkpoint or path to the distcp directory."""
+    output_dir: tyro.conf.Positional[Path]
+    """Output directory to save the converted checkpoints."""
+
+    ema: bool = True
+    """Export EMA weights."""
+
+    s3_args: str | None = None
+    """Additional arguments to pass to s5cmd."""
+
+
+def main():
+    args = tyro.cli(Args, description=__doc__)
+
+    pt_path = args.output_dir / "model.pt"
+    pt_path.unlink(missing_ok=True)
+    pt_ema_fp32_path = args.output_dir / "model_ema_fp32.pt"
+    pt_ema_fp32_path.unlink(missing_ok=True)
+    pt_ema_bf16_path = args.output_dir / "model_ema_bf16.pt"
+    pt_ema_bf16_path.unlink(missing_ok=True)
+
+    if args.input_dir.startswith("s3://"):
+        input_s3 = args.input_dir.rstrip("/")
+        input_s3 = input_s3.removesuffix("/model")
+        distcp_dir = args.output_dir / "model"
+        print(f"Downloading distcp to {distcp_dir}...")
+        # Create the directory if it doesn't exist
+        distcp_dir.mkdir(parents=True, exist_ok=True)
+        # Use sync only if directory exists and has files, otherwise use cp
+        cmd = ["s5cmd"]
+        if args.s3_args:
+            cmd.extend(shlex.split(args.s3_args))
+        if distcp_dir.exists() and any(distcp_dir.iterdir()):
+            cmd.extend(["sync", "--exit-on-error"])
+        else:
+            cmd.extend(["cp", "--show-progress"])
+        cmd.extend(
+            [
+                f"{input_s3}/model/*",
+                f"{distcp_dir}",
+            ]
+        )
+        print(shlex.join(cmd))
+        subprocess.run(cmd, check=True)
+        print(f"Downloaded distcp to '{distcp_dir}'")
+    else:
+        distcp_dir = Path(args.input_dir)
+
+    # Convert distributed checkpoint to torch single checkpoint
+    dcp_to_torch_save(distcp_dir, pt_path)
+    print(f"Converted '{distcp_dir}' to '{pt_path}'")
+
+    if not args.ema:
+        return
+
+    # Drop Reg keys and save EMA weights only in fp32 precision
+    state_dict: dict[str, Any] = torch.load(pt_path, map_location="cpu", weights_only=False)
+    state_dict_ema_fp32: dict[str, Any] = {}
+    for key, value in state_dict.items():
+        if key.startswith("net_ema."):
+            key = key.replace("net_ema.", "net.")
+            state_dict_ema_fp32[key] = value
+    if not state_dict_ema_fp32:
+        raise ValueError("Model doesn't contain EMA weights")
+    torch.save(state_dict_ema_fp32, pt_ema_fp32_path)
+    print(f"Saved EMA fp32 weights from '{pt_path}' to '{pt_ema_fp32_path}'")
+
+    # Save EMA weights only in bf16 precision
+    state_dict_ema_bf16: dict[str, Any] = {}
+    for key, value in state_dict_ema_fp32.items():
+        if isinstance(value, torch.Tensor) and value.dtype == torch.float32:
+            value = value.bfloat16()
+        state_dict_ema_bf16[key] = value
+    torch.save(state_dict_ema_bf16, pt_ema_bf16_path)
+    print(f"fp32 -> bf16: '{pt_ema_fp32_path}' to '{pt_ema_bf16_path}'")
+
+
+if __name__ == "__main__":
+    main()

scripts/eval_gr1_robot.sh

@@ -0,0 +1,245 @@
+#!/bin/bash
+# =============================================================================
+# Cosmos-Predict2.5 Action-Conditioned Evaluation on GR1 Robot Data
+# =============================================================================
+#
+# 一键评估：指定模型路径（DCP 目录或 _ema_bf16.pt），自动完成转换 + 推理。
+#
+# 使用方式 (在 GPU 机器上):
+#   bash scripts/eval_gr1_robot.sh /path/to/checkpoint
+#   bash scripts/eval_gr1_robot.sh /path/to/model_ema_bf16.pt
+#   bash scripts/eval_gr1_robot.sh /path/to/iter_000002000         # DCP 目录
+#   bash scripts/eval_gr1_robot.sh /path/to/iter_000002000 --debug
+#   bash scripts/eval_gr1_robot.sh /path/to/model_ema_bf16.pt --num-episodes 50
+#
+# 选项:
+#   --num-episodes N        评估 episode 数量 (default: 100)
+#   --output-dir PATH       输出目录
+#   --guidance N            CFG guidance 值 (default: 0)
+#   --num-gpus N            使用 GPU 数量 (default: 1)
+#   --experiment NAME       Hydra 实验名 (default: 自动推断)
+#   --debug                 调试模式 (5 episodes, 1 GPU)
+#   -h, --help              显示帮助
+# =============================================================================
+
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# 默认配置
+# ---------------------------------------------------------------------------
+NUM_EPISODES=100
+GUIDANCE=0
+OUTPUT_ROOT="/inspire/ssd/project/security-defense-and-attack/25015/world_model/output"
+DATASET_PATH="/inspire/ssd/project/security-defense-and-attack/25015/world_model/data/PhysicalAI-Robotics-GR00T-Teleop-GR1/GR1_robot"
+COSMOS_ROOT="/inspire/ssd/project/security-defense-and-attack/25015/cosmos-predict2.5"
+NUM_GPUS=1
+DEBUG=0
+EXPERIMENT_NAME="cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes_release_oss"
+CHECKPOINT_INPUT=""
+
+# ---------------------------------------------------------------------------
+# 解析参数 (第一个位置参数为 checkpoint 路径)
+# ---------------------------------------------------------------------------
+if [ $# -eq 0 ]; then
+    echo "用法: bash scripts/eval_gr1_robot.sh <checkpoint_path> [选项]"
+    echo ""
+    echo "checkpoint_path 可以是:"
+    echo "  - DCP 目录 (如 .../checkpoints/iter_000002000)"
+    echo "  - _ema_bf16.pt 文件 (如 .../iter_000002000/model_ema_bf16.pt)"
+    echo ""
+    sed -n '3,18p' "$0"
+    exit 1
+fi
+
+CHECKPOINT_INPUT="$1"
+shift
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --num-episodes)    NUM_EPISODES="$2"; shift 2 ;;
+        --output-dir)      OUTPUT_ROOT="$2"; shift 2 ;;
+        --guidance)        GUIDANCE="$2"; shift 2 ;;
+        --num-gpus)        NUM_GPUS="$2"; shift 2 ;;
+        --experiment)      EXPERIMENT_NAME="$2"; shift 2 ;;
+        --debug)           DEBUG=1; shift ;;
+        -h|--help)
+            sed -n '3,18p' "$0"
+            exit 0
+            ;;
+        *) echo "Unknown option: $1"; exit 1 ;;
+    esac
+done
+
+# 调试模式覆盖
+if [ "$DEBUG" -eq 1 ]; then
+    NUM_EPISODES=5
+    NUM_GPUS=1
+    echo "=== DEBUG MODE ==="
+fi
+
+# ---------------------------------------------------------------------------
+# Step 1: 确定 _ema_bf16.pt 路径 (自动转换 DCP → PT)
+# ---------------------------------------------------------------------------
+if [ ! -e "$CHECKPOINT_INPUT" ]; then
+    echo "ERROR: 路径不存在: ${CHECKPOINT_INPUT}"
+    exit 1
+fi
+
+if [ -f "$CHECKPOINT_INPUT" ]; then
+    # 已经是 .pt 文件，直接使用
+    CHECKPOINT_PATH="$CHECKPOINT_INPUT"
+    echo "使用已有 PT checkpoint: ${CHECKPOINT_PATH}"
+else
+    # 是目录，需要找到或转换 _ema_bf16.pt
+    CKPT_DIR="$CHECKPOINT_INPUT"
+
+    # 检查是否已有 _ema_bf16.pt
+    if [ -f "${CKPT_DIR}/model_ema_bf16.pt" ]; then
+        CHECKPOINT_PATH="${CKPT_DIR}/model_ema_bf16.pt"
+        echo "找到已有 PT checkpoint: ${CHECKPOINT_PATH}"
+    elif [ -d "${CKPT_DIR}/model" ]; then
+        # DCP 格式，需要转换
+        echo "DCP checkpoint，开始转换..."
+        echo "  输入: ${CKPT_DIR}/model/"
+        echo "  输出: ${CKPT_DIR}/model_ema_bf16.pt"
+        echo ""
+
+        cd "${COSMOS_ROOT}"
+        python scripts/convert_distcp_to_pt.py \
+            "${CKPT_DIR}/model" \
+            "${CKPT_DIR}"
+
+        if [ -f "${CKPT_DIR}/model_ema_bf16.pt" ]; then
+            CHECKPOINT_PATH="${CKPT_DIR}/model_ema_bf16.pt"
+            echo ""
+            echo "转换成功: ${CHECKPOINT_PATH}"
+        else
+            echo "ERROR: DCP → PT 转换失败"
+            exit 1
+        fi
+    else
+        echo "ERROR: 目录中既没有 model_ema_bf16.pt 也没有 model/ 子目录"
+        echo "  内容: $(ls "$CKPT_DIR")"
+        exit 1
+    fi
+fi
+
+# ---------------------------------------------------------------------------
+# Step 2: 准备评估数据
+# ---------------------------------------------------------------------------
+EVAL_DATA_DIR="${OUTPUT_ROOT}/eval/gr1_eval_data"
+
+if [ ! -d "${EVAL_DATA_DIR}" ] || [ -z "$(ls -A "${EVAL_DATA_DIR}" 2>/dev/null)" ]; then
+    echo ""
+    echo "准备评估数据..."
+    cd "${COSMOS_ROOT}"
+    python scripts/prepare_gr1_eval_data.py \
+        --dataset-path "${DATASET_PATH}" \
+        --output-dir "${EVAL_DATA_DIR}" \
+        --num-episodes "${NUM_EPISODES}"
+else
+    echo "评估数据已存在: ${EVAL_DATA_DIR}"
+fi
+
+# ---------------------------------------------------------------------------
+# Step 3: 设置环境变量
+# ---------------------------------------------------------------------------
+export IMAGINAIRE_OUTPUT_ROOT="${OUTPUT_ROOT}"
+export HF_HOME="${HF_HOME:-${OUTPUT_ROOT}/hf_cache}"
+
+# 设置 HF 离线模式，防止尝试下载
+export HF_HUB_OFFLINE=1
+export TRANSFORMERS_OFFLINE=1
+
+# ---------------------------------------------------------------------------
+# Step 4: 运行推理
+# ---------------------------------------------------------------------------
+cd "${COSMOS_ROOT}"
+
+PREDICT_DIR="${OUTPUT_ROOT}/eval/gr1_predicted"
+mkdir -p "${PREDICT_DIR}"
+
+# 日志文件
+LOG_DIR="${OUTPUT_ROOT}/logs"
+mkdir -p "${LOG_DIR}"
+LOG_FILE="${LOG_DIR}/eval_$(date +%Y%m%d_%H%M%S).log"
+
+echo ""
+echo "==================================================================="
+echo "Cosmos-Predict2.5  Action-Conditioned Evaluation  (GR1 Robot)"
+echo "==================================================================="
+echo "  Checkpoint:    ${CHECKPOINT_PATH}"
+echo "  Experiment:    ${EXPERIMENT_NAME}"
+echo "  评估数据:      ${EVAL_DATA_DIR}"
+echo "  预测输出:      ${PREDICT_DIR}"
+echo "  GPU 数量:      ${NUM_GPUS}"
+echo "  Episode 数:    ${NUM_EPISODES}"
+echo "  Guidance:      ${GUIDANCE}"
+echo "  日志:          ${LOG_FILE}"
+echo "==================================================================="
+
+# 统计输入文件数
+NUM_MP4=$(ls "${EVAL_DATA_DIR}"/*.mp4 2>/dev/null | wc -l)
+echo "  输入视频数:    ${NUM_MP4}"
+echo ""
+
+# 设置 CUDA 可见设备
+if [ "$NUM_GPUS" -eq 1 ]; then
+    CUDA_DEVICES="0"
+else
+    CUDA_DEVICES=$(seq -s, 0 $((NUM_GPUS - 1)))
+fi
+
+# ---------------------------------------------------------------------------
+# 本地模型路径 (与训练脚本一致，避免从 HF/S3 下载)
+# ---------------------------------------------------------------------------
+VAE_PATH="/inspire/ssd/project/security-defense-and-attack/25015/world_model/Wan2.1-VAE/Wan2.1_VAE.pth"
+REASON1_PATH="/inspire/ssd/project/security-defense-and-attack/25015/world_model/Cosmos-Reason1-7B"
+QWEN_PATH="/inspire/ssd/project/security-defense-and-attack/25015/models/Qwen2.5-VL-7B-Instruct"
+
+# 构建推理命令的本地路径参数
+LOCAL_PATH_ARGS=""
+if [ -f "$VAE_PATH" ]; then
+    LOCAL_PATH_ARGS="${LOCAL_PATH_ARGS} --vae_path ${VAE_PATH}"
+    echo "  使用本地 VAE: ${VAE_PATH}"
+else
+    echo "  WARNING: 本地 VAE 不存在: ${VAE_PATH}"
+fi
+if [ -d "$REASON1_PATH" ]; then
+    LOCAL_PATH_ARGS="${LOCAL_PATH_ARGS} --text_encoder_path ${REASON1_PATH}"
+    echo "  使用本地 Cosmos-Reason1-7B: ${REASON1_PATH}"
+else
+    echo "  WARNING: 本地 Cosmos-Reason1-7B 不存在: ${REASON1_PATH}"
+fi
+if [ -d "$QWEN_PATH" ]; then
+    LOCAL_PATH_ARGS="${LOCAL_PATH_ARGS} --qwen_path ${QWEN_PATH}"
+    echo "  使用本地 Qwen2.5-VL-7B-Instruct: ${QWEN_PATH}"
+else
+    echo "  WARNING: 本地 Qwen2.5-VL-7B-Instruct 不存在: ${QWEN_PATH}"
+fi
+
+echo ""
+echo "启动推理..."
+
+CUDA_VISIBLE_DEVICES=${CUDA_DEVICES} PYTHONPATH=. python \
+    cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
+    --experiment="${EXPERIMENT_NAME}" \
+    --ckpt_path="${CHECKPOINT_PATH}" \
+    --input_video_root="${EVAL_DATA_DIR}" \
+    --save_root="${PREDICT_DIR}" \
+    --resolution 480,832 \
+    --guidance ${GUIDANCE} \
+    --chunk_size 12 \
+    --start 0 \
+    --end ${NUM_EPISODES} \
+    --fps_downsample_ratio 2 \
+    --save_fps 10 \
+    ${LOCAL_PATH_ARGS} \
+    2>&1 | tee "${LOG_FILE}"
+
+echo ""
+echo "==================================================================="
+echo "评估完成!"
+echo "  预测视频: ${PREDICT_DIR}"
+echo "  日志:     ${LOG_FILE}"
+echo "==================================================================="

scripts/prepare_gr1_eval_data.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""
+Prepare GR1 robot evaluation data for Cosmos-Predict2.5 inference.
+
+Extracts episodes from the LeRobot dataset and creates paired
+  episode_XXXXXX.mp4  +  episode_XXXXXX_actions.npy
+files in the format expected by inference_gr00t.py.
+
+Usage:
+  python scripts/prepare_gr1_eval_data.py
+  python scripts/prepare_gr1_eval_data.py --num-episodes 50 --start-episode 100
+  python scripts/prepare_gr1_eval_data.py --output-dir /path/to/output
+"""
+
+import argparse
+import json
+import os
+
+import numpy as np
+import pyarrow.parquet as pq
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare GR1 eval data for inference")
+    parser.add_argument(
+        "--dataset-path",
+        default="/inspire/ssd/project/security-defense-and-attack/25015/world_model/data/PhysicalAI-Robotics-GR00T-Teleop-GR1/GR1_robot",
+    )
+    parser.add_argument("--output-dir", default=None, help="Output directory (default: {dataset_path}/../eval/gr1_eval_data)")
+    parser.add_argument("--num-episodes", type=int, default=100, help="Number of episodes to extract")
+    parser.add_argument("--start-episode", type=int, default=0, help="First episode index")
+    parser.add_argument("--fps-downsample-ratio", type=int, default=2, help="Downsample ratio (20fps -> 10fps)")
+    parser.add_argument("--action-dim", type=int, default=29, help="Action dimension (29 for GR1 selected keys)")
+    args = parser.parse_args()
+
+    dataset_path = args.dataset_path
+    output_dir = args.output_dir or os.path.join(os.path.dirname(dataset_path), "eval", "gr1_eval_data")
+    os.makedirs(output_dir, exist_ok=True)
+
+    # GR1 action key selection: left_arm(0:7), left_hand(7:13), right_arm(22:29), right_hand(29:35), waist(41:44)
+    action_indices = list(range(0, 7)) + list(range(7, 13)) + list(range(22, 29)) + list(range(29, 35)) + list(range(41, 44))
+
+    # Read episodes metadata
+    episodes_path = os.path.join(dataset_path, "meta", "episodes.jsonl")
+    episodes = []
+    with open(episodes_path) as f:
+        for line in f:
+            if line.strip():
+                episodes.append(json.loads(line))
+
+    total_episodes = len(episodes)
+    end_episode = min(args.start_episode + args.num_episodes, total_episodes)
+    selected = episodes[args.start_episode : end_episode]
+
+    print(f"Dataset: {dataset_path}")
+    print(f"Total episodes: {total_episodes}")
+    print(f"Selected: episodes {args.start_episode} - {end_episode - 1} ({len(selected)} episodes)")
+    print(f"Output: {output_dir}")
+    print(f"Action dim: {args.action_dim} (from 44-dim raw, indices {action_indices})")
+    print(f"FPS downsample ratio: {args.fps_downsample_ratio}")
+    print()
+
+    # Read info.json for chunk size
+    with open(os.path.join(dataset_path, "meta", "info.json")) as f:
+        info = json.load(f)
+    chunk_size = info.get("chunks_size", 1000)  # episodes per chunk
+
+    skipped = 0
+    for i, ep in enumerate(selected):
+        ep_idx = ep["episode_index"]
+        chunk_idx = ep_idx // chunk_size
+
+        # Video path
+        video_key = "observation.images.ego_view_freq20"
+        video_src = os.path.join(
+            dataset_path,
+            "videos",
+            f"chunk-{chunk_idx:03d}",
+            video_key,
+            f"episode_{ep_idx:06d}.mp4",
+        )
+
+        # Parquet path
+        parquet_path = os.path.join(
+            dataset_path,
+            "data",
+            f"chunk-{chunk_idx:03d}",
+            f"episode_{ep_idx:06d}.parquet",
+        )
+
+        if not os.path.exists(video_src):
+            print(f"  [SKIP] ep {ep_idx}: video not found: {video_src}")
+            skipped += 1
+            continue
+        if not os.path.exists(parquet_path):
+            print(f"  [SKIP] ep {ep_idx}: parquet not found: {parquet_path}")
+            skipped += 1
+            continue
+
+        # Symlink video
+        video_dst = os.path.join(output_dir, f"episode_{ep_idx:06d}.mp4")
+        if os.path.islink(video_dst) or os.path.exists(video_dst):
+            os.remove(video_dst)
+        os.symlink(os.path.abspath(video_src), video_dst)
+
+        # Extract actions from parquet
+        table = pq.read_table(parquet_path)
+        actions_raw = np.array(table.column("action").to_pylist())  # (T, 44)
+
+        # Select 29 dims and downsample
+        actions_selected = actions_raw[:, action_indices]  # (T, 29)
+        actions_downsampled = actions_selected[:: args.fps_downsample_ratio][:-1]  # drop last, matching inference_gr00t.py
+
+        # Save NPY
+        npy_path = os.path.join(output_dir, f"episode_{ep_idx:06d}_actions.npy")
+        np.save(npy_path, actions_downsampled)
+
+        if (i + 1) % 20 == 0 or i == 0:
+            print(f"  [{i + 1}/{len(selected)}] ep {ep_idx}: {actions_downsampled.shape[0]} steps, {actions_downsampled.shape[1]} dims")
+
+    print()
+    print(f"Done. {len(selected) - skipped} episodes prepared, {skipped} skipped.")
+    print(f"Output directory: {output_dir}")
+
+
+if __name__ == "__main__":
+    main()