CoMind single-ego 2-view training code + README (what to change)

README.md

@@ -0,0 +1,133 @@
+# Training code — CoMind 2-view generation (single-ego setup)
+
+This bundle is our current 2-view egocentric video training stack (Cosmos-Predict2.5, `predict2_multiview`).
+It is set up for a **2-view actor-actor** model with a lot of cross-view conditioning. You are going to
+**strip it down to a single-ego generation setup that still outputs 2 views**, trained **from scratch on the
+CoMind dataset**. This README lists exactly what to change and what to watch out for.
+
+--------------------------------------------------------------------------------
+## 0. What the current model conditions on (the parts you will REMOVE / keep)
+
+Per view, each generated clip is conditioned on:
+| Signal | Key | What it is | Single-ego plan |
+|---|---|---|---|
+| warped RGB | `control_input_warped` | RGB warped into the target view from a **combined pool (both views' past)** | **KEEP the channel, but change the DATA**: warp only from **that view's own past context** |
+| pose | `control_input_pose` | skeletons. `person_pose=True` = identity-colored, **both people** visible | **KEEP the channel, change the DATA**: render **only the wearer's own** pose |
+| in-context references | `reference_frames` (+ `reference_pose`, `reference_cam_w2c`) | clean appearance/pose anchor frames appended as extra tokens | **REMOVE entirely** |
+| camera Plücker | `plucker_map` / `reference_plucker` | per-pixel rays in a **shared cross-view canonical frame** | **REMOVE** (its only purpose is cross-view alignment; single-ego doesn't need it) |
+| view embedding | net `view_embeddings` | per-view identity added to tokens | **REMOVE** (drop `concat_view_embedding`) |
+| cross-view self-attention | token layout `B (V·t) …` | both views attend each other | **Leave as-is** — you do NOT need to separate it; unified attention is fine, and with all cross-view conditioning removed the two views are effectively independent anyway |
+
+So the single-ego model keeps only: **warped_cond (own past) + own pose**, generated per view, no refs / no
+Plücker / no view-embedding, trained from the Cosmos base 2B.
+
+--------------------------------------------------------------------------------
+## 1. Code changes (all in `cosmos_predict2/_src/predict2_multiview/`)
+
+Make a new experiment function in
+`configs/vid2vid/experiment/nymeria_pose_2actor.py` (copy `comind_actoractor_personpose_shared` as a start),
+and flip these toggles. The flags exist already — you're just turning conditioning OFF.
+
+### 1a. model config (`models/multiview_pose_model_rectified_flow.py` fields, set in the experiment)
+```python
+cfg["model"]["config"]["num_reference_frames"]      = 0       # was 4
+cfg["model"]["config"]["enable_reference_plucker"]  = False   # was True
+cfg["model"]["config"]["enable_reference_pose"]     = False   # was True
+cfg["model"]["config"]["enable_plucker"]            = False   # drop cross-view camera Plücker
+```
+
+### 1b. net config (`networks/multiview_pose_dit.py`, set via `cfg["model"]["config"]["net"].update(...)`)
+```python
+cfg["model"]["config"]["net"].update(
+    enable_reference_frames = False,
+    num_reference_frames    = 0,
+    shared_reference        = False,
+    enable_reference_plucker= False,
+    enable_reference_pose   = False,
+    enable_plucker          = False,
+    concat_view_embedding   = False,   # <-- drops the view embedding
+)
+```
+`pose_mode` stays `"vae_concat"` (pose is VAE-encoded and concatenated into `cond_embedder`; keep it). The
+`cond_embedder` in-channels auto-compute as `warped_latent(16) + visibility(1) + pose_latent(16)` — leave that.
+
+### 1c. checkpoint = from scratch
+```python
+cfg["checkpoint"]["load_path"]      = _base_2b_multiview_ckpt()   # Cosmos base 2B, NOT a warm-start
+cfg["checkpoint"]["strict_resume"]  = False
+```
+Because refs/Plücker/view-emb are gone and the pose/warp embedders are zero-init or freshly-shaped, **iter-1
+loss will start HIGH** (~base video prior) — that's correct for a from-scratch run, unlike a warm-start.
+
+### 1d. data loader
+Point `override /data_train` and `override /data_val` at a CoMind data config that:
+- reads **your regenerated** CoMind clips (own-past warp + own-only pose),
+- sets `num_reference_frames=0`, `shared_reference=False`, `emit_reference_pose=False` (do NOT emit any
+  `reference_*`), `person_pose=` your choice (you're rendering own-only pose, so the identity-color flag is moot).
+
+The CoMind loader lives in `datasets/comind_pairs.py` / the actor-actor path in `datasets/nymeria_pairs.py`
+(`get_nymeria_actor_actor_loader`, class `NymeriaActorActorDataset`). Register a new `cs.store(...)` data
+config with the flags above.
+
+### 1e. validation viz
+`callbacks/nymeria_validation_viz.py`: set `num_reference_frames=0, shared_reference=False,
+emit_reference_pose=False` in the viz kwargs too, or the viz dataset build will look for refs that aren't there.
+
+--------------------------------------------------------------------------------
+## 2. DATA you must regenerate (this is the real work — NOT in this code bundle)
+
+The training code just consumes `clip.npz` per view. For the single-ego setup you must rebuild the CoMind
+`clip.npz` so that:
+
+1. **`warped_cond`** = warp into the target view using **only that view's OWN past frames** as the source pool.
+   The current pipeline uses a **combined pool** (both views' past, max-coverage source selection). Change the
+   source pool to per-view own-past only. (Depth + campose per view are already there.)
+2. **`pose` / `pose_person`** = render **only the wearer's own** skeleton in each view (not the partner's).
+   Currently `pose_person` renders both people identity-colored.
+
+Everything else in `clip.npz` (`target_rgb`, `visibility_packed`, per-view `campose`) stays. You do **not**
+need `refs_shared.npz` at all (no references).
+
+--------------------------------------------------------------------------------
+## 3. Things to know / gotchas
+
+- **Env / launch**: `sh/train_nymeria_longer.sh`, run as
+  `CUDA_VISIBLE_DEVICES=0,1,2,3 EXP=<your_exp_name> NPROC=4 bash sh/train_nymeria_longer.sh`.
+  Conda env `ego_dh`. Needs `COSMOS_QWEN_TOKENIZER_DIR=/data/cosmos_reason1_7b` for the online text encoder,
+  `IMAGINAIRE_OUTPUT_ROOT=<out>`.
+- **Online text encoding**: CoMind emits `ai_caption`; the model text-encodes it online each step (Qwen). Keep
+  the tokenizer dir env set or conditioning will crash on `t5_text_embeddings=None`.
+- **DCP checkpoints** are reshardable across NPROC (4↔3↔2 resume works). `save_iter` in the experiment.
+- **register the experiment**: append your function to the `experiments = [...]` list at the bottom of
+  `nymeria_pose_2actor.py`, and confirm it builds:
+  `python -c "..."` overriding `experiment=<name>` (see how the file builds configs) before launching 4-GPU.
+- **`state_t`**: latent temporal length = `1 + (T-1)//4`. For 77 frames it's 20. It appears in the model config;
+  keep it consistent between data (`num_video_frames`) and model.
+- **AR / noise recipe**: `predict2/models/video2world_model_rectified_flow.py` has an optional
+  `noisy_conditioning_*` + `conditional_frames_probs` recipe (teacher-forcing for autoregressive rollout). It is
+  OFF by default (`{1:1.0}`, prob 0). Ignore it unless you want AR — not needed for the single-ego baseline.
+- **Cross-attention**: you asked whether to separate it — you don't need to. The DiT already runs unified
+  self-attention over `B (V·t) H W D`; with no cross-view conditioning the two views just don't share useful
+  signal, which is exactly the single-ego behavior. Leaving it unified is simplest and correct.
+
+--------------------------------------------------------------------------------
+## 4. File map (what's in this bundle)
+
+```
+cosmos_predict2/_src/predict2_multiview/
+  configs/vid2vid/experiment/nymeria_pose_2actor.py   # all experiments (copy comind_* -> your single-ego exp)
+  configs/vid2vid/defaults/{conditioner,data,model}.py
+  datasets/nymeria_pairs.py            # NymeriaActorActorDataset + loaders + cs.store data configs
+  datasets/comind_pairs.py             # CoMind dataset/loader
+  networks/multiview_pose_dit.py       # the DiT: all the enable_* / view-emb / reference embedders
+  models/multiview_pose_model_rectified_flow.py   # conditioning preprocessing (refs/plucker/refpose/depth)
+  models/multiview_vid2vid_model_rectified_flow.py
+  callbacks/nymeria_validation_viz.py  # validation grid
+predict2/models/video2world_model_rectified_flow.py  # base rectified-flow denoise (+optional AR noise)
+sh/train_nymeria_longer.sh            # launch
+```
+
+Summary of the diff you need: **turn OFF** `enable_reference_frames/plucker/pose`, `num_reference_frames=0`,
+`shared_reference=False`, `enable_plucker=False`, `concat_view_embedding=False`; **from base 2B**; and feed
+**CoMind clips rebuilt with own-past warp + own-only pose**. Keep warped_cond + pose channels and unified
+cross-view attention.

cosmos_predict2/_src/predict2/models/video2world_model_rectified_flow.py

@@ -0,0 +1,247 @@
+# 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.
+
+from enum import Enum
+from typing import Callable, Dict, Optional, Tuple
+
+import attrs
+import torch
+from megatron.core import parallel_state
+from torch import Tensor
+
+from cosmos_predict2._src.predict2.conditioner import DataType
+from cosmos_predict2._src.predict2.configs.video2world.defaults.conditioner import Video2WorldCondition
+from cosmos_predict2._src.predict2.models.denoise_prediction import DenoisePrediction
+from cosmos_predict2._src.predict2.models.text2world_model_rectified_flow import (
+    Text2WorldCondition,
+    Text2WorldModelRectifiedFlow,
+    Text2WorldModelRectifiedFlowConfig,
+)
+
+NUM_CONDITIONAL_FRAMES_KEY: str = "num_conditional_frames"
+
+
+class ConditioningStrategy(str, Enum):
+    FRAME_REPLACE = "frame_replace"  # First few frames of the video are replaced with the conditional frames
+
+    def __str__(self) -> str:
+        return self.value
+
+
+def apply_noisy_conditioning(cond_state_B_C_T_H_W, cond_mask_B_T, *, training, prob, scale, min_frames, n_views=1):
+    """Add gaussian noise to the CONDITION latents only (network input), leaving the loss target clean.
+    Counts cond latents PER VIEW and uses the max (uneven view lengths would otherwise average below the
+    min_frames threshold and silently disable noise). See AR handoff recipe."""
+    if not training or prob <= 0.0:
+        return cond_state_B_C_T_H_W
+    B, _, T = cond_state_B_C_T_H_W.shape[:3]
+    if n_views > 1 and T % n_views == 0:
+        L = T // n_views
+        ncf_B = torch.stack([cond_mask_B_T[:, v * L:(v + 1) * L].sum(dim=1) for v in range(n_views)],
+                            dim=0).max(dim=0).values
+    else:
+        ncf_B = cond_mask_B_T.sum(dim=1)
+    for b in range(B):
+        if ncf_B[b].item() >= min_frames and torch.rand(1).item() < prob:
+            cond_state_B_C_T_H_W[b] = cond_state_B_C_T_H_W[b] + scale * torch.randn_like(cond_state_B_C_T_H_W[b])
+    return cond_state_B_C_T_H_W
+
+
+@attrs.define(slots=False)
+class Video2WorldModelRectifiedFlowConfig(Text2WorldModelRectifiedFlowConfig):
+    min_num_conditional_frames: int = 1  # Minimum number of latent conditional frames
+    max_num_conditional_frames: int = 2  # Maximum number of latent conditional frames
+    conditional_frame_timestep: float = (
+        -1.0
+    )  # Noise level used for conditional frames; default is -1 which will not take effective
+    conditioning_strategy: str = str(ConditioningStrategy.FRAME_REPLACE)  # What strategy to use for conditioning
+    denoise_replace_gt_frames: bool = True  # Whether to denoise the ground truth frames
+    conditional_frames_probs: Optional[Dict[int, float]] = None  # Probability distribution for conditional frames
+    # AR: add noise to the CONDITION frames (network input only; loss target stays clean) to close the
+    # train-inference gap — at inference the AR cond frames are the model's own imperfect outputs.
+    noisy_conditioning_prob: float = 0.0    # per-sample prob to apply (0 = off)
+    noisy_conditioning_scale: float = 0.1   # gaussian noise std in latent units
+    noisy_conditioning_min_frames: int = 2  # only when >= this many cond latents (per view) — i.e. the AR case
+
+    def __attrs_post_init__(self):
+        super().__attrs_post_init__()
+        assert self.conditioning_strategy in [
+            str(ConditioningStrategy.FRAME_REPLACE),
+        ]
+
+
+class Video2WorldModelRectifiedFlow(Text2WorldModelRectifiedFlow):
+    def get_data_and_condition(
+        self, data_batch: dict[str, torch.Tensor]
+    ) -> Tuple[Tensor, Tensor, Video2WorldCondition]:
+        # generate random number of conditional frames for training
+        raw_state, latent_state, condition = super().get_data_and_condition(data_batch)
+        condition = condition.set_video_condition(
+            gt_frames=latent_state.to(**self.tensor_kwargs),
+            random_min_num_conditional_frames=self.config.min_num_conditional_frames,
+            random_max_num_conditional_frames=self.config.max_num_conditional_frames,
+            num_conditional_frames=data_batch.get(NUM_CONDITIONAL_FRAMES_KEY, None),
+            conditional_frames_probs=self.config.conditional_frames_probs,
+        )
+        return raw_state, latent_state, condition
+
+    def denoise(
+        self,
+        noise: torch.Tensor,
+        xt_B_C_T_H_W: torch.Tensor,
+        timesteps_B_T: torch.Tensor,
+        condition: Text2WorldCondition,
+    ) -> DenoisePrediction:
+        """
+        Args:
+            xt (torch.Tensor): The input noise data.
+            sigma (torch.Tensor): The noise level.
+            condition (Text2WorldCondition): conditional information, generated from self.conditioner
+
+        Returns:
+            velocity prediction
+        """
+        if condition.is_video:
+            condition_state_in_B_C_T_H_W = condition.gt_frames.type_as(xt_B_C_T_H_W)
+            if not condition.use_video_condition:
+                # When using random dropout, we zero out the ground truth frames
+                condition_state_in_B_C_T_H_W = condition_state_in_B_C_T_H_W * 0
+
+            _, C, _, _, _ = xt_B_C_T_H_W.shape
+            condition_video_mask = condition.condition_video_input_mask_B_C_T_H_W.repeat(1, C, 1, 1, 1).type_as(
+                xt_B_C_T_H_W
+            )
+
+            # AR: noise the condition latents (network input only). The loss-target replacement below uses the
+            # CLEAN gt_frames, so the target is unchanged. Counts cond latents per view (state_t) -> max.
+            if getattr(self.config, "noisy_conditioning_prob", 0.0) > 0.0:
+                cond_mask_B_T = condition.condition_video_input_mask_B_C_T_H_W[:, 0, :, 0, 0]  # [B,T]
+                st = getattr(self.config, "state_t", 0)
+                nv = max(1, (cond_mask_B_T.shape[1] // st)) if st and st > 0 else 1
+                condition_state_in_B_C_T_H_W = apply_noisy_conditioning(
+                    condition_state_in_B_C_T_H_W, cond_mask_B_T, training=self.training,
+                    prob=self.config.noisy_conditioning_prob, scale=self.config.noisy_conditioning_scale,
+                    min_frames=self.config.noisy_conditioning_min_frames, n_views=nv,
+                )
+
+            # Make the first few frames of x_t be the ground truth frames
+            xt_B_C_T_H_W = condition_state_in_B_C_T_H_W * condition_video_mask + xt_B_C_T_H_W * (
+                1 - condition_video_mask
+            )
+
+            if self.config.conditional_frame_timestep >= 0:
+                condition_video_mask_B_1_T_1_1 = condition_video_mask.mean(dim=[1, 3, 4], keepdim=True)
+                timestep_cond_B_1_T_1_1 = (
+                    torch.ones_like(condition_video_mask_B_1_T_1_1) * self.config.conditional_frame_timestep
+                )
+
+                timesteps_B_1_T_1_1 = timestep_cond_B_1_T_1_1 * condition_video_mask_B_1_T_1_1 + timesteps_B_T * (
+                    1 - condition_video_mask_B_1_T_1_1
+                )
+
+                timesteps_B_T = timesteps_B_1_T_1_1.squeeze()
+                timesteps_B_T = (
+                    timesteps_B_T.unsqueeze(0) if timesteps_B_T.ndim == 1 else timesteps_B_T
+                )  # add dimension for batch
+
+        # forward pass through the network
+        net_output_B_C_T_H_W = self.net(
+            x_B_C_T_H_W=xt_B_C_T_H_W.to(**self.tensor_kwargs),  # Eq. 7 of https://arxiv.org/pdf/2206.00364.pdf
+            timesteps_B_T=timesteps_B_T,  # Eq. 7 of https://arxiv.org/pdf/2206.00364.pdf
+            **condition.to_dict(),
+        ).float()
+
+        if condition.is_video and self.config.denoise_replace_gt_frames:
+            gt_frames_x0 = condition.gt_frames.type_as(net_output_B_C_T_H_W)
+            gt_frames_velocity = noise - gt_frames_x0
+            net_output_B_C_T_H_W = gt_frames_velocity * condition_video_mask + net_output_B_C_T_H_W * (
+                1 - condition_video_mask
+            )
+
+        return net_output_B_C_T_H_W
+
+    def get_velocity_fn_from_batch(
+        self,
+        data_batch: Dict,
+        guidance: float = 1.5,
+        is_negative_prompt: bool = False,
+    ) -> Callable:
+        """
+        Generates a callable function `x0_fn` based on the provided data batch and guidance factor.
+
+        This function first processes the input data batch through a conditioning workflow (`conditioner`) to obtain conditioned and unconditioned states. It then defines a nested function `x0_fn` which applies a denoising operation on an input `noise_x` at a given noise level `sigma` using both the conditioned and unconditioned states.
+
+        Args:
+        - data_batch (Dict): A batch of data used for conditioning. The format and content of this dictionary should align with the expectations of the `self.conditioner`
+        - guidance (float, optional): A scalar value that modulates the influence of the conditioned state relative to the unconditioned state in the output. Defaults to 1.5.
+        - is_negative_prompt (bool): use negative prompt t5 in uncondition if true
+
+        Returns:
+        - Callable: A function `x0_fn(noise_x, sigma)` that takes two arguments, `noise_x` and `sigma`, and return velocity predictoin
+
+        The returned function is suitable for use in scenarios where a denoised state is required based on both conditioned and unconditioned inputs, with an adjustable level of guidance influence.
+        """
+
+        if NUM_CONDITIONAL_FRAMES_KEY in data_batch:
+            num_conditional_frames = data_batch[NUM_CONDITIONAL_FRAMES_KEY]
+        else:
+            num_conditional_frames = 1
+
+        if is_negative_prompt:
+            condition, uncondition = self.conditioner.get_condition_with_negative_prompt(data_batch)
+        else:
+            condition, uncondition = self.conditioner.get_condition_uncondition(data_batch)
+
+        is_image_batch = self.is_image_batch(data_batch)
+        condition = condition.edit_data_type(DataType.IMAGE if is_image_batch else DataType.VIDEO)
+        uncondition = uncondition.edit_data_type(DataType.IMAGE if is_image_batch else DataType.VIDEO)
+        _, x0, _ = self.get_data_and_condition(data_batch)
+        # override condition with inference mode; num_conditional_frames used Here!
+        condition = condition.set_video_condition(
+            gt_frames=x0,
+            random_min_num_conditional_frames=self.config.min_num_conditional_frames,
+            random_max_num_conditional_frames=self.config.max_num_conditional_frames,
+            num_conditional_frames=num_conditional_frames,
+            conditional_frames_probs=self.config.conditional_frames_probs,
+        )
+        uncondition = uncondition.set_video_condition(
+            gt_frames=x0,
+            random_min_num_conditional_frames=self.config.min_num_conditional_frames,
+            random_max_num_conditional_frames=self.config.max_num_conditional_frames,
+            num_conditional_frames=num_conditional_frames,
+            conditional_frames_probs=self.config.conditional_frames_probs,
+        )
+        condition = condition.edit_for_inference(is_cfg_conditional=True, num_conditional_frames=num_conditional_frames)
+        uncondition = uncondition.edit_for_inference(
+            is_cfg_conditional=False, num_conditional_frames=num_conditional_frames
+        )
+
+        _, condition, _, _ = self.broadcast_split_for_model_parallelsim(x0, condition, None, None)
+        _, uncondition, _, _ = self.broadcast_split_for_model_parallelsim(x0, uncondition, None, None)
+
+        if parallel_state.is_initialized():
+            pass
+        else:
+            assert not self.net.is_context_parallel_enabled, (
+                "parallel_state is not initialized, context parallel should be turned off."
+            )
+
+        def velocity_fn(noise: torch.Tensor, noise_x: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+            cond_v = self.denoise(noise, noise_x, timestep, condition)
+            uncond_v = self.denoise(noise, noise_x, timestep, uncondition)
+            velocity_pred = cond_v + guidance * (cond_v - uncond_v)
+            return velocity_pred
+
+        return velocity_fn

cosmos_predict2/_src/predict2_multiview/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+

cosmos_predict2/_src/predict2_multiview/callbacks/every_n_draw_sample_multiviewvideo.py

@@ -0,0 +1,485 @@
+# 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.
+
+from contextlib import nullcontext
+from functools import partial
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import torchvision
+import wandb
+from einops import rearrange, repeat
+
+from cosmos_predict2._src.imaginaire.utils import log, misc
+from cosmos_predict2._src.imaginaire.utils.easy_io import easy_io
+from cosmos_predict2._src.imaginaire.utils.parallel_state_helper import is_tp_cp_pp_rank0
+from cosmos_predict2._src.imaginaire.visualize.video import save_img_or_video
+from cosmos_predict2._src.predict2.callbacks.every_n_draw_sample import (
+    EveryNDrawSample,
+    convert_to_primitive,
+    is_primitive,
+    resize_image,
+)
+from cosmos_predict2._src.predict2.models.video2world_model import NUM_CONDITIONAL_FRAMES_KEY
+from cosmos_predict2._src.predict2_multiview.models.multiview_vid2vid_model_rectified_flow import (
+    MultiviewVid2VidModelRectifiedFlow,
+)
+
+try:
+    import ffmpegcv
+except Exception as e:  # ImportError cannot catch all problems
+    log.info(e)
+    ffmpegcv = None
+import cv2
+import numpy as np
+
+try:
+    import imageio
+except Exception as e:  # ImportError cannot catch all problems
+    log.info(e)
+    imageio = None
+
+CONTROL_WEIGHT_KEY = "control_weight"
+
+# view index order for visualization of 7-view autonomous driving dataset
+camera_to_view_id = {
+    "camera_cross_left_120fov": 5,
+    "camera_cross_right_120fov": 1,
+    "camera_front_tele_30fov": 6,
+    "camera_front_wide_120fov": 0,
+    "camera_rear_left_70fov": 4,
+    "camera_rear_right_70fov": 2,
+    "camera_rear_tele_30fov": 3,
+}
+
+visualization_camera_order = [
+    "camera_rear_left_70fov",
+    "camera_cross_left_120fov",
+    "camera_front_wide_120fov",
+    "camera_cross_right_120fov",
+    "camera_rear_right_70fov",
+    "camera_rear_tele_30fov",
+    "camera_front_tele_30fov",
+]
+
+visualization_view_index_order = [camera_to_view_id[camera] for camera in visualization_camera_order]
+
+
+class EveryNDrawSampleMultiviewVideo(EveryNDrawSample):
+    """
+    This class is a modified version of EveryNDrawSample that saves 12 frames instead of 3.
+    """
+
+    def __init__(
+        self,
+        *args,
+        n_view_embed=None,
+        ctrl_hint_keys=None,
+        control_weights=[1.0],
+        num_cond_frames=[0, 1],
+        fix_batch_fp=None,  # For backward compatibility with transfer2 experiments
+        n_x0_level=None,  # For backward compatibility with transfer2 experiments
+        show_all_frames=None,  # For backward compatibility with transfer2 experiments
+        is_sample=None,  # For backward compatibility with transfer2 experiments
+        **kwargs,
+    ):
+        # For backward compatibility with diffusion/v2 experiments that use is_x0 instead of do_x0_prediction
+        if "is_x0" in kwargs:
+            if "do_x0_prediction" in kwargs:
+                assert kwargs["do_x0_prediction"] == kwargs["is_x0"], "do_x0_prediction and is_x0 must be the same"
+            else:
+                kwargs["do_x0_prediction"] = kwargs["is_x0"]
+            del kwargs["is_x0"]
+
+        # For backward compatibility with transfer2 experiments that use n_x0_level instead of n_sigmas_for_x0_prediction
+        if n_x0_level is not None:
+            if "n_sigmas_for_x0_prediction" in kwargs:
+                assert kwargs["n_sigmas_for_x0_prediction"] == n_x0_level, (
+                    "n_sigmas_for_x0_prediction and n_x0_level must be the same"
+                )
+            else:
+                kwargs["n_sigmas_for_x0_prediction"] = n_x0_level
+
+        super().__init__(*args, **kwargs)
+        self.n_view_embed = n_view_embed
+        self.ctrl_hint_keys = ctrl_hint_keys
+        self.control_weights = control_weights
+        self.num_cond_frames = num_cond_frames
+        self.is_x0 = self.do_x0_prediction
+        if not hasattr(self, "fix_batch"):
+            self.fix_batch = None
+        if not hasattr(self, "is_sample"):
+            self.is_sample = True
+
+    def on_train_start(self, model: MultiviewVid2VidModelRectifiedFlow, iteration: int = 0) -> None:
+        return super().on_train_start(model, iteration)
+
+    def _ensure_even_dimensions(self, frame: np.ndarray) -> np.ndarray:
+        """
+        ffmpeg (H.264) requires both H and W to be even.  If either is odd we pad
+        by 1 pixel on the bottom/right using edge-replication.
+        """
+        h, w = frame.shape[:2]
+        pad_h = h % 2
+        pad_w = w % 2
+        if pad_h or pad_w:
+            frame = cv2.copyMakeBorder(frame, 0, pad_h, 0, pad_w, cv2.BORDER_REPLICATE)
+        return frame
+
+    def save_video(self, grid, video_name, fps: int = 30):
+        grid = (grid * 255).astype(np.uint8)
+        grid = np.transpose(grid, (1, 2, 3, 0))  # (T, H, W, C)
+
+        # Convert frames to RGB format and ensure even dimensions
+        processed_frames = []
+        for frame in grid:
+            frame = self._ensure_even_dimensions(frame)
+            processed_frames.append(frame)
+
+        # Use imageio.mimsave instead of ffmpegcv.VideoWriter for better error handling
+        try:
+            if imageio is not None:
+                kwargs = {
+                    "fps": fps,
+                    "quality": 5,  # Good quality
+                    "macro_block_size": 1,
+                    "ffmpeg_params": ["-c:v", "libx264", "-preset", "medium"],
+                }
+                imageio.mimsave(video_name, processed_frames, "mp4", **kwargs)
+            else:
+                raise ImportError("imageio not available")
+        except Exception as e:
+            # Fallback to ffmpegcv if imageio fails
+            if ffmpegcv is not None:
+                try:
+                    with ffmpegcv.VideoWriter(video_name, "h264", fps) as writer:
+                        for frame in processed_frames:
+                            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+                            writer.write(frame)
+                except Exception as ffmpeg_error:
+                    raise RuntimeError(
+                        f"Both imageio and ffmpegcv failed to save video. Imageio error: {e}, FFmpeg error: {ffmpeg_error}"
+                    )
+            else:
+                raise RuntimeError(f"Neither imageio nor ffmpegcv are available. Imageio error: {e}")
+
+    def run_save(self, to_show, batch_size, n_views, base_fp_wo_ext) -> Optional[str]:
+        to_show = (1.0 + torch.stack(to_show, dim=0).clamp(-1, 1)) / 2.0  # [n, b, c, t, h, w]
+        is_single_frame = to_show.shape[3] == 1
+        n_viz_sample = min(self.n_viz_sample, batch_size)
+
+        # ! we only save first n_sample_to_save video!
+        if self.save_s3 and self.data_parallel_id < self.n_sample_to_save:
+            save_img_or_video(
+                rearrange(to_show, "n b c t h w -> c t (n h) (b w)"),
+                f"s3://rundir/{self.name}/{base_fp_wo_ext}",
+                fps=self.fps,
+            )
+
+        file_base_fp = f"{base_fp_wo_ext}_resize.jpg"
+        local_path = f"{self.local_dir}/{file_base_fp}"
+
+        file_base_fp_12frames = f"{base_fp_wo_ext}_12frames.jpg"
+        local_path_12frames = f"{self.local_dir}/{file_base_fp_12frames}"
+
+        if self.rank == 0 and wandb.run:
+            if is_single_frame:  # image case
+                to_show = rearrange(
+                    to_show[:, :n_viz_sample],
+                    "n b c t h w -> t c (n h) (b w)",
+                )
+                image_grid = torchvision.utils.make_grid(to_show, nrow=1, padding=0, normalize=False)
+                # resize so that wandb can handle it
+                torchvision.utils.save_image(resize_image(image_grid, 1024), local_path, nrow=1, scale_each=True)
+            else:
+                to_show = to_show[:, :n_viz_sample]  # [n, b, c, t, h, w]
+                # Select 12 frames for the grid
+                _T = to_show.shape[3]
+                n = 12
+                twelve_frames_list = [round(ix * (_T - 1) / (n - 1)) for ix in range(n)]
+                to_show_12frames = to_show[:, :, :, twelve_frames_list]
+                to_show_12frames = rearrange(to_show_12frames, "n b c t h w -> 1 c (n h) (b t w)")
+                image_grid_12frames = torchvision.utils.make_grid(to_show_12frames, nrow=1, padding=0, normalize=False)
+                torchvision.utils.save_image(
+                    resize_image(image_grid_12frames, 1024), local_path_12frames, nrow=1, scale_each=True
+                )
+                # Create a single stacked video
+                video_tensor = rearrange(to_show, "n b c t h (v w) -> t (n h) (b v w) c", v=n_views)
+
+                # Resize width to 1024 while preserving aspect ratio (keep float to avoid quantization before resize)
+                max_w = 2048
+                T, H, W, C = video_tensor.shape
+                if W > max_w:
+                    scale = max_w / W
+                    new_w = max_w
+                    new_h = int(H * scale)
+                    # video_tensor is currently float in 0-1 range -> convert [T, H, W, C] to [T, C, H, W]
+                    video_tensor_f = video_tensor.permute(0, 3, 1, 2)
+                    video_tensor_f = F.interpolate(
+                        video_tensor_f, size=(new_h, new_w), mode="bilinear", align_corners=False
+                    )
+                    video_tensor = video_tensor_f.permute(0, 2, 3, 1)  # [T, H, W, C]
+
+                video_tensor = rearrange(video_tensor, "T H W C -> C T H W")
+                # Write the video
+                video_fp = f"{self.local_dir}/{base_fp_wo_ext}.mp4"
+                self.save_video(video_tensor.cpu().numpy(), video_fp, fps=self.fps)
+
+            return local_path, local_path_12frames, video_fp
+        return None
+
+    @torch.no_grad()
+    def every_n_impl(self, trainer, model, data_batch, output_batch, loss, iteration):
+        return self.every_n_impl_multiview(
+            trainer, model, None, data_batch, output_batch=output_batch, loss=loss, iteration=iteration
+        )
+
+    @torch.no_grad()
+    def every_n_impl_multiview(
+        self, trainer, model, data_batch_sample_all, data_batch_sample_n, output_batch, loss, iteration
+    ):
+        if self.is_ema:
+            if not model.config.ema.enabled:
+                return
+            context = partial(model.ema_scope, "every_n_sampling")
+        else:
+            context = nullcontext
+
+        tag = "ema" if self.is_ema else "reg"
+        sample_counter = getattr(trainer, "sample_counter", iteration)
+        data_batch_for_info = data_batch_sample_all if data_batch_sample_all is not None else data_batch_sample_n
+        batch_info = {
+            "data": {
+                k: convert_to_primitive(v)
+                for k, v in data_batch_for_info.items()
+                if is_primitive(v) or isinstance(v, (list, dict))
+            },
+            "sample_counter": sample_counter,
+            "iteration": iteration,
+            "sample_n_views": data_batch_for_info["sample_n_views"].cpu().item(),
+            "n_view_embed": self.n_view_embed,
+        }
+        if is_tp_cp_pp_rank0():
+            if self.save_s3 and self.data_parallel_id < self.n_sample_to_save:
+                easy_io.dump(
+                    batch_info,
+                    f"s3://rundir/{self.name}/BatchInfo_ReplicateID{self.data_parallel_id:04d}_Iter{iteration:09d}.json",
+                )
+
+        samples_img_fp = []
+        with context():
+            if self.is_x0:
+                x0_img_fp, mse_loss, sigmas = self.x0_pred(
+                    trainer,
+                    model,
+                    data_batch_for_info,
+                    output_batch,
+                    loss,
+                    iteration,
+                )
+                if self.save_s3 and self.rank == 0:
+                    easy_io.dump(
+                        {
+                            "mse_loss": mse_loss.tolist(),
+                            "sigmas": sigmas.tolist(),
+                            "iteration": iteration,
+                        },
+                        f"s3://rundir/{self.name}/{tag}_MSE_Iter{iteration:09d}.json",
+                    )
+            if self.is_sample:
+                for data_batch in [data_batch_sample_all, data_batch_sample_n]:
+                    if data_batch is None:
+                        samples_img_fp.append(None)
+                        continue
+                    sample_img_fp = self.sample(
+                        trainer,
+                        model,
+                        data_batch,
+                        output_batch,
+                        loss,
+                        iteration,
+                    )
+                    samples_img_fp.append(sample_img_fp)
+            if self.fix_batch is not None:
+                misc.to(self.fix_batch, "cpu")
+
+            dist.barrier()
+        if wandb.run:
+            sample_counter = getattr(trainer, "sample_counter", iteration)
+            data_type = "image" if model.is_image_batch(data_batch) else "video"
+            tag += f"_{data_type}"
+            info = {
+                "trainer/global_step": iteration,
+                "sample_counter": sample_counter,
+            }
+            if self.is_x0:
+                info[f"{self.name}/{tag}_x0"] = wandb.Image(x0_img_fp, caption=f"{sample_counter}")
+                # convert mse_loss to a dict
+                mse_loss = mse_loss.tolist()
+                info.update({f"x0_pred_mse_{tag}/Sigma{sigmas[i]:0.5f}": mse_loss[i] for i in range(len(mse_loss))})
+
+            if self.is_sample:
+                sample_all_img_fp, sample_n_img_fp = samples_img_fp
+                if sample_all_img_fp is not None:
+                    # info[f"{self.name}/{tag}_sample"] = wandb.Image(sample_all_img_fp[0], caption=f"{sample_counter}")
+                    info[f"{self.name}/{tag}_sample_allviews_frames"] = wandb.Image(
+                        sample_all_img_fp[1], caption=f"{sample_counter}"
+                    )
+                    info[f"{self.name}/{tag}_sample_allviews"] = wandb.Video(
+                        sample_all_img_fp[2], caption=f"{sample_counter}"
+                    )
+
+                # info[f"{self.name}/{tag}_sample"] = wandb.Image(sample_n_img_fp[0], caption=f"{sample_counter}")
+                info[f"{self.name}/{tag}_sample_nviews_frames"] = wandb.Image(
+                    sample_n_img_fp[1], caption=f"{sample_counter}"
+                )
+                info[f"{self.name}/{tag}_sample_nviews"] = wandb.Video(sample_n_img_fp[2], caption=f"{sample_counter}")
+            wandb.log(
+                info,
+                step=iteration,
+            )
+        torch.cuda.empty_cache()
+
+    @misc.timer("EveryNDrawSample: sample")
+    def sample(self, trainer, model, data_batch, output_batch, loss, iteration):
+        """
+        Args:
+            skip_save: to make sure FSDP can work, we run forward pass on all ranks even though we only save on rank 0 and 1
+        """
+        n_views = len(data_batch["view_indices_selection"][0])
+        if self.fix_batch is not None:
+            data_batch = misc.to(self.fix_batch, **model.tensor_kwargs)
+        tag = "ema" if self.is_ema else "reg"
+        raw_data, x0, condition = model.get_data_and_condition(data_batch)
+        if self.use_negative_prompt:
+            batch_size = x0.shape[0]
+            data_batch["neg_t5_text_embeddings"] = misc.to(
+                repeat(
+                    self.negative_prompt_data["t5_text_embeddings"],
+                    "l ... -> b (v l) ...",
+                    b=batch_size,
+                    v=n_views,
+                ),
+                **model.tensor_kwargs,
+            )
+            assert data_batch["neg_t5_text_embeddings"].shape == data_batch["t5_text_embeddings"].shape, (
+                f"{data_batch['neg_t5_text_embeddings'].shape} != {data_batch['t5_text_embeddings'].shape}"
+            )
+            data_batch["neg_t5_text_mask"] = data_batch["t5_text_mask"]
+
+        def time_to_width_dimension(mv_video):
+            """
+            Args:
+                mv_video: (B, C, V * T, H, W)
+            Returns:
+                (B, C, T, H, V * W)
+            """
+            current_view_index_order = [i.item() for i in data_batch["view_indices_selection"][0]]
+            expected_view_index_order = visualization_view_index_order
+
+            # Reorder views to match expected visualization order
+            if (
+                len(current_view_index_order) == len(expected_view_index_order)
+                and current_view_index_order != expected_view_index_order
+            ):
+                # Create mapping from current order to expected order
+                reorder_indices = []
+                for expected_view in expected_view_index_order:
+                    if expected_view in current_view_index_order:
+                        reorder_indices.append(current_view_index_order.index(expected_view))
+
+                # Reshape to separate view and time dimensions
+                B, C, VT, H, W = mv_video.shape
+                T = VT // n_views
+                mv_video = rearrange(mv_video, "B C (V T) H W -> B C V T H W", V=n_views)
+
+                # Reorder views according to expected order
+                mv_video = mv_video[:, :, reorder_indices, :, :, :]
+
+                # Reshape back to original format
+                mv_video = rearrange(mv_video, "B C V T H W -> B C (V T) H W")
+
+            return rearrange(mv_video, "B C (V T) H W -> B C T H (V W)", V=n_views)
+
+        # GPU memory management before sampling to avoid OOM
+        if torch.cuda.is_available():
+            mem_allocated_before = torch.cuda.memory_allocated() / 1e9
+            mem_reserved_before = torch.cuda.memory_reserved() / 1e9
+            print(
+                f"[Rank {dist.get_rank() if dist.is_initialized() else 0}] Before sampling - "
+                f"Allocated: {mem_allocated_before:.2f}GB, Reserved: {mem_reserved_before:.2f}GB"
+            )
+
+            # Clear GPU cache to free up fragmented memory
+            torch.cuda.empty_cache()
+
+            mem_allocated_after = torch.cuda.memory_allocated() / 1e9
+            mem_reserved_after = torch.cuda.memory_reserved() / 1e9
+            print(
+                f"[Rank {dist.get_rank() if dist.is_initialized() else 0}] After cleanup - "
+                f"Allocated: {mem_allocated_after:.2f}GB, Reserved: {mem_reserved_after:.2f}GB "
+                f"(freed: Allocated={mem_allocated_before - mem_allocated_after:.2f}GB, "
+                f"Reserved={mem_reserved_before - mem_reserved_after:.2f}GB)"
+            )
+
+        to_show = []
+        # for use_apg in [False, True]:
+        for use_apg in [False]:
+            for num_cond_frames in self.num_cond_frames:
+                for control_weight in self.control_weights:
+                    data_batch[NUM_CONDITIONAL_FRAMES_KEY] = num_cond_frames
+                    data_batch[CONTROL_WEIGHT_KEY] = control_weight
+                    for guidance in self.guidance:
+                        sample = model.generate_samples_from_batch(
+                            data_batch,
+                            guidance=guidance,
+                            # make sure no mismatch and also works for cp
+                            state_shape=x0.shape[1:],
+                            n_sample=x0.shape[0],
+                            num_steps=self.num_sampling_step,
+                            is_negative_prompt=True if self.use_negative_prompt else False,
+                        )
+                        if hasattr(model, "decode"):
+                            sample = model.decode(sample)
+                        to_show.append(sample.float().cpu())
+
+        to_show.append(raw_data.float().cpu())
+
+        # Transfer2-multiview: visualize control input
+        if self.ctrl_hint_keys:
+            # visualize input video
+            if "hint_key" in data_batch:
+                hint = data_batch[data_batch["hint_key"]]
+                for idx in range(0, hint.size(1), 3):
+                    x_rgb = hint[:, idx : idx + 3]
+                    to_show.append(x_rgb.float().cpu())
+            else:
+                for key in self.ctrl_hint_keys:
+                    if key in data_batch and data_batch[key] is not None:
+                        hint = data_batch[key]
+                        log.info(f"hint: {hint.shape}")
+                        to_show.append(hint.float().cpu())
+
+        to_show = [time_to_width_dimension(t) for t in to_show]
+
+        base_fp_wo_ext = f"{tag}_ReplicateID{self.data_parallel_id:04d}_Sample_Iter{iteration:09d}_{n_views}views"
+        batch_size = x0.shape[0]
+
+        if is_tp_cp_pp_rank0():
+            local_path = self.run_save(to_show, batch_size, n_views, base_fp_wo_ext)
+            return local_path
+        return None

cosmos_predict2/_src/predict2_multiview/callbacks/frame_loss_log.py

@@ -0,0 +1,43 @@
+# 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 torch
+
+from cosmos_predict2._src.imaginaire.model import ImaginaireModel
+from cosmos_predict2._src.imaginaire.utils.callback import Callback
+
+"""
+   Dummy FrameLossLog callback used in multiview training with view dropout, where batches don't have the same number of views / frames.
+"""
+
+
+class DummyFrameLossLog(Callback):
+    def __init__(
+        self,
+        logging_iter_multipler: int = 1,
+        save_logging_iter_multipler: int = 1,
+        save_s3: bool = False,
+    ) -> None:
+        pass
+
+    def on_training_step_end(
+        self,
+        model: ImaginaireModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor],
+        loss: torch.Tensor,
+        iteration: int = 0,
+    ):
+        pass

cosmos_predict2/_src/predict2_multiview/callbacks/log_weight.py

@@ -0,0 +1,62 @@
+# 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.
+
+from typing import Optional
+
+import torch
+import wandb
+
+from cosmos_predict2._src.imaginaire.callbacks.every_n import EveryN
+from cosmos_predict2._src.imaginaire.model import ImaginaireModel
+from cosmos_predict2._src.imaginaire.trainer import ImaginaireTrainer
+from cosmos_predict2._src.imaginaire.utils import distributed, log
+
+
+class LogWeight(EveryN):
+    def __init__(
+        self,
+        every_n: Optional[int] = 100,
+        step_size: int = 1,
+        barrier_after_run: bool = True,
+        run_at_start: bool = False,
+    ):
+        super().__init__(
+            every_n=every_n,
+            step_size=step_size,
+            barrier_after_run=barrier_after_run,
+            run_at_start=run_at_start,
+        )
+
+    def every_n_impl(
+        self,
+        trainer: ImaginaireTrainer,
+        model: ImaginaireModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor],
+        loss: torch.Tensor,
+        iteration: int,
+    ) -> None:
+        if "logging_dict" in output_batch:
+            logging_dict = output_batch["logging_dict"]
+
+            if distributed.is_rank0():
+                info = {}
+                for k, v in logging_dict.items():
+                    info[f"model_weight/{k}"] = v
+
+                if info and wandb.run:
+                    wandb.log(info, step=iteration)
+
+                log.info(f"Log weight at iteration {iteration}: {info}")

cosmos_predict2/_src/predict2_multiview/callbacks/nymeria_validation_viz.py

@@ -0,0 +1,335 @@
+# 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.
+
+"""Validation visualization for 2-actor Nymeria generation.
+
+Every ``every_n`` iterations, generates videos for a FIXED set of ``num_train`` train + ``num_val`` val
+samples and saves, per sample, a comparison video laid out as:
+
+         | warped cond | pose input | GT video | generated |
+  egoA   |    ....      |    ....    |   ....    |    ....    |   (row 1)
+  egoB   |    ....      |    ....    |   ....    |    ....    |   (row 2)
+
+i.e. 4 columns (warped / pose / GT / generated) x 2 rows (actor A / actor B). Files are written under
+``<output>/validation/`` and logged to W&B. Samples are deterministic (fixed indices) so progress is
+comparable across iterations.
+"""
+
+import os
+from contextlib import nullcontext
+from functools import partial
+from typing import List, Optional
+
+import cv2
+import numpy as np
+import torch
+import wandb
+
+from cosmos_predict2._src.imaginaire.utils import distributed, log
+from cosmos_predict2._src.imaginaire.utils.callback import Callback
+from cosmos_predict2._src.imaginaire.visualize.video import save_img_or_video
+
+_COLS = ["warped", "pose", "GT", "generated"]
+
+
+class NymeriaValidationViz(Callback):
+    def __init__(
+        self,
+        every_n: int = 2000,
+        num_train: int = 5,
+        num_val: int = 5,
+        num_sampling_step: int = 35,
+        guidance: float = 7.0,
+        fps: int = 10,
+        viz_size: int = 192,
+        is_ema: bool = False,
+        run_at_start: bool = False,
+        # nymeria dataset params (match the training dataloader)
+        root: str = "/data2/nymeria_processed_longer",
+        num_video_frames: int = 77,
+        resolution_hw: tuple = (480, 480),
+        val_num_sessions: int = 4,
+        split_seed: int = 1234,
+        single_view: bool = False,
+        actor_observer: bool = False,  # build the 2-view actor-observer val set (train_split/val_split.csv)
+        actor_actor: bool = False,  # build the 2-view actor-ACTOR val set (longer dataset, per-view real captions)
+        comind: bool = False,  # build the 2-view CoMind actor-actor val set (shared-world clip.npz L_/H_ format)
+        single_view_stage1: bool = False,  # build the STAGE-1 single-view pooled ego-clip val set (real captions)
+        num_reference_frames: int = 0,  # emit `reference_frames` in the val set (match the training dataloader)
+        shared_reference: bool = False,  # use SHARED refs (refs_shared.npz) — match the training dataloader
+        emit_camera_poses: bool = False,  # emit camera_w2c/K in the val set (for Plücker conditioning)
+        emit_depth: bool = False,  # emit control_input_depth (composite depth) -> adds a "depth" viz column
+        person_pose: bool = False,  # use identity-colored pose_person in the val set (match training)
+        emit_reference_pose: bool = False,  # emit reference_pose (ref skeletons) for the reference-pose condition
+        train_fixed_ids: Optional[List[str]] = None,  # pin the viz TRAIN samples to these pair_ids (stable
+        val_fixed_ids: Optional[List[str]] = None,    # across runs / manifest-size changes). Falls back to
+        #                                               evenly-spaced sampling when None.
+        val_manifest: str = "manifest/val_split.csv",  # override the val manifest (e.g. a curated subset for
+        #                                                checkpoint-eval sweeps: interaction-only, synthetic, ...)
+    ):
+        self.train_fixed_ids = train_fixed_ids
+        self.val_fixed_ids = val_fixed_ids
+        self.val_manifest = val_manifest
+        self.actor_observer = actor_observer
+        self.actor_actor = actor_actor
+        self.comind = comind
+        self.single_view_stage1 = single_view_stage1
+        self.num_reference_frames = num_reference_frames
+        self.shared_reference = shared_reference
+        self.emit_camera_poses = emit_camera_poses
+        self.emit_depth = emit_depth
+        self.person_pose = person_pose
+        self.emit_reference_pose = emit_reference_pose
+        self.every_n = every_n
+        self.num_train = num_train
+        self.num_val = num_val
+        self.num_sampling_step = num_sampling_step
+        self.guidance = guidance
+        self.fps = fps
+        self.viz_size = viz_size
+        self.is_ema = is_ema
+        self.run_at_start = run_at_start
+        self._ds_kwargs = dict(
+            root=root, num_video_frames=num_video_frames, resolution_hw=tuple(resolution_hw),
+            val_num_sessions=val_num_sessions, split_seed=split_seed, single_view=single_view,
+            num_reference_frames=num_reference_frames, shared_reference=shared_reference,
+            emit_camera_poses=emit_camera_poses, emit_depth=emit_depth,
+            person_pose=person_pose, emit_reference_pose=emit_reference_pose,
+        )
+        self._samples: List[tuple] = []  # (tag, name, dataset, idx)
+
+    # ------------------------------------------------------------------ setup
+    def on_train_start(self, model, iteration: int = 0) -> None:
+        from cosmos_predict2._src.predict2_multiview.datasets.nymeria_pairs import (
+            NymeriaActorActorDataset,
+            NymeriaActorObserverDataset,
+            NymeriaPairsConfig,
+            NymeriaPairsDataset,
+        )
+
+        self.local_dir = f"{self.config.job.path_local}/validation"
+        if distributed.get_rank() == 0:
+            os.makedirs(self.local_dir, exist_ok=True)
+            log.info(f"NymeriaValidationViz: local_dir={self.local_dir}")
+
+        if self.comind:  # 2-view CoMind actor-actor (shared-world clip.npz L_/H_ format), held-out recs for val
+            from cosmos_predict2._src.predict2_multiview.datasets.comind_pairs import (
+                ComindActorActorDataset, ComindPairsConfig,
+            )
+            cok = dict(root=self._ds_kwargs["root"], num_video_frames=self._ds_kwargs["num_video_frames"],
+                       resolution_hw=self._ds_kwargs["resolution_hw"],
+                       num_reference_frames=self.num_reference_frames, shared_reference=self.shared_reference,
+                       emit_camera_poses=self.emit_camera_poses)
+            train_ds = ComindActorActorDataset(ComindPairsConfig(split="train", role_mix=False, **cok))
+            val_ds = ComindActorActorDataset(ComindPairsConfig(split="val", role_mix=False, **cok))
+        elif self.single_view_stage1:  # STAGE-1 pooled ego clips (V=1) with real per-clip captions
+            from cosmos_predict2._src.predict2_multiview.datasets.nymeria_pairs import (
+                NymeriaSingleViewDataset, _STAGE1_SOURCES,
+            )
+            aok = {k: v for k, v in self._ds_kwargs.items()
+                   if k in ("num_video_frames", "resolution_hw", "num_reference_frames", "emit_camera_poses")}
+            cfg_t = NymeriaPairsConfig(single_view=True, shared_reference=False, **aok)
+            train_ds = NymeriaSingleViewDataset(cfg_t, _STAGE1_SOURCES)
+            val_srcs = [dict(s, manifest_csv=s["manifest_csv"].replace("train_split", "val_split")) for s in _STAGE1_SOURCES]
+            val_ds = NymeriaSingleViewDataset(NymeriaPairsConfig(single_view=True, shared_reference=False, **aok), val_srcs)
+        elif self.actor_observer or self.actor_actor:  # 2-view (real per-view captions, role-mix off for determinism)
+            aok = {k: v for k, v in self._ds_kwargs.items() if k in ("root", "num_video_frames", "resolution_hw", "num_reference_frames", "shared_reference", "emit_camera_poses", "emit_depth", "person_pose", "emit_reference_pose")}
+            _DS = NymeriaActorActorDataset if self.actor_actor else NymeriaActorObserverDataset
+            train_ds = _DS(NymeriaPairsConfig(manifest_csv="manifest/train_split.csv", role_mix=False, **aok))
+            val_ds = _DS(NymeriaPairsConfig(manifest_csv=self.val_manifest, role_mix=False, **aok))
+        else:
+            dummy = getattr(model.config, "text_encoder_config", None) is None
+            train_ds = NymeriaPairsDataset(NymeriaPairsConfig(split="train", dummy_text_embeddings=dummy, **self._ds_kwargs))
+            val_ds = NymeriaPairsDataset(NymeriaPairsConfig(split="val", clean_only=True, dummy_text_embeddings=dummy, **self._ds_kwargs))
+
+        # fixed, evenly-spaced indices (identical on every rank -> consistent collective generation)
+        def pick(ds, n):
+            n = min(n, len(ds))
+            return [int(i * (len(ds) / n)) for i in range(n)] if n else []
+
+        # pin by pair_id when given (stable across runs / manifest resizes); else evenly-spaced fallback.
+        def select(ds, n, fixed_ids):
+            if fixed_ids:
+                id2idx = {p.get("pair_id"): i for i, p in enumerate(ds.pairs)}
+                idxs = [id2idx[pid] for pid in fixed_ids if pid in id2idx]
+                missing = [pid for pid in fixed_ids if pid not in id2idx]
+                if missing:
+                    log.warning(f"NymeriaValidationViz: {len(missing)} fixed pair_ids not in dataset: {missing[:3]}...")
+                if idxs:
+                    return idxs
+                log.warning("NymeriaValidationViz: no fixed pair_ids matched; falling back to even sampling")
+            return pick(ds, n)
+
+        self._samples = [("train", f"train{i:02d}", train_ds, idx)
+                         for i, idx in enumerate(select(train_ds, self.num_train, self.train_fixed_ids))]
+        self._samples += [("val", f"val{i:02d}", val_ds, idx)
+                          for i, idx in enumerate(select(val_ds, self.num_val, self.val_fixed_ids))]
+
+        # run_at_start: generate immediately here (pure inference, before any training step). This lets a
+        # single-GPU checkpoint-eval run (max_iter=0) produce all viz without the training-step backward/optim
+        # memory that would OOM an unsharded model.
+        if self.run_at_start:
+            log.info(f"NymeriaValidationViz: run_at_start -> generating {len(self._samples)} samples now")
+            self._generate(model, iteration)
+
+    # ------------------------------------------------------------------ trigger
+    def on_training_step_end(self, model, data_batch, output_batch, loss, iteration: int = 0) -> None:
+        # run_at_start is handled in on_train_start (pure inference, no training step needed -> lets a
+        # single-GPU checkpoint-eval run with max_iter=0 avoid the training-step OOM). Here only periodic.
+        if iteration > 0 and iteration % self.every_n == 0:
+            self._generate(model, iteration)
+
+    def _generate(self, model, iteration: int) -> None:
+        if not self._samples:
+            return
+        if self.is_ema and not model.config.ema.enabled:
+            return
+        context = partial(model.ema_scope, "validation_viz") if self.is_ema else nullcontext
+        log_payload = {}
+        with context():
+            for tag, name, ds, idx in self._samples:
+                try:
+                    path, caption = self._viz_one(model, ds, idx, iteration, name)  # collective on all ranks
+                except Exception as e:  # don't let viz crash training
+                    log.warning(f"NymeriaValidationViz {name} failed: {type(e).__name__}: {e}")
+                    continue
+                if distributed.get_rank() == 0 and path is not None and wandb.run:
+                    # caption = per-view text fed to this sample (egoA / egoB), shown under the W&B video
+                    log_payload[f"validation/{tag}/{name}"] = wandb.Video(path, fps=self.fps, format="mp4", caption=caption)
+                if distributed.get_rank() == 0 and (len(self._samples) > 20) and ((self._samples.index((tag, name, ds, idx)) + 1) % 25 == 0):
+                    log.info(f"NymeriaValidationViz: generated {self._samples.index((tag, name, ds, idx))+1}/{len(self._samples)}")
+        if distributed.get_rank() == 0 and wandb.run and log_payload:
+            wandb.log(log_payload, step=iteration)
+
+    # ------------------------------------------------------------------ one sample
+    @torch.no_grad()
+    def _viz_one(self, model, ds, idx, iteration, name):
+        from cosmos_predict2._src.predict2_multiview.datasets.multiview import collate_fn as mv_collate
+
+        data_batch = mv_collate([ds[idx]])
+        from cosmos_predict2._src.imaginaire.utils import misc
+
+        # per-view captions actually fed to this sample (before online text-embed consumes them)
+        caps = None
+        _ai = data_batch.get("ai_caption")
+        if _ai and isinstance(_ai[0], (list, tuple)):
+            caps = list(_ai[0])  # [view0_caption, view1_caption]
+
+        data_batch = misc.to(data_batch, device="cuda")
+
+        # real-text runs provide per-view captions (ai_caption), not embeddings -> compute them online first
+        # (mirrors training_step_multiview), otherwise the conditioner has no t5_text_embeddings to read.
+        tec = getattr(model.config, "text_encoder_config", None)
+        if tec is not None and getattr(tec, "compute_online", False):
+            model.inplace_compute_text_embeddings_online(data_batch)
+
+        # generation runs the net in bf16 (training casts inputs explicitly; here we use autocast so the
+        # float32 text + sampling inputs line up with the bf16 weights, mirroring the training forward)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            # GT + latent shape (also runs the model's pose preprocessing into data_batch)
+            raw_data, x0, _ = model.get_data_and_condition(data_batch)
+            sample = model.generate_samples_from_batch(
+                data_batch, guidance=self.guidance, state_shape=x0.shape[1:], n_sample=x0.shape[0],
+                num_steps=self.num_sampling_step, is_negative_prompt=False,
+            )
+            if hasattr(model, "decode"):
+                sample = model.decode(sample)  # (B, 3, V*T, H, W) in [-1, 1]
+
+        if distributed.get_rank() != 0:
+            return None, None
+
+        def to01(x):  # (B,C,V*T,H,W) -> (C,V*T,H,W) in [0,1]
+            return ((x[0].float().cpu().clamp(-1, 1) + 1.0) / 2.0)
+
+        gen = to01(sample)
+        gt = to01(raw_data)
+        warped = data_batch["control_input_warped"][0].float().cpu() / 255.0  # (3,V*T,H,W)
+        pose = data_batch["control_input_pose"][0].float().cpu() / 255.0
+        depth = None
+        if "control_input_depth" in data_batch:  # composite-depth condition (warped scene + human mesh) column
+            depth = data_batch["control_input_depth"][0].float().cpu() / 255.0
+        n_views = int(data_batch["sample_n_views"][0])
+        # shared reference frames fed to this sample (3, R, H, W); normalized to [0,1] for the viz strip
+        refs = None
+        if "reference_frames" in data_batch:
+            rf = data_batch["reference_frames"][0].float().cpu()
+            refs = rf / 255.0 if rf.max() > 1.5 else rf
+        grid = self._compose(warped, pose, gt, gen, n_views, refs=refs, depth=depth)  # (3, T, n_views*s + hs, ncols*s)
+        fp = f"{self.local_dir}/{name}_iter{iteration:09d}"
+        save_img_or_video(grid, fp, fps=self.fps)
+
+        # save the per-view captions locally (sidecar .txt) and return them for W&B logging
+        labels = ["egoA", "egoB", "egoC", "egoD"]
+        caption_str = None
+        if caps is not None:
+            lines = [f"{labels[v] if v < len(labels) else f'view{v}'}: {c}" for v, c in enumerate(caps)]
+            caption_str = "\n".join(lines)
+            with open(f"{fp}.txt", "w") as f:
+                f.write(caption_str + "\n")
+        return f"{fp}.mp4", caption_str
+
+    # ------------------------------------------------------------------ grid + labels
+    def _compose(self, warped, pose, gt, gen, n_views: int = 2, refs=None, depth=None):
+        s = self.viz_size
+        T = gen.shape[1] // n_views  # frames per actor/view
+        titles = list(_COLS) + (["depth"] if depth is not None else [])  # extra depth-condition column when present
+        ncols = len(titles)
+
+        def cell(x, v):  # x:(3,V*T,H,W) -> view v -> (T,s,s,3) numpy uint8
+            a = x[:, v * T : (v + 1) * T]  # (3,T,H,W)
+            a = torch.nn.functional.interpolate(a.permute(1, 0, 2, 3), size=(s, s), mode="bilinear", align_corners=False)
+            a = (a.clamp(0, 1).permute(0, 2, 3, 1).numpy() * 255).astype(np.uint8)  # (T,s,s,3)
+            return a
+
+        rows = []
+        labels = ["egoA", "egoB", "egoC", "egoD"]
+        for v in range(n_views):
+            row_label = labels[v] if v < len(labels) else f"ego{v}"
+            cols = [cell(warped, v), cell(pose, v), cell(gt, v), cell(gen, v)]  # each (T,s,s,3)
+            if depth is not None:
+                cols.append(cell(depth, v))
+            row = np.ascontiguousarray(np.concatenate(cols, axis=2))  # (T,s,ncols*s,3)
+            # column headers (top of actorA row) + row label (left); per-frame contiguous buffer for cv2
+            for t in range(row.shape[0]):
+                frame = np.ascontiguousarray(row[t])
+                if v == 0:
+                    for c, title in enumerate(titles):
+                        cv2.putText(frame, title, (c * s + 4, 16), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1, cv2.LINE_AA)
+                cv2.putText(frame, row_label, (4, s - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1, cv2.LINE_AA)
+                row[t] = frame
+            rows.append(row)
+
+        # shared-reference strip: R thumbnails side-by-side (static across all T frames), one row ABOVE the
+        # views so you can see exactly which reference images conditioned this sample.
+        if refs is not None and refs.shape[1] > 0:
+            R = refs.shape[1]
+            grid_w = ncols * s
+            hs = max(1, grid_w // R)  # thumbnail size so R thumbs span the full grid width
+            thumbs = []
+            for r in range(R):
+                a = torch.nn.functional.interpolate(refs[:, r].unsqueeze(0), size=(hs, hs), mode="bilinear",
+                                                    align_corners=False)[0]
+                a = (a.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8)  # (hs,hs,3)
+                thumbs.append(a)
+            strip = np.ascontiguousarray(np.concatenate(thumbs, axis=1))  # (hs, R*hs, 3)
+            if strip.shape[1] != grid_w:  # pad/trim to exact grid width
+                strip = np.ascontiguousarray(cv2.resize(strip, (grid_w, hs)))
+            cv2.putText(strip, f"shared refs (x{R})", (4, 14), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, cv2.LINE_AA)
+            ref_row = np.repeat(strip[None], T, axis=0)  # (T, hs, 4s, 3) static
+            rows.insert(0, ref_row)  # TOP strip (above the view rows)
+
+        grid = np.concatenate(rows, axis=1)  # (T, hs + n_views*s, 4s, 3)
+        grid = torch.from_numpy(grid).float().div(255).permute(3, 0, 1, 2)  # (3, T, H, 4s)
+        return grid

cosmos_predict2/_src/predict2_multiview/callbacks/sigma_loss_analysis_per_frame.py

@@ -0,0 +1,338 @@
+# 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.
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.distributed as dist
+import wandb
+
+from cosmos_predict2._src.imaginaire.model import ImaginaireModel
+from cosmos_predict2._src.imaginaire.utils import distributed, misc
+from cosmos_predict2._src.imaginaire.utils.callback import Callback
+from cosmos_predict2._src.imaginaire.utils.easy_io import easy_io
+
+
+class DummySigmaLossAnalysisPerFrame(Callback):
+    def __init__(
+        self,
+        logging_iter_multipler: int = 1,
+        logging_viz_iter_multipler: int = 1,
+        save_s3: bool = False,
+    ) -> None:
+        super().__init__()
+        pass
+
+
+def _get_normal_quantile_bins():
+    """Get predefined bins based on exp(N(0,1)) distribution quantiles"""
+    # Using torch.special.ndtri (inverse of standard normal CDF)
+    # and taking exponential to get exponentially spaced bins
+    probs = torch.linspace(0.0, 1.0, 11)  # 11 points gives 10 bins
+    points = torch.special.ndtri(probs)
+    # Take exponential to get exponentially spaced bins
+    points = torch.exp(points)
+    # Replace extreme values at boundaries
+    points[0] = points[1] / (points[2] / points[1])  # Extrapolate left boundary
+    points[-1] = points[-2] * (points[-2] / points[-3])  # Extrapolate right boundary
+    return points.numpy()
+
+
+@dataclass
+class _SigmaLossCache:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.sigma_list: List[torch.Tensor] = []
+        self.loss_list: List[torch.Tensor] = []
+
+    def add(self, sigma: torch.Tensor, loss: torch.Tensor):
+        # Convert to bf16 and store on CPU
+        self.sigma_list.append(sigma.detach().cpu().to(torch.bfloat16))
+        self.loss_list.append(loss.detach().cpu().to(torch.bfloat16))
+
+    def get_arrays(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[int]]:
+        if not self.sigma_list:
+            return torch.tensor([], dtype=torch.bfloat16), torch.tensor([], dtype=torch.bfloat16), None
+
+        sigma_arr = torch.cat(self.sigma_list, dim=0)  # [B*N, T] or [B*N, 1]
+        loss_arr = torch.cat(self.loss_list, dim=0)  # [B*N, T]
+
+        # Handle broadcasting case where sigma is shape [B, 1]
+        if sigma_arr.shape[-1] == 1 and loss_arr.shape[-1] > 1:
+            sigma_arr = sigma_arr.expand(-1, loss_arr.shape[-1])
+
+        num_frames = loss_arr.shape[-1] if len(loss_arr.shape) > 1 else 1
+
+        assert sigma_arr.shape == loss_arr.shape, (sigma_arr.shape, loss_arr.shape)
+
+        return sigma_arr, loss_arr, num_frames
+
+
+class SigmaLossAnalysisPerFrame(Callback):
+    def __init__(
+        self,
+        logging_iter_multipler: int = 1,
+        logging_viz_iter_multipler: int = 1,
+        save_s3: bool = False,
+    ) -> None:
+        super().__init__()
+        self.save_s3 = save_s3
+        self.logging_iter_multipler = logging_iter_multipler
+        assert logging_viz_iter_multipler % logging_iter_multipler == 0
+        self.logging_viz_iter_multipler = logging_viz_iter_multipler
+        self.name = self.__class__.__name__
+
+        self.image_cache = _SigmaLossCache()
+        self.video_cache = _SigmaLossCache()
+
+    def _create_analysis_plots(
+        self,
+        sigma_arr: torch.Tensor,
+        loss_arr: torch.Tensor,
+        frame_idx: Optional[int] = None,  # [N]  # [N]
+    ) -> Optional[wandb.Image]:
+        if len(sigma_arr) == 0:
+            return None
+
+        # Convert to numpy for plotting
+        sigma_np = sigma_arr.cpu().float().numpy()[:800]
+        loss_np = loss_arr.cpu().float().numpy()[:800]
+
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
+
+        # Get predefined bins based on normal distribution quantiles
+        sigma_bins = _get_normal_quantile_bins()
+
+        y_tick_min, y_tick_max = 0, 1.0
+        # 2D histogram with exponential sigma bins and fixed [0,1] loss range
+        loss_bins = np.linspace(y_tick_min, y_tick_max, 20)
+
+        counts, xedges, yedges = np.histogram2d(sigma_np, loss_np, bins=(sigma_bins, loss_bins))
+        if counts.max() < 0.1:
+            return None
+
+        # Plot heatmap with exponential scale colormap
+        im = ax1.imshow(
+            counts.T,
+            origin="lower",
+            aspect="auto",
+            extent=[sigma_bins[0], sigma_bins[-1], y_tick_min, y_tick_max],
+            norm=matplotlib.colors.LogNorm(vmin=1, vmax=counts.max()),
+        )
+        plt.colorbar(im, ax=ax1)
+
+        # Set fixed loss ticks from 0 to 1
+        yticks = np.linspace(y_tick_min, y_tick_max, 6)
+        ax1.set_yticks(yticks)
+        ax1.set_yticklabels([f"{y:.1f}" for y in yticks])
+
+        ax1.set_xlabel("Sigma (Standard Normal Quantiles)")
+        ax1.set_ylabel("Loss")
+        title = "Sigma vs Loss Distribution"
+        if frame_idx is not None:
+            title += f" (Frame {frame_idx})"
+        ax1.set_title(title)
+
+        # Sigma histogram with loss statistics per bin
+        hist_counts, _ = np.histogram(sigma_np, bins=sigma_bins)
+        bin_indices = np.digitize(sigma_np, sigma_bins) - 1
+
+        # Calculate statistics per bin
+        n_bins = len(sigma_bins) - 1
+        means = np.zeros(n_bins)
+        stds = np.zeros(n_bins)
+        for i in range(n_bins):
+            bin_mask = bin_indices == i
+            if bin_mask.any():
+                means[i] = loss_np[bin_mask].mean()
+                stds[i] = loss_np[bin_mask].std()
+            else:
+                means[i] = np.nan
+                stds[i] = np.nan
+
+        # Plot histogram
+        bin_centers = (sigma_bins[:-1] + sigma_bins[1:]) / 2
+        ax2.bar(bin_centers, hist_counts, width=np.diff(sigma_bins), alpha=0.3, align="center")
+
+        # Plot loss statistics on twin axis
+        ax2_twin = ax2.twinx()
+        valid_mask = ~np.isnan(means)
+        ax2_twin.errorbar(
+            bin_centers[valid_mask], means[valid_mask], yerr=stds[valid_mask], color="red", fmt="o-", alpha=0.5
+        )
+
+        ax2.set_xlabel("Sigma (Standard Normal Quantiles)")
+        ax2.set_ylabel("Count")
+        ax2_twin.set_ylabel("Loss (mean ± std)")
+        title = "Sigma Distribution with Loss Statistics"
+        if frame_idx is not None:
+            title += f" (Frame {frame_idx})"
+        ax2.set_title(title)
+
+        # Add grid for better readability
+        ax1.grid(True, alpha=0.3)
+        ax2.grid(True, alpha=0.3)
+
+        # Add quantile labels
+        probs = torch.linspace(0.0, 1.0, 11)  # 10 points for 9 internal quantiles
+        quantile_labels = [f"{p:.1%}" for p in probs]
+        ax1.set_xticks(sigma_bins[1:-1])  # Skip boundary bins
+        ax1.set_xticklabels(quantile_labels[1:-1], rotation=45)
+        ax1.set_xscale("log")
+        ax2.set_xticks(sigma_bins[1:-1])
+        ax2.set_xticklabels(quantile_labels[1:-1], rotation=45)
+        ax2.set_xscale("log")
+
+        plt.tight_layout()
+        fig_img = wandb.Image(fig)
+        plt.close(fig)
+
+        return fig_img
+
+    def _process_frame_stats(self, sigma: torch.Tensor, loss: torch.Tensor, frame_idx: int) -> dict:
+        """Calculate statistics for a specific frame"""
+        return {
+            "sigma_log_mean": float(sigma.log().mean()),
+            "sigma_log_std": float(sigma.log().std()),
+            "loss_mean": float(loss.mean()),
+            "loss_std": float(loss.std()),
+            "loss_min": float(loss.min()),
+            "loss_max": float(loss.max()),
+            "loss_median": float(loss.median()),
+            "loss_q1": float(torch.quantile(loss.float(), 0.25)),
+            "loss_q3": float(torch.quantile(loss.float(), 0.75)),
+        }
+
+    def _gather_and_save(self, cache: _SigmaLossCache, iteration: int, prefix: str, log_viz: bool = True) -> dict:
+        info = {}
+
+        # Gather data from all ranks
+        local_sigma, local_loss, num_frames = cache.get_arrays()
+        world_size = dist.get_world_size()
+
+        if world_size > 1:
+            # Gather sizes first
+            local_size = torch.tensor([len(local_sigma)], dtype=torch.long, device="cuda")
+            sizes = [torch.zeros_like(local_size) for _ in range(world_size)]
+            dist.all_gather(sizes, local_size)
+            sizes = [s.item() for s in sizes]
+
+            # Gather data
+            max_size = max(sizes)
+            if max_size > 0:
+                # Move to GPU for gathering
+                padded_sigma = torch.zeros(max_size, num_frames or 1, dtype=torch.bfloat16, device="cuda")
+                padded_loss = torch.zeros(max_size, num_frames or 1, dtype=torch.bfloat16, device="cuda")
+
+                if len(local_sigma) > 0:
+                    padded_sigma[: len(local_sigma)] = local_sigma.cuda()
+                    padded_loss[: len(local_loss)] = local_loss.cuda()
+
+                all_sigma = [torch.zeros_like(padded_sigma) for _ in range(world_size)]
+                all_loss = [torch.zeros_like(padded_loss) for _ in range(world_size)]
+
+                dist.all_gather(all_sigma, padded_sigma)
+                dist.all_gather(all_loss, padded_loss)
+
+                if distributed.is_rank0():
+                    # Combine data from all ranks
+                    valid_sigma = []
+                    valid_loss = []
+                    for sigma, loss, size in zip(all_sigma, all_loss, sizes):
+                        if size > 0:
+                            valid_sigma.append(sigma[:size])
+                            valid_loss.append(loss[:size])
+
+                    if valid_sigma:
+                        sigma_arr = torch.cat(valid_sigma)
+                        loss_arr = torch.cat(valid_loss)
+
+                        # Overall statistics
+                        info[f"{prefix}/total_samples"] = sigma_arr.shape[0]
+
+                        # Per-frame statistics
+                        if num_frames and num_frames > 1:
+                            for t in range(num_frames):
+                                frame_stats = self._process_frame_stats(sigma_arr[:, t], loss_arr[:, t], t)
+                                frame_prefix = f"{prefix}/frame_{t}"
+                                info.update({f"{frame_prefix}/{k}": v for k, v in frame_stats.items()})
+
+                                # Create per-frame visualization
+                                if log_viz:
+                                    fig_img = self._create_analysis_plots(sigma_arr[:, t], loss_arr[:, t], t)
+                                    if fig_img is not None:
+                                        info[f"{frame_prefix}/distribution_plot"] = fig_img
+                        else:
+                            # Single frame case (images or single-frame stats)
+                            frame_stats = self._process_frame_stats(sigma_arr.squeeze(), loss_arr.squeeze(), None)
+                            info.update({f"{prefix}/{k}": v for k, v in frame_stats.items()})
+
+                            # Create visualization
+                            if log_viz:
+                                fig_img = self._create_analysis_plots(sigma_arr.squeeze(), loss_arr.squeeze())
+                                if fig_img is not None:
+                                    info[f"{prefix}/distribution_plot"] = fig_img
+
+                        if self.save_s3:
+                            save_data = {
+                                "sigma": sigma_arr.cpu(),
+                                "loss": loss_arr.cpu(),
+                                "stats": {k: v for k, v in info.items() if not isinstance(v, wandb.Image)},
+                            }
+                            easy_io.dump(
+                                save_data,
+                                f"s3://rundir/{self.name}/{prefix}_Iter{iteration:09d}.pkl",
+                            )
+
+        cache.reset()
+        return info
+
+    def on_training_step_end(
+        self,
+        model: ImaginaireModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor],
+        loss: torch.Tensor,
+        iteration: int = 0,
+    ):
+        sigma = output_batch["sigma"]
+        loss_per_frame = output_batch["edm_loss_per_frame"]
+
+        if model.is_image_batch(data_batch):
+            self.image_cache.add(sigma, loss_per_frame)
+        else:
+            self.video_cache.add(sigma, loss_per_frame)
+
+        if iteration % (self.config.trainer.logging_iter * self.logging_iter_multipler) == 0:
+            info = {}
+
+            with misc.timer("sigma_loss_analysis"):
+                log_viz = iteration % (self.config.trainer.logging_iter * self.logging_viz_iter_multipler) == 0
+                # Process image data
+                if len(self.image_cache.sigma_list) > 0:
+                    info.update(self._gather_and_save(self.image_cache, iteration, "sigma_loss_image", log_viz=log_viz))
+
+                # Process video data
+                if len(self.video_cache.sigma_list) > 0:
+                    info.update(self._gather_and_save(self.video_cache, iteration, "sigma_loss_video", log_viz=log_viz))
+
+                if distributed.is_rank0() and info and wandb.run:
+                    wandb.log(info, step=iteration)

cosmos_predict2/_src/predict2_multiview/conditioner.py

@@ -0,0 +1,103 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+import torch
+
+from cosmos_predict2._src.imaginaire.utils.easy_io import easy_io
+from cosmos_predict2._src.predict2.conditioner import AbstractEmbModel
+
+
+class MVTextAttr(AbstractEmbModel):
+    def __init__(
+        self,
+        input_key: List[str],
+        dropout_rate: Optional[float] = 0.0,
+        use_empty_string: bool = False,
+        empty_string_embeddings_path: str = "s3://bucket/predict2_assets/reason1_empty_string_embeddings.pt",
+        credential_path: str = "credentials/s3_training.secret",
+        single_caption_length: int = 512,
+    ):
+        super().__init__()
+        self._input_key = input_key
+        self._dropout_rate = dropout_rate
+        # if True, will use empty string embeddings
+        # otherwise use zero tensor embeddings
+        self.use_empty_string = use_empty_string
+        self._empty_string_embeddings_cache = None
+        self.empty_string_embeddings_path = empty_string_embeddings_path
+        self.credential_path = credential_path
+        self.single_caption_length = single_caption_length
+
+    def forward(self, token: torch.Tensor):
+        return {"crossattn_emb": token}
+
+    def _get_empty_string_embeddings(self) -> torch.Tensor:
+        """Lazy load and cache empty string embeddings."""
+        if self._empty_string_embeddings_cache is None:
+            self._empty_string_embeddings_cache = easy_io.load(
+                self.empty_string_embeddings_path,
+                backend_args={"backend": "s3", "s3_credential_path": self.credential_path},
+            )
+        return self._empty_string_embeddings_cache
+
+    def random_dropout_input(
+        self, in_tensor: torch.Tensor, dropout_rate: Optional[float] = None, key: Optional[str] = None
+    ) -> torch.Tensor:
+        if key is not None and "mask" in key:
+            return in_tensor
+
+        dropout_rate = dropout_rate if dropout_rate is not None else self._dropout_rate
+
+        B = in_tensor.shape[0]  # batch size
+        S_per_view = self.single_caption_length  # sequence length per view
+        if in_tensor.shape[1] % S_per_view != 0:
+            raise ValueError(
+                f"in_tensor sequence length {in_tensor.shape[1]} is not divisible by single_caption_length {S_per_view}"
+            )
+        V = in_tensor.shape[1] // S_per_view  # number of views
+
+        # reshape input tensor to [B, V, S, C]
+        in_tensor_reshaped = in_tensor.view(B, V, S_per_view, -1)
+
+        # create independent dropout mask for each view: [B, V]
+        dropout_rates = torch.ones(B, V, device=in_tensor.device) * dropout_rate
+        keep_mask = torch.bernoulli(1.0 - dropout_rates).type_as(in_tensor)
+        keep_mask = keep_mask.view(B, V, 1, 1)  # broadcastable shape
+
+        # prepare empty prompt data
+        if not self.use_empty_string:
+            empty_string_embeddings = torch.zeros(
+                1, 1, S_per_view, in_tensor_reshaped.shape[-1], dtype=in_tensor.dtype, device=in_tensor.device
+            )
+        else:
+            empty_string_embeddings = (
+                self._get_empty_string_embeddings().to(dtype=in_tensor.dtype, device=in_tensor.device).unsqueeze(1)
+            )  # [1, 1, 512, C]
+
+        # expand empty prompt data to all views: [B, V, S, C]
+        empty_string_embeddings = empty_string_embeddings.expand(B, V, S_per_view, -1)
+
+        # apply dropout independently for each view
+        output_reshaped = keep_mask * in_tensor_reshaped + (1.0 - keep_mask) * empty_string_embeddings
+
+        # reshape back to original shape: [B, V * S, C]
+        return output_reshaped.view(B, V * S_per_view, -1)
+
+    def details(self) -> str:
+        return "Output key: [crossattn_emb]"

cosmos_predict2/_src/predict2_multiview/configs/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/config.py

@@ -0,0 +1,47 @@
+# 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.
+
+from cosmos_predict2._src.imaginaire.flags import INTERNAL
+from cosmos_predict2._src.imaginaire.utils.config_helper import import_all_modules_from_package
+from cosmos_predict2._src.predict2.configs.video2world.config import make_config as vid2vid_make_config
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.callbacks import register_callbacks
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.conditioner import register_conditioner
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.dataloader import (
+    register_multiview_dataloader,
+)
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.dataloader_local import register_waymo_dataloader
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.model import register_model
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.net import register_net
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.optimizer import register_optimizer
+from cosmos_predict2._src.predict2_multiview.datasets.nymeria_pairs import register_nymeria_pairs_dataloader
+from cosmos_predict2._src.predict2_multiview.datasets.comind_pairs import register_comind_data
+
+
+def make_config():
+    c = vid2vid_make_config()
+    c.job.project = "cosmos_predict2_multiview"
+    register_conditioner()
+    register_model()
+    register_net()
+    register_multiview_dataloader()
+    register_waymo_dataloader()
+    register_nymeria_pairs_dataloader()
+    register_comind_data()
+    register_callbacks()
+    register_optimizer()
+    import_all_modules_from_package("cosmos_predict2._src.predict2_multiview.configs.vid2vid.experiment", reload=True)
+    if not INTERNAL:
+        import_all_modules_from_package("cosmos_predict2.experiments.multiview", reload=True)
+    return c

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/defaults/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/defaults/callbacks.py

@@ -0,0 +1,50 @@
+# 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.
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.predict2_multiview.callbacks.log_weight import LogWeight
+from cosmos_predict2._src.predict2_multiview.callbacks.sigma_loss_analysis_per_frame import SigmaLossAnalysisPerFrame
+
+LOG_SIGMA_LOSS_CALLBACKS = dict(
+    sigma_loss_log=L(SigmaLossAnalysisPerFrame)(
+        save_s3="${upload_reproducible_setup}",
+        logging_iter_multipler=2,
+        logging_viz_iter_multipler=10,
+    ),
+)
+
+LOG_WEIGHT_CALLBACKS = dict(
+    log_weight=L(LogWeight)(
+        every_n=100,
+    ),
+)
+
+
+def register_callbacks():
+    cs = ConfigStore.instance()
+    cs.store(
+        group="callbacks",
+        package="trainer.callbacks",
+        name="log_sigma_loss",
+        node=LOG_SIGMA_LOSS_CALLBACKS,
+    )
+    cs.store(
+        group="callbacks",
+        package="trainer.callbacks",
+        name="log_weight",
+        node=LOG_WEIGHT_CALLBACKS,
+    )

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/defaults/conditioner.py

@@ -0,0 +1,720 @@
+# 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 copy
+import random
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from einops import rearrange
+from hydra.core.config_store import ConfigStore
+from omegaconf import ListConfig
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.imaginaire.lazy_config import LazyDict
+from cosmos_predict2._src.imaginaire.utils import log
+from cosmos_predict2._src.imaginaire.utils.context_parallel import broadcast_split_tensor
+from cosmos_predict2._src.imaginaire.utils.validator import Validator
+from cosmos_predict2._src.predict2.conditioner import Text2WorldCondition, TextAttr
+from cosmos_predict2._src.predict2.configs.video2world.defaults.conditioner import (
+    _SHARED_CONFIG,
+    GeneralConditioner,
+    ReMapkey,
+    Video2WorldCondition,
+)
+from cosmos_predict2._src.predict2_multiview.conditioner import MVTextAttr
+
+
+class ConditionLocation(Enum):
+    """
+    Enum representing different camera condition locations for anymulti-to-multiview video generation.
+
+    Attributes:
+        NO_CAM: Indicates no camera is used for conditioning (i.e text2world)
+        REF_CAM: Indicates a reference camera is used for conditioning. (i.e single-to-multiview-text2world)
+        ANY_CAM: Indicates any camera can be used for conditioning. (i.e any-to-multiview-text2world)
+        FIRST_RANDOM_N: Indicates a random number of frames from all cameras are used for conditioning. (i.e video2world-multiview)
+
+    Note: Multiple locations can be set together when compatible.
+        - NO_CAM cannot be set with any other location.
+        - ANY_CAM and REF_CAM cannot be set simultaneously.
+        - FIRST_RANDOM_N can be set with ANY_CAM or REF_CAM.
+    """
+
+    NO_CAM = "no_cam"
+    REF_CAM = "ref_cam"
+    ANY_CAM = "any_cam"
+    FIRST_RANDOM_N = "first_random_n"
+
+
+class ConditionLocationListValidator(Validator):
+    """
+    Validator for a list of ConditionLocation objects.
+    Validates that:
+        - NO_CAM is not set with any other location
+        - ANY_CAM and REF_CAM are not set together
+    """
+
+    def __init__(self, default: List[ConditionLocation], hidden=False, tooltip=None):
+        self.default = default
+        self.hidden = hidden
+        self.tooltip = tooltip
+
+    def validate(self, value: List[ConditionLocation]):
+        for v in value:
+            if not isinstance(v, ConditionLocation):
+                raise TypeError(f"All elements must be ConditionLocation enums, got {type(v)}: {v}")
+        if ConditionLocation.NO_CAM in value:
+            assert len(value) == 1, f"Cannot set ConditionLocation.NO_CAM and other locations together. Got {value=}"
+        elif ConditionLocation.ANY_CAM in value and ConditionLocation.REF_CAM in value:
+            raise ValueError("ConditionLocation.ANY_CAM and ConditionLocation.REF_CAM cannot be set together.")
+        return value
+
+    def __repr__(self) -> str:
+        return f"ConditionLocationValidator({self.default=}, {self.hidden=})"
+
+    def json(self):
+        return {
+            "type": ConditionLocationListValidator.__name__,
+            "default": self.default,
+            "tooltip": self.tooltip,
+        }
+
+
+class ConditionLocationList(list):
+    def __init__(self, locations: List[ConditionLocation]):
+        enum_locations = []
+        for loc in locations:
+            if not isinstance(loc, ConditionLocation):
+                loc = ConditionLocation(loc)  # Will raise ValueError if invalid
+            enum_locations.append(loc)
+        super().__init__(enum_locations)
+        self.validator = ConditionLocationListValidator(default=[])
+        self.validator.validate(self)
+
+    def __repr__(self) -> str:
+        return f"ConditionLocationList({super().__repr__()})"
+
+    def to_json(self):
+        return {
+            "type": ConditionLocationList.__name__,
+            "locations": [location.value for location in self],
+        }
+
+
+@dataclass(frozen=True)
+class MultiViewCondition(Video2WorldCondition):
+    state_t: Optional[int] = None
+    view_indices_B_T: Optional[torch.Tensor] = None
+    ref_cam_view_idx_sample_position: Optional[torch.Tensor] = None
+    # Pose / warped-frame conditioning for 2-actor joint generation. All laid out (B, C, V*T, H, W):
+    #   pose_map_B_C_T_H_W      : projected 2D pose RGB at pixel resolution (encoded by the net's pose encoder)
+    #   pose_latent_B_C_T_H_W   : frozen-VAE pose latent (16 ch, latent resolution) for vae_concat/vae_mlp_add
+    #   warped_latent_B_C_T_H_W : warped past-frame VAE latent (16 ch, latent resolution)
+    #   visibility_mask_B_C_T_H_W: warp-validity mask (1 ch, latent resolution)
+    pose_map_B_C_T_H_W: Optional[torch.Tensor] = None
+    pose_latent_B_C_T_H_W: Optional[torch.Tensor] = None
+    warped_latent_B_C_T_H_W: Optional[torch.Tensor] = None
+    visibility_mask_B_C_T_H_W: Optional[torch.Tensor] = None
+    #   reference_latent_B_C_R_H_W: clean reference-frame VAE latent (16 ch, latent res), (B, C, V*R, H, W),
+    #                               appended as extra in-context frames inside the net for appearance conditioning
+    reference_latent_B_C_R_H_W: Optional[torch.Tensor] = None
+    #   plucker_map_B_C_T_H_W: per-pixel Plücker ray map (6 ch = dir 3 + moment 3, latent res), (B, 6, V*T, H, W),
+    #                          in a per-pair canonical frame (view0 frame0) for cross-view shared-space grounding
+    plucker_map_B_C_T_H_W: Optional[torch.Tensor] = None
+    #   reference_plucker_map_B_C_R_H_W: Plücker rays of the reference frames' past poses, same canonical frame
+    reference_plucker_map_B_C_R_H_W: Optional[torch.Tensor] = None
+    #   reference_pose_latent_B_C_R_H_W: VAE latent (16 ch) of each reference frame's per-person skeleton render,
+    #   (B, C, V*R, H, W), added to the reference tokens by the net's zero-init reference_pose_embedder.
+    reference_pose_latent_B_C_R_H_W: Optional[torch.Tensor] = None
+    #   depth_latent_B_C_T_H_W: composite-depth VAE latent (16 ch, latent res) = warped scene depth + human mesh depth
+    depth_latent_B_C_T_H_W: Optional[torch.Tensor] = None
+
+    def set_video_condition(
+        self,
+        state_t: int,
+        gt_frames: torch.Tensor,
+        condition_locations: Union[ConditionLocationList, ListConfig] = field(
+            default_factory=lambda: ConditionLocationList([])
+        ),
+        random_min_num_conditional_frames_per_view: Optional[int] = None,
+        random_max_num_conditional_frames_per_view: Optional[int] = None,
+        num_conditional_frames_per_view: Optional[int | List[int]] = None,
+        condition_cam_idx: Optional[int] = None,
+        view_condition_dropout_max: int = 0,
+        conditional_frames_probs: Optional[Dict[int, float]] = None,
+    ) -> "MultiViewCondition":
+        """
+        Sets the video conditioning frames for anymulti-to-multiview generation.
+
+        This method creates a conditioning mask for the input video frames that determines
+        which frames will be used as context frames for generating new frames. The method
+        handles video batches (T>1) and does not support images (T=1).
+
+        Args:
+            gt_frames: A tensor of ground truth frames with shape [B, C, T, H, W], where:
+                B = batch size
+                C = number of channels
+                T = number of frames per view * self.sample_n_views
+                H = height
+                W = width
+
+            random_min_num_conditional_frames_per_view: Minimum number of frames per view to use for conditioning
+                when randomly selecting a number of conditioning frames.
+
+            random_max_num_conditional_frames_per_view: Maximum number of frames per view to use for conditioning
+                when randomly selecting a number of conditioning frames.
+
+            num_conditional_frames_per_view: Optional[int | List[int]]; If provided, all examples in the batch will use
+                exactly this many frames per view for conditioning. If None, a random number of frames per view
+                between random_min_num_conditional_frames_per_view and random_max_num_conditional_frames_per_view
+                will be selected for each example in the batch. Can also be a list of integers, one for each view.
+
+            condition_cam_idx: Optional; Used only if ConditionLocation.ANY_CAM is in condition_locations.
+                If provided, all examples in the batch will use the same cam_idx for conditioning. If None,
+                a random cam_idx will be selected for each example in the batch.
+            view_condition_dropout_max: Optional; If provided and > 0, then a random number of views will be dropped from the conditioning.
+
+            conditional_frames_probs: Optional; Dictionary mapping number of frames to probabilities.
+                If provided, overrides the random_min/max_num_conditional_frames with weighted sampling.
+                Example: {0: 0.5, 1: 0.25, 2: 0.25} for 50% chance of 0 frames, 25% for 1, 25% for 2.
+
+        Returns:
+            A new MultiViewCondition object with the gt_frames and conditioning mask set.
+            The conditioning mask (condition_video_input_mask_B_C_T_H_W) is a binary tensor
+            of shape [B, 1, T, H, W] where 1 indicates frames used for conditioning and 0
+            indicates frames to be generated.
+
+        Notes:
+            - Image batches are not supported.
+            - For video batches multiple condition_locations can be provided and combined:
+                - If num_conditional_frames_per_view is provided and "random_n" is in condition_locations,
+                then all examples will use the same number of frames per view for conditioning,
+                otherwise, if num_conditional_frames_per_view is not provided,
+                then each example will randomly uses between random_min_num_conditional_frames_per_view
+                and random_max_num_conditional_frames_per_view frames per view.
+                - If "ref_cam" is in condition_locations, then for each example,
+                all frames of the first view will be used for conditioning.
+        """
+        kwargs = self.to_dict(skip_underscore=False)
+        kwargs["state_t"] = state_t
+        kwargs["gt_frames"] = gt_frames
+        B, _, T, H, W = gt_frames.shape
+
+        if not isinstance(condition_locations, ConditionLocationList):
+            condition_locations = ConditionLocationList(condition_locations)
+        assert len(condition_locations) > 0, "condition_locations must be provided."
+        assert state_t is not None, "state_t must be provided."
+        assert T > 1, "Image batches are not supported."
+        assert T % state_t == 0, f"T must be a multiple of state_t. Got T={T} and state_t={state_t}."
+        sample_n_views = T // state_t
+        condition_video_input_mask_B_C_V_T_H_W = torch.zeros(
+            B, 1, sample_n_views, state_t, H, W, dtype=gt_frames.dtype, device=gt_frames.device
+        )
+        views_eligible_for_dropout = list(range(sample_n_views))
+
+        if ConditionLocation.REF_CAM in condition_locations:
+            ref_cam_view_idx_sample_position = kwargs["ref_cam_view_idx_sample_position"]
+            ref_cam_idx_B = (
+                torch.ones(B, dtype=torch.int32, device=ref_cam_view_idx_sample_position.device)
+                * ref_cam_view_idx_sample_position
+            )
+            condition_video_input_mask_B_C_V_T_H_W = self.enable_ref_cam_condition(
+                ref_cam_idx_B, condition_video_input_mask_B_C_V_T_H_W
+            )
+            assert (ref_cam_view_idx_sample_position == ref_cam_view_idx_sample_position[0]).all(), (
+                f"ref_cam_view_idx_sample_position must be the same for all examples. Got {ref_cam_view_idx_sample_position=}"
+            )
+            ref_cam_view_idx_sample_position_int = ref_cam_view_idx_sample_position[0].cpu().item()
+            views_eligible_for_dropout.remove(ref_cam_view_idx_sample_position_int)
+        elif ConditionLocation.ANY_CAM in condition_locations:
+            if condition_cam_idx is None:
+                assert kwargs["view_indices_B_T"].shape[-1] % sample_n_views == 0, (
+                    f"view_indices_B_T last dimension must be a multiple of sample_n_views. Got view_indices_B_T.shape={kwargs['view_indices_B_T'].shape}, sample_n_views={sample_n_views}"
+                )
+                view_indices = kwargs["view_indices_B_T"]
+                selected_cam_latent_t_index = torch.randint(0, state_t, size=(B,))
+                any_cam_idx_B = view_indices[torch.arange(B), selected_cam_latent_t_index]
+            else:
+                any_cam_idx_B = torch.full((B,), condition_cam_idx, dtype=torch.int32)
+            condition_video_input_mask_B_C_V_T_H_W = self.enable_ref_cam_condition(
+                any_cam_idx_B, condition_video_input_mask_B_C_V_T_H_W
+            )
+            assert (any_cam_idx_B == any_cam_idx_B[0]).all(), (
+                f"any_cam_idx_B must be the same for all examples. Got {any_cam_idx_B=}"
+            )
+            any_cam_idx_B_int = any_cam_idx_B[0].cpu().item()
+            views_eligible_for_dropout.remove(any_cam_idx_B_int)
+        if ConditionLocation.FIRST_RANDOM_N in condition_locations:
+            if (
+                num_conditional_frames_per_view is None
+                and random_min_num_conditional_frames_per_view == random_max_num_conditional_frames_per_view
+            ):
+                num_conditional_frames_per_view = random_min_num_conditional_frames_per_view
+            if num_conditional_frames_per_view is not None:
+                if isinstance(num_conditional_frames_per_view, list):
+                    assert len(num_conditional_frames_per_view) == sample_n_views, (
+                        f"num_conditional_frames_per_view must be a list of length {sample_n_views}. Got {num_conditional_frames_per_view=}"
+                    )
+                    log.info(
+                        f"Setting num_conditional_frames_per_view_B_V explicitly from list: {num_conditional_frames_per_view}"
+                    )
+                    num_conditional_frames_per_view_B_V = torch.tensor(
+                        num_conditional_frames_per_view, dtype=torch.int32
+                    ).repeat(B, 1)
+                else:
+                    num_conditional_frames_per_view_B_V = (
+                        torch.ones((B, sample_n_views), dtype=torch.int32) * num_conditional_frames_per_view
+                    )
+            elif conditional_frames_probs is not None:
+                # Use weighted sampling based on provided probabilities
+                frames_options = list(conditional_frames_probs.keys())
+                weights = list(conditional_frames_probs.values())
+                num_conditional_frames_per_view_B_V = (
+                    torch.tensor(random.choices(frames_options, weights=weights, k=B), dtype=torch.int32)
+                    .view(B, 1)
+                    .repeat(1, sample_n_views)
+                )
+            else:
+                assert (
+                    random_min_num_conditional_frames_per_view is not None
+                    and random_max_num_conditional_frames_per_view is not None
+                ), (
+                    f"random_min_num_conditional_frames_per_view and random_max_num_conditional_frames_per_view must be provided if num_conditional_frames_per_view is None. Got {random_min_num_conditional_frames_per_view=}, {random_max_num_conditional_frames_per_view=}, {num_conditional_frames_per_view=}"
+                )
+                num_conditional_frames_per_view_B_V = torch.randint(
+                    random_min_num_conditional_frames_per_view,
+                    random_max_num_conditional_frames_per_view + 1,
+                    size=(B, 1),
+                ).repeat(1, sample_n_views)
+            condition_video_input_mask_B_C_V_T_H_W = self.enable_first_random_n_condition(
+                condition_video_input_mask_B_C_V_T_H_W, num_conditional_frames_per_view_B_V
+            )
+        if view_condition_dropout_max > 0:
+            random.shuffle(views_eligible_for_dropout)
+            n_views_to_dropout = random.randint(0, view_condition_dropout_max)
+            views_to_dropout = views_eligible_for_dropout[:n_views_to_dropout]
+            for view_idx in views_to_dropout:
+                condition_video_input_mask_B_C_V_T_H_W[:, :, view_idx] = 0
+
+        condition_video_input_mask_B_C_T_H_W = rearrange(
+            condition_video_input_mask_B_C_V_T_H_W, "B C V T H W -> B C (V T) H W", V=sample_n_views
+        )
+        kwargs["condition_video_input_mask_B_C_T_H_W"] = condition_video_input_mask_B_C_T_H_W
+        return type(self)(**kwargs)
+
+    def enable_ref_cam_condition(self, cam_idx_B: torch.Tensor, condition_video_input_mask_B_C_V_T_H_W: torch.Tensor):
+        """
+        Sets condition video input mask to 1 for all frames of the cam_idx[i] view in each example i
+        Args:
+            cam_idx_B: A tensor of shape [B]
+            condition_video_input_mask_B_C_V_T_H_W: A tensor of shape [B, 1, V, T, H, W]
+            where V is the number of views, T is the number of frames per view, H is the height, and W is the width
+        Returns:
+            A copy of the condition video input mask with the cam_idx[i] view set to 1 for example i
+        """
+        assert condition_video_input_mask_B_C_V_T_H_W.ndim == 6, (
+            f"condition_video_input_mask_B_C_V_T_H_W must have 6 dimensions. Got {condition_video_input_mask_B_C_V_T_H_W.shape=}"
+        )
+        assert cam_idx_B.ndim == 1, f"cam_idx_B must have 1 dimension. Got {cam_idx_B.shape=}"
+        copy_condition_video_input_mask_B_C_V_T_H_W = condition_video_input_mask_B_C_V_T_H_W.clone()
+        for i in range(copy_condition_video_input_mask_B_C_V_T_H_W.shape[0]):
+            copy_condition_video_input_mask_B_C_V_T_H_W[i, :, cam_idx_B[i]] = 1
+        return copy_condition_video_input_mask_B_C_V_T_H_W
+
+    def enable_first_random_n_condition(
+        self, condition_video_input_mask_B_C_V_T_H_W: torch.Tensor, num_conditional_frames_per_view_B_V: torch.Tensor
+    ):
+        """
+        Sets condition video input mask to 1 for the first num_conditional_frames_per_view_B frames of each view
+        Args:
+            condition_video_input_mask_B_C_V_T_H_W: A tensor of shape [B, 1, V, T, H, W]
+            num_conditional_frames_per_view_B_V: A tensor of shape [B, V]
+        Returns:
+            A copy of the condition video input mask with the first num_conditional_frames_per_view_B_V frames of each view set to 1
+        """
+        assert condition_video_input_mask_B_C_V_T_H_W.ndim == 6, (
+            "condition_video_input_mask_B_C_V_T_H_W must have 6 dimensions"
+        )
+        B, _, _, _, _, _ = condition_video_input_mask_B_C_V_T_H_W.shape
+        copy_condition_video_input_mask_B_C_V_T_H_W = condition_video_input_mask_B_C_V_T_H_W.clone()
+        for idx in range(B):
+            for view_idx in range(num_conditional_frames_per_view_B_V.shape[1]):
+                copy_condition_video_input_mask_B_C_V_T_H_W[
+                    idx, :, view_idx, : num_conditional_frames_per_view_B_V[idx, view_idx]
+                ] = 1
+        return copy_condition_video_input_mask_B_C_V_T_H_W
+
+    def edit_for_inference(
+        self,
+        condition_locations: Union[ConditionLocationList, ListConfig] = field(
+            default_factory=lambda: ConditionLocationList([])
+        ),
+        is_cfg_conditional: bool = True,
+        num_conditional_frames_per_view: int = 1,
+    ) -> "MultiViewCondition":
+        _condition = self.set_video_condition(
+            state_t=self.state_t,
+            gt_frames=self.gt_frames,
+            condition_locations=condition_locations,
+            random_min_num_conditional_frames_per_view=0,
+            random_max_num_conditional_frames_per_view=0,
+            num_conditional_frames_per_view=num_conditional_frames_per_view,
+            view_condition_dropout_max=0,
+        )
+        if not is_cfg_conditional:
+            # Do not use classifier free guidance on conditional frames.
+            # YB found that it leads to worse results.
+            _condition.use_video_condition.fill_(True)
+        return _condition
+
+    def broadcast(self, process_group: torch.distributed.ProcessGroup) -> "MultiViewCondition":
+        if self.is_broadcasted:
+            return self
+        gt_frames_B_C_T_H_W = self.gt_frames
+        view_indices_B_T = self.view_indices_B_T
+        condition_video_input_mask_B_C_T_H_W = self.condition_video_input_mask_B_C_T_H_W
+        # Pose conditioning tensors are laid out (V*T) like gt_frames. Keep them out of the generic
+        # broadcast (which would .cuda()+broadcast from rank0) and re-attach unchanged. For CP==1 (our
+        # target) this pass-through is correct; CP>1 would need the same seq-split as gt_frames below.
+        pose_map_B_C_T_H_W = self.pose_map_B_C_T_H_W
+        pose_latent_B_C_T_H_W = self.pose_latent_B_C_T_H_W
+        warped_latent_B_C_T_H_W = self.warped_latent_B_C_T_H_W
+        visibility_mask_B_C_T_H_W = self.visibility_mask_B_C_T_H_W
+        reference_latent_B_C_R_H_W = self.reference_latent_B_C_R_H_W
+        plucker_map_B_C_T_H_W = self.plucker_map_B_C_T_H_W
+        reference_plucker_map_B_C_R_H_W = self.reference_plucker_map_B_C_R_H_W
+        reference_pose_latent_B_C_R_H_W = self.reference_pose_latent_B_C_R_H_W
+        depth_latent_B_C_T_H_W = self.depth_latent_B_C_T_H_W
+        kwargs = self.to_dict(skip_underscore=False)
+        kwargs["gt_frames"] = None
+        kwargs["condition_video_input_mask_B_C_T_H_W"] = None
+        kwargs["view_indices_B_T"] = None
+        kwargs["pose_map_B_C_T_H_W"] = None
+        kwargs["pose_latent_B_C_T_H_W"] = None
+        kwargs["warped_latent_B_C_T_H_W"] = None
+        kwargs["visibility_mask_B_C_T_H_W"] = None
+        kwargs["reference_latent_B_C_R_H_W"] = None
+        kwargs["plucker_map_B_C_T_H_W"] = None
+        kwargs["reference_plucker_map_B_C_R_H_W"] = None
+        kwargs["reference_pose_latent_B_C_R_H_W"] = None
+        kwargs["depth_latent_B_C_T_H_W"] = None
+        new_condition = Text2WorldCondition.broadcast(
+            type(self)(**kwargs),
+            process_group,
+        )
+
+        kwargs = new_condition.to_dict(skip_underscore=False)
+        _, _, T, _, _ = gt_frames_B_C_T_H_W.shape
+        n_views = T // self.state_t
+        assert T % self.state_t == 0, f"T must be a multiple of state_t. Got T={T} and state_t={self.state_t}."
+        if process_group is not None:
+            if T > 1 and process_group.size() > 1:
+                log.debug(f"Broadcasting {gt_frames_B_C_T_H_W.shape=} to {n_views=} views")
+                gt_frames_B_C_V_T_H_W = rearrange(gt_frames_B_C_T_H_W, "B C (V T) H W -> B C V T H W", V=n_views)
+                condition_video_input_mask_B_C_V_T_H_W = rearrange(
+                    condition_video_input_mask_B_C_T_H_W, "B C (V T) H W -> B C V T H W", V=n_views
+                )
+                view_indices_B_V_T = rearrange(view_indices_B_T, "B (V T) -> B V T", V=n_views)
+
+                gt_frames_B_C_V_T_H_W = broadcast_split_tensor(
+                    gt_frames_B_C_V_T_H_W, seq_dim=3, process_group=process_group
+                )
+                condition_video_input_mask_B_C_V_T_H_W = broadcast_split_tensor(
+                    condition_video_input_mask_B_C_V_T_H_W, seq_dim=3, process_group=process_group
+                )
+                view_indices_B_V_T = broadcast_split_tensor(view_indices_B_V_T, seq_dim=2, process_group=process_group)
+
+                gt_frames_B_C_T_H_W = rearrange(gt_frames_B_C_V_T_H_W, "B C V T H W -> B C (V T) H W", V=n_views)
+                condition_video_input_mask_B_C_T_H_W = rearrange(
+                    condition_video_input_mask_B_C_V_T_H_W, "B C V T H W -> B C (V T) H W", V=n_views
+                )
+                view_indices_B_T = rearrange(view_indices_B_V_T, "B V T -> B (V T)", V=n_views)
+
+        kwargs["gt_frames"] = gt_frames_B_C_T_H_W
+        kwargs["condition_video_input_mask_B_C_T_H_W"] = condition_video_input_mask_B_C_T_H_W
+        kwargs["view_indices_B_T"] = view_indices_B_T
+        kwargs["pose_map_B_C_T_H_W"] = pose_map_B_C_T_H_W
+        kwargs["pose_latent_B_C_T_H_W"] = pose_latent_B_C_T_H_W
+        kwargs["warped_latent_B_C_T_H_W"] = warped_latent_B_C_T_H_W
+        kwargs["visibility_mask_B_C_T_H_W"] = visibility_mask_B_C_T_H_W
+        kwargs["reference_latent_B_C_R_H_W"] = reference_latent_B_C_R_H_W
+        kwargs["plucker_map_B_C_T_H_W"] = plucker_map_B_C_T_H_W
+        kwargs["reference_plucker_map_B_C_R_H_W"] = reference_plucker_map_B_C_R_H_W
+        kwargs["reference_pose_latent_B_C_R_H_W"] = reference_pose_latent_B_C_R_H_W
+        kwargs["depth_latent_B_C_T_H_W"] = depth_latent_B_C_T_H_W
+        return type(self)(**kwargs)
+
+
+class MultiViewConditioner(GeneralConditioner):
+    def _forward(self, batch: Dict, override_dropout_rate: Optional[Dict[str, float]] = None) -> Dict:
+        """Like GeneralConditioner._forward but SKIPS any embedder whose (string) input_key is absent from the
+        batch. This lets optional conditioning (reference_latent / plucker_map) be present in the shared
+        conditioner config while individual experiments produce only the subset they enable."""
+        from collections import defaultdict
+        from contextlib import nullcontext
+
+        output = defaultdict(list)
+        override_dropout_rate = override_dropout_rate or {}
+        for emb_name in override_dropout_rate.keys():
+            assert emb_name in self.embedders, f"invalid name found {emb_name}"
+        for emb_name, embedder in self.embedders.items():
+            if isinstance(embedder.input_key, str) and embedder.input_key not in batch:
+                continue  # optional condition not produced by this experiment -> skip (no KeyError)
+            context = nullcontext if embedder.is_trainable else torch.no_grad
+            with context():
+                if isinstance(embedder.input_key, str):
+                    emb_out = embedder(
+                        embedder.random_dropout_input(
+                            batch[embedder.input_key], override_dropout_rate.get(emb_name, None)
+                        )
+                    )
+                else:
+                    emb_out = embedder(
+                        *[
+                            embedder.random_dropout_input(batch.get(k), override_dropout_rate.get(emb_name, None), k)
+                            for k in embedder.input_key
+                        ]
+                    )
+            for k, v in emb_out.items():
+                output[k].append(v)
+        return {k: torch.cat(v, dim=self.KEY2DIM.get(k, -1)) for k, v in output.items()}
+
+    def forward(self, batch: Dict, override_dropout_rate: Optional[Dict[str, float]] = None) -> MultiViewCondition:
+        output = self._forward(batch, override_dropout_rate)
+        return MultiViewCondition(**output)
+
+    def get_condition_with_negative_prompt(
+        self,
+        data_batch: Dict,
+    ) -> Tuple[Any, Any]:
+        """
+        Similar functionality as get_condition_uncondition
+        But use negative prompts for unconditon
+        """
+        cond_dropout_rates, uncond_dropout_rates = {}, {}
+        for emb_name, embedder in self.embedders.items():
+            cond_dropout_rates[emb_name] = 0.0
+            if isinstance(embedder, TextAttr) or isinstance(embedder, MVTextAttr):
+                uncond_dropout_rates[emb_name] = 0.0
+            else:
+                uncond_dropout_rates[emb_name] = 1.0 if embedder.dropout_rate > 1e-4 else 0.0
+
+        data_batch_neg_prompt = copy.deepcopy(data_batch)
+        if "neg_t5_text_embeddings" in data_batch_neg_prompt:
+            if isinstance(data_batch_neg_prompt["neg_t5_text_embeddings"], torch.Tensor):
+                data_batch_neg_prompt["t5_text_embeddings"] = data_batch_neg_prompt["neg_t5_text_embeddings"]
+
+        condition: Any = self(data_batch, override_dropout_rate=cond_dropout_rates)
+        un_condition: Any = self(data_batch_neg_prompt, override_dropout_rate=uncond_dropout_rates)
+
+        return condition, un_condition
+
+
+MultiViewConditionerConfig: LazyDict = L(MultiViewConditioner)(
+    **_SHARED_CONFIG,
+    view_indices_B_T=L(ReMapkey)(
+        input_key="latent_view_indices_B_T",
+        output_key="view_indices_B_T",
+        dropout_rate=0.0,
+        dtype=None,
+    ),
+    ref_cam_view_idx_sample_position=L(ReMapkey)(
+        input_key="ref_cam_view_idx_sample_position",
+        output_key="ref_cam_view_idx_sample_position",
+        dropout_rate=0.0,
+        dtype=None,
+    ),
+)
+
+
+class TextAttrEmptyStringDropout(TextAttr):
+    def __init__(
+        self,
+        input_key: str,
+        pos_input_key: str,
+        dropout_input_key: str,
+        dropout_rate: Optional[float] = 0.0,
+        use_empty_string: bool = False,
+        **kwargs,
+    ):
+        self._input_key = input_key
+        self._pos_input_key = pos_input_key
+        self._dropout_input_key = dropout_input_key
+        self._dropout_rate = dropout_rate
+        self._use_empty_string = use_empty_string
+        super().__init__(input_key, dropout_rate)
+
+    def forward(self, tensor: torch.Tensor):
+        return {"crossattn_emb": tensor}
+
+    def random_dropout_input(
+        self,
+        in_tensor_dict: torch.Tensor | Dict[str, torch.Tensor],
+        dropout_rate: Optional[float] = None,
+        key: Optional[str] = None,
+    ) -> torch.Tensor:
+        if key is not None and "mask" in key:
+            return in_tensor_dict
+        del key
+        assert isinstance(in_tensor_dict, dict), f"in_tensor_dict must be a dict. Got {type(in_tensor_dict)}"
+        in_tensor = in_tensor_dict[self._pos_input_key]
+        B = in_tensor.shape[0]
+        dropout_rate = dropout_rate if dropout_rate is not None else self.dropout_rate
+        keep_mask = torch.bernoulli((1.0 - dropout_rate) * torch.ones(B)).type_as(in_tensor)
+        if self._use_empty_string:
+            empty_prompt = in_tensor_dict[self._dropout_input_key]
+            if empty_prompt.shape[0] != B:
+                empty_prompt = empty_prompt.repeat(B, 1, 1)
+        else:
+            empty_prompt = torch.zeros_like(in_tensor)
+
+        return keep_mask * in_tensor + (1 - keep_mask) * empty_prompt
+
+    def details(self) -> str:
+        return "Output key: [crossattn_emb]"
+
+
+_SHARED_CONFIG_PER_VIEW_DROPOUT = copy.deepcopy(_SHARED_CONFIG)
+_SHARED_CONFIG_PER_VIEW_DROPOUT["text"] = L(MVTextAttr)(
+    input_key=["t5_text_embeddings"],
+    dropout_rate=0.2,
+    use_empty_string=False,
+)
+
+MultiViewConditionerPerViewDropoutConfig: LazyDict = L(MultiViewConditioner)(
+    **_SHARED_CONFIG_PER_VIEW_DROPOUT,
+    view_indices_B_T=L(ReMapkey)(
+        input_key="latent_view_indices_B_T",
+        output_key="view_indices_B_T",
+        dropout_rate=0.0,
+        dtype=None,
+    ),
+    ref_cam_view_idx_sample_position=L(ReMapkey)(
+        input_key="ref_cam_view_idx_sample_position",
+        output_key="ref_cam_view_idx_sample_position",
+        dropout_rate=0.0,
+        dtype=None,
+    ),
+)
+
+
+# Conditioner for 2-actor pose-conditioned joint generation. Adds three ReMapkey embedders that thread the
+# preprocessed pose / warped-latent / visibility tensors (written into the data_batch by the model's
+# get_data_and_condition) into the MultiViewCondition fields consumed by MultiViewPoseDiT.forward.
+#   - pose_map        : dropout 0.1 -> classifier-free-guidance on pose (zeroed in the uncond branch).
+#   - warped_latent   : dropout 0.0 -> geometric conditioning always present (no guidance).
+#   - visibility_mask : dropout 0.0 -> always present.
+MultiViewPoseConditionerConfig: LazyDict = L(MultiViewConditioner)(
+    **_SHARED_CONFIG,
+    view_indices_B_T=L(ReMapkey)(
+        input_key="latent_view_indices_B_T",
+        output_key="view_indices_B_T",
+        dropout_rate=0.0,
+        dtype=None,
+    ),
+    ref_cam_view_idx_sample_position=L(ReMapkey)(
+        input_key="ref_cam_view_idx_sample_position",
+        output_key="ref_cam_view_idx_sample_position",
+        dropout_rate=0.0,
+        dtype=None,
+    ),
+    pose_map=L(ReMapkey)(
+        input_key="pose_map",
+        output_key="pose_map_B_C_T_H_W",
+        dropout_rate=0.1,
+        dtype=None,
+    ),
+    warped_latent=L(ReMapkey)(
+        input_key="warped_latent",
+        output_key="warped_latent_B_C_T_H_W",
+        dropout_rate=0.0,
+        dtype=None,
+    ),
+    visibility_mask=L(ReMapkey)(
+        input_key="visibility_mask",
+        output_key="visibility_mask_B_C_T_H_W",
+        dropout_rate=0.0,
+        dtype=None,
+    ),
+)
+
+# Variant that ALSO remaps the frozen-VAE pose latent (for net pose_mode = vae_concat / vae_mlp_add).
+# pose_latent carries the classifier-free-guidance dropout (0.1) for pose; pose_map is still present but the
+# net ignores it in the VAE pose modes.
+MultiViewPoseLatentConditionerConfig: LazyDict = L(MultiViewConditioner)(
+    **_SHARED_CONFIG,
+    view_indices_B_T=L(ReMapkey)(
+        input_key="latent_view_indices_B_T", output_key="view_indices_B_T", dropout_rate=0.0, dtype=None,
+    ),
+    ref_cam_view_idx_sample_position=L(ReMapkey)(
+        input_key="ref_cam_view_idx_sample_position", output_key="ref_cam_view_idx_sample_position",
+        dropout_rate=0.0, dtype=None,
+    ),
+    pose_latent=L(ReMapkey)(
+        input_key="pose_latent", output_key="pose_latent_B_C_T_H_W", dropout_rate=0.1, dtype=None,
+    ),
+    warped_latent=L(ReMapkey)(
+        input_key="warped_latent", output_key="warped_latent_B_C_T_H_W", dropout_rate=0.0, dtype=None,
+    ),
+    visibility_mask=L(ReMapkey)(
+        input_key="visibility_mask", output_key="visibility_mask_B_C_T_H_W", dropout_rate=0.0, dtype=None,
+    ),
+    reference_latent=L(ReMapkey)(
+        input_key="reference_latent", output_key="reference_latent_B_C_R_H_W", dropout_rate=0.0, dtype=None,
+    ),
+    plucker_map=L(ReMapkey)(
+        input_key="plucker_map", output_key="plucker_map_B_C_T_H_W", dropout_rate=0.0, dtype=None,
+    ),
+    reference_plucker_map=L(ReMapkey)(
+        input_key="reference_plucker_map", output_key="reference_plucker_map_B_C_R_H_W", dropout_rate=0.0, dtype=None,
+    ),
+    reference_pose_latent=L(ReMapkey)(
+        input_key="reference_pose_latent", output_key="reference_pose_latent_B_C_R_H_W", dropout_rate=0.0, dtype=None,
+    ),
+    depth_latent=L(ReMapkey)(
+        input_key="depth_latent", output_key="depth_latent_B_C_T_H_W", dropout_rate=0.0, dtype=None,
+    ),
+)
+
+
+def register_conditioner():
+    cs = ConfigStore.instance()
+    cs.store(
+        group="conditioner",
+        package="model.config.conditioner",
+        name="video_prediction_multiview_conditioner",
+        node=MultiViewConditionerConfig,
+    )
+    cs.store(
+        group="conditioner",
+        package="model.config.conditioner",
+        name="video_prediction_multiview_pose_conditioner",
+        node=MultiViewPoseConditionerConfig,
+    )
+    cs.store(
+        group="conditioner",
+        package="model.config.conditioner",
+        name="video_prediction_multiview_pose_latent_conditioner",
+        node=MultiViewPoseLatentConditionerConfig,
+    )
+    cs.store(
+        group="conditioner",
+        package="model.config.conditioner",
+        name="video_prediction_multiview_conditioner_per_view_dropout",
+        node=MultiViewConditionerPerViewDropoutConfig,
+    )

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/defaults/dataloader.py

@@ -0,0 +1,101 @@
+# 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.
+
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.predict2_multiview.datasets.multiview import (
+    DEFAULT_CAMERAS,
+    AugmentationConfig,
+    get_multiview_video_loader,
+)
+
+DEFAULT_CAMERA_VIEW_CONFIGS = {
+    "7views": DEFAULT_CAMERAS,
+    "4views": [
+        "camera_front_wide_120fov",
+        "camera_cross_right_120fov",
+        "camera_rear_tele_30fov",
+        "camera_cross_left_120fov",
+    ],
+}
+
+
+def register_multiview_dataloader() -> None:
+    """Register multiview video dataloader configurations."""
+
+    cs = ConfigStore.instance()
+
+    # alpamayo
+    datasets = ["alpamayo_dec2024"]
+    object_stores = ["gcs", "s3"]
+    resolutions = [
+        ("480p", (480, 832)),
+        ("720p", (720, 1280)),
+        ("1080p", (1080, 1920)),
+    ]
+    fps = [
+        ("10fps", 3),
+        ("15fps", 2),
+        ("30fps", 1),
+    ]
+    num_video_frames = [
+        ("29frames", 29),
+        ("61frames", 61),
+        ("93frames", 93),
+    ]
+    cs.store(
+        group="data_val",
+        package="dataloader_val",
+        name="mock",
+        node=L(get_multiview_video_loader)(
+            dataset_name=datasets[0],
+            is_train=False,
+            object_store="s3",
+            augmentation_config=L(AugmentationConfig)(
+                resolution_hw=resolutions[0][1],
+                fps_downsample_factor=fps[0][1],
+                num_video_frames=num_video_frames[0][1],
+                camera_keys=DEFAULT_CAMERA_VIEW_CONFIGS["7views"],
+            ),
+            batch_size=1,
+            num_workers=2,
+        ),
+    )
+
+    for dataset in datasets:
+        for object_store in object_stores:
+            for resolution_str, resolution_hw in resolutions:
+                for fps_str, downsample_factor in fps:
+                    for num_video_frames_str, num_frames in num_video_frames:
+                        for views_str, camera_keys in DEFAULT_CAMERA_VIEW_CONFIGS.items():
+                            name = f"video_{dataset}_{object_store}_{resolution_str}_{fps_str}_{num_video_frames_str}_{views_str}"
+                            cs.store(
+                                group="data_train",
+                                package="dataloader_train",
+                                name=name,
+                                node=L(get_multiview_video_loader)(
+                                    dataset_name=dataset,
+                                    is_train=True,
+                                    object_store=object_store,
+                                    augmentation_config=L(AugmentationConfig)(
+                                        resolution_hw=resolution_hw,
+                                        fps_downsample_factor=downsample_factor,
+                                        num_video_frames=num_frames,
+                                        camera_keys=camera_keys,
+                                    ),
+                                ),
+                            )

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/defaults/dataloader_local.py

@@ -0,0 +1,111 @@
+# 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.
+
+"""Local file-based dataloader configurations."""
+
+import torch.distributed as dist
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.predict2.datasets.local_datasets.dataset_video import get_generic_dataloader, get_sampler
+from cosmos_predict2._src.predict2_multiview.datasets.local import WaymoLocalDataset
+from cosmos_predict2._src.predict2_multiview.datasets.multiview import AugmentationConfig, collate_fn
+
+
+def register_waymo_dataloader() -> None:
+    """Register local file-based dataloader configurations."""
+
+    cs = ConfigStore.instance()
+
+    waymo_dataset = L(WaymoLocalDataset)(
+        video_file_dirs=["datasets/multiview/waymo/input"],
+        augmentation_config=L(AugmentationConfig)(
+            resolution_hw=(720, 1280),
+            fps_downsample_factor=1,
+            num_video_frames=29,
+            camera_keys=[
+                "pinhole_front",
+                "pinhole_front_right",
+                "pinhole_side_right",
+                "pinhole_side_left",
+                "pinhole_front_left",
+            ],
+            camera_view_mapping={
+                "pinhole_front": 0,
+                "pinhole_front_right": 1,
+                "pinhole_side_right": 2,
+                # no pinhole_back in the dataset, so skip ID 3
+                "pinhole_side_left": 4,
+                "pinhole_front_left": 5,
+                # no pinehole_front_tele in the dataset skip ID 6
+            },
+            camera_video_key_mapping={
+                "pinhole_front": "video_pinhole_front",
+                "pinhole_front_right": "video_pinhole_front_right",
+                "pinhole_side_right": "video_pinhole_side_right",
+                "pinhole_side_left": "video_pinhole_side_left",
+                "pinhole_front_left": "video_pinhole_front_left",
+            },
+            camera_caption_key_mapping={
+                "pinhole_front": "caption_pinhole_front",
+                "pinhole_front_right": "caption_pinhole_front_right",
+                "pinhole_side_right": "caption_pinhole_side_right",
+                "pinhole_side_left": "caption_pinhole_side_left",
+                "pinhole_front_left": "caption_pinhole_front_left",
+            },
+            caption_probability={
+                "long": 1.0,
+            },
+            single_caption_camera_name="pinhole_front",
+            add_view_prefix_to_caption=True,
+            camera_prefix_mapping={
+                "pinhole_front": "The video is captured from a camera mounted on a car. The camera is facing forward.",
+                "pinhole_front_right": "The video is captured from a camera mounted on a car. The camera is facing to the front right.",
+                "pinhole_side_right": "The video is captured from a camera mounted on a car. The camera is facing to the side right.",
+                "pinhole_side_left": "The video is captured from a camera mounted on a car. The camera is facing to the side left.",
+                "pinhole_front_left": "The video is captured from a camera mounted on a car. The camera is facing to the front left.",
+            },
+        ),
+    )
+
+    cs.store(
+        group="data_train",
+        package="dataloader_train",
+        name=f"waymo",
+        node=L(get_generic_dataloader)(
+            dataset=waymo_dataset,
+            sampler=L(get_sampler)(dataset=waymo_dataset) if dist.is_initialized() else None,
+            collate_fn=collate_fn,
+            batch_size=1,
+            drop_last=True,
+            num_workers=4,
+            pin_memory=True,
+        ),
+    )
+
+    cs.store(
+        group="data_val",
+        package="dataloader_val",
+        name=f"waymo",
+        node=L(get_generic_dataloader)(
+            dataset=waymo_dataset,
+            sampler=L(get_sampler)(dataset=waymo_dataset) if dist.is_initialized() else None,
+            collate_fn=collate_fn,
+            batch_size=1,
+            drop_last=True,
+            num_workers=4,
+            pin_memory=True,
+        ),
+    )

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/defaults/model.py

@@ -0,0 +1,57 @@
+# 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.
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.predict2_multiview.models.multiview_pose_model_rectified_flow import (
+    MultiviewVid2VidPoseModelRectifiedFlow,
+    MultiviewVid2VidPoseModelRectifiedFlowConfig,
+)
+from cosmos_predict2._src.predict2_multiview.models.multiview_vid2vid_model_rectified_flow import (
+    MultiviewVid2VidModelRectifiedFlow,
+    MultiviewVid2VidModelRectifiedFlowConfig,
+)
+
+FSDP_RECTIFIED_FLOW_CONFIG = dict(
+    trainer=dict(
+        distributed_parallelism="fsdp",
+    ),
+    model=L(MultiviewVid2VidModelRectifiedFlow)(
+        config=MultiviewVid2VidModelRectifiedFlowConfig(),
+        _recursive_=False,
+    ),
+)
+
+FSDP_RECTIFIED_FLOW_POSE_CONFIG = dict(
+    trainer=dict(
+        distributed_parallelism="fsdp",
+    ),
+    model=L(MultiviewVid2VidPoseModelRectifiedFlow)(
+        config=MultiviewVid2VidPoseModelRectifiedFlowConfig(),
+        _recursive_=False,
+    ),
+)
+
+
+def register_model():
+    cs = ConfigStore.instance()
+    cs.store(group="model", package="_global_", name="fsdp_rectified_flow_multiview", node=FSDP_RECTIFIED_FLOW_CONFIG)
+    cs.store(
+        group="model",
+        package="_global_",
+        name="fsdp_rectified_flow_multiview_pose",
+        node=FSDP_RECTIFIED_FLOW_POSE_CONFIG,
+    )

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/defaults/net.py

@@ -0,0 +1,190 @@
+# 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 copy
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.imaginaire.lazy_config import LazyDict
+from cosmos_predict2._src.predict2.networks.minimal_v4_dit import SACConfig
+from cosmos_predict2._src.predict2_multiview.networks.multiview_cross_dit import MultiViewCrossDiT, MultiViewSACConfig
+from cosmos_predict2._src.predict2_multiview.networks.multiview_dit import MultiViewDiT
+from cosmos_predict2._src.predict2_multiview.networks.multiview_pose_dit import MultiViewPoseDiT
+
+COSMOS_V1_7B_MULTIVIEW_NET: LazyDict = L(MultiViewDiT)(
+    max_img_h=240,
+    max_img_w=240,
+    max_frames=128,
+    in_channels=16,
+    out_channels=16,
+    patch_spatial=2,
+    patch_temporal=1,
+    model_channels=4096,
+    num_blocks=28,
+    num_heads=32,
+    concat_padding_mask=True,
+    pos_emb_cls="rope3d",
+    pos_emb_learnable=True,
+    pos_emb_interpolation="crop",
+    use_adaln_lora=True,
+    adaln_lora_dim=256,
+    atten_backend="minimal_a2a",
+    extra_per_block_abs_pos_emb=True,
+    rope_h_extrapolation_ratio=1.0,
+    rope_w_extrapolation_ratio=1.0,
+    rope_t_extrapolation_ratio=2.0,
+    sac_config=SACConfig(),
+    n_cameras_emb=7,
+    view_condition_dim=6,
+    concat_view_embedding=True,
+    use_wan_fp32_strategy=False,
+    layer_mask=None,
+)
+
+COSMOS_V1_2B_MULTIVIEW_NET = copy.deepcopy(COSMOS_V1_7B_MULTIVIEW_NET)
+COSMOS_V1_2B_MULTIVIEW_NET.model_channels = 2048
+COSMOS_V1_2B_MULTIVIEW_NET.num_blocks = 28
+COSMOS_V1_2B_MULTIVIEW_NET.num_heads = 16
+COSMOS_V1_2B_MULTIVIEW_NET.extra_per_block_abs_pos_emb = False
+COSMOS_V1_2B_MULTIVIEW_NET.rope_t_extrapolation_ratio = 1.0
+
+COSMOS_V1_14B_MULTIVIEW_NET = copy.deepcopy(COSMOS_V1_7B_MULTIVIEW_NET)
+COSMOS_V1_14B_MULTIVIEW_NET.model_channels = 5120
+COSMOS_V1_14B_MULTIVIEW_NET.num_blocks = 36
+COSMOS_V1_14B_MULTIVIEW_NET.num_heads = 40
+COSMOS_V1_14B_MULTIVIEW_NET.extra_per_block_abs_pos_emb = False
+COSMOS_V1_14B_MULTIVIEW_NET.rope_t_extrapolation_ratio = 1.0
+
+mini_net = copy.deepcopy(COSMOS_V1_7B_MULTIVIEW_NET)
+mini_net.model_channels = 1024
+mini_net.num_heads = 8
+mini_net.num_blocks = 2
+mini_net.rope_t_extrapolation_ratio = 1.0
+
+# 2B multi-view network + pose encoder + warped/visibility conditioning (2-actor joint generation).
+# Derived from the 2B multiview net; only the target class and the pose-specific kwargs change.
+COSMOS_V1_2B_MULTIVIEW_POSE_NET = copy.deepcopy(COSMOS_V1_2B_MULTIVIEW_NET)
+COSMOS_V1_2B_MULTIVIEW_POSE_NET._target_ = MultiViewPoseDiT
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.warped_latent_channels = 16
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.visibility_channels = 1
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.pose_in_channels = 3
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.pose_hidden_dim = 16
+# in-context reference-frame appearance conditioning (default OFF -> other experiments unchanged)
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.enable_reference_frames = False
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.num_reference_frames = 0
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.ref_rope_offset = 50
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.ref_rope_stride = 5
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.shared_reference = False  # refs = one shared set, no per-view view-embedding
+# per-pixel Plücker ray conditioning (default OFF)
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.enable_plucker = False
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.plucker_channels = 6
+COSMOS_V1_2B_MULTIVIEW_POSE_NET.enable_reference_plucker = False
+
+# tiny variant for fast shape/smoke tests
+mini_pose_net = copy.deepcopy(COSMOS_V1_2B_MULTIVIEW_POSE_NET)
+mini_pose_net.model_channels = 1024
+mini_pose_net.num_heads = 8
+mini_pose_net.num_blocks = 2
+
+# modified according to COSMOS_V1_2B_MULTIVIEW_NET
+COSMOS_V1_2B_MULTIVIEW_CROSSVIEW_NET = LazyDict = L(MultiViewCrossDiT)(
+    max_img_h=240,
+    max_img_w=240,
+    max_frames=128,
+    in_channels=16,
+    out_channels=16,
+    patch_spatial=2,
+    patch_temporal=1,
+    model_channels=2048,
+    num_blocks=28,
+    num_heads=16,
+    concat_padding_mask=True,
+    pos_emb_cls="rope3d",
+    pos_emb_learnable=True,
+    pos_emb_interpolation="crop",
+    use_adaln_lora=True,
+    adaln_lora_dim=256,
+    atten_backend="minimal_a2a",
+    extra_per_block_abs_pos_emb=False,
+    rope_h_extrapolation_ratio=1.0,
+    rope_w_extrapolation_ratio=1.0,
+    rope_t_extrapolation_ratio=1.0,
+    sac_config=MultiViewSACConfig(),
+    n_cameras_emb=7,
+    view_condition_dim=6,
+    concat_view_embedding=False,
+    adaln_view_embedding=True,
+    enable_cross_view_attn=True,
+    use_wan_fp32_strategy=False,
+    layer_mask=None,
+)
+
+# modify according to COSMOS_V1_14B_MULTIVIEW_NET
+COSMOS_V1_14B_MULTIVIEW_CROSSVIEW_NET = LazyDict = L(MultiViewCrossDiT)(
+    max_img_h=240,
+    max_img_w=240,
+    max_frames=128,
+    in_channels=16,
+    out_channels=16,
+    patch_spatial=2,
+    patch_temporal=1,
+    model_channels=5120,
+    num_blocks=36,
+    num_heads=40,
+    concat_padding_mask=True,
+    pos_emb_cls="rope3d",
+    pos_emb_learnable=True,
+    pos_emb_interpolation="crop",
+    use_adaln_lora=True,
+    adaln_lora_dim=256,
+    atten_backend="minimal_a2a",
+    extra_per_block_abs_pos_emb=False,
+    rope_h_extrapolation_ratio=1.0,
+    rope_w_extrapolation_ratio=1.0,
+    rope_t_extrapolation_ratio=1.0,
+    sac_config=MultiViewSACConfig(),
+    n_cameras_emb=7,
+    view_condition_dim=6,
+    concat_view_embedding=False,
+    adaln_view_embedding=True,
+    enable_cross_view_attn=True,
+    use_wan_fp32_strategy=False,
+    layer_mask=None,
+)
+
+
+def register_net():
+    cs = ConfigStore.instance()
+    cs.store(group="net", package="model.config.net", name="mini_net", node=mini_net)
+    cs.store(group="net", package="model.config.net", name="cosmos_v1_2B_multiview", node=COSMOS_V1_2B_MULTIVIEW_NET)
+    cs.store(group="net", package="model.config.net", name="cosmos_v1_7B_multiview", node=COSMOS_V1_7B_MULTIVIEW_NET)
+    cs.store(group="net", package="model.config.net", name="cosmos_v1_14B_multiview", node=COSMOS_V1_14B_MULTIVIEW_NET)
+    cs.store(
+        group="net",
+        package="model.config.net",
+        name="cosmos_v1_2B_multiview_crossview",
+        node=COSMOS_V1_2B_MULTIVIEW_CROSSVIEW_NET,
+    )
+    cs.store(
+        group="net",
+        package="model.config.net",
+        name="cosmos_v1_14B_multiview_crossview",
+        node=COSMOS_V1_14B_MULTIVIEW_CROSSVIEW_NET,
+    )
+    cs.store(
+        group="net", package="model.config.net", name="cosmos_v1_2B_multiview_pose", node=COSMOS_V1_2B_MULTIVIEW_POSE_NET
+    )
+    cs.store(group="net", package="model.config.net", name="mini_pose_net", node=mini_pose_net)

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/defaults/optimizer.py

@@ -0,0 +1,72 @@
+# 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.
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import PLACEHOLDER, LazyDict
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.predict2.utils.optim_instantiate import get_base_optimizer
+from cosmos_predict2._src.predict2_multiview.utils.optim_instantiate import get_multiple_optimizer
+
+AdamWConfig = L(get_base_optimizer)(
+    model=PLACEHOLDER,
+    lr=1e-4,
+    weight_decay=0.1,
+    betas=[0.9, 0.99],
+    optim_type="adamw",
+    eps=1e-8,
+    fused=True,
+)
+
+FusedAdamWConfig: LazyDict = L(get_base_optimizer)(
+    model=PLACEHOLDER,
+    lr=1e-4,
+    weight_decay=0.1,
+    betas=[0.9, 0.99],
+    optim_type="fusedadam",
+    eps=1e-8,
+    master_weights=True,
+    capturable=True,
+)
+
+MultipleAdamWConfig = L(get_multiple_optimizer)(
+    model=PLACEHOLDER,
+    lr=1e-4,
+    weight_decay=1e-3,
+    betas=[0.9, 0.999],
+    optim_type="adamw",
+    eps=1e-8,
+    fused=True,
+    lr_overrides=[],  # New format: list of dicts with 'pattern', 'lr', and optional 'match_type'
+)
+
+
+MultipleFusedAdamWConfig = L(get_multiple_optimizer)(
+    model=PLACEHOLDER,
+    lr=1e-4,
+    weight_decay=1e-3,
+    betas=[0.9, 0.999],
+    optim_type="fusedadam",
+    eps=1e-8,
+    lr_overrides=[],  # New format: list of dicts with 'pattern', 'lr', and optional 'match_type'
+)
+
+
+def register_optimizer():
+    cs = ConfigStore.instance()
+    cs.store(group="optimizer", package="optimizer", name="fusedadamw", node=FusedAdamWConfig)
+    cs.store(group="optimizer", package="optimizer", name="adamw", node=AdamWConfig)
+    cs.store(group="optimizer", package="optimizer", name="multipleadamw", node=MultipleAdamWConfig)
+    cs.store(group="optimizer", package="optimizer", name="multiplefusedadamw", node=MultipleFusedAdamWConfig)

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/experiment/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/experiment/buttercup2p5_rectified_flow.py

@@ -0,0 +1,229 @@
+# 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.
+
+"""Experiments for vanilla attention training on 2B model.
+
+All experiments are using datasets located on AWS S3 unless specified.
+If you want to use datasets on GCS, change the name of the override /data_train from e.g.
+video_alpamayo_dec2024_s3_720p_10fps_93frames_7views -> video_alpamayo_dec2024_gcs_720p_10fps_93frames_7views
+"""
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.predict2.text_encoders.text_encoder import EmbeddingConcatStrategy
+from cosmos_predict2._src.predict2_multiview.callbacks.every_n_draw_sample_multiviewvideo import (
+    EveryNDrawSampleMultiviewVideo,
+)
+
+
+def buttercup_predict2p5_2b_7views_res720p_fps30_t8_joint_alpamayo1capviewprefix_allcapsviewprefix_29frames_nofps_uniform_dropoutt0() -> (
+    dict
+):
+    state_t = 8
+    return dict(
+        defaults=[
+            {"override /ckpt_type": "dcp"},
+            {"override /optimizer": "adamw"},
+            {"override /callbacks": ["basic", "viz_online_sampling", "wandb", "cluster_speed"]},
+            {"override /checkpoint": "s3"},
+            {"override /tokenizer": "wan2pt1_tokenizer"},
+            {"override /data_train": "video_alpamayo_dec2024_s3_720p_30fps_29frames_7views"},
+            {"override /conditioner": "video_prediction_multiview_conditioner"},
+            {"override /model": "fsdp_rectified_flow_multiview"},
+            {"override /net": "cosmos_v1_2B_multiview"},
+            "_self_",
+        ],
+        job=dict(
+            group="cosmos2_mv",
+            name="buttercup_predict2p5_2b_7views_res720p_fps30_t8_joint_alpamayo1capviewprefix_allcapsviewprefix_29frames_nofps_uniform_dropoutt0",
+        ),
+        optimizer=dict(
+            lr=3e-5,  # 2**(-14.5) = 3.0517578125e-05
+            weight_decay=1e-3,
+            betas=[0.9, 0.999],
+        ),
+        scheduler=dict(
+            f_max=[0.99],
+            f_min=[0.4],
+            warm_up_steps=[100],
+            cycle_lengths=[400_000],
+        ),
+        checkpoint=dict(
+            load_from_object_store=dict(
+                enabled=True,
+            ),
+            save_to_object_store=dict(
+                enabled=True,
+            ),
+            save_iter=500,
+            load_path="cosmos_predict2_multiview/cosmos2_mv/buttercup_predict2p5_2b_7views_res720p_fps30_t8_from48kfps30mv_condprobs0442_joint_alpamayo1capnoviewprefix_allcapsviewprefix_29frames_nofps-0/checkpoints/iter_000005000",
+        ),
+        trainer=dict(
+            callbacks=dict(
+                compile_tokenizer=dict(
+                    enabled=False,
+                ),
+                iter_speed=dict(
+                    hit_thres=50,
+                    every_n=100,
+                ),
+                every_n_sample_reg=L(EveryNDrawSampleMultiviewVideo)(
+                    every_n=2_000,
+                    do_x0_prediction=False,
+                    is_ema=False,
+                    num_sampling_step=35,
+                    guidance=[0, 3, 7],
+                    fps=30,
+                ),
+                every_n_sample_ema=L(EveryNDrawSampleMultiviewVideo)(
+                    every_n=2_000,
+                    do_x0_prediction=False,
+                    is_ema=True,
+                    num_sampling_step=35,
+                    guidance=[0, 3, 7],
+                    fps=30,
+                ),
+            ),
+        ),
+        model_parallel=dict(
+            context_parallel_size=8,
+        ),
+        model=dict(
+            config=dict(
+                min_num_conditional_frames=0,
+                max_num_conditional_frames=2,
+                conditional_frames_probs={0: 0.5, 1: 0.25, 2: 0.25},
+                condition_locations=["first_random_n"],
+                fsdp_shard_size=8,
+                resolution="720p",
+                state_t=state_t,
+                shift=5,
+                use_dynamic_shift=False,
+                train_time_weight="uniform",
+                train_time_distribution="logitnormal",
+                online_text_embeddings_as_dict=False,
+                net=dict(
+                    concat_view_embedding=True,
+                    view_condition_dim=7,
+                    state_t=8,
+                    n_cameras_emb=7,
+                    rope_enable_fps_modulation=False,
+                    rope_h_extrapolation_ratio=3.0,
+                    rope_w_extrapolation_ratio=3.0,
+                    rope_t_extrapolation_ratio=float(state_t) / 24,
+                    timestep_scale=0.001,
+                    sac_config=dict(
+                        mode="predict2_2b_720",
+                    ),
+                    use_crossattn_projection=True,
+                    crossattn_proj_in_channels=100352,
+                    crossattn_emb_channels=1024,
+                    use_wan_fp32_strategy=True,
+                ),
+                conditioner=dict(
+                    use_video_condition=dict(
+                        dropout_rate=0.0,
+                    ),
+                    text=dict(
+                        dropout_rate=0.0,
+                        use_empty_string=False,
+                    ),
+                ),
+                tokenizer=dict(
+                    temporal_window=16,
+                ),
+                text_encoder_class="reason1p1_7B",
+                text_encoder_config=dict(
+                    embedding_concat_strategy=str(EmbeddingConcatStrategy.FULL_CONCAT),
+                    compute_online=True,
+                    ckpt_path="s3://bucket/cosmos_reasoning1/sft_exp700/sft_exp721-1_qwen7b_tl_721_5vs5_s3_balanced_n32_resume_16k/checkpoints/iter_000016000/model/",
+                ),
+            ),
+        ),
+        dataloader_train=dict(
+            augmentation_config=dict(
+                single_caption_camera_name="camera_front_wide_120fov",
+                add_view_prefix_to_caption=True,
+            )
+        ),
+    )
+
+
+# Note, use GB200 cluster for this config, as it OOMs on H100 clusters.
+# For context parallelism of 24, use multiples of 24 GPUs or 6 nodes (since 4 GB200 per node).
+def buttercup_predict2p5_2b_mv_7views_res720p_fps30_t24_joint_alpamayo1capviewprefix_allcapsviewprefix_93frames_nofps_uniform_dropoutt0() -> (
+    dict
+):
+    state_t = 24
+    return dict(
+        defaults=[
+            "/experiment/buttercup_predict2p5_2b_7views_res720p_fps30_t8_joint_alpamayo1capviewprefix_allcapsviewprefix_29frames_nofps_uniform_dropoutt0",
+            {"override /data_train": "video_alpamayo_dec2024_gcs_720p_30fps_93frames_7views"},
+            "_self_",
+        ],
+        job=dict(
+            group="cosmos2_mv",
+            name="buttercup_predict2p5_2b_mv_7views_res720p_fps30_t24_joint_alpamayo1capviewprefix_allcapsviewprefix_93frames_nofps_uniform_dropoutt0",
+        ),
+        checkpoint=dict(
+            load_path="cosmos_predict2_multiview/cosmos2_mv/buttercup_predict2p5_2b_7views_res720p_fps30_t8_from48kfps30mv_condprobs0442_joint_alpamayo1capnoviewprefix_allcapsviewprefix_29frames_nofps-0/checkpoints/iter_000005000",
+        ),
+        model_parallel=dict(
+            context_parallel_size=8,
+        ),
+        model=dict(
+            config=dict(
+                state_t=state_t,
+                net=dict(
+                    state_t=state_t,
+                    rope_enable_fps_modulation=False,
+                    rope_h_extrapolation_ratio=3.0,
+                    rope_w_extrapolation_ratio=3.0,
+                    rope_t_extrapolation_ratio=state_t / 24.0,
+                    sac_config=dict(
+                        mode="predict2_2b_720",
+                    ),
+                ),
+            ),
+        ),
+        trainer=dict(
+            straggler_detection=dict(
+                enabled=False,
+            ),
+        ),
+        dataloader_train=dict(
+            augmentation_config=dict(
+                single_caption_camera_name="camera_front_wide_120fov",
+                add_view_prefix_to_caption=True,
+            ),
+        ),
+    )
+
+
+experiments = [
+    buttercup_predict2p5_2b_7views_res720p_fps30_t8_joint_alpamayo1capviewprefix_allcapsviewprefix_29frames_nofps_uniform_dropoutt0(),
+    buttercup_predict2p5_2b_mv_7views_res720p_fps30_t24_joint_alpamayo1capviewprefix_allcapsviewprefix_93frames_nofps_uniform_dropoutt0(),
+]
+
+cs = ConfigStore.instance()
+
+for _item in experiments:
+    cs.store(
+        group="experiment",
+        package="_global_",
+        name=_item["job"]["name"],
+        node=_item,
+    )

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/experiment/buttercup2p5_rectified_flow_14b.py

@@ -0,0 +1,243 @@
+# 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.
+
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.predict2.text_encoders.text_encoder import EmbeddingConcatStrategy
+from cosmos_predict2._src.predict2_multiview.callbacks.every_n_draw_sample_multiviewvideo import (
+    EveryNDrawSampleMultiviewVideo,
+)
+from cosmos_predict2._src.predict2_multiview.callbacks.log_weight import LogWeight
+
+"""
+torchrun --nproc_per_node=8 --master_port=12341 -m scripts.train --config=cosmos_predict2/_src/predict2_multiview/configs/vid2vid/config.py -- experiment=buttercup_predict2p5_14b_7views_res720p_fps30_t8_frombase2p5_condprobs0442_joint_alpamayo1capnoviewprefix_allcapsviewprefix_29frames_nofps_uniform
+"""
+
+
+def buttercup_predict2p5_14b_7views_res720p_fps10_t8_fromavfinetune_allcaps_29frames_nofps_uniform() -> dict:
+    state_t = 8
+    return dict(
+        defaults=[
+            {"override /data_train": "video_alpamayo_dec2024_gcs_720p_10fps_29frames_7views"},
+            {"override /conditioner": "video_prediction_multiview_conditioner"},
+            {"override /model": "fsdp_rectified_flow_multiview"},
+            {"override /net": "cosmos_v1_14B_multiview"},
+            {"override /ckpt_type": "dcp"},
+            {"override /optimizer": "adamw"},
+            {
+                "override /callbacks": [
+                    "basic",
+                    "viz_online_sampling",
+                    "wandb",
+                    "cluster_speed",
+                ]
+            },
+            {"override /checkpoint": "s3"},
+            {"override /tokenizer": "wan2pt1_tokenizer"},
+            "_self_",
+        ],
+        job=dict(
+            group="cosmos2p5_mv",
+            name="buttercup_predict2p5_14b_7views_res720p_fps10_t8_fromavfinetune_allcaps_29frames_nofps_uniform",
+        ),
+        checkpoint=dict(
+            save_to_object_store=dict(
+                enabled=True,
+            ),
+            load_from_object_store=dict(
+                enabled=True,
+            ),
+            save_iter=250,
+            load_path="cosmos_diffusion_v2/official_runs_vid2vid/Stage-c_pt_4-Index-43-Size-14B-Res-720-Fps-16-Note-rf_av_high_sigma_uniform/checkpoints/iter_000022000/",
+        ),
+        model_parallel=dict(
+            context_parallel_size=8,
+        ),
+        optimizer=dict(
+            lr=2 ** (-14.5),
+            weight_decay=0.001,
+            betas=[0.9, 0.999],
+        ),
+        scheduler=dict(
+            f_max=[0.3],
+            f_min=[0.1],
+            warm_up_steps=[2_000],
+            cycle_lengths=[200_000],
+        ),
+        model=dict(
+            config=dict(
+                min_num_conditional_frames=0,
+                max_num_conditional_frames=2,
+                conditional_frames_probs={0: 0.6, 1: 0.2, 2: 0.2},
+                condition_locations=["first_random_n"],
+                fsdp_shard_size=32,
+                resolution="720p",
+                online_text_embeddings_as_dict=False,
+                state_t=state_t,
+                shift=5,
+                use_dynamic_shift=False,
+                train_time_weight="uniform",
+                train_time_distribution="logitnormal",
+                use_high_sigma_strategy=True,
+                net=dict(
+                    concat_view_embedding=True,
+                    view_condition_dim=7,
+                    n_cameras_emb=7,
+                    state_t=state_t,
+                    rope_enable_fps_modulation=False,
+                    rope_h_extrapolation_ratio=3.0,
+                    rope_w_extrapolation_ratio=3.0,
+                    rope_t_extrapolation_ratio=float(state_t) / 24,
+                    timestep_scale=0.001,
+                    sac_config=dict(
+                        mode="predict2_14b_720_aggressive",
+                    ),
+                    use_crossattn_projection=True,
+                    crossattn_proj_in_channels=100352,
+                    crossattn_emb_channels=1024,
+                    use_wan_fp32_strategy=True,
+                ),
+                conditioner=dict(
+                    use_video_condition=dict(
+                        dropout_rate=0.0,
+                    ),
+                    text=dict(
+                        dropout_rate=0.2,
+                        use_empty_string=False,
+                    ),
+                ),
+                tokenizer=dict(
+                    temporal_window=16,
+                ),
+                text_encoder_class="reason1p1_7B",
+                text_encoder_config=dict(
+                    embedding_concat_strategy=str(EmbeddingConcatStrategy.FULL_CONCAT),
+                    compute_online=True,
+                    ckpt_path="s3://bucket/cosmos_reasoning1/sft_exp700/sft_exp721-1_qwen7b_tl_721_5vs5_s3_balanced_n32_resume_16k/checkpoints/iter_000016000/model/",
+                    s3_credential_path="credentials/s3_checkpoint.secret",
+                ),
+            )
+        ),
+        trainer=dict(
+            straggler_detection=dict(
+                enabled=False,
+            ),
+            callbacks=dict(
+                compile_tokenizer=dict(
+                    enabled=False,
+                ),
+                iter_speed=dict(
+                    hit_thres=50,
+                    every_n=100,
+                ),
+                every_n_sample_ema=L(EveryNDrawSampleMultiviewVideo)(
+                    every_n=2_000,
+                    do_x0_prediction=False,
+                    is_ema=True,
+                    num_sampling_step=35,
+                    guidance=[0, 3],
+                    fps=10,
+                    run_at_start=True,
+                ),
+            ),
+        ),
+        dataloader_train=dict(
+            augmentation_config=dict(
+                caption_probability={
+                    "qwen2p5_7b_caption": 0.7,
+                    "qwen2p5_7b_caption_medium": 0.2,
+                    "qwen2p5_7b_caption_short": 0.1,
+                },
+            )
+        ),
+    )
+
+
+def buttercup_predict2p5_14b_crossview_7views_res720p_fps10_t8_fromavfinetune_allcaps_29frames_nofps_uniform():
+    return dict(
+        defaults=[
+            "/experiment/buttercup_predict2p5_14b_7views_res720p_fps10_t8_fromavfinetune_allcaps_29frames_nofps_uniform",
+            {"override /net": "cosmos_v1_14B_multiview_crossview"},
+            {"override /optimizer": "multiplefusedadamw"},
+            "_self_",
+        ],
+        job=dict(
+            group="cosmos2_mv2",
+            name="buttercup_predict2p5_14b_crossview_7views_res720p_fps10_t8_fromavfinetune_allcaps_29frames_nofps_uniform",
+        ),
+        model=dict(
+            config=dict(
+                net=dict(
+                    cross_view_attn_map_str={
+                        "camera_front_wide_120fov": [
+                            "camera_cross_left_120fov",
+                            "camera_cross_right_120fov",
+                            "camera_front_tele_30fov",
+                        ],
+                        "camera_cross_left_120fov": ["camera_front_wide_120fov", "camera_rear_left_70fov"],
+                        "camera_cross_right_120fov": ["camera_front_wide_120fov", "camera_rear_right_70fov"],
+                        "camera_rear_left_70fov": ["camera_cross_left_120fov", "camera_rear_tele_30fov"],
+                        "camera_rear_right_70fov": ["camera_cross_right_120fov", "camera_rear_tele_30fov"],
+                        "camera_rear_tele_30fov": ["camera_rear_left_70fov", "camera_rear_right_70fov"],
+                        "camera_front_tele_30fov": ["camera_front_wide_120fov"],
+                    },
+                    camera_to_view_id={
+                        "camera_front_wide_120fov": 0,
+                        "camera_cross_left_120fov": 5,
+                        "camera_cross_right_120fov": 1,
+                        "camera_rear_left_70fov": 4,
+                        "camera_rear_right_70fov": 2,
+                        "camera_rear_tele_30fov": 3,
+                        "camera_front_tele_30fov": 6,
+                    },
+                ),
+            ),
+        ),
+        optimizer=dict(
+            lr=3e-5,
+            lr_overrides=[
+                {"pattern": "cross_view_attn", "lr": 1e-4, "match_type": "contains"},
+            ],
+        ),
+        trainer=dict(
+            logging_iter=50,
+            callbacks=dict(
+                log_weight=L(LogWeight)(
+                    every_n=50,
+                ),
+                every_n_sample_reg=dict(
+                    every_n=1500,
+                ),
+            ),
+        ),
+    )
+
+
+experiments = [
+    buttercup_predict2p5_14b_7views_res720p_fps10_t8_fromavfinetune_allcaps_29frames_nofps_uniform(),
+    buttercup_predict2p5_14b_crossview_7views_res720p_fps10_t8_fromavfinetune_allcaps_29frames_nofps_uniform(),
+]
+
+cs = ConfigStore.instance()
+
+for _item in experiments:
+    cs.store(
+        group="experiment",
+        package="_global_",
+        name=_item["job"]["name"],
+        node=_item,
+    )

cosmos_predict2/_src/predict2_multiview/configs/vid2vid/experiment/nymeria_pose_2actor.py

@@ -0,0 +1,1005 @@
+# 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.
+
+"""2-actor joint video generation on Nymeria, full fine-tune of Cosmos Predict 2.5 (2B) on 4 GPUs with FSDP.
+
+Actor 1 / Actor 2 are the two "views" (V=2). Per actor: frame 0 is given, frames 1.. are generated; the
+warped past-frame (pixel -> VAE latent) + visibility mask are channel-concat conditioning (additive zero-init
+PatchEmbed), and the projected 2D pose RGB map is fed through the 3D-conv pose encoder and added to the patch
+tokens. VAE + text encoder are frozen; the DiT backbone + pose encoder + view embeddings are fully trained.
+
+Architecture mirrors the pretrained ``buttercup_predict2p5_2b ... t8`` 2B multiview net so the base checkpoint
+warm-starts cleanly (the new pose/cond modules load as missing keys).
+
+state_t = 9 latent frames  (= 1 + (33 - 1) // 4 for the 33-frame nymeria_processed clips).
+
+Launch (4 GPUs):
+  torchrun --nproc_per_node=4 -m scripts.train \
+      --config=cosmos_predict2/_src/predict2_multiview/configs/vid2vid/config.py -- \
+      experiment=nymeria_pose_2actor_2b_t9_fsdp4
+
+NOTE (deployment): set ``checkpoint.load_path`` to your local 2B multiview checkpoint dir for warm start, and
+ensure the frozen text-encoder checkpoint (``model.config.text_encoder_config.ckpt_path``) is reachable. Both
+are inherited from the buttercup recipe below and may point at S3 by default.
+"""
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.predict2.text_encoders.text_encoder import EmbeddingConcatStrategy
+from cosmos_predict2._src.predict2_multiview.callbacks.every_n_draw_sample_multiviewvideo import (
+    EveryNDrawSampleMultiviewVideo,
+)
+from cosmos_predict2._src.predict2_multiview.callbacks.nymeria_validation_viz import NymeriaValidationViz
+
+
+def nymeria_pose_2actor_2b_t9_fsdp4() -> dict:
+    state_t = 9  # 1 + (33 - 1) // 4 ; matches 33 target frames per actor @10fps
+    return dict(
+        defaults=[
+            {"override /ckpt_type": "dcp"},
+            {"override /optimizer": "adamw"},
+            {"override /callbacks": ["basic", "viz_online_sampling"]},
+            {"override /checkpoint": "s3"},  # used with object-store disabled below -> writes to local job dir
+            {"override /tokenizer": "wan2pt1_tokenizer"},
+            {"override /data_train": "nymeria_pairs_480p_10fps_33frames"},
+            {"override /data_val": "nymeria_pairs_480p_10fps_33frames"},
+            {"override /conditioner": "video_prediction_multiview_pose_conditioner"},
+            {"override /model": "fsdp_rectified_flow_multiview_pose"},
+            {"override /net": "cosmos_v1_2B_multiview_pose"},
+            "_self_",
+        ],
+        job=dict(
+            group="nymeria_pose",
+            name="nymeria_pose_2actor_2b_t9_fsdp4",
+        ),
+        optimizer=dict(
+            lr=3e-5,
+            weight_decay=1e-3,
+            betas=[0.9, 0.999],
+        ),
+        scheduler=dict(
+            f_max=[0.99],
+            f_min=[0.4],
+            warm_up_steps=[100],
+            cycle_lengths=[400_000],
+        ),
+        checkpoint=dict(
+            # local training: keep checkpoints on the local filesystem (no object store).
+            load_from_object_store=dict(enabled=False),
+            save_to_object_store=dict(enabled=False),
+            save_iter=500,
+            # TODO(deploy): point at your local 2B multiview checkpoint dir to warm-start the DiT backbone.
+            # Leave as "" to train the backbone from its random init (pose/cond modules are always fresh).
+            load_path="",
+            strict_resume=False,  # tolerate the new pose/cond parameters missing from the base checkpoint
+        ),
+        trainer=dict(
+            max_iter=400_000,
+            callbacks=dict(
+                every_n_sample_reg=L(EveryNDrawSampleMultiviewVideo)(
+                    every_n=2_000,
+                    do_x0_prediction=False,
+                    is_ema=False,
+                    num_sampling_step=35,
+                    guidance=[0, 3, 7],
+                    fps=10,
+                ),
+                every_n_sample_ema=L(EveryNDrawSampleMultiviewVideo)(
+                    every_n=2_000,
+                    do_x0_prediction=False,
+                    is_ema=True,
+                    num_sampling_step=35,
+                    guidance=[0, 3, 7],
+                    fps=10,
+                ),
+            ),
+        ),
+        model_parallel=dict(
+            context_parallel_size=1,
+        ),
+        model=dict(
+            config=dict(
+                # frame 0 given, the rest generated (1 conditional frame per actor)
+                min_num_conditional_frames_per_view=1,
+                max_num_conditional_frames_per_view=1,
+                conditional_frames_probs={1: 1.0},
+                condition_locations=["first_random_n"],
+                fsdp_shard_size=4,  # FSDP across the 4 GPUs
+                resolution="480",
+                state_t=state_t,
+                shift=5,
+                use_dynamic_shift=False,
+                train_time_weight="uniform",
+                train_time_distribution="logitnormal",
+                online_text_embeddings_as_dict=False,
+                net=dict(
+                    concat_view_embedding=True,
+                    view_condition_dim=7,
+                    state_t=state_t,
+                    n_cameras_emb=7,
+                    rope_enable_fps_modulation=False,
+                    rope_h_extrapolation_ratio=3.0,
+                    rope_w_extrapolation_ratio=3.0,
+                    rope_t_extrapolation_ratio=float(state_t) / 24.0,
+                    timestep_scale=0.001,
+                    sac_config=dict(mode="predict2_2b_720"),
+                    use_crossattn_projection=True,
+                    crossattn_proj_in_channels=100352,
+                    crossattn_emb_channels=1024,
+                    use_wan_fp32_strategy=True,
+                    # pose / warped / visibility conditioning channels
+                    warped_latent_channels=16,
+                    visibility_channels=1,
+                    pose_in_channels=3,
+                    pose_hidden_dim=16,
+                ),
+                conditioner=dict(
+                    use_video_condition=dict(dropout_rate=0.0),
+                    text=dict(dropout_rate=0.0, use_empty_string=False),
+                ),
+                tokenizer=dict(temporal_window=16),
+                text_encoder_class="reason1p1_7B",
+                text_encoder_config=dict(
+                    embedding_concat_strategy=str(EmbeddingConcatStrategy.FULL_CONCAT),
+                    compute_online=True,
+                    ckpt_path="s3://bucket/cosmos_reasoning1/sft_exp700/sft_exp721-1_qwen7b_tl_721_5vs5_s3_balanced_n32_resume_16k/checkpoints/iter_000016000/model/",
+                ),
+            ),
+        ),
+    )
+
+
+def _local_wan_vae_path() -> str:
+    """Resolve the locally-cached Wan2.1 VAE (downloaded from HF); used for offline runs."""
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download("Wan-AI/Wan2.1-T2V-1.3B", "Wan2.1_VAE.pth")
+
+
+def _base_2b_multiview_ckpt() -> str:
+    """Local path to the pretrained 2B multiview checkpoint (HF mirror of MODEL_CHECKPOINTS auto/multiview).
+
+    Returns "" if it cannot be resolved/downloaded (then the backbone trains from random init).
+    """
+    try:
+        from huggingface_hub import hf_hub_download
+
+        return hf_hub_download(
+            "nvidia/Cosmos-Predict2.5-2B",
+            "auto/multiview/524af350-2e43-496c-8590-3646ae1325da_ema_bf16.pt",
+            revision="865baf084d4c9e850eac59a021277d5a9b9e8b63",
+        )
+    except Exception:
+        return ""
+
+
+def nymeria_pose_2actor_2b_smoke() -> dict:
+    """Offline smoke test: real Wan VAE (local) + zero dummy text (no S3 text encoder), random-init backbone.
+
+    Exercises the full pipeline end-to-end: dataloader -> VAE encode (video + warped) -> pose preprocessing
+    -> MultiViewPoseDiT forward -> rectified-flow loss -> backward -> FSDP -> optimizer step -> checkpoint.
+
+    Single GPU:   torchrun --nproc_per_node=1 -m scripts.train --config=<cfg> -- experiment=nymeria_pose_2actor_2b_smoke model.config.fsdp_shard_size=1
+    4-GPU FSDP:   torchrun --nproc_per_node=4 -m scripts.train --config=<cfg> -- experiment=nymeria_pose_2actor_2b_smoke
+    """
+    state_t = 9
+    return dict(
+        defaults=[
+            {"override /ckpt_type": "dcp"},
+            {"override /optimizer": "adamw"},
+            {"override /callbacks": ["basic"]},
+            {"override /checkpoint": "s3"},
+            {"override /tokenizer": "wan2pt1_tokenizer"},
+            {"override /data_train": "nymeria_pairs_smoke"},
+            {"override /data_val": "nymeria_pairs_smoke"},
+            {"override /conditioner": "video_prediction_multiview_pose_conditioner"},
+            {"override /model": "fsdp_rectified_flow_multiview_pose"},
+            {"override /net": "cosmos_v1_2B_multiview_pose"},
+            "_self_",
+        ],
+        job=dict(group="nymeria_pose", name="nymeria_pose_2actor_2b_smoke"),
+        optimizer=dict(lr=3e-5, weight_decay=1e-3, betas=[0.9, 0.999]),
+        scheduler=dict(f_max=[0.99], f_min=[0.4], warm_up_steps=[1], cycle_lengths=[1000]),
+        checkpoint=dict(
+            load_from_object_store=dict(enabled=False),
+            save_to_object_store=dict(enabled=False),
+            save_iter=1000,
+            load_path="",
+            strict_resume=False,
+        ),
+        trainer=dict(max_iter=3, logging_iter=1),
+        model_parallel=dict(context_parallel_size=1),
+        model=dict(
+            config=dict(
+                min_num_conditional_frames_per_view=1,
+                max_num_conditional_frames_per_view=1,
+                conditional_frames_probs={1: 1.0},
+                condition_locations=["first_random_n"],
+                fsdp_shard_size=4,
+                resolution="480",
+                state_t=state_t,
+                shift=5,
+                use_dynamic_shift=False,
+                train_time_weight="uniform",
+                train_time_distribution="logitnormal",
+                online_text_embeddings_as_dict=False,
+                ema=dict(enabled=False),
+                net=dict(
+                    concat_view_embedding=True,
+                    view_condition_dim=7,
+                    state_t=state_t,
+                    n_cameras_emb=7,
+                    rope_enable_fps_modulation=False,
+                    rope_h_extrapolation_ratio=3.0,
+                    rope_w_extrapolation_ratio=3.0,
+                    rope_t_extrapolation_ratio=float(state_t) / 24.0,
+                    timestep_scale=0.001,
+                    sac_config=dict(mode="predict2_2b_720"),
+                    use_crossattn_projection=False,  # consume dummy crossattn (1024-d) directly
+                    crossattn_emb_channels=1024,
+                    use_wan_fp32_strategy=True,
+                    warped_latent_channels=16,
+                    visibility_channels=1,
+                    pose_in_channels=3,
+                    pose_hidden_dim=16,
+                ),
+                conditioner=dict(
+                    use_video_condition=dict(dropout_rate=0.0),
+                    text=dict(dropout_rate=0.0, use_empty_string=False),
+                ),
+                tokenizer=dict(vae_pth=_local_wan_vae_path(), temporal_window=16),
+                text_encoder_class="reason1p1_7B",  # valid label; encoder is NOT built (config below is None)
+                text_encoder_config=None,  # disable the (S3-only) text encoder; dataset provides dummy embeddings
+            ),
+        ),
+    )
+
+
+def nymeria_pose_2actor_2b_longer() -> dict:
+    """Production run on the LONGER dataset (200 past + 77 target frames), 4-GPU FSDP full fine-tune.
+
+    - Warm-starts the 2B multiview backbone from the HF checkpoint (auto/multiview ema_bf16.pt); the new
+      pose encoder + warped/visibility embedders load as missing keys (kept at zero-init) under strict_resume=False.
+    - VAE = local Wan2.1 VAE. Text is disabled (Nymeria has no captions) -> dummy zero crossattn + no crossattn
+      projection. To use real text instead: set data_train/val to `nymeria_pairs_longer_77frames`, restore
+      `use_crossattn_projection=True` / `crossattn_proj_in_channels=100352`, and provide the reason1 text encoder.
+    - Outputs (checkpoints + validation videos) go under $IMAGINAIRE_OUTPUT_ROOT/multi-ego/nymeria_pose/<name>,
+      and W&B logs to entity/project from $WANDB_ENTITY / job.project with run name == <name> (== the folder).
+    """
+    state_t = 20  # 1 + (77 - 1) // 4 ; matches 77 target frames per actor @10fps
+    return dict(
+        defaults=[
+            {"override /ckpt_type": "dcp"},
+            {"override /optimizer": "adamw"},
+            {"override /callbacks": ["basic", "wandb"]},
+            {"override /checkpoint": "s3"},
+            {"override /tokenizer": "wan2pt1_tokenizer"},
+            {"override /data_train": "nymeria_train_ready_smoke"},
+            {"override /data_val": "nymeria_train_ready_smoke"},
+            {"override /conditioner": "video_prediction_multiview_pose_conditioner"},
+            {"override /model": "fsdp_rectified_flow_multiview_pose"},
+            {"override /net": "cosmos_v1_2B_multiview_pose"},
+            "_self_",
+        ],
+        # folder -> $IMAGINAIRE_OUTPUT_ROOT/train/<name> (project="train", empty group). W&B project/entity are
+        # decoupled via $WANDB_PROJECT / $WANDB_ENTITY (set in sh/train_nymeria_longer.sh).
+        job=dict(group="", name="nymeria_pose_2actor_2b_longer", project="train"),
+        optimizer=dict(lr=3e-5, weight_decay=1e-3, betas=[0.9, 0.999]),
+        scheduler=dict(f_max=[0.99], f_min=[0.4], warm_up_steps=[100], cycle_lengths=[400_000]),
+        checkpoint=dict(
+            load_from_object_store=dict(enabled=False),
+            save_to_object_store=dict(enabled=False),
+            save_iter=500,
+            load_path=_base_2b_multiview_ckpt(),  # warm-start backbone; "" -> random init
+            strict_resume=False,
+        ),
+        trainer=dict(
+            max_iter=400_000,
+            logging_iter=50,
+            callbacks=dict(
+                # custom 2-actor validation viz: [warped | pose | GT | generated] x [egoA; egoB], 5 train + 5 val.
+                validation_viz=L(NymeriaValidationViz)(
+                    every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+                    num_video_frames=77,
+                ),
+            ),
+        ),
+        model_parallel=dict(context_parallel_size=1),
+        model=dict(
+            config=dict(
+                min_num_conditional_frames_per_view=1,
+                max_num_conditional_frames_per_view=1,
+                conditional_frames_probs={1: 1.0},
+                condition_locations=["first_random_n"],
+                fsdp_shard_size=4,
+                resolution="480",
+                state_t=state_t,
+                shift=5,
+                use_dynamic_shift=False,
+                train_time_weight="uniform",
+                train_time_distribution="logitnormal",
+                online_text_embeddings_as_dict=False,
+                net=dict(
+                    concat_view_embedding=True,
+                    view_condition_dim=7,
+                    state_t=state_t,
+                    n_cameras_emb=7,
+                    rope_enable_fps_modulation=False,
+                    rope_h_extrapolation_ratio=3.0,
+                    rope_w_extrapolation_ratio=3.0,
+                    rope_t_extrapolation_ratio=float(state_t) / 24.0,
+                    timestep_scale=0.001,
+                    sac_config=dict(mode="predict2_2b_720"),
+                    use_crossattn_projection=False,
+                    crossattn_emb_channels=1024,
+                    use_wan_fp32_strategy=True,
+                    warped_latent_channels=16,
+                    visibility_channels=1,
+                    pose_in_channels=3,
+                    pose_hidden_dim=16,
+                ),
+                conditioner=dict(
+                    use_video_condition=dict(dropout_rate=0.0),
+                    text=dict(dropout_rate=0.0, use_empty_string=False),
+                ),
+                tokenizer=dict(vae_pth=_local_wan_vae_path(), temporal_window=16),
+                text_encoder_class="reason1p1_7B",
+                text_encoder_config=None,
+            ),
+        ),
+    )
+
+
+def nymeria_pose_2b_stage1_singleview() -> dict:
+    """Stage 1: single-view (V=1) pose-encoder pretraining, 2-GPU FSDP.
+
+    Each actor is an independent V=1 sample (simpler + 2x data + faster iters). Warm-starts from base 2B.
+    The pose encoder uses a SMALL-RANDOM final-proj init (pose_proj_init_std>0) instead of zero-init so its
+    conv feature extractor gets gradient from step 0 (zero-init starves the convs until the proj grows, which
+    is a key cause of slow pose learning). Warped conditioning is kept as-is. After this warms up the pose
+    encoder, fine-tune the joint 2-actor model (nymeria_pose_2actor_2b_longer) warm-started from this stage.
+    """
+    return dict(
+        defaults=[
+            "/experiment/nymeria_pose_2actor_2b_longer",
+            {"override /data_train": "nymeria_single_view"},
+            {"override /data_val": "nymeria_single_view"},
+            "_self_",
+        ],
+        job=dict(group="", name="nymeria_pose_2b_stage1_singleview", project="train"),
+        model=dict(
+            config=dict(
+                fsdp_shard_size=2,  # 2-GPU run
+                net=dict(
+                    pose_proj_init_std=0.02,  # small-random pose proj init -> faster pose learning
+                    freeze_view_embedding=True,  # view embedding is meaningless with V=1 -> don't train it
+                ),
+            )
+        ),
+        trainer=dict(
+            callbacks=dict(
+                validation_viz=L(NymeriaValidationViz)(
+                    every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+                    num_video_frames=77, single_view=True,
+                ),
+            ),
+        ),
+    )
+
+
+def nymeria_pose_2b_stage1_sv_vaeconcat() -> dict:
+    """Stage-1 single-view, VAE pose encoder + CHANNEL CONCAT (Setting 1).
+
+    The pixel pose RGB is passed through the FROZEN Wan2.1 VAE -> a 16ch latent-resolution pose latent, which
+    is CHANNEL-CONCATENATED with the warped latent + visibility mask into the single zero-init cond_embedder
+    (one linear patch projection mixes all conditioning latents). The Conv3d PoseEncoder is NOT built.
+    Warm-starts from base 2B; the cond_embedder loads as a (zero-init) missing key. Run on GPU 0,1.
+    """
+    return dict(
+        defaults=[
+            "/experiment/nymeria_pose_2b_stage1_singleview",
+            {"override /conditioner": "video_prediction_multiview_pose_latent_conditioner"},
+            "_self_",
+        ],
+        job=dict(group="", name="nymeria_pose_2b_stage1_sv_vaeconcat", project="train"),
+        model=dict(
+            config=dict(
+                pose_via_vae=True,  # preprocess: pose RGB -> frozen VAE -> pose_latent
+                net=dict(pose_mode="vae_concat"),
+            )
+        ),
+    )
+
+
+def nymeria_pose_2b_stage1_sv_vaemlp() -> dict:
+    """Stage-1 single-view, VAE pose encoder + trainable MLP + ADD (Setting 2).
+
+    Pixel pose RGB -> FROZEN Wan2.1 VAE -> 16ch pose latent -> trainable PatchEmbed + MLP -> ADDED to the
+    patch tokens (MLP last layer zero-init -> day-0 == base). The Conv3d PoseEncoder is NOT built; the VAE
+    stays frozen and only the small pose MLP/embedder (+ backbone) train. Run on GPU 2,3.
+    """
+    return dict(
+        defaults=[
+            "/experiment/nymeria_pose_2b_stage1_singleview",
+            {"override /conditioner": "video_prediction_multiview_pose_latent_conditioner"},
+            "_self_",
+        ],
+        job=dict(group="", name="nymeria_pose_2b_stage1_sv_vaemlp", project="train"),
+        model=dict(
+            config=dict(
+                pose_via_vae=True,
+                net=dict(pose_mode="vae_mlp_add"),
+            )
+        ),
+    )
+
+
+def nymeria_pose_2view_actorobs() -> dict:
+    """2-view actor-observer joint generation: vae_concat pose + REAL per-view text (reason1), 4-GPU FSDP.
+
+    - Data: /data2/nymeria_processed_single actor+observer pairs (train_split/val_split.csv), role-mix 50% so
+      the view embedding can't memorize actor/observer roles.
+    - Text: HF Cosmos-Reason1-7B (FULL_CONCAT 100352) online, per-view captions (actor caption on the actor
+      view; "C is observing the partner. The partner ..." on the observer view). crossattn projection restored
+      (100352->1024) -> warm-starts from base 2B.
+    - Pose: vae_concat (pose RGB -> frozen VAE -> pose_latent, channel-concat with warped+vis in cond_embedder).
+      Uses the pose_latent conditioner. View embedding TRAINABLE (role-mix handles actor/observer symmetry).
+    - Warm-start: base 2B multiview. Requires env COSMOS_QWEN_TOKENIZER_DIR=/data/cosmos_reason1_7b.
+    """
+    state_t = 20
+    return dict(
+        defaults=[
+            {"override /ckpt_type": "dcp"},
+            {"override /optimizer": "adamw"},
+            {"override /callbacks": ["basic", "wandb"]},
+            {"override /checkpoint": "s3"},
+            {"override /tokenizer": "wan2pt1_tokenizer"},
+            {"override /data_train": "nymeria_actor_observer"},
+            {"override /data_val": "nymeria_actor_observer"},
+            {"override /conditioner": "video_prediction_multiview_pose_latent_conditioner"},
+            {"override /model": "fsdp_rectified_flow_multiview_pose"},
+            {"override /net": "cosmos_v1_2B_multiview_pose"},
+            "_self_",
+        ],
+        job=dict(group="", name="nymeria_pose_2view_actorobs", project="train"),
+        optimizer=dict(lr=3e-5, weight_decay=1e-3, betas=[0.9, 0.999]),
+        scheduler=dict(f_max=[0.99], f_min=[0.4], warm_up_steps=[100], cycle_lengths=[400_000]),
+        checkpoint=dict(
+            load_from_object_store=dict(enabled=False),
+            save_to_object_store=dict(enabled=False),
+            save_iter=500,
+            load_path=_base_2b_multiview_ckpt(),
+            strict_resume=False,
+        ),
+        trainer=dict(
+            max_iter=400_000,
+            logging_iter=50,
+            callbacks=dict(
+                validation_viz=L(NymeriaValidationViz)(
+                    every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+                    num_video_frames=77, root="/data2/nymeria_processed_single", actor_observer=True,
+                ),
+            ),
+        ),
+        model_parallel=dict(context_parallel_size=1),
+        model=dict(
+            config=dict(
+                min_num_conditional_frames_per_view=1,
+                max_num_conditional_frames_per_view=1,
+                conditional_frames_probs={1: 1.0},
+                condition_locations=["first_random_n"],
+                fsdp_shard_size=4,
+                resolution="480",
+                state_t=state_t,
+                shift=5,
+                use_dynamic_shift=False,
+                train_time_weight="uniform",
+                train_time_distribution="logitnormal",
+                online_text_embeddings_as_dict=False,
+                pose_via_vae=True,  # pose RGB -> frozen VAE -> pose_latent (for vae_concat)
+                net=dict(
+                    concat_view_embedding=True,
+                    view_condition_dim=7,
+                    state_t=state_t,
+                    n_cameras_emb=7,
+                    rope_enable_fps_modulation=False,
+                    rope_h_extrapolation_ratio=3.0,
+                    rope_w_extrapolation_ratio=3.0,
+                    rope_t_extrapolation_ratio=float(state_t) / 24.0,
+                    timestep_scale=0.001,
+                    sac_config=dict(mode="predict2_2b_720"),
+                    use_crossattn_projection=True,
+                    crossattn_proj_in_channels=100352,
+                    crossattn_emb_channels=1024,
+                    use_wan_fp32_strategy=True,
+                    warped_latent_channels=16,
+                    visibility_channels=1,
+                    pose_in_channels=3,
+                    pose_hidden_dim=16,
+                    pose_mode="vae_concat",
+                ),
+                conditioner=dict(
+                    use_video_condition=dict(dropout_rate=0.0),
+                    text=dict(dropout_rate=0.2, use_empty_string=False),
+                ),
+                tokenizer=dict(vae_pth=_local_wan_vae_path(), temporal_window=16),
+                text_encoder_class="reason1p1_7B",
+                text_encoder_config=dict(
+                    embedding_concat_strategy=str(EmbeddingConcatStrategy.FULL_CONCAT),
+                    compute_online=True,
+                    ckpt_path="/data/cosmos_reason1_7b",
+                ),
+            ),
+        ),
+    )
+
+
+def nymeria_pose_2view_actorobs_refs() -> dict:
+    """nymeria_pose_2view_actorobs + in-context REFERENCE-FRAME appearance conditioning.
+
+    Appends R=6 clean source frames per view as extra self-attention tokens at fixed non-contiguous temporal
+    RoPE positions [50,55,60,65,70,75], to sharpen appearance detail (complements the geometry-only warped_cond).
+    Only new parameter is a zero-init ref gate -> day-0 output ~= base; warm-starts from base 2B. Fresh run.
+    """
+    cfg = nymeria_pose_2view_actorobs()
+    cfg["job"]["name"] = "nymeria_pose_2view_actorobs_refs"
+    # dataset variants that emit `reference_frames` (R=6 clean source frames/view)
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_observer_refs"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_actor_observer_refs"
+    # model: VAE-encode the reference frames
+    cfg["model"]["config"]["num_reference_frames"] = 4
+    # net: enable the in-context reference tokens
+    cfg["model"]["config"]["net"].update(
+        enable_reference_frames=True, num_reference_frames=4, ref_rope_offset=50, ref_rope_stride=5,
+    )
+    # validation viz must also emit reference frames so generation sees them
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data2/nymeria_processed_single", actor_observer=True,
+        num_reference_frames=4,
+    )
+    return cfg
+
+
+def nymeria_pose_2view_actorobs_refs_campose() -> dict:
+    """refs (R=4) + per-pixel Plücker camera-ray conditioning, for cross-view shared-space awareness.
+
+    Adds a 6-ch Plücker ray map (dir+moment, per-pair canonical frame = view0 frame0) via a zero-init PatchEmbed
+    added to the tokens (same additive pattern as warped/pose; day-0 == base). Targets the two views not
+    recognizing they're the same space / cross-view consistency. Fresh from base 2B.
+    """
+    cfg = nymeria_pose_2view_actorobs_refs()  # inherits refs (R=4) config
+    cfg["job"]["name"] = "nymeria_pose_2view_actorobs_refs_campose"
+    # dataloaders that also emit camera poses (campose.npz)
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_observer_refs_campose"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_actor_observer_refs_campose"
+    # model: compute Plücker maps
+    cfg["model"]["config"]["enable_plucker"] = True
+    # net: enable the Plücker embedder
+    cfg["model"]["config"]["net"].update(enable_plucker=True, plucker_channels=6)
+    # validation viz must emit camera poses too
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data2/nymeria_processed_single", actor_observer=True,
+        num_reference_frames=4, emit_camera_poses=True,
+    )
+    return cfg
+
+
+def nymeria_pose_2view_actorobs_campose() -> dict:
+    """Ablation of `..._refs_campose` with the REFERENCE-frame condition REMOVED: Plücker camera conditioning
+    only (no in-context reference frames). A/B partner to isolate what the reference frames contribute vs the
+    Plücker shared-space signal. Same split/val, fresh from base 2B."""
+    cfg = nymeria_pose_2view_actorobs_refs_campose()
+    cfg["job"]["name"] = "nymeria_pose_2view_actorobs_campose"
+    # camera-pose-only dataloaders (no reference frames)
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_observer_campose"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_actor_observer_campose"
+    # disable reference frames (keep Plücker on)
+    cfg["model"]["config"]["num_reference_frames"] = 0
+    cfg["model"]["config"]["net"].update(enable_reference_frames=False, num_reference_frames=0)
+    # validation viz: no reference frames, keep camera poses
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data2/nymeria_processed_single", actor_observer=True,
+        num_reference_frames=0, emit_camera_poses=True,
+    )
+    return cfg
+
+
+def nymeria_pose_2view_actorobs_refs_campose_refpose() -> dict:
+    """`..._refs_campose` + POSED references: the in-context reference frames also get a Plücker ray map from
+    their OWN past camera poses (in the same per-pair canonical frame), so references are geometrically grounded
+    (posed reference views), not just floating appearance. Fresh from base 2B."""
+    cfg = nymeria_pose_2view_actorobs_refs_campose()
+    cfg["job"]["name"] = "nymeria_pose_2view_actorobs_refs_campose_refpose"
+    cfg["model"]["config"]["enable_reference_plucker"] = True
+    cfg["model"]["config"]["net"].update(enable_reference_plucker=True)
+    # (dataloaders already emit reference frames + camera poses; the dataset also emits reference_cam_w2c)
+    return cfg
+
+
+def nymeria_pose_2view_actorobs_refs_campose_refpose_shared() -> dict:
+    """`..._refs_campose_refpose` but with SHARED references: instead of per-view reference pools, one greedy
+    set-cover selection over the COMBINED pool (both views' past) is fed to both view slots; the per-view
+    view-embedding on reference tokens is DROPPED (shared_reference=True). Refs are geometrically identified
+    only by their posed Plücker; both views attend all of them via cross-view self-attention. Uses
+    refs_shared.npz (backfill_shared_refs.py). Fresh from base 2B."""
+    cfg = nymeria_pose_2view_actorobs_refs_campose_refpose()
+    cfg["job"]["name"] = "nymeria_pose_2view_actorobs_refs_campose_refpose_shared"
+    # net + (net-nested) share flag: drop the ref view-embedding
+    cfg["model"]["config"]["net"].update(shared_reference=True)
+    # shared-reference dataloaders (refs_shared.npz)
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_observer_refs_campose_shared"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_actor_observer_refs_campose_shared"
+    # validation viz: shared refs on the actor-observer val set
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data2/nymeria_processed_single", actor_observer=True,
+        num_reference_frames=4, shared_reference=True, emit_camera_poses=True,
+    )
+    return cfg
+
+
+def nymeria_pose_2view_actoractor_refs_campose_refpose() -> dict:
+    """Same structure as `..._refs_campose_refpose` (refs + Plücker + posed refs) but on the ACTOR-ACTOR longer
+    dataset (/data2/nymeria_processed_longer): two ego actors generated jointly, each with its own real action
+    caption. Fresh from base 2B."""
+    cfg = nymeria_pose_2view_actorobs_refs_campose_refpose()
+    cfg["job"]["name"] = "nymeria_pose_2view_actoractor_refs_campose_refpose"
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_actor_refs_campose"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_actor_actor_refs_campose"
+    # validation viz on the actor-actor longer val set
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data2/nymeria_processed_longer", actor_actor=True,
+        num_reference_frames=4, emit_camera_poses=True,
+    )
+    return cfg
+
+
+# fixed viz samples (real longer): 5 val (Loc_46+Loc_BX held-out) + 5 train (6 train locations) -> stable
+# wandb comparison across all actor-actor runs regardless of manifest size. See manifest/viz_fixed_ids.json.
+_AA_VAL_FIXED_IDS = [
+    "paul_act1_t2fpwg_003927__justin_act1_q1vpto_003941",
+    "paul_act2_pryhwf_031416__justin_act2_bw0jz0_031395",
+    "patricia_act2_c993gv_003234__greg_act2_jc6wnc_003243",
+    "patricia_act4_ifc9zd_005775__greg_act4_chhp0x_004682",
+    "patricia_act5_swxtds_038808__greg_act5_hu4348_040953",
+]
+_AA_TRAIN_FIXED_IDS = [
+    "carly_act0_et4azi_025410__andrew_act0_us908q_025439",
+    "sarah_act0_v24gb2_034881__alexis_act0_0e9t84_034200",
+    "jodi_act4_iplp4x_018018__katie_act4_ie44mj_017665",
+    "shawn_act7_onxbbz_019404__brenda_act7_mnqj9y_019417",
+    "randall_act2_nwllyf_016632__erica_act2_nbdcde_016609",
+]
+
+
+def nymeria_pose_2view_actoractor_refs_campose_refpose_shared() -> dict:
+    """SHARED-reference ACTOR-ACTOR, stage-2 of the stagewise plan: warm-start from the actor-observer SHARED
+    checkpoint, then fine-tune on the REAL longer (location-holdout) + SYNTHETIC combined set. Both views are
+    ego actors with their own captions; shared refs (no view-embedding) + Plücker + posed refs. Validation runs
+    on the REAL longer location-holdout val with FIXED viz pair_ids (stable wandb comparison)."""
+    cfg = nymeria_pose_2view_actorobs_refs_campose_refpose_shared()  # shared net + refs+plucker+refpose
+    cfg["job"]["name"] = "nymeria_pose_2view_actoractor_refs_campose_refpose_shared"
+    # warm-start from the actor-observer SHARED checkpoint (model-only staged dir -> fresh optimizer, iter 0)
+    cfg["checkpoint"]["load_path"] = "/data/model_output/warmstart/actorobs_shared_iter5500"
+    cfg["checkpoint"]["strict_resume"] = False
+    # data: REAL longer + SYNTHETIC combined (train); val = REAL longer location-holdout only
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_actor_refs_campose_shared_plus_synth"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_actor_actor_refs_campose_shared"
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data2/nymeria_processed_longer", actor_actor=True,
+        num_reference_frames=4, shared_reference=True, emit_camera_poses=True,
+        train_fixed_ids=_AA_TRAIN_FIXED_IDS, val_fixed_ids=_AA_VAL_FIXED_IDS,
+    )
+    return cfg
+
+
+def nymeria_pose_2b_stage1_lora() -> dict:
+    """STAGE-1 of the motion curriculum: single-view (V=1) LoRA pretraining on the MAX pooled ego-clip set
+    (real single actor + longer + synthetic, ~41k clips). Full conditioning minus cross-view: real per-clip
+    caption + pose + warping + plucker(self-frame0) + OWN refs (R=4). Backbone = LoRA(rank 64, prior-preserving);
+    new zero-init conditioning modules (pose/plucker/ref embedders, gates) stay fully trainable. Fresh from base
+    2B; warm-starts the 2-view stages 2/3. 2-GPU."""
+    cfg = nymeria_pose_2view_actorobs_refs_campose_refpose()  # refs(4)+plucker+refpose, per-view (own refs)
+    cfg["job"]["name"] = "nymeria_pose_2b_stage1_lora"
+    cfg["checkpoint"]["load_path"] = _base_2b_multiview_ckpt()  # start of the curriculum chain
+    cfg["checkpoint"]["strict_resume"] = False
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_single_view_stage1"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_single_view_stage1"
+    mc = cfg["model"]["config"]
+    mc["fsdp_shard_size"] = 2
+    mc["use_lora"] = True
+    mc["lora_rank"] = 64
+    mc["lora_alpha"] = 64
+    mc["net"].update(freeze_view_embedding=True)  # V=1 -> view embedding meaningless
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, single_view_stage1=True,
+        num_reference_frames=4, shared_reference=False, emit_camera_poses=True,
+    )
+    return cfg
+
+
+def nymeria_pose_2view_actoractor_refs_campose_refpose_shared_synthonly() -> dict:
+    """DIAGNOSTIC: SYNTHETIC-ONLY actor-actor (fresh from base 2B, no real-data influence). Isolates whether the
+    black/broken human appearance seen in synthetic val of the real+synth run comes from the REAL dataset
+    (masking prior) or is inherent to the synthetic data/masking. Val viz on synthetic val (no real fixed ids)."""
+    cfg = nymeria_pose_2view_actoractor_refs_campose_refpose_shared()
+    cfg["job"]["name"] = "nymeria_pose_2view_actoractor_refs_campose_refpose_shared_synthonly"
+    cfg["checkpoint"]["load_path"] = _base_2b_multiview_ckpt()  # fresh from base -> zero real-data influence
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_actor_refs_campose_shared_synthonly"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_actor_actor_refs_campose_shared_synthonly"
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data3/synthetic_processed_multi", actor_actor=True,
+        num_reference_frames=4, shared_reference=True, emit_camera_poses=True,
+        train_fixed_ids=None, val_fixed_ids=None,  # synthetic viz (real longer fixed ids don't apply)
+    )
+    return cfg
+
+
+# CoMind val handover/collaboration clips (from val holdout recs 1326e688/313483e1) — pinned so every validation
+# shows whether the model generates person-to-person object transfer. Verified loadable (refs+campose present).
+_COMIND_HANDOVER_VAL_IDS = [
+    "313483e1_000831",  # partner washes zucchini -> HANDS IT TO the wearer, who dries it
+    "1326e688_003624",  # partner rinses carrot; wearer's hands enter to TAKE THE CARROT FROM the partner
+    "1326e688_052008",  # partner washes dishes -> HANDS THE WEARER A PAIR OF CHOPSTICKS
+    "1326e688_040200",  # partner stirs pan; wearer takes ... (near GT handover: bowl @frame 40292)
+    "1326e688_028392",  # partner holds bowl+spoon, places in sink; wearer reaches (bowl transfer)
+    "1326e688_055608",  # pan handling near GT handover (pan @frame 55781)
+]
+
+
+def comind_actoractor_personpose_shared() -> dict:
+    """CoMind (Aria 2-person kitchen) 2-view actor-actor. Dataset = CoMind ONLY (32 multislam shared-world recs,
+    ~12k pairs). Skeleton pose is IDENTITY-colored (person_pose: leader=EGO / helper=PARTNER, view-invariant),
+    NOT wearer/partner role-colored. Conditioning = combined-pool warp RGB + person pose + GREEDY shared refs +
+    cross-view Plücker (campose/refpose). DEPTH channels are stored in clip.npz but NOT loaded (no depth condition
+    for now). Warm-start from the Nymeria actor-observer SHARED checkpoint. Val = CoMind held-out recs."""
+    cfg = nymeria_pose_2view_actoractor_refs_campose_refpose_shared()  # shared net + refs+plucker+refpose
+    cfg["job"]["name"] = "comind_actoractor_personpose_shared"
+    cfg["checkpoint"]["load_path"] = "/data/model_output/warmstart/actorobs_shared_iter5500"  # nymeria actor-observer
+    cfg["checkpoint"]["strict_resume"] = False
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "comind_actoractor_personpose_shared"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "comind_actoractor_personpose_shared"
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=len(_COMIND_HANDOVER_VAL_IDS), num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data4/comind_dataset", comind=True,
+        num_reference_frames=4, shared_reference=True, emit_camera_poses=True,
+        train_fixed_ids=None, val_fixed_ids=_COMIND_HANDOVER_VAL_IDS,  # pin handover/collab clips
+    )
+    return cfg
+
+
+def nymeria_synth_actoractor_depth_shared() -> dict:
+    """Nymeria longer + SYNTHETIC actor-actor + composite DEPTH condition. Depth = warped scene depth + human
+    mesh depth (mesh_cond depth_comp, RGB) -> frozen VAE -> depth_latent(16) -> zero-init depth_embedder ADDED to
+    tokens (== VAE-channel-concat, warm-start clean). Same net (refs+plucker+refpose) as the actor-actor shared
+    run; warm-start from the Nymeria actor-observer SHARED checkpoint. Val = nymeria longer holdout.
+    NOTE: synthetic depth is 100% ready; nymeria longer mesh Stage-B is in progress -> missing clips get ZERO
+    depth (coverage grows as Stage-B completes)."""
+    cfg = nymeria_pose_2view_actoractor_refs_campose_refpose_shared()  # refs+plucker+refpose, warm-start actorobs
+    cfg["job"]["name"] = "nymeria_synth_actoractor_depth_shared"
+    # warm-start from the DEPTH actor-observer run (nymeria_actorobs_depth_shared iter 5500) — pose+refs+plucker
+    # +DEPTH embedder are ALL already trained there (net matches exactly), so this continues from a depth-aware
+    # actor-observer rather than the depth-less actorobs_shared. Model-only staged dir -> fresh optimizer.
+    cfg["checkpoint"]["load_path"] = "/data/model_output/warmstart/actorobs_depth_iter5500"
+    cfg["checkpoint"]["strict_resume"] = False
+    cfg["model"]["config"]["depth_via_vae"] = True             # RGB depth -> frozen VAE -> depth_latent
+    cfg["model"]["config"]["net"].update(enable_depth=True)    # zero-init depth_embedder (additive)
+    cfg["trainer"]["callbacks"]["validation_viz"]["emit_depth"] = True  # add a "depth" column to the viz grid
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_actor_refs_campose_shared_plus_synth_depth"
+    return cfg
+
+
+def nymeria_actorobs_depth_shared() -> dict:
+    """Nymeria ACTOR-OBSERVER (single) + composite DEPTH condition — the stage-1 depth run (train this FIRST,
+    then warm-start the actor-actor depth run from it). Actor-observer SHARED model (refs+plucker+refpose) with
+    warped-scene + human-mesh depth (mesh_cond depth_comp) -> frozen VAE -> depth_latent(16) -> zero-init
+    depth_embedder ADDED to tokens (== VAE-channel-concat, day-0 contribution 0 so warm-start is clean).
+    Warm-start from the trained actor-observer SHARED checkpoint (actorobs_shared_iter5500) so only the depth
+    branch is new. single depth (mesh Stage-B) ~100% ready. Val = nymeria single holdout + a depth viz column."""
+    cfg = nymeria_pose_2view_actorobs_refs_campose_refpose_shared()  # actor-observer shared: refs+plucker+refpose
+    cfg["job"]["name"] = "nymeria_actorobs_depth_shared"
+    # FRESH from Cosmos base 2B (NOT warm-started from a trained actor-observer ckpt): the pose/refs/plucker/
+    # refpose/depth modules are all new (zero-init or missing keys) and train from scratch on top of the base
+    # video backbone. So iter-1 loss starts HIGH (~base video prior), unlike a warm-start.
+    cfg["checkpoint"]["load_path"] = _base_2b_multiview_ckpt()
+    cfg["checkpoint"]["strict_resume"] = False
+    cfg["model"]["config"]["depth_via_vae"] = True             # RGB depth -> frozen VAE -> depth_latent
+    cfg["model"]["config"]["net"].update(enable_depth=True)    # zero-init depth_embedder (additive)
+    cfg["trainer"]["callbacks"]["validation_viz"]["emit_depth"] = True  # add a "depth" column to the viz grid
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_actor_observer_refs_campose_shared_depth"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_actor_observer_refs_campose_shared_depth"
+    return cfg
+
+
+def comind_synth_actoractor_personpose_shared() -> dict:
+    """CoMind + SYNTHETIC combined 2-view actor-actor (both person-pose). Train = CoMind (32 shared-world recs,
+    ~10.8k) + synthetic (/data3/synthetic_processed_multi, ~17.4k) pooled. Validation stays on the CoMind val
+    (handover/collab clips pinned, same as comind_actoractor_personpose_shared). Warm-start from the Nymeria
+    actor-observer SHARED checkpoint. Depth NOT loaded."""
+    cfg = comind_actoractor_personpose_shared()  # comind net/conditioning/viz + actor-observer warm-start
+    cfg["job"]["name"] = "comind_synth_actoractor_personpose_shared"
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "comind_synth_actoractor_personpose_shared"  # comind + synthetic
+        # data_val stays "comind_actoractor_personpose_shared" (CoMind val) — set by the parent
+    return cfg
+
+
+def vroid_actoractor_refpose() -> dict:
+    """Synthetic vroid_batch (/data4/vroid_batch) 2-view actor-actor, FRESH from Cosmos base 2B (no warm-start).
+    Same net as the shared actor-actor runs (person-pose + shared refs + Plücker + posed refs) PLUS a new
+    REFERENCE-POSE condition: each reference frame's per-person skeleton render (refs_shared `reference_pose`)
+    -> frozen VAE -> zero-init `reference_pose_embedder` ADDED to the reference tokens (mirrors reference-Plücker,
+    so the in-context references now carry BOTH RGB appearance and pose). 4-GPU."""
+    cfg = nymeria_pose_2view_actoractor_refs_campose_refpose_shared()  # shared refs + plucker + refpose net
+    cfg["job"]["name"] = "vroid_actoractor_refpose_ar_generaltext"  # AR recipe (2-cond + noise) + general captions
+    cfg["checkpoint"]["load_path"] = _base_2b_multiview_ckpt()  # FRESH from Cosmos base 2B
+    cfg["checkpoint"]["strict_resume"] = False
+    # NEW: reference-pose conditioning (model VAE-encodes reference_pose -> net zero-init reference_pose_embedder)
+    cfg["model"]["config"]["enable_reference_pose"] = True
+    cfg["model"]["config"]["net"].update(enable_reference_pose=True)
+    # AR training recipe (from the handoff): condition on 1 OR 2 latents (50/50 per view) + noise the cond
+    # frames (prob 0.5) so the model is robust to imperfect (self-generated) cond frames at AR inference.
+    # `0: 0.0` is load-bearing (explicit so it can't be picked). See AR_HANDOFF/README.
+    cfg["model"]["config"]["conditional_frames_probs"] = {0: 0.0, 1: 0.5, 2: 0.5}
+    cfg["model"]["config"]["max_num_conditional_frames_per_view"] = 2
+    cfg["model"]["config"]["noisy_conditioning_prob"] = 0.5
+    cfg["model"]["config"]["noisy_conditioning_scale"] = 0.1
+    cfg["model"]["config"]["noisy_conditioning_min_frames"] = 2
+    # data: vroid synthetic (person-pose + reference-pose emitted)
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "vroid_actoractor_refpose"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "vroid_actoractor_refpose"
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data4/vroid_batch", actor_actor=True,
+        num_reference_frames=4, shared_reference=True, emit_camera_poses=True,
+        person_pose=True, emit_reference_pose=True, val_manifest="manifest/val_avatar_disjoint.csv",
+    )
+    return cfg
+
+
+def vroid_overfit5() -> dict:
+    """OVERFIT test: warm-start from the vroid ar_generaltext iter9000 checkpoint and deliberately overfit on
+    5 hand-picked val pairs (manifest/overfit5.csv). Train AND validate on the SAME 5 pairs so we can watch the
+    model memorize them. Same net/recipe as vroid_actoractor_refpose_ar_generaltext (AR 2-cond+noise, refpose)."""
+    cfg = vroid_actoractor_refpose()  # inherits AR recipe + refpose + general captions
+    cfg["job"]["name"] = "vroid_overfit5"
+    # warm-start from the iter9000 checkpoint of the main run (NOT base 2B)
+    cfg["checkpoint"]["load_path"] = (
+        "/data/model_output/train/vroid_actoractor_refpose_ar_generaltext/checkpoints/iter_000009000"
+    )
+    cfg["checkpoint"]["strict_resume"] = False  # load model weights; fresh optimizer/iteration for the overfit
+    cfg["checkpoint"]["save_iter"] = 200         # save often so we can inspect the memorization curve
+    cfg["trainer"]["max_iter"] = 4000
+    # data: the 5-pair overfit set for BOTH train and val
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "vroid_overfit5"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "vroid_overfit5"
+    # validation viz on the SAME 5 pairs, frequently (num_val=5 -> all of them)
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=200, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data4/vroid_batch", actor_actor=True,
+        num_reference_frames=4, shared_reference=True, emit_camera_poses=True,
+        person_pose=True, emit_reference_pose=True, val_manifest="manifest/overfit5.csv",
+    )
+    return cfg
+
+
+def nymeria_vroid_mixed_refpose_pretrain() -> dict:
+    """PRETRAIN on a MIX of nymeria actor-OBSERVER (real /data2/nymeria_processed_single) + vroid actor-ACTOR
+    (synthetic /data4/vroid_batch), pooled uniformly. Full conditioning structure INCLUDING reference-HUMAN-POSE
+    (skeleton): shared refs + Plücker campose + person-pose + reference_pose. Fresh from Cosmos base 2B. Standard
+    training recipe (no AR teacher-forcing/noise). 4-GPU."""
+    cfg = nymeria_pose_2view_actoractor_refs_campose_refpose_shared()  # shared refs + plucker + refpose net
+    cfg["job"]["name"] = "nymeria_vroid_mixed_refpose_pretrain"
+    cfg["checkpoint"]["load_path"] = _base_2b_multiview_ckpt()  # FRESH from Cosmos base 2B
+    cfg["checkpoint"]["strict_resume"] = False
+    # reference-HUMAN-POSE conditioning (VAE-encode reference_pose -> net zero-init reference_pose_embedder)
+    cfg["model"]["config"]["enable_reference_pose"] = True
+    cfg["model"]["config"]["net"].update(enable_reference_pose=True)
+    # data: the mixed obs + vroid loader (person-pose + reference-pose emitted for both)
+    for d in cfg["defaults"]:
+        if isinstance(d, dict) and "override /data_train" in d:
+            d["override /data_train"] = "nymeria_obs_vroid_mixed_refpose"
+        if isinstance(d, dict) and "override /data_val" in d:
+            d["override /data_val"] = "nymeria_obs_vroid_mixed_refpose"
+    # validation viz on the vroid actor-actor val set (person-pose + reference-pose)
+    cfg["trainer"]["callbacks"]["validation_viz"] = L(NymeriaValidationViz)(
+        every_n=500, num_train=5, num_val=5, num_sampling_step=35, guidance=7.0, fps=10,
+        num_video_frames=77, root="/data4/vroid_batch", actor_actor=True,
+        num_reference_frames=4, shared_reference=True, emit_camera_poses=True,
+        person_pose=True, emit_reference_pose=True, val_manifest="manifest/val_avatar_disjoint.csv",
+    )
+    return cfg
+
+
+experiments = [
+    nymeria_pose_2actor_2b_smoke(),
+    nymeria_pose_2actor_2b_longer(),
+    nymeria_pose_2b_stage1_singleview(),
+    nymeria_pose_2b_stage1_sv_vaeconcat(),
+    nymeria_pose_2b_stage1_sv_vaemlp(),
+    nymeria_pose_2view_actorobs(),
+    nymeria_pose_2view_actorobs_refs(),
+    nymeria_pose_2view_actorobs_refs_campose(),
+    nymeria_pose_2view_actorobs_campose(),
+    nymeria_pose_2view_actorobs_refs_campose_refpose(),
+    nymeria_pose_2view_actorobs_refs_campose_refpose_shared(),
+    nymeria_pose_2view_actoractor_refs_campose_refpose(),
+    nymeria_pose_2view_actoractor_refs_campose_refpose_shared(),
+    nymeria_pose_2view_actoractor_refs_campose_refpose_shared_synthonly(),
+    nymeria_pose_2b_stage1_lora(),
+    comind_actoractor_personpose_shared(),
+    comind_synth_actoractor_personpose_shared(),
+    nymeria_synth_actoractor_depth_shared(),
+    nymeria_actorobs_depth_shared(),
+    vroid_actoractor_refpose(),
+    vroid_overfit5(),
+    nymeria_vroid_mixed_refpose_pretrain(),
+]
+
+cs = ConfigStore.instance()
+for _item in experiments:
+    cs.store(
+        group="experiment",
+        package="_global_",
+        name=_item["job"]["name"],
+        node=_item,
+    )

cosmos_predict2/_src/predict2_multiview/datasets/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+

cosmos_predict2/_src/predict2_multiview/datasets/comind_pairs.py

@@ -0,0 +1,302 @@
+"""CoMind 2-view actor-actor dataset (shared-world, multislam-eligible 32 recs).
+
+Unlike Nymeria (two separate per-clip-dir clip.npz per pair), ONE CoMind clip.npz holds BOTH ego views —
+``L_`` (leader) and ``H_`` (helper) — as JPEG byte-array streams (warped_cond / pose / target) + packed
+visibility + Sapiens keypoints. Sidecars produced by the shared-world rebuild:
+  * ``<clip>.refs.npz``    — 8 GREEDY shared reference frames (combined leader+helper pool) + shared-world w2c/K.
+  * ``<clip>.campose.npz`` — both views' TARGET-frame camera poses (77) + K, in the multislam shared world.
+Captions: ``manifest/captions/<clip_id>.json`` -> {"leader": ..., "helper": ...}.
+
+Emits the SAME sample dict as ``NymeriaActorActorDataset`` (shared refs / camera poses / role-mix) so it drops
+straight into the existing 2-view experiment + net. Combine with synthetic via ``get_comind_actor_actor_combined_loader``.
+"""
+import csv
+import json
+import os
+import random
+from typing import Any, Optional
+
+import attrs
+import cv2
+import numpy as np
+import torch
+from hydra.core.config_store import ConfigStore
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.imaginaire.utils import log
+from cosmos_predict2._src.predict2.datasets.cached_replay_dataloader import get_cached_replay_dataloader
+from cosmos_predict2._src.predict2_multiview.datasets.multiview import collate_fn
+from cosmos_predict2._src.predict2_multiview.datasets.nymeria_pairs import (
+    NymeriaActorActorDataset,
+    NymeriaPairsConfig,
+    _resize_thwc_uint8,
+)
+
+cv2.setNumThreads(1)
+
+
+@attrs.define(slots=False)
+class ComindPairsConfig:
+    root: str = "/data4/comind_dataset"
+    resolution_hw: tuple[int, int] = (480, 480)
+    num_video_frames: int = 77
+    fps: float = 10.0
+    num_reference_frames: int = 4          # per-view slot; shared set = 2*R (=8) greedy refs
+    shared_reference: bool = True          # refs carry no view identity (net drops ref view-emb)
+    emit_camera_poses: bool = True         # cross-view Plücker (campose.npz) + posed refs (refs.npz w2c)
+    role_mix: bool = True                  # swap view0<->view1 (+captions) 50%
+    person_pose: bool = True               # use IDENTITY-colored pose (L/H_pose_person: leader=EGO, helper=PARTNER
+    #                                        view-invariant) instead of role-colored pose (wearer/partner)
+    reference_pose: bool = True             # emit reference_pose (Sapiens skeleton on each reference frame, person-ID
+    #                                        palette) alongside reference_frames — human-pose condition for the refs
+    split: str = "all"                     # "train" | "val" | "all" — held out by RECORDING
+    val_num_sessions: int = 3              # recordings held out for val
+    split_seed: int = 1234
+
+
+def _jdec(b: bytes) -> np.ndarray:
+    return cv2.cvtColor(cv2.imdecode(np.frombuffer(b, np.uint8), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
+
+
+def _stack_jpeg(arr: np.ndarray) -> np.ndarray:  # object array of JPEG bytes -> (T,H,W,3) uint8
+    return np.stack([_jdec(b) for b in arr])
+
+
+def _fit_poses(w2c: np.ndarray, n: int) -> np.ndarray:
+    m = w2c.shape[0]
+    if m < n:
+        return np.concatenate([w2c, np.repeat(w2c[-1:], n - m, 0)], 0)
+    return w2c[:n]
+
+
+class ComindActorActorDataset(torch.utils.data.Dataset):
+    def __init__(self, config: ComindPairsConfig):
+        self.config = config
+        self.root = config.root
+        self.clips_dir = os.path.join(self.root, "clips")
+        self.resolution_hw = tuple(config.resolution_hw)
+        self.num_frames = int(config.num_video_frames)
+        self.R = int(config.num_reference_frames)
+        self.emit_cam = bool(config.emit_camera_poses)
+        self.cap_dir = os.path.join(self.root, "manifest", "captions")
+
+        elig = sorted(json.load(open(os.path.join(self.root, "manifest", "multislam_paths.json"))))
+        rng = random.Random(config.split_seed)
+        shuffled = list(elig)
+        rng.shuffle(shuffled)
+        val_recs = set(shuffled[: config.val_num_sessions])
+        elig = set(elig)
+
+        self.pairs: list[dict[str, str]] = []
+        with open(os.path.join(self.root, "manifest", "clip_pairs.csv"), newline="") as f:
+            for row in csv.DictReader(f):
+                rec8 = row["recording"][:8]
+                if rec8 not in elig:
+                    continue
+                cid = row["clip_id"]
+                npz = os.path.join(self.clips_dir, rec8, f"{cid}.npz")
+                refs = npz[:-4] + ".refs.npz"
+                camp = npz[:-4] + ".campose.npz"
+                if not (os.path.isfile(npz) and os.path.isfile(refs)):
+                    continue
+                if self.emit_cam and not os.path.isfile(camp):
+                    continue
+                in_val = rec8 in val_recs
+                if config.split == "train" and in_val:
+                    continue
+                if config.split == "val" and not in_val:
+                    continue
+                self.pairs.append({"pair_id": cid, "rec8": rec8, "npz": npz, "refs": refs, "camp": camp})
+        log.info(
+            f"ComindActorActorDataset[{config.split}]: {len(self.pairs)} clips "
+            f"({len(elig)} eligible recs, {len(val_recs)} held out for val, R={self.R}, emit_cam={self.emit_cam})"
+        )
+        if not self.pairs:
+            raise RuntimeError(f"No usable CoMind clips for split={config.split} under {self.clips_dir}")
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def _resize(self, thwc: np.ndarray, n: int) -> torch.Tensor:
+        return _resize_thwc_uint8(thwc, n, self.resolution_hw)
+
+    def _load_view(self, z, refs, camp, side: str, slot: int) -> dict[str, torch.Tensor]:
+        """side: 'L'(leader) | 'H'(helper). slot: final view position (0/1) -> shared-ref R-slice."""
+        out: dict[str, torch.Tensor] = {}
+        out["video"] = self._resize(_stack_jpeg(z[f"{side}_target"]), self.num_frames)
+        pose_key = f"{side}_pose_person" if (self.config.person_pose and f"{side}_pose_person" in z.files) else f"{side}_pose"
+        out["control_input_pose"] = self._resize(_stack_jpeg(z[pose_key]), self.num_frames)
+        out["control_input_warped"] = self._resize(_stack_jpeg(z[f"{side}_warped"]), self.num_frames)
+        vt, vh, vw = (int(x) for x in z["vis_shape"])
+        vis = np.unpackbits(z[f"{side}_vis_packed"])[: vt * vh * vw].reshape(vt, vh, vw)
+        vis = (vis.astype(np.uint8) * 255)[..., None]  # (T,H,W,1)
+        out["control_input_visibility"] = self._resize(vis, self.num_frames)
+        if self.R > 0:
+            allf = refs["reference_frames"]  # (2R, h, h, 3) greedy shared set (top-gain first)
+            lo = slot * self.R
+            out["reference_frames"] = self._resize(allf[lo : lo + self.R], self.R)
+            if self.config.reference_pose and "reference_pose" in refs.files:  # Sapiens skeleton per ref frame (person-ID palette)
+                out["reference_pose"] = self._resize(refs["reference_pose"][lo : lo + self.R], self.R)
+            if self.emit_cam:
+                out["cam_ref_w2c"] = torch.from_numpy(refs["reference_w2c"].astype(np.float32)[lo : lo + self.R])
+        if self.emit_cam:
+            out["cam_w2c"] = torch.from_numpy(_fit_poses(camp[f"{side}_w2c"].astype(np.float32), self.num_frames))
+            out["cam_K"] = torch.from_numpy(camp[f"K_{'leader' if side == 'L' else 'helper'}"].astype(np.float32))
+            out["cam_src_res"] = torch.tensor(504, dtype=torch.int64)
+        return out
+
+    def _caption(self, cid: str) -> tuple[str, str]:
+        p = os.path.join(self.cap_dir, f"{cid}.json")
+        if os.path.isfile(p):
+            d = json.load(open(p))
+            return d.get("leader", ""), d.get("helper", "")
+        return "", ""
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        for _ in range(16):
+            try:
+                return self._build_sample(idx)
+            except Exception as e:
+                log.warning(f"ComindActorActorDataset: bad sample {self.pairs[idx]['pair_id']} ({type(e).__name__}: {str(e)[:80]}); retrying")
+                idx = random.randrange(len(self.pairs))
+        raise RuntimeError("ComindActorActorDataset: too many unreadable samples in a row")
+
+    def _build_sample(self, idx: int) -> dict[str, Any]:
+        p = self.pairs[idx]
+        z = np.load(p["npz"], allow_pickle=True)
+        refs = np.load(p["refs"], allow_pickle=True)
+        camp = np.load(p["camp"], allow_pickle=True) if self.emit_cam else None
+        cap_l, cap_h = self._caption(p["pair_id"])
+
+        views = [("L", cap_l), ("H", cap_h)]  # view0=leader, view1=helper
+        if self.config.role_mix and random.random() < 0.5:
+            views = views[::-1]
+        actors = [self._load_view(z, refs, camp, side, slot=i) for i, (side, _) in enumerate(views)]
+        captions = [c for _, c in views]
+        n_views = 2
+
+        sample: dict[str, Any] = {}
+        for key in actors[0]:
+            if key.startswith("cam_"):
+                continue
+            sample[key] = torch.cat([a[key] for a in actors], dim=1).contiguous()
+        if "cam_w2c" in actors[0]:
+            sample["camera_w2c"] = torch.stack([a["cam_w2c"] for a in actors], dim=0).contiguous()
+            sample["camera_K"] = torch.stack([a["cam_K"] for a in actors], dim=0).contiguous()
+            sample["camera_src_res"] = torch.stack([a["cam_src_res"] for a in actors], dim=0).contiguous()
+        if "cam_ref_w2c" in actors[0]:
+            sample["reference_cam_w2c"] = torch.stack([a["cam_ref_w2c"] for a in actors], dim=0).contiguous()
+
+        T = self.num_frames
+        sample.update(
+            {
+                "__key__": p["pair_id"],
+                "__url__": p["pair_id"],
+                "ai_caption": captions,
+                "view_indices": torch.tensor([0] * T + [1] * T, dtype=torch.int64),
+                "fps": torch.tensor(self.config.fps, dtype=torch.float64),
+                "num_video_frames_per_view": torch.tensor(T, dtype=torch.int64),
+                "view_indices_selection": torch.tensor([0, 1], dtype=torch.int64),
+                "camera_keys_selection": ["view0", "view1"],
+                "sample_n_views": torch.tensor(n_views, dtype=torch.int64),
+                "padding_mask": torch.zeros((1, *self.resolution_hw), dtype=torch.float32),
+                "ref_cam_view_idx_sample_position": torch.tensor(-1, dtype=torch.int64),
+                "front_cam_view_idx_sample_position": torch.tensor(0, dtype=torch.int64),
+            }
+        )
+        return sample
+
+
+def get_comind_actor_actor_combined_loader(
+    *,
+    comind_root: str = "/data4/comind_dataset",
+    synth_root: str = "/data3/synthetic_processed_multi",
+    synth_manifest: str = "manifest/train_split.csv",
+    resolution_hw: tuple[int, int] = (480, 480),
+    num_video_frames: int = 77,
+    fps: float = 10.0,
+    role_mix: bool = True,
+    num_reference_frames: int = 4,
+    shared_reference: bool = True,
+    emit_camera_poses: bool = True,
+    person_pose: bool = True,
+    split: str = "train",
+    batch_size: int = 1,
+    num_workers: int = 4,
+    prefetch_factor: Optional[int] = 2,
+    is_train: bool = True,
+    **kwargs: Any,
+):
+    """ConcatDataset over CoMind (shared-world 32 recs) [+ synthetic]. Both emit the same sample dict; pooled
+    uniformly. synth_root="" -> CoMind-only. Val (is_train=False) uses each set's held-out split."""
+    parts: list[torch.utils.data.Dataset] = [
+        ComindActorActorDataset(
+            ComindPairsConfig(
+                root=comind_root, resolution_hw=tuple(resolution_hw), num_video_frames=num_video_frames, fps=fps,
+                num_reference_frames=num_reference_frames, shared_reference=shared_reference,
+                emit_camera_poses=emit_camera_poses, role_mix=role_mix, person_pose=person_pose, split=split,
+            )
+        )
+    ]
+    if synth_root:
+        parts.append(
+            NymeriaActorActorDataset(
+                NymeriaPairsConfig(
+                    root=synth_root, manifest_csv=synth_manifest, resolution_hw=tuple(resolution_hw),
+                    num_video_frames=num_video_frames, fps=fps, role_mix=role_mix,
+                    num_reference_frames=num_reference_frames, shared_reference=shared_reference,
+                    emit_camera_poses=emit_camera_poses, person_pose=person_pose,
+                )
+            )
+        )
+    dataset = torch.utils.data.ConcatDataset(parts) if len(parts) > 1 else parts[0]
+    log.info(f"comind+synth combined[{split}]: {[len(p) for p in parts]} -> {len(dataset)} clips")
+    sampler = None
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=is_train, drop_last=True)
+    return get_cached_replay_dataloader(
+        webdataset=False, use_cache=False, dataset=dataset, batch_size=batch_size, num_workers=num_workers,
+        sampler=sampler, shuffle=(sampler is None and is_train), drop_last=True,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None, persistent_workers=num_workers > 0,
+        pin_memory=False, collate_fn=collate_fn, cache_replay_name="comind_actor_actor_combined_dataloader",
+    )
+
+
+def register_comind_data() -> None:
+    cs = ConfigStore.instance()
+    # CoMind-ONLY (shared-world 32 recs), 2-view actor-actor, IDENTITY-colored pose (person_pose), greedy shared
+    # refs + cross-view Plücker. Depth channels stored but NOT loaded (person_pose pose, no depth condition).
+    cs.store(
+        group="data_train", package="dataloader_train", name="comind_actoractor_personpose_shared",
+        node=L(get_comind_actor_actor_combined_loader)(
+            synth_root="", split="train", num_video_frames=77, role_mix=True, num_reference_frames=4,
+            shared_reference=True, emit_camera_poses=True, person_pose=True, is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="comind_actoractor_personpose_shared",
+        node=L(get_comind_actor_actor_combined_loader)(
+            synth_root="", split="val", num_video_frames=77, role_mix=False, num_reference_frames=4,
+            shared_reference=True, emit_camera_poses=True, person_pose=True, is_train=False, batch_size=1, num_workers=2),
+    )
+    # CoMind + SYNTHETIC combined, 2-view actor-actor, person-pose (both use pose_person). Train mixes both;
+    # validation stays on the CoMind val (data_val below = comind-only).
+    cs.store(
+        group="data_train", package="dataloader_train", name="comind_synth_actoractor_personpose_shared",
+        node=L(get_comind_actor_actor_combined_loader)(
+            synth_root="/data3/synthetic_processed_multi", split="train", num_video_frames=77, role_mix=True,
+            num_reference_frames=4, shared_reference=True, emit_camera_poses=True, person_pose=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+    # (kept) CoMind + synthetic combined variant, role pose — earlier config.
+    cs.store(
+        group="data_train", package="dataloader_train", name="comind_actor_actor_refs_campose_shared_plus_synth",
+        node=L(get_comind_actor_actor_combined_loader)(
+            split="train", num_video_frames=77, role_mix=True, num_reference_frames=4,
+            shared_reference=True, emit_camera_poses=True, person_pose=True, is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="comind_actor_actor_refs_campose_shared",
+        node=L(get_comind_actor_actor_combined_loader)(
+            synth_root="", split="val", num_video_frames=77, role_mix=False, num_reference_frames=4,
+            shared_reference=True, emit_camera_poses=True, person_pose=True, is_train=False, batch_size=1, num_workers=2),
+    )

cosmos_predict2/_src/predict2_multiview/datasets/local.py

@@ -0,0 +1,173 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Local file-based datasets.
+"""
+
+import glob
+import os
+import random
+from pathlib import Path
+
+import pandas as pd
+from torch.utils.data import Dataset
+
+from cosmos_predict2._src.predict2_multiview.datasets.multiview import (
+    AugmentationConfig,
+    make_augmentations,
+)
+
+
+class WaymoLocalDataset(Dataset):
+    def __init__(
+        self,
+        video_file_dirs: list[str],
+        augmentation_config: AugmentationConfig,
+        shuffle: bool = True,
+        gc_every_n: int = 100,
+    ) -> None:
+        self.video_file_dirs = video_file_dirs
+        self.augmentation_config = augmentation_config
+        self.shuffle = shuffle
+        self.gc_every_n = gc_every_n
+
+        self.augmentations, self.dataset_keys = make_augmentations(augmentation_config)
+        self.sample_dirs = self.build_sample_dirs(self.video_file_dirs)
+
+        if self.shuffle:
+            random.shuffle(self.sample_dirs)
+
+    def build_sample_dirs(self, video_file_dirs: list[str]) -> list[str]:
+        sample_dirs = []
+        for video_file_dir in video_file_dirs:
+            for sample_dir in glob.glob(os.path.join(video_file_dir, "**")):
+                if os.path.isdir(sample_dir):
+                    sample_dirs.append(sample_dir)
+        return sample_dirs
+
+    def load_data(self, sample_dir: str) -> dict:
+        sample_id = sample_dir.split("/")[-1]
+        data_dict = dict()
+        for filename in glob.glob(os.path.join(sample_dir, "*.mp4")):
+            with open(filename, "rb") as f:
+                view = filename.split("/")[-1].split(".")[0]
+                data_dict[f"video_{view}"] = f.read()
+        return data_dict
+
+    def load_caption(self, sample_dir: str) -> dict:
+        caption_path = os.path.join(sample_dir, "caption.jsonl")
+        with open(caption_path, "r") as f:
+            caption_df = pd.read_json(f, lines=True, orient="records")
+
+        caption_dict = dict()
+        for view_name, view_df in caption_df.groupby("view"):
+            caption_styles = dict()
+            for row in view_df.itertuples():
+                caption = row.caption
+                tag = None if pd.isna(row.tag) else row.tag
+                caption_styles[tag or "long"] = caption
+            caption_dict[f"caption_{view_name}"] = {
+                "t2w_windows": [
+                    {"start_frame": 0, "end_frame": self.augmentation_config.num_video_frames, **caption_styles}
+                ]
+            }
+        return caption_dict
+
+    def __len__(self) -> int:
+        return len(self.sample_dirs)
+
+    def __getitem__(self, index: int) -> dict:
+        sample_dir = self.sample_dirs[index]
+        data_dict = {
+            "__key__": str(index),
+            "__url__": str(sample_dir[index]),
+        }
+        data_dict.update(self.load_data(sample_dir))
+        data_dict.update(self.load_caption(sample_dir))
+        for k, aug in self.augmentations.items():
+            data_dict = aug(data_dict)
+        return data_dict
+
+
+class LocalMultiViewDataset(Dataset):
+    """Dataset wrapper for local multiview sample."""
+
+    def __init__(
+        self,
+        video_file_dicts: list[dict[str, bytes | Path | None]],
+        prompts: list[str],
+        augmentation_config: AugmentationConfig,
+        camera_key_adapter: dict[str, str] | None = None,
+        control_file_dicts: list[dict[str, bytes | Path | None]] | None = None,
+    ) -> None:
+        self.video_file_dicts = video_file_dicts
+        self.prompts = prompts
+        self.augmentation_config = augmentation_config
+        self.camera_key_adapter = camera_key_adapter
+        self.control_file_dicts = control_file_dicts
+
+        if self.control_file_dicts is not None and len(self.video_file_dicts) != len(self.control_file_dicts):
+            raise ValueError("Number of video file dicts and control file dicts must be the same!")
+
+        if len(self.prompts) != len(self.video_file_dicts):
+            raise ValueError("Number of prompts and video file dicts must be the same!")
+
+        if self.augmentation_config.single_caption_camera_name is None:
+            raise ValueError(
+                "`single_caption_camera_name` must be set since only single prompt is provided by dataset!"
+            )
+
+        self.augmentations, self.dataset_keys = make_augmentations(augmentation_config)
+
+    def __len__(self) -> int:
+        return len(self.video_file_dicts)
+
+    def __getitem__(self, index: int) -> dict:
+        data_dict = {
+            "__key__": str(index),
+            "__url__": "local_dataset",
+        }
+
+        for view_key, filepath in self.video_file_dicts[index].items():
+            if filepath is None:
+                raise ValueError(f"view_key {view_key} has null filepath!")
+            default_key = self.camera_key_adapter[view_key] if self.camera_key_adapter else view_key
+            video_key = self.augmentation_config.camera_video_key_mapping[default_key]
+            if isinstance(filepath, bytes):
+                data_dict[video_key] = filepath
+            else:
+                with open(filepath, "rb") as f:
+                    data_dict[video_key] = f.read()
+
+        if self.control_file_dicts is not None:
+            for view_key, filepath in self.control_file_dicts[index].items():
+                if filepath is None:
+                    raise ValueError(f"view_key {view_key} has null filepath!")
+                default_key = self.camera_key_adapter[view_key] if self.camera_key_adapter else view_key
+                control_key = self.augmentation_config.camera_control_key_mapping[default_key]
+                if isinstance(filepath, bytes):
+                    data_dict[control_key] = filepath
+                else:
+                    with open(filepath, "rb") as f:
+                        data_dict[control_key] = f.read()
+
+        caption_styles = dict(
+            zip(
+                self.augmentation_config.caption_probability.keys(),
+                [self.prompts[index] for _ in range(len(self.augmentation_config.caption_probability))],
+            )
+        )
+
+        caption_key = self.augmentation_config.camera_caption_key_mapping[
+            self.augmentation_config.single_caption_camera_name
+        ]
+        data_dict[caption_key] = {
+            "t2w_windows": [
+                {"start_frame": 0, "end_frame": self.augmentation_config.num_video_frames, **caption_styles}
+            ]
+        }
+
+        for k, aug in self.augmentations.items():
+            data_dict = aug(data_dict)
+        return data_dict

cosmos_predict2/_src/predict2_multiview/datasets/multiview.py

@@ -0,0 +1,547 @@
+# 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.
+
+"""
+Webloaders of datasets and augmentations for visual-text multiview dataset for AV
+"""
+
+try:
+    from megatron.core import parallel_state
+
+    USE_MEGATRON = True
+except ImportError:
+    USE_MEGATRON = False
+import io
+import random
+from typing import Any, Final, Literal, Optional, TypeAlias
+
+import attrs
+import torch
+import webdataset as wds
+from einops import rearrange
+from torchvision.transforms import InterpolationMode, Resize
+
+import cosmos_predict2._src.predict2.datasets.distributor.parallel_sync_multi_aspect_ratio as parallel_sync_multi_aspect_ratio
+from cosmos_predict2._src.imaginaire.datasets.decoders.json_loader import json_decoder
+from cosmos_predict2._src.imaginaire.datasets.decoders.video_decoder import video_naive_bytes
+from cosmos_predict2._src.imaginaire.datasets.webdataset.augmentors.augmentor import Augmentor
+from cosmos_predict2._src.imaginaire.datasets.webdataset.config.schema import DatasetConfig
+from cosmos_predict2._src.imaginaire.datasets.webdataset.distributors import ShardlistBasic
+from cosmos_predict2._src.imaginaire.datasets.webdataset.webdataset_ext import Dataset
+from cosmos_predict2._src.imaginaire.utils import log
+from cosmos_predict2._src.predict2.datasets.cached_replay_dataloader import get_cached_replay_dataloader
+from cosmos_predict2._src.predict2_multiview.datasets.wdinfo_utils import DEFAULT_CATALOG, get_video_dataset_info
+
+CameraKeyType: TypeAlias = str
+
+# for Autonomous Driving Dataset (Alpamayo and Mads)
+
+DEFAULT_CAMERAS: Final[tuple[CameraKeyType, ...]] = (
+    "camera_front_wide_120fov",
+    "camera_cross_right_120fov",
+    "camera_rear_right_70fov",
+    "camera_rear_tele_30fov",
+    "camera_rear_left_70fov",
+    "camera_cross_left_120fov",
+    "camera_front_tele_30fov",
+)
+
+DEFAULT_CAMERA_VIEW_MAPPING: Final = dict(zip(DEFAULT_CAMERAS, range(len(DEFAULT_CAMERAS))))
+
+DEFAULT_CAPTION_PREFIXES: Final = {
+    "camera_front_wide_120fov": "The video is captured from a camera mounted on a car. The camera is facing forward.",
+    "camera_cross_right_120fov": "The video is captured from a camera mounted on a car. The camera is facing to the right.",
+    "camera_rear_right_70fov": "The video is captured from a camera mounted on a car. The camera is facing the rear right side.",
+    "camera_rear_tele_30fov": "The video is captured from a camera mounted on a car. The camera is facing backwards.",
+    "camera_rear_left_70fov": "The video is captured from a camera mounted on a car. The camera is facing the rear left side.",
+    "camera_cross_left_120fov": "The video is captured from a camera mounted on a car. The camera is facing to the left.",
+    "camera_front_tele_30fov": "The video is captured from a telephoto camera mounted on a car. The camera is facing forward.",
+}
+
+DEFAULT_CAPTION_KEY_MAPPING: Final = dict(
+    zip(DEFAULT_CAMERAS, [f"metas_{k}_10s_chunks_qwen2p5_vl_32b" for k in DEFAULT_CAMERAS])
+)
+DEFAULT_VIDEO_KEY_MAPPING: Final = dict(zip(DEFAULT_CAMERAS, [f"video_{k}" for k in DEFAULT_CAMERAS]))
+
+# Agibot 3-view (head_color, hand_left, hand_right) multiview multicontrol
+AGIBOT_VIEWS: Final[tuple[CameraKeyType, ...]] = ("head_color", "hand_left", "hand_right")
+AGIBOT_VIEW_MAPPING: Final = dict(zip(AGIBOT_VIEWS, range(len(AGIBOT_VIEWS))))
+AGIBOT_VIDEO_KEY_MAPPING: Final = dict(zip(AGIBOT_VIEWS, [f"video_{k}" for k in AGIBOT_VIEWS]))
+AGIBOT_CAPTION_KEY_MAPPING: Final = dict(zip(AGIBOT_VIEWS, [f"metas_{k}" for k in AGIBOT_VIEWS]))
+AGIBOT_CONTROL_KEY_MAPPING: Final = dict(zip(AGIBOT_VIEWS, [f"control_{k}" for k in AGIBOT_VIEWS]))
+AGIBOT_CAPTION_PREFIXES: Final = {
+    "head_color": "The video is captured from a camera mounted on the head of the subject, facing forward.",
+    "hand_left": "The video is captured from a camera mounted on the left hand of the subject.",
+    "hand_right": "The video is captured from a camera mounted on the right hand of the subject.",
+}
+
+
+class UnpackMetas(Augmentor):
+    """Unpack metas from single meta dicts list into per-camera meta dicts."""
+
+    def __init__(
+        self, position_to_camera_mapping: dict[int, str], input_key: str = "metas", output_prefix: str = "metas_"
+    ) -> None:
+        super().__init__([], {})
+        self.position_to_camera_mapping = position_to_camera_mapping
+        self.input_key = input_key
+        self.output_prefix = output_prefix
+
+    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
+        metas = data.pop(self.input_key)
+        for i, meta in enumerate(metas):
+            camera_name = self.position_to_camera_mapping[i]
+            data[f"{self.output_prefix}{camera_name}"] = meta
+        return data
+
+
+class ExtractFramesAndCaptions(Augmentor):
+    """Extract frames from a videos."""
+
+    def __init__(
+        self,
+        camera_order: list[CameraKeyType],
+        num_frames: int,
+        resolution_hw: tuple[int, int],
+        fps_downsample_factor: int,
+        caption_probability: dict[str, float],
+        camera_view_mapping: dict[CameraKeyType, int],
+        camera_caption_key_mapping: dict[CameraKeyType, str],
+        camera_video_key_mapping: dict[CameraKeyType, str],
+        camera_control_key_mapping: Optional[dict[CameraKeyType, str]] = None,
+        add_view_prefix_to_caption: bool = False,
+        camera_prefix_mapping: Optional[dict[CameraKeyType, str]] = None,
+        single_caption_camera_name: Optional[CameraKeyType] = None,
+        window_random_frame_offset_range: Optional[tuple[int, int]] = None,
+    ) -> None:
+        """Extracts frames and captions from video/metadata dicts.
+
+        Args:
+            camera_order: Order of cameras to extract
+            num_frames: Number of frames to extract
+            resolution_hw: Resolution of the extracted frames
+            fps_downsample_factor: FPS downsample factor
+            caption_probability: Probability of each caption type in t2w window
+            camera_view_mapping: Mapping of camera keys to view indices
+            camera_caption_key_mapping: Mapping of camera keys to caption keys
+            camera_video_key_mapping: Mapping of camera keys to video keys
+            camera_control_key_mapping: Mapping of camera keys to control keys
+            add_view_prefix_to_caption: Whether to add caption prefix for all views
+            camera_prefix_mapping: Mapping of camera keys to prefixes
+            single_caption_camera_name: Name of the camera key to use for single caption conditioning.
+                If `add_view_prefix_to_caption` is True, will still provide prefixes for other views.
+            window_random_frame_offset_range: Optional range of random offset to add to the start frame of the extracted window.
+
+        Returns:
+            data: Dictionary with resized tensors of frames and captions
+        """
+        super().__init__([], {})
+        self.camera_order = camera_order
+        self.num_frames = num_frames
+        self.resolution_hw = resolution_hw
+        self.fps_downsample_factor = fps_downsample_factor
+        self.caption_probability = caption_probability
+        self.camera_view_mapping = camera_view_mapping
+        self.camera_caption_key_mapping = camera_caption_key_mapping
+        self.camera_video_key_mapping = camera_video_key_mapping
+        self.camera_control_key_mapping = camera_control_key_mapping
+        self.add_view_prefix_to_caption = add_view_prefix_to_caption
+        self.camera_prefix_mapping = camera_prefix_mapping
+        self.single_caption_camera_name = single_caption_camera_name
+        self.window_random_frame_offset_range = window_random_frame_offset_range
+
+        if self.add_view_prefix_to_caption and self.camera_prefix_mapping is None:
+            raise ValueError("camera_prefix_mapping is required when add_view_prefix_to_caption is True")
+
+        if set(self.camera_caption_key_mapping.keys()) != set(self.camera_video_key_mapping.keys()):
+            raise ValueError(
+                f"Mismatching keys {set(self.camera_caption_key_mapping.keys())} != {set(self.camera_video_key_mapping.keys())}"
+            )
+        if self.camera_control_key_mapping is not None:
+            if set(self.camera_control_key_mapping.keys()) != set(self.camera_caption_key_mapping.keys()):
+                raise ValueError(
+                    f"Mismatching keys {set(self.camera_control_key_mapping.keys())} != {set(self.camera_caption_key_mapping.keys())}"
+                )
+        for camera_name in self.camera_caption_key_mapping.keys():
+            if camera_name not in self.camera_view_mapping:
+                raise ValueError(f"Camera name {camera_name} not found in camera view mapping")
+        if self.single_caption_camera_name and self.single_caption_camera_name not in self.camera_order:
+            raise ValueError(
+                f"Single caption camera name {self.single_caption_camera_name} must appear in selected cameras"
+            )
+
+        if self.window_random_frame_offset_range is not None:
+            start_range, end_range = self.window_random_frame_offset_range
+            if start_range < 0 or end_range < 0:
+                raise ValueError("`window_random_frame_offset_range` must be non-negative")
+            if start_range > end_range:
+                raise ValueError("`window_random_frame_offset_range` start must be less than end")
+
+    def __call__(self, data: dict[str, Any]) -> dict[str, Any] | None:
+        """Extract frames from a video."""
+
+        chunk_index, extracted_frame_ids, video_fps = None, None, None
+        (
+            captions,
+            multiview_frames,
+            multiview_control,
+            view_indices,
+            view_indices_selection,
+            camera_keys_selection,
+            original_sizes,
+        ) = (
+            [],
+            [],
+            [],
+            [],
+            [],
+            [],
+            [],
+        )
+
+        for camera_name in self.camera_order:
+            video_key = self.camera_video_key_mapping[camera_name]
+            if self.single_caption_camera_name:
+                meta_key = self.camera_caption_key_mapping[self.single_caption_camera_name]
+            else:
+                meta_key = self.camera_caption_key_mapping[camera_name]
+
+            t2w_windows = data[meta_key]["t2w_windows"]
+            if chunk_index is None:
+                chunk_index = random.choice(list(range(len(t2w_windows))))
+            window = t2w_windows[chunk_index]
+
+            # extract caption
+            choices = list(self.caption_probability.keys())
+            weights = list(self.caption_probability.values())
+            caption_style = random.choices(choices, weights=weights)[0]
+            caption = ""
+            if self.single_caption_camera_name:
+                if camera_name == self.single_caption_camera_name:
+                    caption = window[caption_style]
+            else:
+                caption = window[caption_style]
+
+            assert isinstance(caption, str), f"Caption is not a string: {caption}"
+            if self.add_view_prefix_to_caption:
+                caption = f"{self.camera_prefix_mapping[camera_name]} {caption}"
+            captions.append(caption)
+
+            # extract frames
+            random_offset = 0
+            if self.window_random_frame_offset_range is not None:
+                random_offset = random.randint(*self.window_random_frame_offset_range)
+            frame_start = window["start_frame"] + random_offset
+            frame_end = frame_start + self.num_frames * self.fps_downsample_factor
+            frame_indices = list(range(frame_start, frame_end, self.fps_downsample_factor))
+            try:
+                frames, original_fps, original_hw = self.extract_frames(
+                    data[video_key], frame_indices, self.resolution_hw
+                )
+            except Exception as e:
+                log.error(f"Error extracting frames for camera {camera_name}: {e}")
+                return None
+            assert len(frames) == self.num_frames, f"Expected {self.num_frames} frames, got {len(frames)}"
+            multiview_frames.append(frames)
+
+            # check consistency between videos
+            if extracted_frame_ids is None:
+                extracted_frame_ids = frame_indices
+            elif frame_indices != extracted_frame_ids:
+                raise ValueError("Extracted frame IDs do not match")
+
+            if video_fps is None:
+                video_fps = original_fps
+            elif video_fps != original_fps:
+                raise ValueError("Video FPS does not match")
+            original_sizes.append(list(original_hw))
+
+            # extract control frames if available
+            if self.camera_control_key_mapping is not None:
+                control_key = self.camera_control_key_mapping[camera_name]
+                try:
+                    control_frames, control_fps, _ = self.extract_frames(
+                        data[control_key], frame_indices, self.resolution_hw
+                    )
+                except Exception as e:
+                    log.error(f"Error extracting control frames for camera {camera_name}: {e}")
+                    return None
+                if len(control_frames) != self.num_frames:
+                    raise ValueError(f"Expected {self.num_frames} frames, got {len(control_frames)}")
+                if control_fps != original_fps:
+                    raise ValueError(f"Control FPS {control_fps} does not match video FPS {original_fps}")
+                multiview_control.append(control_frames)
+
+            view_indices.extend([self.camera_view_mapping[camera_name]] * self.num_frames)
+            view_indices_selection.append(self.camera_view_mapping[camera_name])
+            camera_keys_selection.append(camera_name)
+
+        front_cam_view_idx_sample_position = (
+            torch.tensor(self.camera_order.index(self.single_caption_camera_name), dtype=torch.int64)
+            if self.single_caption_camera_name
+            else None
+        )
+        if self.single_caption_camera_name and not self.add_view_prefix_to_caption:
+            captions = [captions[front_cam_view_idx_sample_position]]
+
+        if video_fps % self.fps_downsample_factor != 0:
+            raise ValueError("Original FPS is not divisible by FPS downsample factor")
+        fps = video_fps / self.fps_downsample_factor
+
+        sample = {
+            "__key__": data["__key__"],
+            "__url__": data["__url__"],
+            "video": rearrange(torch.cat(multiview_frames, dim=0), "t c h w -> c t h w"),
+            "ai_caption": captions,
+            "view_indices": torch.tensor(view_indices, dtype=torch.int64),
+            "fps": torch.tensor(fps, dtype=torch.float64),
+            "chunk_index": torch.tensor(chunk_index, dtype=torch.int64),
+            "frame_indices": torch.tensor(extracted_frame_ids, dtype=torch.int64),
+            "num_video_frames_per_view": torch.tensor(len(extracted_frame_ids), dtype=torch.int64),
+            "view_indices_selection": torch.tensor(view_indices_selection, dtype=torch.int64),
+            "camera_keys_selection": camera_keys_selection,
+            "sample_n_views": torch.tensor(len(camera_keys_selection), dtype=torch.int64),
+            "padding_mask": torch.zeros((1, *self.resolution_hw), dtype=torch.float32),
+            "ref_cam_view_idx_sample_position": torch.tensor(-1, dtype=torch.int64),
+            "front_cam_view_idx_sample_position": front_cam_view_idx_sample_position,
+            "original_hw": torch.tensor(original_sizes, dtype=torch.int64),
+        }
+        if self.camera_control_key_mapping is not None:
+            sample["control_input_hdmap_bbox"] = rearrange(torch.cat(multiview_control, dim=0), "t c h w -> c t h w")
+        return sample
+
+    @staticmethod
+    def extract_frames(
+        video: bytes, frame_indices: list[int], resolution_hw: tuple[int, int]
+    ) -> tuple[torch.Tensor, float, tuple[int, int]]:
+        """Extract frames from a video given start and end frame range."""
+
+        from decord import VideoReader
+
+        video_reader = VideoReader(io.BytesIO(video))
+        fps = video_reader.get_avg_fps()
+        frames = video_reader.get_batch(frame_indices).asnumpy()
+        frames = rearrange(torch.from_numpy(frames), "t h w c -> t c h w")
+        original_h, original_w = frames.shape[-2:]
+        return (
+            Resize(resolution_hw, interpolation=InterpolationMode.BILINEAR, antialias=True)(frames),
+            fps,
+            (original_h, original_w),
+        )
+
+
+def get_multiview_dataset(
+    *,
+    dataset_name: str,
+    is_train: bool,
+    object_store: Literal["gcs", "s3"],
+    dataset_keys: list[str],
+    augmentations: dict[str, Augmentor],
+    dataset_catalog: dict[str, dict[str, list[str]]],
+) -> Dataset:
+    """Get video-text dataset with optional custom augmentation factory.
+
+    Args:
+        is_train: Whether this is for training
+        dataset_name: Name of dataset to use for loading wdinfo files
+        object_store: Object store to use ("gcs" or "s3")
+        dataset_keys: List of keys to use for loading dataset
+        augmentations: Augmentations map to apply to dataset
+        dataset_catalog: Dataset catalog to use for loading dataset
+    """
+
+    dataset_info = get_video_dataset_info(
+        dataset_name,
+        object_store=object_store,
+        dataset_keys=dataset_keys,
+        dataset_catalog=dataset_catalog,
+    )
+
+    if (
+        USE_MEGATRON
+        and parallel_state.is_initialized()
+        and (
+            parallel_state.get_context_parallel_world_size() > 1
+            or parallel_state.get_tensor_model_parallel_world_size() > 1
+        )
+    ):
+        distributor_fn = parallel_sync_multi_aspect_ratio.ShardlistMultiAspectRatioParallelSync
+    else:
+        distributor_fn = ShardlistBasic
+
+    video_data_config = DatasetConfig(
+        keys=[],  # keys are defined per dataset
+        buffer_size=1,
+        streaming_download=True,
+        dataset_info=dataset_info,
+        distributor=distributor_fn(
+            shuffle=is_train,
+            split_by_node=True,
+            split_by_worker=True,
+            resume_flag=True,
+            verbose=False,
+            is_infinite_loader=is_train,
+        ),
+        decoders=[
+            video_naive_bytes(),
+            json_decoder,
+        ],
+        augmentation=augmentations,
+        remove_extension_from_keys=True,
+    )
+
+    return Dataset(
+        config=video_data_config,
+        handler=wds.warn_and_continue,
+        decoder_handler=wds.warn_and_continue,
+        detshuffle=False,
+    )
+
+
+def collate_fn(batch: list[dict[str, Any]]) -> dict[str, Any]:
+    merged = dict()
+    is_tensor = dict()
+    for row in batch:
+        for key, value in row.items():
+            if key not in merged:
+                merged[key] = []
+            if isinstance(value, torch.Tensor):
+                is_tensor[key] = True
+            merged[key].append(value)
+    for key, value in merged.items():
+        if is_tensor.get(key, False):
+            merged[key] = torch.stack(value, dim=0)
+    return merged
+
+
+@attrs.define(slots=False)
+class AugmentationConfig:
+    """Configuration for video augmentation."""
+
+    resolution_hw: tuple[int, int] = (1080, 1920)
+    fps_downsample_factor: int = 1
+    num_video_frames: int = 93
+    caption_probability: dict[str, float] = {
+        "qwen2p5_7b_caption": 0.7,
+        "qwen2p5_7b_caption_medium": 0.2,
+        "qwen2p5_7b_caption_short": 0.1,
+    }
+    camera_keys: tuple[CameraKeyType, ...] = DEFAULT_CAMERAS
+    camera_view_mapping: dict[CameraKeyType, int] = DEFAULT_CAMERA_VIEW_MAPPING
+    camera_caption_key_mapping: dict[CameraKeyType, str] = DEFAULT_CAPTION_KEY_MAPPING
+    camera_video_key_mapping: dict[CameraKeyType, str] = DEFAULT_VIDEO_KEY_MAPPING
+    camera_control_key_mapping: Optional[dict[CameraKeyType, str]] = None
+    position_to_camera_mapping: Optional[dict[int, CameraKeyType]] = None
+    add_view_prefix_to_caption: bool = False
+    camera_prefix_mapping: Optional[dict[CameraKeyType, str]] = DEFAULT_CAPTION_PREFIXES
+    single_caption_camera_name: Optional[CameraKeyType] = None
+    window_random_frame_offset_range: Optional[tuple[int, int]] = None
+
+    def __attrs_post_init__(self) -> None:
+        """Post initialization checks for camera keys consistency."""
+
+        for camera_key in self.camera_keys:
+            for attr_name in [
+                "camera_view_mapping",
+                "camera_caption_key_mapping",
+                "camera_video_key_mapping",
+                "camera_control_key_mapping",
+                "camera_prefix_mapping",
+            ]:
+                attr = getattr(self, attr_name)
+                if attr is not None:
+                    if camera_key not in attr:
+                        raise ValueError(f"Camera key {camera_key} not found in `{attr_name}` mapping!")
+        if self.single_caption_camera_name is not None:
+            if self.single_caption_camera_name not in self.camera_keys:
+                raise ValueError(
+                    f"Single caption camera key {self.single_caption_camera_name} not found in camera keys!"
+                )
+
+
+def make_augmentations(augmentation_config: AugmentationConfig) -> tuple[dict[str, Augmentor], list[str]]:
+    """Make augmentations for multiview video dataset."""
+
+    augmentations = dict()
+    if augmentation_config.position_to_camera_mapping is not None:
+        augmentations["unpack_metas"] = UnpackMetas(
+            position_to_camera_mapping=augmentation_config.position_to_camera_mapping
+        )
+
+    augmentations["extract_frames_and_captions"] = ExtractFramesAndCaptions(
+        camera_order=augmentation_config.camera_keys,
+        num_frames=augmentation_config.num_video_frames,
+        resolution_hw=augmentation_config.resolution_hw,
+        fps_downsample_factor=augmentation_config.fps_downsample_factor,
+        caption_probability=augmentation_config.caption_probability,
+        camera_view_mapping=augmentation_config.camera_view_mapping,
+        camera_caption_key_mapping=augmentation_config.camera_caption_key_mapping,
+        camera_video_key_mapping=augmentation_config.camera_video_key_mapping,
+        camera_control_key_mapping=augmentation_config.camera_control_key_mapping,
+        add_view_prefix_to_caption=augmentation_config.add_view_prefix_to_caption,
+        camera_prefix_mapping=augmentation_config.camera_prefix_mapping,
+        single_caption_camera_name=augmentation_config.single_caption_camera_name,
+        window_random_frame_offset_range=augmentation_config.window_random_frame_offset_range,
+    )
+
+    # define dataset keys to load
+    dataset_keys = list(augmentation_config.camera_video_key_mapping.values())
+    if augmentation_config.position_to_camera_mapping is not None:
+        dataset_keys.append("metas")
+    else:
+        dataset_keys.extend(augmentation_config.camera_caption_key_mapping.values())
+    if augmentation_config.camera_control_key_mapping is not None:
+        dataset_keys.extend(augmentation_config.camera_control_key_mapping.values())
+
+    return augmentations, dataset_keys
+
+
+def get_multiview_video_loader(
+    *,
+    dataset_name: str,
+    is_train: bool,
+    object_store: Literal["gcs", "s3"] = "s3",
+    augmentation_config: AugmentationConfig = AugmentationConfig(),
+    batch_size: int = 1,
+    num_workers: int = 4,
+    prefetch_factor: int | None = 1,
+    **kwargs: Any,
+):
+    """Get video loader for alpamayo multiview dataset
+    pass kwargs to tolerate `dataloaders` from inheritance
+    """
+
+    # make augmentations
+    augmentations, dataset_keys = make_augmentations(augmentation_config)
+
+    # get dataloader
+    return get_cached_replay_dataloader(
+        dataset=get_multiview_dataset(
+            is_train=is_train,
+            object_store=object_store,
+            dataset_name=dataset_name,
+            dataset_keys=dataset_keys,
+            dataset_catalog=DEFAULT_CATALOG,
+            augmentations=augmentations,
+        ),
+        num_workers=num_workers,
+        batch_size=batch_size,
+        sampler=None,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None,
+        persistent_workers=num_workers > 0,
+        pin_memory=False,
+        collate_fn=collate_fn,
+        cache_replay_name="video_dataloader",
+    )

cosmos_predict2/_src/predict2_multiview/datasets/nymeria_pairs.py

@@ -0,0 +1,1074 @@
+# 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.
+
+"""Nymeria 2-actor paired clip dataset for joint multi-view video generation.
+
+Driven by ``manifest/train_ready.csv`` (only pairs whose BOTH clips finished person+hands warping). Each row
+references two per-actor ``clip.npz`` files (``actor1_npz`` / ``actor2_npz``); each npz holds raw arrays:
+  * ``target_rgb``        (T, H, W, 3) uint8 -- ground-truth video to generate
+  * ``pose``              (T, H, W, 3) uint8 -- projected 2D pose RGB (driving pose)
+  * ``warped_cond``       (T, H, W, 3) uint8 -- warped past-frame conditioning (pixel RGB)
+  * ``visibility_packed`` (packbits) + ``vis_shape`` (T, H, W) -- warp-validity mask (bool -> 1ch)
+  * ``keypoints`` (T, 2, 23, 3), ``src_frame`` (T,)  -- unused here
+
+(Falls back to per-clip mp4 files if a CSV row has no ``*_npz`` columns, e.g. the old ``clip_pairs.csv``.)
+
+Actor 1 / Actor 2 are emitted as the two "views" (V=2): every tensor is laid out ``(C, V*T, H, W)`` with the
+two actors concatenated along the temporal axis. ``actor{1,2}_n_extra`` count untracked 3rd parties; set
+``clean_only=True`` to keep only pairs with no 3rd party (n_extra==0 for both actors).
+"""
+
+import csv
+import os
+import random
+import re
+from typing import Any, Optional
+
+import attrs
+import numpy as np
+import torch
+from einops import rearrange
+from hydra.core.config_store import ConfigStore
+from torchvision.transforms import InterpolationMode, Resize
+
+from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L
+from cosmos_predict2._src.imaginaire.utils import log
+from cosmos_predict2._src.predict2.datasets.cached_replay_dataloader import get_cached_replay_dataloader
+from cosmos_predict2._src.predict2_multiview.datasets.multiview import collate_fn
+
+# npz array keys
+_NPZ_KEYS = {"video": "target_rgb", "control_input_pose": "pose", "control_input_warped": "warped_cond"}
+# mp4 fallback files
+_VIDEO_FILES = {
+    "video": "target.mp4",
+    "control_input_pose": "pose.mp4",
+    "control_input_warped": "warped_cond.mp4",
+    "control_input_visibility": "visibility.mp4",
+}
+
+
+def _resize_thwc_uint8(frames_thwc: np.ndarray, num_frames: int, resolution_hw: tuple[int, int]) -> torch.Tensor:
+    """(T, H, W, C) uint8 -> resized (C, T, H', W') uint8, padded/truncated to num_frames."""
+    t = torch.from_numpy(np.ascontiguousarray(frames_thwc))  # (T,H,W,C)
+    t = rearrange(t, "t h w c -> t c h w")
+    n = t.shape[0]
+    if n < num_frames:
+        t = torch.cat([t, t[-1:].repeat(num_frames - n, 1, 1, 1)], dim=0)
+    elif n > num_frames:
+        t = t[:num_frames]
+    t = Resize(resolution_hw, interpolation=InterpolationMode.BILINEAR, antialias=True)(t.float()).to(torch.uint8)
+    return rearrange(t, "t c h w -> c t h w").contiguous()
+
+
+def _read_video_frames(path: str, num_frames: int, resolution_hw: tuple[int, int]) -> torch.Tensor:
+    """mp4 fallback: first num_frames frames -> (C, T, H', W') uint8."""
+    from decord import VideoReader
+
+    vr = VideoReader(path)
+    n = min(num_frames, len(vr))
+    frames = vr.get_batch(list(range(n))).asnumpy()  # (n,H,W,C)
+    return _resize_thwc_uint8(frames, num_frames, resolution_hw)
+
+
+@attrs.define(slots=False)
+class NymeriaPairsConfig:
+    root: str = "/data2/nymeria_processed_longer"
+    manifest_csv: str = "manifest/train_ready.csv"
+    resolution_hw: tuple[int, int] = (480, 480)
+    num_video_frames: int = 77  # target frames per actor (77 for longer/train_ready)
+    fps: float = 10.0
+    caption: str = "a person performing an everyday activity"  # placeholder; text encoder is frozen
+    clean_only: bool = False  # keep only pairs with no untracked 3rd party (n_extra==0 for both actors)
+    single_view: bool = False  # stage-1 pose pretraining: emit each actor as an independent V=1 sample
+    num_reference_frames: int = 0  # >0: emit `reference_frames` (R clean source frames/view) for appearance
+    #                                conditioning (in-context reference tokens). Loaded from clip.npz
+    #                                `reference_frames` or a sibling refs.npz (see backfill_refs.py).
+    shared_reference: bool = False  # SHARED refs: instead of per-view pools, load ONE greedy-selected set from
+    #                                 refs_shared.npz (identical in both clip dirs; V*R frames + w2c) and hand
+    #                                 view-slot v its contiguous R-slice (top-gain first). Refs carry no view
+    #                                 identity (net drops the ref view-embedding); both views attend all via
+    #                                 cross-view self-attn. Pairs with net.shared_reference=True.
+    emit_camera_poses: bool = False  # emit per-view camera_w2c/K for Plücker ray conditioning (campose.npz)
+    emit_depth: bool = False  # emit `control_input_depth` = composite depth RGB (mesh_cond.npz `depth_comp`:
+    #                           warped scene depth + human mesh depth) for the net's VAE depth conditioning
+    person_pose: bool = False  # use IDENTITY-colored pose (npz `pose_person`) instead of role-colored `pose`
+    emit_reference_pose: bool = False  # emit `reference_pose` = per-person skeleton of each shared ref (refs_shared.npz)
+    role_mix: bool = False  # actor-observer 2-view: randomly swap view0<->view1 (+captions) 50% so the
+    #                         view embedding can't memorize actor/observer roles
+    # train/val split — held out by SESSION (=(actor1_seq, actor2_seq)) so adjacent tiled windows of the same
+    # recording never leak across the split. "all" uses everything (no split).
+    split: str = "all"  # "train" | "val" | "all"
+    val_num_sessions: int = 4  # number of sessions held out for validation (~10% of 36 sessions)
+    split_seed: int = 1234  # deterministic session assignment
+    val_max_pairs: int = 40  # cap the val set size for fast periodic validation (0 = uncapped)
+    # smoke/offline mode: emit zero text embeddings so the (S3-only) text encoder can be disabled.
+    # text_tokens_per_view MUST be 512 (MultiViewCrossAttention derives n_views as context_len // 512).
+    dummy_text_embeddings: bool = False
+    text_embed_dim: int = 1024
+    text_tokens_per_view: int = 512
+
+
+class NymeriaPairsDataset(torch.utils.data.Dataset):
+    def __init__(self, config: NymeriaPairsConfig):
+        self.config = config
+        self.root = config.root
+        self.clips_dir = os.path.join(self.root, "clips")
+        self.resolution_hw = tuple(config.resolution_hw)
+        self.num_frames = int(config.num_video_frames)
+        manifest_path = os.path.join(self.root, config.manifest_csv)
+
+        # 1) collect ALL usable candidate rows (with session + clean flag). The clean filter is applied AFTER
+        #    the session split so val_sessions are computed from the same full universe for train and val.
+        candidates: list[dict[str, str]] = []
+        with open(manifest_path, newline="") as f:
+            for row in csv.DictReader(f):
+                session = f"{row.get('actor1_seq', '')}__{row.get('actor2_seq', '')}"
+                clean = int(row.get("actor1_n_extra", 0)) == 0 and int(row.get("actor2_n_extra", 0)) == 0
+                a1_npz, a2_npz = row.get("actor1_npz", ""), row.get("actor2_npz", "")
+                if a1_npz and a2_npz:  # npz path (train_ready.csv)
+                    if os.path.isfile(a1_npz) and os.path.isfile(a2_npz):
+                        candidates.append({"pair_id": row["pair_id"], "a1": a1_npz, "a2": a2_npz, "mode": "npz", "session": session, "clean": clean})
+                else:  # mp4 fallback (clip_pairs.csv)
+                    a1 = os.path.join(self.clips_dir, row["actor1_clip"])
+                    a2 = os.path.join(self.clips_dir, row["actor2_clip"])
+                    if self._mp4_complete(a1) and self._mp4_complete(a2):
+                        candidates.append({"pair_id": row["pair_id"], "a1": a1, "a2": a2, "mode": "mp4", "session": session, "clean": clean})
+
+        # 2) deterministic session-level holdout from the FULL session universe (no temporal-window leakage,
+        #    and train/val are guaranteed disjoint by session regardless of clean_only)
+        import random as _random
+
+        sessions = sorted({c["session"] for c in candidates})
+        rng = _random.Random(config.split_seed)
+        rng.shuffle(sessions)
+        val_sessions = set(sessions[: config.val_num_sessions])
+        if config.split == "train":
+            candidates = [c for c in candidates if c["session"] not in val_sessions]
+        elif config.split == "val":
+            candidates = [c for c in candidates if c["session"] in val_sessions]
+
+        # 3) clean filter (after split)
+        n_extra_dropped = 0
+        if config.clean_only:
+            before = len(candidates)
+            candidates = [c for c in candidates if c["clean"]]
+            n_extra_dropped = before - len(candidates)
+
+        # 3b) single-view (stage-1): expand each pair into two independent V=1 samples (one per actor)
+        if config.single_view:
+            singles = []
+            for c in candidates:
+                for k in ("a1", "a2"):
+                    singles.append({"pair_id": f"{c['pair_id']}__{k}", "a1": c[k], "mode": c["mode"], "session": c["session"]})
+            candidates = singles
+
+        # 4) val cap (after expansion)
+        if config.split == "val" and config.val_max_pairs and len(candidates) > config.val_max_pairs:
+            candidates.sort(key=lambda c: c["pair_id"])
+            step = len(candidates) / config.val_max_pairs
+            candidates = [candidates[int(i * step)] for i in range(config.val_max_pairs)]
+
+        self.pairs = candidates
+        log.info(
+            f"NymeriaPairsDataset[{config.split}]: {len(self.pairs)} pairs from {manifest_path} "
+            f"(clean_only={config.clean_only}, {len(sessions)} sessions, {len(val_sessions)} held out for val, "
+            f"dropped {n_extra_dropped} for 3rd-party)"
+        )
+        if len(self.pairs) == 0:
+            raise RuntimeError(f"No usable clip pairs for split={config.split} in {manifest_path}")
+
+    @staticmethod
+    def _mp4_complete(clip_dir: str) -> bool:
+        return os.path.isdir(clip_dir) and all(
+            os.path.isfile(os.path.join(clip_dir, fn)) for fn in _VIDEO_FILES.values()
+        )
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def _load_actor(self, ref: str, mode: str, slot: int = 0) -> dict[str, torch.Tensor]:
+        out: dict[str, torch.Tensor] = {}
+        if mode == "npz":
+            d = np.load(ref, allow_pickle=True)
+            for key, npz_key in _NPZ_KEYS.items():
+                out[key] = _resize_thwc_uint8(d[npz_key], self.num_frames, self.resolution_hw)
+            if getattr(self.config, "person_pose", False) and "pose_person" in d.files:  # identity-colored pose
+                out["control_input_pose"] = _resize_thwc_uint8(d["pose_person"], self.num_frames, self.resolution_hw)
+            if getattr(self.config, "emit_depth", False):  # composite depth RGB (mesh_cond.npz depth_comp)
+                mc = os.path.join(os.path.dirname(ref), "mesh_cond.npz")
+                if os.path.isfile(mc):
+                    dc = np.load(mc)["depth_comp"]
+                else:  # clips still lacking mesh_cond (Stage-B in progress) -> zero depth (shape-consistent for concat)
+                    dc = np.zeros((self.num_frames, self.resolution_hw[0], self.resolution_hw[1], 3), np.uint8)
+                out["control_input_depth"] = _resize_thwc_uint8(dc, self.num_frames, self.resolution_hw)
+            # visibility: unpack bits -> (T,H,W) -> (T,H,W,1)
+            vt, vh, vw = (int(x) for x in d["vis_shape"])
+            vis = np.unpackbits(d["visibility_packed"])[: vt * vh * vw].reshape(vt, vh, vw)
+            vis = (vis.astype(np.uint8) * 255)[..., None]  # (T,H,W,1) in {0,255}
+            out["control_input_visibility"] = _resize_thwc_uint8(vis, self.num_frames, self.resolution_hw)
+            # reference frames (clean source frames) for appearance conditioning -> (3, R, H, W) uint8
+            R = int(getattr(self.config, "num_reference_frames", 0))
+            if R > 0 and getattr(self.config, "shared_reference", False):
+                # SHARED refs: one greedy set (V*R frames, top-gain first) identical in both clip dirs; hand
+                # THIS view-slot its contiguous R-slice. Refs carry no view identity (net drops ref view-emb).
+                sib = os.path.join(os.path.dirname(ref), "refs_shared.npz")
+                if not os.path.isfile(sib):
+                    raise KeyError(f"no refs_shared for {ref} (run backfill_shared_refs.py)")
+                z = np.load(sib)
+                allf = z["reference_frames"]  # (V*R, h, h, 3)
+                lo = slot * R
+                if lo + R > allf.shape[0]:
+                    raise KeyError(f"refs_shared has {allf.shape[0]} < {lo+R} frames for slot {slot} R={R}")
+                out["reference_frames"] = _resize_thwc_uint8(allf[lo : lo + R], R, self.resolution_hw)
+                if getattr(self.config, "emit_camera_poses", False):
+                    rw2c = z["reference_w2c"].astype(np.float32)[lo : lo + R]  # (R,4,4) source-frame world->cam
+                    out["cam_ref_w2c"] = torch.from_numpy(rw2c)
+                if getattr(self.config, "emit_reference_pose", False) and "reference_pose" in z.files:
+                    # per-person skeleton render of each shared ref (this view-slot's R-slice), same layout as refs
+                    out["reference_pose"] = _resize_thwc_uint8(z["reference_pose"][lo : lo + R], R, self.resolution_hw)
+            elif R > 0:
+                if "reference_frames" in d.files:  # clip.npz native (new warps)
+                    ref_frames = d["reference_frames"]
+                else:  # sibling refs.npz (backfilled)
+                    sib = os.path.join(os.path.dirname(ref), "refs.npz")
+                    if not os.path.isfile(sib):
+                        raise KeyError(f"no reference_frames for {ref} (run backfill_refs.py)")
+                    ref_frames = np.load(sib)["reference_frames"]
+                # stored refs may hold != R frames; pick R (coverage-spaced if more, pad-last if fewer). The
+                # SAME `sel` is reused for the reference camera poses so frames <-> poses stay aligned.
+                k_stored = ref_frames.shape[0]
+                if k_stored >= R:
+                    sel = np.linspace(0, k_stored - 1, R).round().astype(int)
+                else:
+                    sel = np.concatenate([np.arange(k_stored), np.full(R - k_stored, k_stored - 1, dtype=int)])
+                out["reference_frames"] = _resize_thwc_uint8(ref_frames[sel], R, self.resolution_hw)
+                # posed-reference Plücker: past camera poses of the reference frames (aligned via `sel`)
+                if getattr(self.config, "emit_camera_poses", False):
+                    rp = os.path.join(os.path.dirname(ref), "refpose.npz")
+                    if not os.path.isfile(rp):
+                        raise KeyError(f"no refpose for {ref} (run backfill_refpose.py)")
+                    rw2c = np.load(rp)["reference_w2c"].astype(np.float32)[sel]  # (R,4,4)
+                    out["cam_ref_w2c"] = torch.from_numpy(rw2c)
+            # camera poses (target-frame w2c + K, shared SLAM world) for Plücker ray conditioning
+            if getattr(self.config, "emit_camera_poses", False):
+                cp = os.path.join(os.path.dirname(ref), "campose.npz")
+                if not os.path.isfile(cp):
+                    raise KeyError(f"no campose for {ref} (run backfill_campose.py)")
+                cpz = np.load(cp)
+                w2c = cpz["w2c"].astype(np.float32)  # (T,4,4) target-frame world->cam
+                n = w2c.shape[0]
+                if n < self.num_frames:  # pad by repeating last frame's pose
+                    w2c = np.concatenate([w2c, np.repeat(w2c[-1:], self.num_frames - n, 0)], 0)
+                elif n > self.num_frames:
+                    w2c = w2c[: self.num_frames]
+                out["cam_w2c"] = torch.from_numpy(w2c)  # (T,4,4)
+                out["cam_K"] = torch.from_numpy(cpz["K"].astype(np.float32))  # (3,3)
+                out["cam_src_res"] = torch.tensor(int(cpz["src_res"]), dtype=torch.int64)
+        else:
+            for key, fn in _VIDEO_FILES.items():
+                out[key] = _read_video_frames(os.path.join(ref, fn), self.num_frames, self.resolution_hw)
+        return out
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        # robust to corrupt / mid-write clip.npz (warping may still be running): skip + retry another sample
+        for _ in range(16):
+            try:
+                return self._build_sample(idx)
+            except Exception as e:
+                log.warning(f"NymeriaPairsDataset: bad sample {self.pairs[idx]['pair_id']} ({type(e).__name__}); retrying")
+                idx = random.randrange(len(self.pairs))
+        raise RuntimeError("NymeriaPairsDataset: too many unreadable samples in a row")
+
+    def _build_sample(self, idx: int) -> dict[str, Any]:
+        pair = self.pairs[idx]
+        # single-view: one actor (V=1); paired: both actors stacked on the temporal axis (V=2)
+        actors = [self._load_actor(pair["a1"], pair["mode"], slot=0)]
+        if not self.config.single_view:
+            actors.append(self._load_actor(pair["a2"], pair["mode"], slot=1))
+        n_views = len(actors)
+
+        sample: dict[str, Any] = {}
+        for key in actors[0]:  # concat actors along temporal axis -> (C, V*T, H, W)
+            if key.startswith("cam_"):  # per-view metadata, stacked on the view dim below
+                continue
+            sample[key] = torch.cat([a[key] for a in actors], dim=1).contiguous()
+        if "cam_w2c" in actors[0]:
+            sample["camera_w2c"] = torch.stack([a["cam_w2c"] for a in actors], dim=0).contiguous()
+            sample["camera_K"] = torch.stack([a["cam_K"] for a in actors], dim=0).contiguous()
+            sample["camera_src_res"] = torch.stack([a["cam_src_res"] for a in actors], dim=0).contiguous()
+        if "cam_ref_w2c" in actors[0]:  # reference-frame past poses: (V, R, 4, 4)
+            sample["reference_cam_w2c"] = torch.stack([a["cam_ref_w2c"] for a in actors], dim=0).contiguous()
+
+        T = self.num_frames
+        view_indices = [v for v in range(n_views) for _ in range(T)]
+        sample.update(
+            {
+                "__key__": pair["pair_id"],
+                "__url__": pair["pair_id"],
+                "ai_caption": [self.config.caption for _ in range(n_views)],
+                "view_indices": torch.tensor(view_indices, dtype=torch.int64),
+                "fps": torch.tensor(self.config.fps, dtype=torch.float64),
+                "num_video_frames_per_view": torch.tensor(T, dtype=torch.int64),
+                "view_indices_selection": torch.tensor(list(range(n_views)), dtype=torch.int64),
+                "camera_keys_selection": [f"actor{v + 1}" for v in range(n_views)],
+                "sample_n_views": torch.tensor(n_views, dtype=torch.int64),
+                "padding_mask": torch.zeros((1, *self.resolution_hw), dtype=torch.float32),
+                "ref_cam_view_idx_sample_position": torch.tensor(-1, dtype=torch.int64),
+                "front_cam_view_idx_sample_position": torch.tensor(0, dtype=torch.int64),
+            }
+        )
+        if self.config.dummy_text_embeddings:
+            n_tok = n_views * self.config.text_tokens_per_view
+            sample["t5_text_embeddings"] = torch.zeros(n_tok, self.config.text_embed_dim, dtype=torch.float32)
+            sample["t5_text_mask"] = torch.ones(n_tok, dtype=torch.float32)
+        return sample
+
+
+class NymeriaActorObserverDataset(NymeriaPairsDataset):
+    """2-view actor-observer dataset (e.g. /data2/nymeria_processed_single).
+
+    Each row of the manifest = one actor(head) clip + one observer clip + a `text` caption of the actor's
+    action. Reads a pre-split manifest (train_split.csv / val_split.csv) directly. Emits V=2 with:
+      * role mixing (``role_mix``): 50% of the time view0/view1 (and their captions) are swapped, so the
+        view embedding cannot encode "view0=actor / view1=observer".
+      * per-view captions for per-view text cross-attention:
+          actor view    -> the CSV `text`
+          observer view -> "C is observing the partner. The partner <text with 'C' removed>"
+    Real captions are emitted in ``ai_caption`` (no dummy embeddings): the model's online multiview text
+    encoder computes one 512-token embedding per view.
+    """
+
+    def __init__(self, config: NymeriaPairsConfig):
+        self.config = config
+        self.root = config.root
+        self.clips_dir = os.path.join(self.root, "clips")
+        self.resolution_hw = tuple(config.resolution_hw)
+        self.num_frames = int(config.num_video_frames)
+        manifest_path = os.path.join(self.root, config.manifest_csv)
+        self.pairs: list[dict[str, str]] = []
+        with open(manifest_path, newline="") as f:
+            for row in csv.DictReader(f):
+                actor_npz = os.path.join(self.clips_dir, row["actor1_clip"], "clip.npz")  # head/actor
+                obs_npz = os.path.join(self.clips_dir, row["actor2_clip"], "clip.npz")  # observer
+                if os.path.isfile(actor_npz) and os.path.isfile(obs_npz):
+                    self.pairs.append(
+                        {"pair_id": row["pair_id"], "actor": actor_npz, "observer": obs_npz, "text": row.get("text", "")}
+                    )
+        log.info(f"NymeriaActorObserverDataset: {len(self.pairs)} pairs from {manifest_path} (role_mix={config.role_mix})")
+        if len(self.pairs) == 0:
+            raise RuntimeError(f"No usable actor-observer pairs in {manifest_path}")
+
+    @staticmethod
+    def _observer_caption(actor_text: str) -> str:
+        # frame the actor action from the observer's view: subsequent "C" subjects become "the partner"
+        # (kept as subjects so multi-sentence actions stay grammatical); the first is dropped since the
+        # template already prepends "The partner".
+        action = re.sub(r"\bC\b", "the partner", actor_text).strip()
+        action = re.sub(r"^the partner\s+", "", action)  # drop the leading one (prepended below)
+        caption = f"C is observing the partner. The partner {action}"
+        caption = re.sub(r"([.!?]\s+)the partner", r"\1The partner", caption)  # capitalize after sentence end
+        return caption
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        # robust to corrupt / mid-write clip.npz (warping may still be running): skip + retry another sample
+        for _ in range(16):
+            try:
+                return self._build_sample(idx)
+            except Exception as e:
+                log.warning(
+                    f"NymeriaActorObserverDataset: bad sample {self.pairs[idx]['pair_id']} "
+                    f"({type(e).__name__}: {str(e)[:80]}); retrying another"
+                )
+                idx = random.randrange(len(self.pairs))
+        raise RuntimeError("NymeriaActorObserverDataset: too many unreadable samples in a row")
+
+    def _get_roles(self, p: dict) -> list:
+        """(clip_npz, caption) per view. Actor-observer: view0=actor(real text), view1=observer(transformed)."""
+        return [(p["actor"], p["text"]), (p["observer"], self._observer_caption(p["text"]))]
+
+    def _build_sample(self, idx: int) -> dict[str, Any]:
+        p = self.pairs[idx]
+        roles = self._get_roles(p)
+        if self.config.role_mix and random.random() < 0.5:
+            roles = roles[::-1]  # swap the two view slots (caption follows the clip)
+
+        actors = [self._load_actor(npz, "npz", slot=i) for i, (npz, _) in enumerate(roles)]
+        captions = [cap for _, cap in roles]
+        n_views = 2
+
+        sample: dict[str, Any] = {}
+        for key in actors[0]:  # concat the two views along the temporal axis -> (C, V*T, H, W)
+            if key.startswith("cam_"):  # per-view metadata, stacked on the view dim below (not temporal-concat)
+                continue
+            sample[key] = torch.cat([a[key] for a in actors], dim=1).contiguous()
+        if "cam_w2c" in actors[0]:  # camera poses: (V, T, 4, 4) / (V, 3, 3) / (V,)
+            sample["camera_w2c"] = torch.stack([a["cam_w2c"] for a in actors], dim=0).contiguous()
+            sample["camera_K"] = torch.stack([a["cam_K"] for a in actors], dim=0).contiguous()
+            sample["camera_src_res"] = torch.stack([a["cam_src_res"] for a in actors], dim=0).contiguous()
+        if "cam_ref_w2c" in actors[0]:  # reference-frame past poses: (V, R, 4, 4)
+            sample["reference_cam_w2c"] = torch.stack([a["cam_ref_w2c"] for a in actors], dim=0).contiguous()
+
+        T = self.num_frames
+        sample.update(
+            {
+                "__key__": p["pair_id"],
+                "__url__": p["pair_id"],
+                "ai_caption": captions,  # [view0_caption, view1_caption] -> per-view cross-attention
+                "view_indices": torch.tensor([0] * T + [1] * T, dtype=torch.int64),
+                "fps": torch.tensor(self.config.fps, dtype=torch.float64),
+                "num_video_frames_per_view": torch.tensor(T, dtype=torch.int64),
+                "view_indices_selection": torch.tensor([0, 1], dtype=torch.int64),
+                "camera_keys_selection": ["view0", "view1"],
+                "sample_n_views": torch.tensor(n_views, dtype=torch.int64),
+                "padding_mask": torch.zeros((1, *self.resolution_hw), dtype=torch.float32),
+                "ref_cam_view_idx_sample_position": torch.tensor(-1, dtype=torch.int64),
+                "front_cam_view_idx_sample_position": torch.tensor(0, dtype=torch.int64),
+            }
+        )
+        return sample
+
+
+class NymeriaActorActorDataset(NymeriaActorObserverDataset):
+    """2-view actor-ACTOR dataset (/data2/nymeria_processed_longer): two ego actors generated jointly, each with
+    its OWN real action caption (from captions/<seq>.csv, keyed by clip_id). Same conditioning structure as the
+    actor-observer setup (refs / camera poses / role-mix) but BOTH views are actors (no observer transform).
+    Manifest columns: actor1_clip / actor2_clip (clip dir names)."""
+
+    def __init__(self, config: NymeriaPairsConfig):
+        import glob as _glob
+
+        self.config = config
+        self.root = config.root
+        self.clips_dir = os.path.join(self.root, "clips")
+        self.resolution_hw = tuple(config.resolution_hw)
+        self.num_frames = int(config.num_video_frames)
+        # caption lookup: clip_id -> caption (per-seq CSVs)
+        captions: dict[str, str] = {}
+        for f in _glob.glob(os.path.join(self.root, "captions", "*.csv")):
+            for row in csv.DictReader(open(f)):
+                captions[row["clip_id"]] = row.get("caption", "")
+        manifest_path = os.path.join(self.root, config.manifest_csv)
+        self.pairs = []
+        with open(manifest_path, newline="") as fp:
+            for row in csv.DictReader(fp):
+                a1, a2 = row["actor1_clip"], row["actor2_clip"]
+                # prefer the manifest's absolute npz path (clips may live outside {root}/clips, e.g. vroid groups/)
+                n1 = row.get("actor1_npz") or os.path.join(self.clips_dir, a1, "clip.npz")
+                n2 = row.get("actor2_npz") or os.path.join(self.clips_dir, a2, "clip.npz")
+                c1, c2 = captions.get(a1, ""), captions.get(a2, "")
+                if os.path.isfile(n1) and os.path.isfile(n2) and c1 and c2:
+                    self.pairs.append({"pair_id": row["pair_id"], "actor": n1, "observer": n2, "cap0": c1, "cap1": c2})
+        log.info(f"NymeriaActorActorDataset: {len(self.pairs)} pairs from {manifest_path} (role_mix={config.role_mix})")
+        if not self.pairs:
+            raise RuntimeError(f"No usable actor-actor pairs in {manifest_path}")
+
+    def _get_roles(self, p: dict) -> list:
+        # both views are actors -> each gets its own real caption (no observer transform)
+        return [(p["actor"], p["cap0"]), (p["observer"], p["cap1"])]
+
+
+def get_nymeria_actor_actor_loader(
+    *,
+    root: str = "/data2/nymeria_processed_longer",
+    manifest_csv: str = "manifest/train_split.csv",
+    resolution_hw: tuple[int, int] = (480, 480),
+    num_video_frames: int = 77,
+    fps: float = 10.0,
+    role_mix: bool = True,
+    num_reference_frames: int = 0,
+    shared_reference: bool = False,
+    emit_camera_poses: bool = False,
+    person_pose: bool = False,
+    emit_reference_pose: bool = False,
+    batch_size: int = 1,
+    num_workers: int = 4,
+    prefetch_factor: Optional[int] = 2,
+    is_train: bool = True,
+    **kwargs: Any,
+):
+    dataset = NymeriaActorActorDataset(
+        NymeriaPairsConfig(
+            root=root, manifest_csv=manifest_csv, resolution_hw=tuple(resolution_hw),
+            num_video_frames=num_video_frames, fps=fps, role_mix=role_mix,
+            num_reference_frames=num_reference_frames, shared_reference=shared_reference,
+            emit_camera_poses=emit_camera_poses, person_pose=person_pose,
+            emit_reference_pose=emit_reference_pose,
+        )
+    )
+    sampler = None
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=is_train, drop_last=True)
+    return get_cached_replay_dataloader(
+        webdataset=False, use_cache=False, dataset=dataset, batch_size=batch_size, num_workers=num_workers,
+        sampler=sampler, shuffle=(sampler is None and is_train), drop_last=True,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None, persistent_workers=num_workers > 0,
+        pin_memory=False, collate_fn=collate_fn, cache_replay_name="nymeria_actor_actor_dataloader",
+    )
+
+
+class NymeriaSingleViewDataset(NymeriaPairsDataset):
+    """STAGE-1 single-view: pool individual EGO clips across roots (real single actor, longer both egos,
+    synthetic both egos) as independent V=1 samples, each with its OWN real action caption. Conditions =
+    pose + warping + plucker(self-frame0) + own refs (shared_reference=False -> per-clip refs.npz). Motion-
+    focused pretraining before the 2-view stages. `sources` = list of dicts {root, manifest_csv, ego_cols,
+    caption_kind ("text_col"|"csv_glob")}."""
+
+    def __init__(self, config: NymeriaPairsConfig, sources: list):
+        import glob as _glob
+        self.config = config
+        self.resolution_hw = tuple(config.resolution_hw)
+        self.num_frames = int(config.num_video_frames)
+        self.pairs = []
+        for src in sources:
+            root = src["root"]
+            caps = {}
+            if src["caption_kind"] == "csv_glob":
+                for f in _glob.glob(os.path.join(root, "captions", "*.csv")):
+                    for r in csv.DictReader(open(f)):
+                        caps[r["clip_id"]] = r.get("caption", "")
+            with open(os.path.join(root, src["manifest_csv"]), newline="") as fp:
+                for row in csv.DictReader(fp):
+                    for col in src["ego_cols"]:
+                        clip = row.get(col, "")
+                        if not clip:
+                            continue
+                        npz = os.path.join(root, "clips", clip, "clip.npz")
+                        cap = row.get("text", "") if src["caption_kind"] == "text_col" else caps.get(clip, "")
+                        if cap and os.path.isfile(npz):
+                            self.pairs.append({"a1": npz, "mode": "npz", "pair_id": clip, "caption": cap,
+                                               "session": row.get("actor1_seq", clip)})
+        log.info(f"NymeriaSingleViewDataset: {len(self.pairs)} ego clips from {len(sources)} sources")
+        if not self.pairs:
+            raise RuntimeError("NymeriaSingleViewDataset: no ego clips found")
+
+    def _build_sample(self, idx: int) -> dict[str, Any]:
+        sample = super()._build_sample(idx)  # V=1 (config.single_view=True); placeholder caption
+        sample["ai_caption"] = [self.pairs[idx]["caption"]]  # real per-clip action caption
+        return sample
+
+
+_STAGE1_SOURCES = [
+    dict(root="/data2/nymeria_processed_single", manifest_csv="manifest/train_split.csv",
+         ego_cols=["actor1_clip"], caption_kind="text_col"),   # actor(head)=ego; observer(exo) excluded
+    dict(root="/data2/nymeria_processed_longer", manifest_csv="manifest/train_split.csv",
+         ego_cols=["actor1_clip", "actor2_clip"], caption_kind="csv_glob"),
+    dict(root="/data3/synthetic_processed_multi", manifest_csv="manifest/train_split.csv",
+         ego_cols=["actor1_clip", "actor2_clip"], caption_kind="csv_glob"),
+]
+
+
+def get_nymeria_single_view_loader(
+    *, sources=None, val: bool = False, resolution_hw: tuple[int, int] = (480, 480), num_video_frames: int = 77,
+    fps: float = 10.0, num_reference_frames: int = 4, emit_camera_poses: bool = True,
+    batch_size: int = 1, num_workers: int = 4, prefetch_factor: Optional[int] = 2, is_train: bool = True,
+    **kwargs: Any,
+):
+    srcs = sources if sources is not None else _STAGE1_SOURCES
+    if val:  # swap to each source's val split
+        srcs = [dict(s, manifest_csv=s["manifest_csv"].replace("train_split", "val_split")) for s in srcs]
+    cfg = NymeriaPairsConfig(
+        resolution_hw=tuple(resolution_hw), num_video_frames=num_video_frames, fps=fps, single_view=True,
+        num_reference_frames=num_reference_frames, shared_reference=False, emit_camera_poses=emit_camera_poses,
+    )
+    dataset = NymeriaSingleViewDataset(cfg, srcs)
+    sampler = None
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=is_train, drop_last=True)
+    return get_cached_replay_dataloader(
+        webdataset=False, use_cache=False, dataset=dataset, batch_size=batch_size, num_workers=num_workers,
+        sampler=sampler, shuffle=(sampler is None and is_train), drop_last=True,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None, persistent_workers=num_workers > 0,
+        pin_memory=False, collate_fn=collate_fn, cache_replay_name="nymeria_single_view_dataloader",
+    )
+
+
+def get_nymeria_actor_actor_combined_loader(
+    *,
+    roots: tuple = ("/data2/nymeria_processed_longer", "/data3/synthetic_processed_multi"),
+    manifest_csv: str = "manifest/train_split.csv",
+    resolution_hw: tuple[int, int] = (480, 480),
+    num_video_frames: int = 77,
+    fps: float = 10.0,
+    role_mix: bool = True,
+    num_reference_frames: int = 0,
+    shared_reference: bool = False,
+    emit_camera_poses: bool = False,
+    emit_depth: bool = False,
+    person_pose: bool = False,
+    batch_size: int = 1,
+    num_workers: int = 4,
+    prefetch_factor: Optional[int] = 2,
+    is_train: bool = True,
+    **kwargs: Any,
+):
+    """ConcatDataset over MULTIPLE actor-actor roots (e.g. real longer + synthetic), all sharing the same
+    conditioning config. Each root is a NymeriaActorActorDataset; samples are pooled uniformly (no reweighting)."""
+    parts = [
+        NymeriaActorActorDataset(
+            NymeriaPairsConfig(
+                root=r, manifest_csv=manifest_csv, resolution_hw=tuple(resolution_hw),
+                num_video_frames=num_video_frames, fps=fps, role_mix=role_mix,
+                num_reference_frames=num_reference_frames, shared_reference=shared_reference,
+                emit_camera_poses=emit_camera_poses, emit_depth=emit_depth, person_pose=person_pose,
+            )
+        )
+        for r in roots
+    ]
+    dataset = torch.utils.data.ConcatDataset(parts)
+    log.info(f"actor-actor combined: {[len(p) for p in parts]} -> {len(dataset)} pairs from {list(roots)}")
+    sampler = None
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=is_train, drop_last=True)
+    return get_cached_replay_dataloader(
+        webdataset=False, use_cache=False, dataset=dataset, batch_size=batch_size, num_workers=num_workers,
+        sampler=sampler, shuffle=(sampler is None and is_train), drop_last=True,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None, persistent_workers=num_workers > 0,
+        pin_memory=False, collate_fn=collate_fn, cache_replay_name="nymeria_actor_actor_combined_dataloader",
+    )
+
+
+def get_nymeria_obs_vroid_mixed_loader(
+    *,
+    obs_root: str = "/data2/nymeria_processed_single",
+    vroid_root: str = "/data4/vroid_batch",
+    obs_manifest: str = "manifest/train_split.csv",
+    vroid_manifest: str = "manifest/train_split.csv",
+    resolution_hw: tuple[int, int] = (480, 480),
+    num_video_frames: int = 77,
+    fps: float = 10.0,
+    role_mix: bool = True,
+    num_reference_frames: int = 4,
+    shared_reference: bool = True,
+    emit_camera_poses: bool = True,
+    person_pose: bool = True,
+    emit_reference_pose: bool = True,
+    batch_size: int = 1,
+    num_workers: int = 4,
+    prefetch_factor: Optional[int] = 2,
+    is_train: bool = True,
+    **kwargs: Any,
+):
+    """PRETRAIN mix: nymeria actor-OBSERVER (real) + vroid actor-ACTOR (synthetic), pooled uniformly via a
+    ConcatDataset. Both share the SAME conditioning structure (shared refs + Plücker campose + person-pose +
+    REFERENCE-POSE), so batches are interchangeable to the model; each dataset applies its own role handling
+    (observer transform vs actor-actor) and loads its own captions."""
+    common = dict(
+        resolution_hw=tuple(resolution_hw), num_video_frames=num_video_frames, fps=fps, role_mix=role_mix,
+        num_reference_frames=num_reference_frames, shared_reference=shared_reference,
+        emit_camera_poses=emit_camera_poses, person_pose=person_pose, emit_reference_pose=emit_reference_pose,
+    )
+    obs = NymeriaActorObserverDataset(NymeriaPairsConfig(root=obs_root, manifest_csv=obs_manifest, **common))
+    vro = NymeriaActorActorDataset(NymeriaPairsConfig(root=vroid_root, manifest_csv=vroid_manifest, **common))
+    dataset = torch.utils.data.ConcatDataset([obs, vro])
+    log.info(f"obs+vroid mixed: obs={len(obs)} + vroid={len(vro)} -> {len(dataset)} pairs")
+    sampler = None
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=is_train, drop_last=True)
+    return get_cached_replay_dataloader(
+        webdataset=False, use_cache=False, dataset=dataset, batch_size=batch_size, num_workers=num_workers,
+        sampler=sampler, shuffle=(sampler is None and is_train), drop_last=True,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None, persistent_workers=num_workers > 0,
+        pin_memory=False, collate_fn=collate_fn, cache_replay_name="nymeria_obs_vroid_mixed_dataloader",
+    )
+
+
+def get_nymeria_actor_observer_loader(
+    *,
+    root: str = "/data2/nymeria_processed_single",
+    manifest_csv: str = "manifest/train_split.csv",
+    resolution_hw: tuple[int, int] = (480, 480),
+    num_video_frames: int = 77,
+    fps: float = 10.0,
+    role_mix: bool = True,
+    num_reference_frames: int = 0,
+    shared_reference: bool = False,
+    emit_camera_poses: bool = False,
+    emit_depth: bool = False,
+    person_pose: bool = False,
+    emit_reference_pose: bool = False,
+    batch_size: int = 1,
+    num_workers: int = 4,
+    prefetch_factor: Optional[int] = 2,
+    is_train: bool = True,
+    **kwargs: Any,
+):
+    dataset = NymeriaActorObserverDataset(
+        NymeriaPairsConfig(
+            root=root,
+            manifest_csv=manifest_csv,
+            resolution_hw=tuple(resolution_hw),
+            num_video_frames=num_video_frames,
+            fps=fps,
+            role_mix=role_mix,
+            num_reference_frames=num_reference_frames,
+            shared_reference=shared_reference,
+            emit_camera_poses=emit_camera_poses,
+            emit_depth=emit_depth,
+            person_pose=person_pose,
+            emit_reference_pose=emit_reference_pose,
+        )
+    )
+    sampler = None
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=is_train, drop_last=True)
+    return get_cached_replay_dataloader(
+        webdataset=False,
+        use_cache=False,
+        dataset=dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        sampler=sampler,
+        shuffle=(sampler is None and is_train),
+        drop_last=True,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None,
+        persistent_workers=num_workers > 0,
+        pin_memory=False,
+        collate_fn=collate_fn,
+        cache_replay_name="nymeria_actor_observer_dataloader",
+    )
+
+
+def get_nymeria_pairs_loader(
+    *,
+    root: str = "/data2/nymeria_processed_longer",
+    manifest_csv: str = "manifest/train_ready.csv",
+    resolution_hw: tuple[int, int] = (480, 480),
+    num_video_frames: int = 77,
+    fps: float = 10.0,
+    caption: str = "a person performing an everyday activity",
+    clean_only: bool = False,
+    single_view: bool = False,
+    split: str = "all",
+    val_num_sessions: int = 4,
+    split_seed: int = 1234,
+    val_max_pairs: int = 40,
+    dummy_text_embeddings: bool = False,
+    batch_size: int = 1,
+    num_workers: int = 4,
+    prefetch_factor: Optional[int] = 2,
+    is_train: bool = True,
+    **kwargs: Any,
+):
+    dataset = NymeriaPairsDataset(
+        NymeriaPairsConfig(
+            root=root,
+            manifest_csv=manifest_csv,
+            resolution_hw=tuple(resolution_hw),
+            num_video_frames=num_video_frames,
+            fps=fps,
+            caption=caption,
+            clean_only=clean_only,
+            single_view=single_view,
+            split=split,
+            val_num_sessions=val_num_sessions,
+            split_seed=split_seed,
+            val_max_pairs=val_max_pairs,
+            dummy_text_embeddings=dummy_text_embeddings,
+        )
+    )
+
+    sampler = None
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=is_train, drop_last=True)
+
+    return get_cached_replay_dataloader(
+        webdataset=False,
+        use_cache=False,
+        dataset=dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        sampler=sampler,
+        shuffle=(sampler is None and is_train),
+        drop_last=True,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None,
+        persistent_workers=num_workers > 0,
+        pin_memory=False,
+        collate_fn=collate_fn,
+        cache_replay_name="nymeria_pairs_dataloader",
+    )
+
+
+def register_nymeria_pairs_dataloader() -> None:
+    cs = ConfigStore.instance()
+
+    def _store(name: str, **kw):
+        """Register a name with a session-held-out train split (data_train) and val split (data_val).
+
+        data_train -> split="train" (all sessions except the held-out ones).
+        data_val   -> split="val"   (only held-out sessions, clean subset, capped to val_max_pairs).
+        """
+        cs.store(
+            group="data_train", package="dataloader_train", name=name,
+            node=L(get_nymeria_pairs_loader)(is_train=True, split="train", **kw),
+        )
+        val_kw = dict(kw)
+        val_kw["clean_only"] = True  # cleaner validation regardless of the train clean setting
+        val_kw["num_workers"] = 2  # validation set is small
+        cs.store(
+            group="data_val", package="dataloader_val", name=name,
+            node=L(get_nymeria_pairs_loader)(is_train=False, split="val", **val_kw),
+        )
+
+    # LONGER (200+77, state_t=20) from train_ready.csv (clip.npz). 77 target frames -> state_t = 1+(77-1)//4 = 20.
+    _store("nymeria_train_ready", root="/data2/nymeria_processed_longer", num_video_frames=77, batch_size=1, num_workers=4)
+    _store("nymeria_train_ready_smoke", root="/data2/nymeria_processed_longer", num_video_frames=77,
+           dummy_text_embeddings=True, batch_size=1, num_workers=4)
+    _store("nymeria_train_ready_clean_smoke", root="/data2/nymeria_processed_longer", num_video_frames=77,
+           clean_only=True, dummy_text_embeddings=True, batch_size=1, num_workers=4)
+    # stage-1 single-view (V=1) pose pretraining: each actor an independent sample
+    _store("nymeria_single_view", root="/data2/nymeria_processed_longer", num_video_frames=77,
+           single_view=True, dummy_text_embeddings=True, batch_size=1, num_workers=4)
+
+    # 2-view actor-observer (real per-view text). train: role-mix on; val: role-mix off (deterministic).
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_observer",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_observer",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, is_train=False, batch_size=1, num_workers=2),
+    )
+    # 2-view actor-observer WITH reference-frame appearance conditioning (R=6 clean source frames/view)
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_observer_refs",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_observer_refs",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, is_train=False, batch_size=1, num_workers=2),
+    )
+    # 2-view actor-observer WITH reference frames (R=4) + camera-pose (Plücker) conditioning
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_observer_refs_campose",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, emit_camera_poses=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_observer_refs_campose",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, emit_camera_poses=True,
+            is_train=False, batch_size=1, num_workers=2),
+    )
+    # 2-view ACTOR-ACTOR (longer dataset) with reference frames (R=4) + camera-pose (Plücker) conditioning.
+    # Both views are ego actors, each with its own real caption (captions/<seq>.csv).
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_actor_refs_campose",
+        node=L(get_nymeria_actor_actor_loader)(
+            root="/data2/nymeria_processed_longer", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, emit_camera_poses=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_actor_refs_campose",
+        node=L(get_nymeria_actor_actor_loader)(
+            root="/data2/nymeria_processed_longer", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, emit_camera_poses=True,
+            is_train=False, batch_size=1, num_workers=2),
+    )
+
+    # 2-view actor-observer with SHARED greedy references (R=4/slot -> 8 shared) + camera-pose Plücker.
+    # Pairs with net/model shared_reference=True (ref view-embedding dropped). Uses refs_shared.npz.
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_observer_refs_campose_shared",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_observer_refs_campose_shared",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, is_train=False, batch_size=1, num_workers=2),
+    )
+    # + DEPTH condition: actor-observer SHARED refs with composite-depth (mesh_cond depth_comp) emitted as
+    # control_input_depth -> net VAE depth_embedder. single depth (mesh Stage-B) is ~100% ready.
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_observer_refs_campose_shared_depth",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, emit_depth=True, is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_observer_refs_campose_shared_depth",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, emit_depth=True, is_train=False, batch_size=1, num_workers=2),
+    )
+    # 2-view ACTOR-ACTOR with SHARED greedy references + camera-pose Plücker.
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_actor_refs_campose_shared",
+        node=L(get_nymeria_actor_actor_loader)(
+            root="/data2/nymeria_processed_longer", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_actor_refs_campose_shared",
+        node=L(get_nymeria_actor_actor_loader)(
+            root="/data2/nymeria_processed_longer", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, is_train=False, batch_size=1, num_workers=2),
+    )
+
+    # VROID synthetic actor-actor: person-pose + shared refs + Plücker + REFERENCE-POSE condition
+    # (refs_shared `reference_pose`). /data4/vroid_batch. Captions from captions/core4d_all.csv (general:
+    # "A first-person view of C interacting with a partner."). detail_all.csv (per-action) is kept as .bak.
+    cs.store(
+        group="data_train", package="dataloader_train", name="vroid_actoractor_refpose",
+        node=L(get_nymeria_actor_actor_loader)(
+            root="/data4/vroid_batch", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, person_pose=True, emit_reference_pose=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="vroid_actoractor_refpose",
+        node=L(get_nymeria_actor_actor_loader)(
+            root="/data4/vroid_batch", manifest_csv="manifest/val_avatar_disjoint.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, person_pose=True, emit_reference_pose=True,
+            is_train=False, batch_size=1, num_workers=2),
+    )
+
+    # PRETRAIN MIX: nymeria actor-OBSERVER (real /data2/nymeria_processed_single) + vroid actor-ACTOR
+    # (synthetic /data4/vroid_batch), pooled uniformly. Shared refs + Plücker campose + person-pose + refpose.
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_obs_vroid_mixed_refpose",
+        node=L(get_nymeria_obs_vroid_mixed_loader)(
+            obs_root="/data2/nymeria_processed_single", vroid_root="/data4/vroid_batch",
+            obs_manifest="manifest/train_split.csv", vroid_manifest="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, person_pose=True, emit_reference_pose=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_obs_vroid_mixed_refpose",
+        node=L(get_nymeria_obs_vroid_mixed_loader)(
+            obs_root="/data2/nymeria_processed_single", vroid_root="/data4/vroid_batch",
+            obs_manifest="manifest/val_split.csv", vroid_manifest="manifest/val_avatar_disjoint.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, person_pose=True, emit_reference_pose=True,
+            is_train=False, batch_size=1, num_workers=2),
+    )
+
+    # OVERFIT: train AND val on the SAME 5 hand-picked pairs (manifest/overfit5.csv) — deliberate overfit test.
+    # role_mix=False on both so train and val see the identical pairs/orientation.
+    for grp, pkg in [("data_train", "dataloader_train"), ("data_val", "dataloader_val")]:
+        cs.store(
+            group=grp, package=pkg, name="vroid_overfit5",
+            node=L(get_nymeria_actor_actor_loader)(
+                root="/data4/vroid_batch", manifest_csv="manifest/overfit5.csv",
+                num_video_frames=77, role_mix=False, num_reference_frames=4, shared_reference=True,
+                emit_camera_poses=True, person_pose=True, emit_reference_pose=True,
+                is_train=(grp == "data_train"), batch_size=1, num_workers=2),
+        )
+
+    # 2-view ACTOR-ACTOR SHARED refs, REAL longer + SYNTHETIC combined (all pairs, no reweighting). Val stays
+    # on the REAL longer location-holdout set only (nymeria_actor_actor_refs_campose_shared val).
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_actor_refs_campose_shared_plus_synth",
+        node=L(get_nymeria_actor_actor_combined_loader)(
+            roots=["/data2/nymeria_processed_longer", "/data3/synthetic_processed_multi"],
+            manifest_csv="manifest/train_split.csv", num_video_frames=77, role_mix=True,
+            num_reference_frames=4, shared_reference=True, emit_camera_poses=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+
+    # + DEPTH condition: REAL longer + SYNTHETIC combined actor-actor with composite-depth (mesh_cond depth_comp)
+    # emitted as control_input_depth -> net VAE depth_embedder. Synthetic depth 100% ready; nymeria mesh Stage-B
+    # in progress (missing -> zero depth, coverage grows). Val stays on nymeria longer holdout.
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_actor_refs_campose_shared_plus_synth_depth",
+        node=L(get_nymeria_actor_actor_combined_loader)(
+            roots=["/data2/nymeria_processed_longer", "/data3/synthetic_processed_multi"],
+            manifest_csv="manifest/train_split.csv", num_video_frames=77, role_mix=True,
+            num_reference_frames=4, shared_reference=True, emit_camera_poses=True, emit_depth=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+
+    # STAGE-1 SINGLE-VIEW (motion pretraining): pooled ego clips (real single actor + longer + synthetic),
+    # per-clip real caption + pose + warping + plucker + own refs (R=4). LoRA-friendly, max data.
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_single_view_stage1",
+        node=L(get_nymeria_single_view_loader)(
+            val=False, num_video_frames=77, num_reference_frames=4, emit_camera_poses=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_single_view_stage1",
+        node=L(get_nymeria_single_view_loader)(
+            val=True, num_video_frames=77, num_reference_frames=4, emit_camera_poses=True,
+            is_train=False, batch_size=1, num_workers=2),
+    )
+
+    # 2-view ACTOR-ACTOR SHARED refs, SYNTHETIC ONLY (diagnostic: isolate synthetic human appearance / masking
+    # from the real-data influence in the combined run). root = synthetic; val = synthetic val_split.
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_actor_refs_campose_shared_synthonly",
+        node=L(get_nymeria_actor_actor_loader)(
+            root="/data3/synthetic_processed_multi", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_actor_refs_campose_shared_synthonly",
+        node=L(get_nymeria_actor_actor_loader)(
+            root="/data3/synthetic_processed_multi", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=4, shared_reference=True,
+            emit_camera_poses=True, is_train=False, batch_size=1, num_workers=2),
+    )
+
+    # 2-view actor-observer with camera-pose (Plücker) conditioning ONLY (no reference frames) — ablation
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_actor_observer_campose",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/train_split.csv",
+            num_video_frames=77, role_mix=True, num_reference_frames=0, emit_camera_poses=True,
+            is_train=True, batch_size=1, num_workers=4),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_actor_observer_campose",
+        node=L(get_nymeria_actor_observer_loader)(
+            root="/data2/nymeria_processed_single", manifest_csv="manifest/val_split.csv",
+            num_video_frames=77, role_mix=False, num_reference_frames=0, emit_camera_poses=True,
+            is_train=False, batch_size=1, num_workers=2),
+    )
+
+    # legacy 33-frame (mp4) smoke kept for quick tests (no split)
+    cs.store(
+        group="data_train", package="dataloader_train", name="nymeria_pairs_smoke",
+        node=L(get_nymeria_pairs_loader)(root="/data2/nymeria_processed", manifest_csv="manifest/clip_pairs.csv",
+                                         num_video_frames=33, dummy_text_embeddings=True, batch_size=1, num_workers=2, is_train=True),
+    )
+    cs.store(
+        group="data_val", package="dataloader_val", name="nymeria_pairs_smoke",
+        node=L(get_nymeria_pairs_loader)(root="/data2/nymeria_processed", manifest_csv="manifest/clip_pairs.csv",
+                                         num_video_frames=33, dummy_text_embeddings=True, batch_size=1, num_workers=2, is_train=False),
+    )

cosmos_predict2/_src/predict2_multiview/datasets/wdinfo_utils.py

@@ -0,0 +1,79 @@
+# 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.
+
+"""Utility functions for handling wdinfo files."""
+
+from typing import Literal, Mapping
+
+from cosmos_predict2._src.imaginaire import config
+from cosmos_predict2._src.imaginaire.datasets.webdataset.config.schema import DatasetInfo
+
+DEFAULT_CATALOG: Mapping = {
+    "alpamayo_dec2024": {
+        "sensitive": [
+            "wdinfo/alpamayo_dec2024/v0/resolution_1080/aspect_ratio_16_9/duration_10_30/wdinfo_test.json",
+        ],
+    },
+    "mads_multiview_0823": {
+        "sensitive": [
+            "wdinfo/mads/cosmos-mads-dataset-transfer2-multiview-0823/v0/driving/resolution_720/aspect_ratio_16_9/duration_5_10/wdinfo_08232025.json",
+        ],
+    },
+}
+
+
+def get_video_dataset_info(
+    source_name: str,
+    *,
+    dataset_keys: list[str] | None = None,
+    object_store: Literal["gcs", "s3"] = "gcs",
+    dataset_catalog: dict = DEFAULT_CATALOG,
+) -> list[DatasetInfo]:
+    if source_name not in dataset_catalog:
+        raise KeyError(
+            f"Source {source_name} not found in dataset catalog. Available keys are {dataset_catalog.keys()}"
+        )
+
+    # Create the wdinfo files here
+    dataset_infos = []
+    for sensitive_type, wdinfos in dataset_catalog[source_name].items():
+        if object_store == "gcs":
+            bucket = "bucket" if sensitive_type == "nonsensitive" else "bucket-s"
+        elif object_store == "s3":
+            bucket = "bucket" if sensitive_type == "nonsensitive" else "bucket-sensitive"
+        else:
+            raise ValueError("Cosmos data: only support gcs or s3 for object store")
+
+        if not wdinfos:
+            continue
+
+        dataset_infos.append(
+            DatasetInfo(
+                object_store_config=config.ObjectStoreConfig(
+                    enabled=True,
+                    credentials=(
+                        "credentials/s3_training.secret" if object_store == "s3" else "credentials/gcs_training.secret"
+                    ),
+                    bucket=bucket,
+                ),
+                wdinfo=wdinfos,
+                per_dataset_keys=dataset_keys,
+                source=source_name,
+                opts={
+                    "aspect_ratio": "16,9",
+                },
+            )
+        )
+    return dataset_infos

cosmos_predict2/_src/predict2_multiview/models/multiview_pose_model_rectified_flow.py

@@ -0,0 +1,261 @@
+# 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.
+
+"""Multiview rectified-flow model with pose + warped-frame conditioning for 2-actor joint generation.
+
+This subclass only adds *preprocessing* of the three conditioning inputs and leaves the training loop, FSDP,
+EMA and checkpointing of the base ``MultiviewVid2VidModelRectifiedFlow`` untouched:
+  * ``control_input_warped``     (pixel RGB) -> normalized to [-1,1] -> VAE-encoded (frozen) -> ``warped_latent``
+  * ``control_input_visibility`` (pixel)     -> 1ch [0,1] -> down-sampled to latent (T_lat,H_lat,W_lat) -> ``visibility_mask``
+  * ``control_input_pose``       (pixel RGB) -> passed through at pixel res -> ``pose_map`` (the net's pose
+                                                encoder normalizes /255 and down-samples /16 spatial, /4 temporal)
+
+The preprocessed tensors are written into the data_batch under keys read by the pose conditioner's ReMapkey
+embedders (``pose_map`` / ``warped_latent`` / ``visibility_mask``), so they flow into the MultiViewCondition
+and ultimately into ``MultiViewPoseDiT.forward`` with classifier-free-guidance dropout handled by the
+conditioner. No checkpoint weight surgery is needed: the new pose/cond modules are additive + zero-initialized
+and load as missing keys from the base 2B checkpoint.
+"""
+
+from typing import Callable, Dict, Tuple
+
+import attrs
+import torch
+import torch.nn.functional as F
+
+from cosmos_predict2._src.imaginaire.utils import log
+from einops import rearrange
+from torch import Tensor
+
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.conditioner import MultiViewCondition
+from cosmos_predict2._src.predict2_multiview.models.multiview_vid2vid_model_rectified_flow import (
+    MultiviewVid2VidModelRectifiedFlow,
+    MultiviewVid2VidModelRectifiedFlowConfig,
+)
+
+_POSE_PREPROCESSED_KEY = "_pose_conditioning_preprocessed"
+
+
+@attrs.define(slots=False)
+class MultiviewVid2VidPoseModelRectifiedFlowConfig(MultiviewVid2VidModelRectifiedFlowConfig):
+    # data_batch keys carrying the raw pixel conditioning videos (one per actor, laid out (B,C,V*T,H,W))
+    pose_input_key: str = "control_input_pose"
+    warped_input_key: str = "control_input_warped"
+    visibility_input_key: str = "control_input_visibility"
+    # when True, pass the pose RGB through the frozen VAE -> `pose_latent` (16ch latent res) instead of
+    # (in addition to) the pixel `pose_map`. The net's pose_mode ("vae_concat"/"vae_mlp_add") consumes it.
+    pose_via_vae: bool = False
+    # when True, VAE-encode the composite-depth RGB (`control_input_depth`) into `depth_latent` (16ch latent res)
+    # for the net's zero-init depth_embedder (net.enable_depth). Depth = warped scene depth + human mesh depth.
+    depth_via_vae: bool = False
+    depth_input_key: str = "control_input_depth"
+    # >0: VAE-encode `reference_frames` (R clean source frames/view) into `reference_latent` for the net's
+    # in-context reference appearance conditioning. Must match the net's num_reference_frames.
+    num_reference_frames: int = 0
+    reference_input_key: str = "reference_frames"
+    # compute per-pixel Plücker ray maps (dir+moment, canonical view0-frame0 frame) from camera_w2c/K into
+    # `plucker_map` for the net's cross-view shared-space conditioning.
+    enable_plucker: bool = False
+    # also compute posed-reference Plücker (reference frames' past poses in the same canonical frame) into
+    # `reference_plucker_map`, so the in-context reference frames are geometrically grounded too.
+    enable_reference_plucker: bool = False
+    # VAE-encode the per-person skeleton render of each reference frame (refs_shared `reference_pose`) into
+    # `reference_pose_latent` for the net's zero-init reference_pose_embedder (net.enable_reference_pose).
+    enable_reference_pose: bool = False
+    reference_pose_input_key: str = "reference_pose"
+
+
+class MultiviewVid2VidPoseModelRectifiedFlow(MultiviewVid2VidModelRectifiedFlow):
+    def __init__(self, config: MultiviewVid2VidPoseModelRectifiedFlowConfig):
+        super().__init__(config)
+
+    def set_up_model(self):
+        # super() builds the net + (if config.use_lora) runs add_lora, which FREEZES the whole base network
+        # and leaves only LoRA adapters trainable. Our NEW zero-init conditioning modules are NOT LoRA targets
+        # and must stay FULLY trainable, so un-freeze them here (backbone LoRA + full new-modules hybrid).
+        super().set_up_model()
+        if not getattr(self.config, "use_lora", False):
+            return
+        base = self.net.base_model.model if hasattr(self.net, "base_model") else self.net
+        trainable = ["cond_embedder", "plucker_embedder", "depth_embedder", "reference_pose_embedder", "pose_encoder", "pose_latent_embedder",
+                     "pose_mlp", "ref_gate"]
+        if not getattr(base, "freeze_view_embedding", False):
+            trainable.append("view_embeddings")
+        n_new = 0
+        for name, param in self.net.named_parameters():
+            n = name.replace("base_model.model.", "")
+            if any(n == p or n.startswith(p + ".") for p in trainable):
+                param.requires_grad_(True)
+                param.data = param.data.float()  # match add_lora's fp32 upcast for trainable params
+                n_new += param.numel()
+        tot = sum(p.numel() for p in self.net.parameters())
+        tr = sum(p.numel() for p in self.net.parameters() if p.requires_grad)
+        log.info(f"[LoRA+new-modules] unfroze {n_new:,} new-module params -> trainable {tr:,}/{tot:,} "
+                 f"({100 * tr / max(1, tot):.3f}%); targets={trainable}")
+
+    @torch.no_grad()
+    def _preprocess_pose_conditioning(self, data_batch: Dict[str, torch.Tensor]) -> None:
+        """Encode/down-sample the raw pixel conditioning videos into the keys read by the conditioner.
+
+        Idempotent: guarded so it runs once per data_batch even though both training (get_data_and_condition)
+        and inference (get_velocity_fn_from_batch) call it.
+        """
+        if data_batch.get(_POSE_PREPROCESSED_KEY, False):
+            return
+        pose_key = self.config.pose_input_key
+        warped_key = self.config.warped_input_key
+        vis_key = self.config.visibility_input_key
+        if pose_key not in data_batch or warped_key not in data_batch or vis_key not in data_batch:
+            # nothing to do (e.g. a batch that does not carry pose conditioning)
+            return
+
+        state_t = self.config.state_t
+        num_pixel_frames_per_view = self.tokenizer.get_pixel_num_frames(state_t)
+
+        # --- warped past-frame: pixel RGB -> [-1,1] -> frozen VAE encode (multiview-aware reshape inside) ---
+        warped = data_batch[warped_key].to(**self.tensor_kwargs) / 127.5 - 1.0  # (B,C,V*Tpix,H,W)
+        warped_latent = self.encode(warped).to(**self.tensor_kwargs)  # (B,16,V*state_t,H_lat,W_lat)
+        n_views = warped_latent.shape[2] // state_t
+        _, _, _, h_lat, w_lat = warped_latent.shape
+        data_batch["warped_latent"] = warped_latent
+
+        # --- visibility mask: pixel -> 1ch [0,1] -> down-sample to (state_t, H_lat, W_lat) per view ---
+        vis = data_batch[vis_key].to(**self.tensor_kwargs)
+        if vis.shape[1] > 1:
+            vis = vis.mean(dim=1, keepdim=True)
+        vis = vis / 255.0
+        vis = rearrange(vis, "B C (V T) H W -> (B V) C T H W", V=n_views)
+        vis = F.interpolate(vis.float(), size=(state_t, h_lat, w_lat), mode="trilinear", align_corners=False)
+        vis = rearrange(vis, "(B V) C T H W -> B C (V T) H W", V=n_views).to(**self.tensor_kwargs)
+        data_batch["visibility_mask"] = vis
+
+        # --- pose map: keep pixel resolution & full temporal; the net's pose encoder handles the rest ---
+        data_batch["pose_map"] = data_batch[pose_key].float()  # (B,3,V*Tpix,H,W) in [0,255]
+
+        # --- pose latent (optional): pixel RGB pose -> [-1,1] -> frozen VAE encode (same path as warped) ---
+        # gives a 16ch latent-res pose representation for the "vae_concat"/"vae_mlp_add" net pose modes.
+        if getattr(self.config, "pose_via_vae", False):
+            pose = data_batch[pose_key].to(**self.tensor_kwargs) / 127.5 - 1.0  # (B,3,V*Tpix,H,W)
+            data_batch["pose_latent"] = self.encode(pose).to(**self.tensor_kwargs)  # (B,16,V*state_t,H_lat,W_lat)
+
+        # --- composite DEPTH latent (optional): RGB-encoded depth (warped scene + human mesh) -> [-1,1] -> frozen
+        #     VAE encode (same path as warped/pose) -> depth_latent (B,16,V*state_t,H_lat,W_lat). ---
+        if getattr(self.config, "depth_via_vae", False):
+            dkey = getattr(self.config, "depth_input_key", "control_input_depth")
+            if dkey in data_batch:
+                depth = data_batch[dkey].to(**self.tensor_kwargs) / 127.5 - 1.0  # (B,3,V*Tpix,H,W)
+                data_batch["depth_latent"] = self.encode(depth).to(**self.tensor_kwargs)
+
+        # --- reference frames (optional): R clean source frames/view, VAE-encoded PER FRAME (1-frame clips) so
+        #     each of the R frames yields exactly one latent frame -> (B,16,V*R,H_lat,W_lat) reference_latent. ---
+        ref_key = getattr(self.config, "reference_input_key", "reference_frames")
+        if getattr(self.config, "num_reference_frames", 0) > 0 and ref_key in data_batch:
+            refs = data_batch[ref_key].to(**self.tensor_kwargs) / 127.5 - 1.0  # (B,3,V*R,H,W)
+            data_batch["reference_latent"] = self._encode_reference_frames(refs).to(**self.tensor_kwargs)
+
+        # --- reference POSE (optional): per-person skeleton render of each ref frame, VAE-encoded PER FRAME
+        #     (same path as reference_latent) -> reference_pose_latent (B,16,V*R,H_lat,W_lat). ---
+        if getattr(self.config, "enable_reference_pose", False):
+            rpkey = getattr(self.config, "reference_pose_input_key", "reference_pose")
+            if rpkey in data_batch:
+                rpose = data_batch[rpkey].to(**self.tensor_kwargs) / 127.5 - 1.0  # (B,3,V*R,H,W)
+                data_batch["reference_pose_latent"] = self._encode_reference_frames(rpose).to(**self.tensor_kwargs)
+
+        # --- Plücker ray map (optional): per-pixel camera rays in a per-pair canonical frame (view0 frame0),
+        #     at latent res, for cross-view shared-space grounding. -> plucker_map (B,6,V*state_t,H_lat,W_lat) ---
+        if getattr(self.config, "enable_plucker", False) and "camera_w2c" in data_batch:
+            src = int(data_batch["camera_src_res"][0, 0]) if "camera_src_res" in data_batch else 504
+            # force fp32 (disable autocast): torch.linalg.inv rejects bf16, and validation calls this inside an
+            # autocast(bf16) block where einsum/matmul would otherwise downcast the poses to bf16.
+            with torch.autocast("cuda", enabled=False):
+                w2c_t = data_batch["camera_w2c"].to(torch.float32)  # (B,V,Tpix,4,4)
+                K = data_batch["camera_K"]
+                # per-pair canonical frame = TARGET view0 frame0 (same transform to all views AND references)
+                c2w_ref = torch.linalg.inv(w2c_t[:, 0, 0])  # (B,4,4)
+                Tp = w2c_t.shape[2]
+                idx = torch.linspace(0, Tp - 1, state_t, device=w2c_t.device).round().long()
+                data_batch["plucker_map"] = self._plucker_from_w2c(w2c_t[:, :, idx], K, c2w_ref, src, h_lat, w_lat)
+                # posed-reference Plücker (optional): reference frames' PAST poses in the SAME canonical frame
+                if getattr(self.config, "enable_reference_plucker", False) and "reference_cam_w2c" in data_batch:
+                    data_batch["reference_plucker_map"] = self._plucker_from_w2c(
+                        data_batch["reference_cam_w2c"].to(torch.float32), K, c2w_ref, src, h_lat, w_lat
+                    )
+
+        data_batch[_POSE_PREPROCESSED_KEY] = True
+
+    @torch.no_grad()
+    def _plucker_from_w2c(self, w2c_BVN, camera_K, c2w_ref, src_res, h_lat, w_lat):
+        """Per-pixel Plücker rays (dir 3 + moment 3) at latent res for a set of poses `w2c_BVN` (B,V,N,4,4),
+        expressed in a GIVEN canonical frame `c2w_ref` (B,4,4) = inv(view0-frame0). The SAME c2w_ref is used for
+        the real frames AND the references, so both live in one shared frame. camera_K (B,V,3,3) ->
+        plucker (B, 6, V*N, H_lat, W_lat)."""
+        dt = torch.float32
+        dev = w2c_BVN.device
+        B, V, N = w2c_BVN.shape[:3]
+        w2c = w2c_BVN.to(dt)
+        w2c_canon = torch.einsum("bvnij,bjk->bvnik", w2c, c2w_ref.to(dt))  # (B,V,N,4,4)
+        c2w = torch.linalg.inv(w2c_canon.reshape(-1, 4, 4)).reshape(B, V, N, 4, 4)
+        Rc2w = c2w[..., :3, :3]  # (B,V,N,3,3)
+        C = c2w[..., :3, 3]  # (B,V,N,3)  camera centers in canonical frame
+        # intrinsics scaled from src_res to the latent grid
+        K = camera_K.to(dt).clone()  # (B,V,3,3)
+        K[..., 0, :] *= w_lat / float(src_res)
+        K[..., 1, :] *= h_lat / float(src_res)
+        Kinv = torch.linalg.inv(K)  # (B,V,3,3)
+        ys, xs = torch.meshgrid(
+            torch.arange(h_lat, device=dev, dtype=dt), torch.arange(w_lat, device=dev, dtype=dt), indexing="ij"
+        )
+        pix = torch.stack(
+            [xs.reshape(-1) + 0.5, ys.reshape(-1) + 0.5, torch.ones(h_lat * w_lat, device=dev, dtype=dt)], 0
+        )  # (3, HW)
+        d_cam = torch.einsum("bvij,jp->bvip", Kinv, pix)  # (B,V,3,HW)
+        d_cam = d_cam / d_cam.norm(dim=2, keepdim=True)
+        d_world = torch.einsum("bvnij,bvjp->bvnip", Rc2w, d_cam)  # (B,V,N,3,HW)
+        d_world = d_world / d_world.norm(dim=3, keepdim=True)
+        C_exp = C.unsqueeze(-1).expand(-1, -1, -1, -1, h_lat * w_lat)  # (B,V,N,3,HW)
+        m = torch.cross(C_exp, d_world, dim=3)  # moment C x d
+        plk = torch.cat([d_world, m], dim=3)  # (B,V,N,6,HW)
+        plk = plk.reshape(B, V, N, 6, h_lat, w_lat).permute(0, 3, 1, 2, 4, 5).reshape(
+            B, 6, V * N, h_lat, w_lat
+        )
+        return plk.to(**self.tensor_kwargs)
+
+    @torch.no_grad()
+    def _encode_reference_frames(self, refs_B_C_VR_H_W: torch.Tensor) -> torch.Tensor:
+        """Encode each reference frame as an INDEPENDENT 1-frame clip through the frozen VAE.
+        A plain multiview encode maps V*R pixel frames via get_pixel_num_frames(state_t) and would mis-infer
+        n_views; per-frame encoding gives exactly R latent frames/view. (B,3,V*R,H,W) -> (B,16,V*R,H_lat,W_lat).
+        """
+        B = refs_B_C_VR_H_W.shape[0]
+        x = rearrange(refs_B_C_VR_H_W, "B C VR H W -> (B VR) C H W").unsqueeze(2)  # (B*VR, C, 1, H, W)
+        z = self.tokenizer.encode(x)  # (B*VR, 16, 1, H_lat, W_lat)
+        return rearrange(z.squeeze(2), "(B VR) C H W -> B C VR H W", B=B)
+
+    def get_data_and_condition(
+        self, data_batch: dict[str, torch.Tensor]
+    ) -> Tuple[Tensor, Tensor, MultiViewCondition]:
+        # training path: conditioner runs inside super().get_data_and_condition, so preprocess first
+        self._preprocess_pose_conditioning(data_batch)
+        return super().get_data_and_condition(data_batch)
+
+    def get_velocity_fn_from_batch(
+        self,
+        data_batch: Dict,
+        guidance: float = 1.5,
+        is_negative_prompt: bool = False,
+    ) -> Callable:
+        # inference path: conditioner is invoked before get_data_and_condition, so preprocess up front
+        self._preprocess_pose_conditioning(data_batch)
+        return super().get_velocity_fn_from_batch(data_batch, guidance, is_negative_prompt)

cosmos_predict2/_src/predict2_multiview/models/multiview_vid2vid_model_rectified_flow.py

@@ -0,0 +1,655 @@
+# 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 random
+from dataclasses import field
+from typing import Any, Callable, Dict, Optional, Tuple, cast
+
+import attrs
+import torch
+import torch.distributed as dist
+from einops import rearrange
+from megatron.core import parallel_state
+from torch import Tensor
+from torch.distributed import get_process_group_ranks
+
+from cosmos_predict2._src.imaginaire.flags import SMOKE
+from cosmos_predict2._src.imaginaire.utils import log
+from cosmos_predict2._src.imaginaire.utils.context_parallel import broadcast, broadcast_split_tensor
+from cosmos_predict2._src.predict2.conditioner import DataType
+from cosmos_predict2._src.predict2.models.text2world_model_rectified_flow import IS_PREPROCESSED_KEY
+from cosmos_predict2._src.predict2.models.video2world_model_rectified_flow import (
+    NUM_CONDITIONAL_FRAMES_KEY,
+    Video2WorldModelRectifiedFlow,
+    Video2WorldModelRectifiedFlowConfig,
+)
+from cosmos_predict2._src.predict2.utils.dtensor_helper import broadcast_dtensor_model_states
+from cosmos_predict2._src.predict2_multiview.configs.vid2vid.defaults.conditioner import (
+    ConditionLocationList,
+    MultiViewCondition,
+)
+from cosmos_predict2._src.predict2_multiview.models.view_sampling import sample_n_views_from_data_batch
+
+TRAIN_SAMPLE_N_VIEWS_KEY = "train_sample_n_views"
+TRAIN_SAMPLING_APPLIED_KEY = "train_sampling_applied"
+_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."
+
+
+@attrs.define(slots=False)
+class MultiviewVid2VidModelRectifiedFlowConfig(Video2WorldModelRectifiedFlowConfig):
+    min_num_conditional_frames_per_view: int = 1
+    max_num_conditional_frames_per_view: int = 2
+    train_sample_views_range: Tuple[int, int] | None = None
+    condition_locations: ConditionLocationList = field(default_factory=lambda: ConditionLocationList([]))
+    state_t: int = 0
+    view_condition_dropout_max: int = 0
+    online_text_embeddings_as_dict: bool = True  # For backward compatibility with old experiments
+    conditional_frames_probs: Optional[Dict[int, float]] = None  # Probability distribution for conditional frames
+
+
+class MultiviewVid2VidModelRectifiedFlow(Video2WorldModelRectifiedFlow):
+    def __init__(self, config: MultiviewVid2VidModelRectifiedFlowConfig):
+        super().__init__(config)
+        self.state_t = config.state_t
+        self.empty_string_text_embeddings = None
+        self.neg_text_embeddings = None
+        if self.config.text_encoder_config is not None and self.config.text_encoder_config.compute_online:
+            compute_empty_and_negative_text_embeddings(self)
+
+    @torch.no_grad()
+    def encode(self, state: torch.Tensor) -> torch.Tensor:
+        n_views = state.shape[2] // self.tokenizer.get_pixel_num_frames(self.state_t)
+        cp_size = len(get_process_group_ranks(parallel_state.get_context_parallel_group()))
+        # let n_views 2 also cp-encoded
+        if n_views > 1 and n_views <= cp_size:
+            return self.encode_cp(state)
+        state = rearrange(state, "B C (V T) H W -> (B V) C T H W", V=n_views)
+        encoded_state = super().encode(state)
+        encoded_state = rearrange(encoded_state, "(B V) C T H W -> B C (V T) H W", V=n_views)
+        return encoded_state
+
+    @torch.no_grad()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        n_views = latent.shape[2] // self.state_t
+        cp_size = len(get_process_group_ranks(parallel_state.get_context_parallel_group()))
+        # let n_views 2 also cp-decoded
+        if n_views > 1 and n_views <= cp_size:
+            return self.decode_cp(latent)
+        latent = rearrange(latent, "B C (V T) H W -> (B V) C T H W", V=n_views)
+        decoded_state = super().decode(latent)
+        decoded_state = rearrange(decoded_state, "(B V) C T H W -> B C (V T) H W", V=n_views)
+        return decoded_state
+
+    @torch.no_grad()
+    def encode_cp(self, state: torch.Tensor) -> torch.Tensor:
+        cp_size = len(get_process_group_ranks(parallel_state.get_context_parallel_group()))
+        cp_group = parallel_state.get_context_parallel_group()
+        n_views = state.shape[2] // self.tokenizer.get_pixel_num_frames(self.state_t)
+        assert n_views <= cp_size, f"n_views must be less than cp_size, got n_views={n_views} and cp_size={cp_size}"
+        state_V_B_C_T_H_W = rearrange(state, "B C (V T) H W -> V B C T H W", V=n_views)
+        state_input = torch.zeros((cp_size, *state_V_B_C_T_H_W.shape[1:]), **self.tensor_kwargs)
+        state_input[0:n_views] = state_V_B_C_T_H_W
+        local_state_V_B_C_T_H_W = broadcast_split_tensor(state_input, seq_dim=0, process_group=cp_group)
+        local_state = rearrange(local_state_V_B_C_T_H_W, "V B C T H W -> (B V) C T H W")
+        encoded_state = super().encode(local_state)
+        encoded_state_list = [torch.empty_like(encoded_state) for _ in range(cp_size)]
+        dist.all_gather(encoded_state_list, encoded_state, group=cp_group)
+        encoded_state = torch.cat(encoded_state_list[0:n_views], dim=2)  # [B, C, V * T, H, W]
+        return encoded_state
+
+    @torch.no_grad()
+    def decode_cp(self, latent: torch.Tensor) -> torch.Tensor:
+        cp_size = len(get_process_group_ranks(parallel_state.get_context_parallel_group()))
+        cp_group = parallel_state.get_context_parallel_group()
+        n_views = latent.shape[2] // self.state_t
+        assert n_views <= cp_size, f"n_views must be less than cp_size, got n_views={n_views} and cp_size={cp_size}"
+        latent_V_B_C_T_H_W = rearrange(latent, "B C (V T) H W -> V B C T H W", V=n_views)
+        latent_input = torch.zeros((cp_size, *latent_V_B_C_T_H_W.shape[1:]), **self.tensor_kwargs)
+        latent_input[0:n_views] = latent_V_B_C_T_H_W
+        local_latent_V_B_C_T_H_W = broadcast_split_tensor(latent_input, seq_dim=0, process_group=cp_group)
+        local_latent = rearrange(local_latent_V_B_C_T_H_W, "V B C T H W -> (B V) C T H W")
+        decoded_state = super().decode(local_latent)
+        decoded_state_list = [torch.empty_like(decoded_state) for _ in range(cp_size)]
+        dist.all_gather(decoded_state_list, decoded_state, group=cp_group)
+        decoded_state = torch.cat(decoded_state_list[0:n_views], dim=2)  # [B, C, V * T, H, W]
+        return decoded_state
+
+    def training_step(
+        self, data_batch: dict[str, torch.Tensor], iteration: int
+    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        return training_step_multiview(self, data_batch, iteration)
+
+    def inplace_compute_text_embeddings_online(self, data_batch: dict[str, torch.Tensor]) -> None:
+        inplace_compute_text_embeddings_online_multiview(self, data_batch)
+
+    def broadcast_split_for_model_parallelsim(
+        self,
+        x0_B_C_T_H_W: torch.Tensor,
+        condition: MultiViewCondition,
+        epsilon_B_C_T_H_W: torch.Tensor,
+        sigma_B_T: torch.Tensor,
+    ):
+        n_views = x0_B_C_T_H_W.shape[2] // self.state_t
+        x0_B_C_T_H_W = rearrange(x0_B_C_T_H_W, "B C (V T) H W -> (B V) C T H W", V=n_views).contiguous()
+        if epsilon_B_C_T_H_W is not None:
+            epsilon_B_C_T_H_W = rearrange(epsilon_B_C_T_H_W, "B C (V T) H W -> (B V) C T H W", V=n_views).contiguous()
+        reshape_sigma_B_T = False
+        if sigma_B_T is not None:
+            assert sigma_B_T.ndim == 2, "sigma_B_T should be 2D tensor"
+            if sigma_B_T.shape[-1] != 1:
+                assert sigma_B_T.shape[-1] % n_views == 0, (
+                    f"sigma_B_T temporal dimension T must either be 1 or a multiple of sample_n_views. Got T={sigma_B_T.shape[-1]} and sample_n_views={n_views}"
+                )
+                sigma_B_T = rearrange(sigma_B_T, "B (V T) -> (B V) T", V=n_views).contiguous()
+                reshape_sigma_B_T = True
+        x0_B_C_T_H_W, condition, epsilon_B_C_T_H_W, sigma_B_T = super().broadcast_split_for_model_parallelsim(
+            x0_B_C_T_H_W, condition, epsilon_B_C_T_H_W, sigma_B_T
+        )
+        x0_B_C_T_H_W = rearrange(x0_B_C_T_H_W, "(B V) C T H W -> B C (V T) H W", V=n_views)
+        if epsilon_B_C_T_H_W is not None:
+            epsilon_B_C_T_H_W = rearrange(epsilon_B_C_T_H_W, "(B V) C T H W -> B C (V T) H W", V=n_views)
+        if reshape_sigma_B_T:
+            sigma_B_T = rearrange(sigma_B_T, "(B V) T -> B (V T)", V=n_views)
+        return x0_B_C_T_H_W, condition, epsilon_B_C_T_H_W, sigma_B_T
+
+    def get_data_batch_with_latent_view_indices(self, data_batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        num_video_frames_per_view = int(data_batch["num_video_frames_per_view"].cpu().item())
+        n_views = data_batch["view_indices"].shape[1] // num_video_frames_per_view
+        view_indices_B_V_T = rearrange(data_batch["view_indices"], "B (V T) -> B V T", V=n_views)
+
+        latent_view_indices_B_V_T = view_indices_B_V_T[:, :, 0 : self.config.state_t]
+        latent_view_indices_B_T = rearrange(latent_view_indices_B_V_T, "B V T -> B (V T)")
+        data_batch_with_latent_view_indices = data_batch.copy()
+        data_batch_with_latent_view_indices["latent_view_indices_B_T"] = latent_view_indices_B_T
+        return data_batch_with_latent_view_indices
+
+    def _normalize_video_databatch_inplace(self, data_batch: dict[str, Tensor], input_key: str = None) -> None:
+        input_key = self.input_data_key if input_key is None else input_key
+        is_preprocessed = IS_PREPROCESSED_KEY in data_batch and data_batch[IS_PREPROCESSED_KEY] is True
+
+        num_video_frames_per_view = (
+            self.tokenizer.get_pixel_num_frames(self.state_t)
+            if is_preprocessed
+            else data_batch["num_video_frames_per_view"]
+        )
+        if isinstance(num_video_frames_per_view, torch.Tensor):
+            num_video_frames_per_view = int(num_video_frames_per_view.cpu().item())
+        n_views = data_batch[input_key].shape[2] // num_video_frames_per_view
+        if input_key in data_batch:
+            data_batch[input_key] = rearrange(data_batch[input_key], "B C (V T) H W -> (B V) C T H W", V=n_views)
+            super()._normalize_video_databatch_inplace(data_batch, input_key)
+            data_batch[input_key] = rearrange(data_batch[input_key], "(B V) C T H W -> B C (V T) H W", V=n_views)
+
+    def get_data_and_condition(self, data_batch: dict[str, torch.Tensor]) -> Tuple[Tensor, Tensor, MultiViewCondition]:
+        data_batch_with_latent_view_indices = self.get_data_batch_with_latent_view_indices(data_batch)
+        raw_state, latent_state, condition = super(Video2WorldModelRectifiedFlow, self).get_data_and_condition(
+            data_batch_with_latent_view_indices
+        )
+        condition = cast(MultiViewCondition, condition)
+        condition = condition.set_video_condition(
+            state_t=self.config.state_t,
+            gt_frames=latent_state.to(**self.tensor_kwargs),
+            condition_locations=self.config.condition_locations,
+            random_min_num_conditional_frames_per_view=self.config.min_num_conditional_frames_per_view,
+            random_max_num_conditional_frames_per_view=self.config.max_num_conditional_frames_per_view,
+            num_conditional_frames_per_view=None,
+            view_condition_dropout_max=self.config.view_condition_dropout_max,
+            conditional_frames_probs=self.config.conditional_frames_probs,
+        )
+        return raw_state, latent_state, condition
+
+    def get_velocity_fn_from_batch(
+        self,
+        data_batch: Dict,
+        guidance: float = 1.5,
+        is_negative_prompt: bool = False,
+    ) -> Callable:
+        """
+        Generates a callable function `x0_fn` based on the provided data batch and guidance factor.
+
+        This function first processes the input data batch through a conditioning workflow (`conditioner`) to obtain conditioned and unconditioned states. It then defines a nested function `x0_fn` which applies a denoising operation on an input `noise_x` at a given noise level `sigma` using both the conditioned and unconditioned states.
+
+        Args:
+        - data_batch (Dict): A batch of data used for conditioning. The format and content of this dictionary should align with the expectations of the `self.conditioner`
+        - guidance (float, optional): A scalar value that modulates the influence of the conditioned state relative to the unconditioned state in the output. Defaults to 1.5.
+        - is_negative_prompt (bool): use negative prompt t5 in uncondition if true
+
+        Returns:
+        - Callable: A function `x0_fn(noise_x, sigma)` that takes two arguments, `noise_x` and `sigma`, and return velocity predictoin
+
+        The returned function is suitable for use in scenarios where a denoised state is required based on both conditioned and unconditioned inputs, with an adjustable level of guidance influence.
+        """
+
+        data_batch_with_latent_view_indices = self.get_data_batch_with_latent_view_indices(data_batch)
+        if NUM_CONDITIONAL_FRAMES_KEY in data_batch_with_latent_view_indices:
+            num_conditional_frames = data_batch_with_latent_view_indices[NUM_CONDITIONAL_FRAMES_KEY]
+            log.debug(f"Using {num_conditional_frames=} from data batch")
+        else:
+            num_conditional_frames = 1
+
+        if is_negative_prompt:
+            condition, uncondition = self.conditioner.get_condition_with_negative_prompt(
+                data_batch_with_latent_view_indices
+            )
+        else:
+            condition, uncondition = self.conditioner.get_condition_uncondition(data_batch_with_latent_view_indices)
+
+        is_image_batch = self.is_image_batch(data_batch_with_latent_view_indices)
+        condition = condition.edit_data_type(DataType.IMAGE if is_image_batch else DataType.VIDEO)
+        uncondition = uncondition.edit_data_type(DataType.IMAGE if is_image_batch else DataType.VIDEO)
+        _, x0, _ = self.get_data_and_condition(data_batch_with_latent_view_indices)
+        # override condition with inference mode; num_conditional_frames used Here!
+        condition = condition.set_video_condition(
+            state_t=self.config.state_t,
+            gt_frames=x0,
+            condition_locations=self.config.condition_locations,
+            random_min_num_conditional_frames_per_view=self.config.min_num_conditional_frames_per_view,
+            random_max_num_conditional_frames_per_view=self.config.max_num_conditional_frames_per_view,
+            num_conditional_frames_per_view=num_conditional_frames,
+            view_condition_dropout_max=0,
+            conditional_frames_probs=self.config.conditional_frames_probs,
+        )
+        uncondition = uncondition.set_video_condition(
+            state_t=self.config.state_t,
+            gt_frames=x0,
+            condition_locations=self.config.condition_locations,
+            random_min_num_conditional_frames_per_view=self.config.min_num_conditional_frames_per_view,
+            random_max_num_conditional_frames_per_view=self.config.max_num_conditional_frames_per_view,
+            num_conditional_frames_per_view=num_conditional_frames,
+            view_condition_dropout_max=0,
+            conditional_frames_probs=self.config.conditional_frames_probs,
+        )
+        condition = condition.edit_for_inference(
+            is_cfg_conditional=True,
+            condition_locations=self.config.condition_locations,
+            num_conditional_frames_per_view=num_conditional_frames,
+        )
+        uncondition = uncondition.edit_for_inference(
+            is_cfg_conditional=False,
+            condition_locations=self.config.condition_locations,
+            num_conditional_frames_per_view=num_conditional_frames,
+        )
+        _, condition, _, _ = self.broadcast_split_for_model_parallelsim(x0, condition, None, None)
+        _, uncondition, _, _ = self.broadcast_split_for_model_parallelsim(x0, uncondition, None, None)
+        if parallel_state.is_initialized():
+            pass
+        else:
+            assert not self.net.is_context_parallel_enabled, (
+                "parallel_state is not initialized, context parallel should be turned off."
+            )
+
+        def velocity_fn(noise: torch.Tensor, noise_x: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+            cond_v = self.denoise(noise, noise_x, timestep, condition)
+            uncond_v = self.denoise(noise, noise_x, timestep, uncondition)
+            velocity_pred = uncond_v + guidance * (cond_v - uncond_v)  # align with pred2
+            return velocity_pred
+
+        return velocity_fn
+
+    @torch.no_grad()
+    def generate_samples_from_batch(
+        self,
+        data_batch: dict[str, torch.Tensor],
+        guidance: float = 1.5,
+        seed: int = 1,
+        state_shape: Tuple | None = None,
+        n_sample: int | None = None,
+        is_negative_prompt: bool = False,
+        num_steps: int = 35,
+        shift: float = 5.0,
+        **kwargs,
+    ) -> torch.Tensor:
+        data_batch_with_latent_view_indices = self.get_data_batch_with_latent_view_indices(data_batch)
+        process_group = parallel_state.get_context_parallel_group()
+        cp_size = len(get_process_group_ranks(process_group))
+        samples_B_C_T_H_W = super().generate_samples_from_batch(
+            data_batch_with_latent_view_indices,
+            guidance,
+            seed,
+            state_shape,
+            n_sample,
+            is_negative_prompt,
+            num_steps,
+            shift,
+            **kwargs,
+        )
+        if cp_size > 1:
+            samples_B_C_T_H_W = rearrange(
+                samples_B_C_T_H_W, "B C (c V T) H W -> B C (V c T) H W", c=cp_size, T=self.state_t // cp_size
+            )
+        return samples_B_C_T_H_W
+
+    def set_up_model(self):
+        """Override set_up_model to initialize cross-view attention from base model."""
+        # Call parent's set_up_model first
+        super().set_up_model()
+
+        # Load base model and initialize cross-view attention if enabled
+        if (
+            hasattr(self.net, "enable_cross_view_attn")
+            and self.net.enable_cross_view_attn
+            and self.net.init_cross_view_attn_weight_from is not None
+        ):
+            log.info("Loading base model for cross-view attention initialization")
+            self.net.init_cross_view_attn_with_self_attn_weights(is_ema=False)
+            if self.config.ema.enabled:
+                self.net_ema.init_cross_view_attn_with_self_attn_weights(is_ema=True)
+
+            # Broadcast to ensure all ranks have consistent weights
+            if self.fsdp_device_mesh is not None:
+                log.info("Broadcasting cross-view attention weights for consistency")
+                broadcast_dtensor_model_states(self.net, self.fsdp_device_mesh)
+                if self.config.ema.enabled:
+                    broadcast_dtensor_model_states(self.net_ema, self.fsdp_device_mesh)
+
+
+def compute_text_embeddings_online_multiview_single_caption(
+    model, data_batch: dict[str, torch.Tensor]
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    is_preprocessed = IS_PREPROCESSED_KEY in data_batch and data_batch[IS_PREPROCESSED_KEY] is True
+    num_video_frames_per_view = (
+        model.tokenizer.get_pixel_num_frames(model.state_t)
+        if is_preprocessed
+        else data_batch["num_video_frames_per_view"]
+    )
+    if isinstance(num_video_frames_per_view, torch.Tensor):
+        num_video_frames_per_view = int(num_video_frames_per_view.cpu().item())
+    n_views = data_batch[model.input_data_key].shape[2] // num_video_frames_per_view
+    B, _, _, _, _ = data_batch[model.input_data_key].shape
+
+    # compute prompt embeddings
+    if len(data_batch["ai_caption"]) != 1:
+        raise NotImplementedError(f"Expected batch size of 1, got {len(data_batch['ai_caption'])}")
+
+    if len(data_batch["ai_caption"][0]) != 1:
+        raise ValueError(f"Expected a single caption, got {len(data_batch['ai_caption'][0])}")
+
+    caption = data_batch["ai_caption"][0][0]
+    assert isinstance(caption, str)
+    view0_text_embeddings_B_L_D = model.text_encoder.compute_text_embeddings_online(
+        data_batch={model.input_caption_key: [caption]},
+        input_caption_key=model.input_caption_key,
+    )
+    assert view0_text_embeddings_B_L_D.shape[0] == 1
+    assert view0_text_embeddings_B_L_D.shape[1] == 512, (
+        f"view0_text_embeddings should be of shape (B, 512, D), got {view0_text_embeddings_B_L_D.shape}"
+    )
+    output_text_embeddings = model.empty_string_text_embeddings.clone().repeat(B, n_views, 1)
+    output_neg_text_embeddings = model.empty_string_text_embeddings.clone().repeat(B, n_views, 1)
+    output_text_embeddings = rearrange(output_text_embeddings, "B (V L) D -> V B L D", V=n_views)
+    output_neg_text_embeddings = rearrange(output_neg_text_embeddings, "B (V L) D -> V B L D", V=n_views)
+    # Assign prompt embeddings to the front camera view
+    for i_b in range(B):
+        front_cam_view_idx_sample_position = data_batch["front_cam_view_idx_sample_position"][i_b]
+        output_text_embeddings[front_cam_view_idx_sample_position, i_b] = view0_text_embeddings_B_L_D[i_b]
+        output_neg_text_embeddings[front_cam_view_idx_sample_position, i_b] = model.neg_text_embeddings[0]
+    output_text_embeddings = rearrange(output_text_embeddings, "V B L D -> B (V L) D")
+    output_neg_text_embeddings = rearrange(output_neg_text_embeddings, "V B L D -> B (V L) D")
+
+    dropout_text_embeddings = model.empty_string_text_embeddings.clone().repeat(B, n_views, 1)
+
+    if not model.config.conditioner.text.use_empty_string:
+        dropout_text_embeddings *= 0.0
+
+    return output_text_embeddings, output_neg_text_embeddings, dropout_text_embeddings
+
+
+def compute_text_embeddings_online_multiview_multiple_captions(
+    model, data_batch: dict[str, torch.Tensor]
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    is_preprocessed = IS_PREPROCESSED_KEY in data_batch and data_batch[IS_PREPROCESSED_KEY] is True
+    num_video_frames_per_view = (
+        model.tokenizer.get_pixel_num_frames(model.state_t)
+        if is_preprocessed
+        else data_batch["num_video_frames_per_view"]
+    )
+    if isinstance(num_video_frames_per_view, torch.Tensor):
+        num_video_frames_per_view = int(num_video_frames_per_view.cpu().item())
+    n_views = data_batch[model.input_data_key].shape[2] // num_video_frames_per_view
+    B, _, _, _, _ = data_batch[model.input_data_key].shape
+
+    # compute each view's caption separately
+    if not len(data_batch["ai_caption"]) == 1:
+        raise NotImplementedError(f"Expected batch size of 1, got {len(data_batch['ai_caption'])}")
+
+    captions = data_batch["ai_caption"][0]
+    if len(captions) != n_views:
+        raise ValueError(f"Expected {n_views} captions, got {len(captions)}: {captions}")
+    view_text_embeddings = []
+    for caption in captions:
+        data_batch_per_view = {model.input_caption_key: [caption]}
+        view_text_embedding = model.text_encoder.compute_text_embeddings_online(
+            data_batch_per_view, model.input_caption_key
+        )
+        view_text_embeddings.append(view_text_embedding)
+
+    view_text_embeddings_B_V_L_D = torch.stack(view_text_embeddings, dim=1)
+    assert view_text_embeddings_B_V_L_D.shape[:3] == (
+        B,
+        n_views,
+        512,
+    ), f"view_text_embeddings_B_V_L_D should be of shape (B, n_views, 512, D), got {view_text_embeddings_B_V_L_D.shape}"
+    output_text_embeddings = rearrange(view_text_embeddings_B_V_L_D, "B V L D -> B (V L) D")
+
+    # repeat negative embedding for each per view
+    output_neg_text_embeddings_L_D = model.neg_text_embeddings[0].clone()
+    output_neg_text_embeddings_B_V_L_D = rearrange(output_neg_text_embeddings_L_D, "(B V L) D -> B V L D", B=1, V=1)
+    output_neg_text_embeddings_B_V_L_D = output_neg_text_embeddings_B_V_L_D.repeat(B, n_views, 1, 1)
+    assert output_neg_text_embeddings_B_V_L_D.shape[:3] == (
+        B,
+        n_views,
+        512,
+    ), (
+        f"output_neg_text_embeddings_B_V_L_D should be of shape (B, n_views, 512, D), got {output_neg_text_embeddings_B_V_L_D.shape}"
+    )
+    output_neg_text_embeddings = rearrange(output_neg_text_embeddings_B_V_L_D, "B V L D -> B (V L) D")
+
+    dropout_text_embeddings = model.empty_string_text_embeddings.clone().repeat(B, n_views, 1)
+    if not model.config.conditioner.text.use_empty_string:
+        dropout_text_embeddings *= 0.0
+
+    return output_text_embeddings, output_neg_text_embeddings, dropout_text_embeddings
+
+
+def compute_text_embeddings_online_multiview(
+    model, data_batch: dict[str, torch.Tensor]
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    captions = data_batch["ai_caption"][0]
+    is_preprocessed = IS_PREPROCESSED_KEY in data_batch and data_batch[IS_PREPROCESSED_KEY] is True
+    num_video_frames_per_view = (
+        model.tokenizer.get_pixel_num_frames(model.state_t)
+        if is_preprocessed
+        else data_batch["num_video_frames_per_view"]
+    )
+    if isinstance(num_video_frames_per_view, torch.Tensor):
+        num_video_frames_per_view = int(num_video_frames_per_view.cpu().item())
+    n_views = data_batch[model.input_data_key].shape[2] // num_video_frames_per_view
+    assert len(captions) == 1 or len(captions) == n_views, f"Expected 1 or {n_views} captions, got {len(captions)}"
+    if len(captions) == 1:
+        return compute_text_embeddings_online_multiview_single_caption(model, data_batch)
+    else:
+        return compute_text_embeddings_online_multiview_multiple_captions(model, data_batch)
+
+
+def inplace_compute_text_embeddings_online_multiview(model, data_batch: dict[str, torch.Tensor]) -> None:
+    output_text_embeddings, output_neg_text_embeddings, dropout_text_embeddings = (
+        compute_text_embeddings_online_multiview(model, data_batch)
+    )
+    t5_text_embeddings = {
+        "text_embeddings": output_text_embeddings,
+        "dropout_text_embeddings": dropout_text_embeddings,
+    }
+    neg_t5_text_embeddings = {
+        "text_embeddings": output_neg_text_embeddings,
+        "dropout_text_embeddings": dropout_text_embeddings,
+    }
+    data_batch["t5_text_embeddings"] = (
+        t5_text_embeddings if model.config.online_text_embeddings_as_dict else t5_text_embeddings["text_embeddings"]
+    )
+    data_batch["neg_t5_text_embeddings"] = (
+        neg_t5_text_embeddings
+        if model.config.online_text_embeddings_as_dict
+        else neg_t5_text_embeddings["text_embeddings"]
+    )
+    data_batch["t5_text_mask"] = torch.ones(
+        output_text_embeddings.shape[0], output_text_embeddings.shape[1], device="cuda"
+    )
+
+
+def compute_empty_and_negative_text_embeddings(model):
+    # Compute empty string embeddings for text embedding dropout
+    if model.empty_string_text_embeddings is None:
+        empty_string_data_batch = {
+            model.input_caption_key: [" "],
+        }
+        model.empty_string_text_embeddings = model.text_encoder.compute_text_embeddings_online(
+            empty_string_data_batch, model.input_caption_key
+        )
+
+    # compute negative prompt embeddings for sampling
+    if model.neg_text_embeddings is None:
+        neg_promt_data_batch = {
+            model.input_caption_key: [_DEFAULT_NEGATIVE_PROMPT],
+        }
+        model.neg_text_embeddings = model.text_encoder.compute_text_embeddings_online(
+            neg_promt_data_batch, model.input_caption_key
+        )
+
+
+def preprocess_databatch(
+    data_batch: dict[str, Any],
+    train_sample_views_range: Optional[tuple[int, int]],
+) -> dict[str, Any]:
+    """Preprocess data batch with dynamic view sampling."""
+
+    if TRAIN_SAMPLING_APPLIED_KEY in data_batch and data_batch[TRAIN_SAMPLING_APPLIED_KEY] is True:
+        return data_batch
+    if train_sample_views_range is not None:
+        min_views, max_views = train_sample_views_range
+        log.debug(f"Randomly sampling {min_views} to {max_views} views")
+        if SMOKE:
+            train_sample_n_views = 1
+        elif TRAIN_SAMPLE_N_VIEWS_KEY in data_batch:
+            train_sample_n_views = data_batch[TRAIN_SAMPLE_N_VIEWS_KEY]
+            log.debug(f"Using {TRAIN_SAMPLE_N_VIEWS_KEY} from data batch: {train_sample_n_views}")
+            if train_sample_n_views < 1:  # No sampling is applied
+                return data_batch
+        else:
+            # Sample n_views to keep
+            train_sample_n_views = random.randint(min_views, max_views)
+            if dist.is_initialized():
+                # Have all ranks globally sample the same number of views such that tensors are same shape for efficiency
+                train_sample_n_views_tensor = torch.tensor(
+                    [train_sample_n_views], device=data_batch["sample_n_views"].device, dtype=torch.int64
+                )
+                dist.broadcast(train_sample_n_views_tensor, src=0)
+                train_sample_n_views = train_sample_n_views_tensor[0].cpu().item()
+
+        n_views = data_batch["sample_n_views"].cpu().item()
+        log.debug(f"Sampling {train_sample_n_views} views out of {n_views}")
+        available_view_indices = list(range(n_views))
+        if data_batch["front_cam_view_idx_sample_position"][0] is not None and train_sample_n_views < n_views:
+            # In cases where we only have a single caption, we shouldn't drop the front camera view
+            front_cam_view_idx_sample_position = data_batch["front_cam_view_idx_sample_position"][0].cpu().item()
+            available_view_indices.pop(front_cam_view_idx_sample_position)
+        keep_view_indices = sorted(random.sample(available_view_indices, train_sample_n_views))
+        if parallel_state.get_context_parallel_group() is not None:
+            # Broadcast the exact views to keep for each context parallel group
+            keep_view_indices = broadcast(keep_view_indices, parallel_state.get_context_parallel_group())
+        log.debug(f"Sampled and broadcasted keep_view_indices={keep_view_indices}")
+        data_batch = sample_n_views_from_data_batch(data_batch, keep_view_indices)
+        data_batch[TRAIN_SAMPLING_APPLIED_KEY] = True
+    return data_batch
+
+
+def training_step_multiview(
+    model, data_batch: dict[str, torch.Tensor], iteration: int
+) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+    """
+    Performs a single training step for the diffusion model.
+
+    This method is responsible for executing one iteration of the model's training. It involves:
+    1. Adding noise to the input data using the SDE process.
+    2. Passing the noisy data through the network to generate predictions.
+    3. Computing the loss based on the difference between the predictions and the original data, \
+        considering any configured loss weighting.
+
+    Args:
+        data_batch (dict): raw data batch draw from the training data loader.
+        iteration (int): Current iteration number.
+
+    Returns:
+        tuple: A tuple containing two elements:
+            - dict: additional data that used to debug / logging / callbacks
+            - Tensor: The computed loss for the training step as a PyTorch Tensor.
+
+    Raises:
+        AssertionError: If the class is conditional, \
+            but no number of classes is specified in the network configuration.
+
+    Notes:
+        - The method handles different types of conditioning
+        - The method also supports Kendall's loss
+    """
+    model._update_train_stats(data_batch)
+
+    # only happens in training
+    data_batch = preprocess_databatch(data_batch, model.config.train_sample_views_range)
+
+    if model.config.text_encoder_config is not None and model.config.text_encoder_config.compute_online:
+        model.inplace_compute_text_embeddings_online(data_batch)
+
+    # Get the input data to noise and denoise~(image, video) and the corresponding conditioner.
+    _, x0_B_C_T_H_W, condition = model.get_data_and_condition(data_batch)
+
+    # Sample pertubation noise levels and N(0, 1) noises
+    epsilon_B_C_T_H_W = torch.randn(x0_B_C_T_H_W.size(), **model.tensor_kwargs_fp32)
+    batch_size = x0_B_C_T_H_W.size()[0]
+    t_B = model.rectified_flow.sample_train_time(batch_size).to(**model.tensor_kwargs_fp32)
+    t_B = rearrange(t_B, "b -> b 1")  # add a dimension for T, all frames share the same sigma
+    x0_B_C_T_H_W, condition, epsilon_B_C_T_H_W, t_B = model.broadcast_split_for_model_parallelsim(
+        x0_B_C_T_H_W, condition, epsilon_B_C_T_H_W, t_B
+    )
+    timesteps = model.rectified_flow.get_discrete_timestamp(t_B, model.tensor_kwargs_fp32)
+    sigmas = model.rectified_flow.get_sigmas(
+        timesteps,
+        model.tensor_kwargs_fp32,
+    )
+    timesteps = rearrange(timesteps, "b -> b 1")
+    sigmas = rearrange(sigmas, "b -> b 1")
+    xt_B_C_T_H_W, vt_B_C_T_H_W = model.rectified_flow.get_interpolation(epsilon_B_C_T_H_W, x0_B_C_T_H_W, sigmas)
+
+    vt_pred_B_C_T_H_W = model.denoise(
+        noise=epsilon_B_C_T_H_W,
+        xt_B_C_T_H_W=xt_B_C_T_H_W.to(**model.tensor_kwargs),
+        timesteps_B_T=timesteps,
+        condition=condition,
+    )
+
+    time_weights_B = model.rectified_flow.train_time_weight(timesteps, model.tensor_kwargs_fp32)
+    per_instance_loss = torch.mean((vt_pred_B_C_T_H_W - vt_B_C_T_H_W) ** 2, dim=list(range(1, vt_pred_B_C_T_H_W.dim())))
+
+    loss = torch.mean(time_weights_B * per_instance_loss)
+
+    output_batch = {
+        "x0": x0_B_C_T_H_W,
+        "xt": xt_B_C_T_H_W,
+        "sigma": sigmas,
+        "condition": condition,
+        "model_pred": vt_pred_B_C_T_H_W,
+        "edm_loss": loss,
+    }
+
+    return output_batch, loss

cosmos_predict2/_src/predict2_multiview/models/view_sampling.py

@@ -0,0 +1,65 @@
+# 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.
+
+"""Functions for sampling views from a multiview control video data batch."""
+
+import torch
+
+from cosmos_predict2._src.imaginaire.utils import log
+
+
+def sample_n_views_from_data_batch(data_batch: dict, keep_view_indices: list[int]) -> dict:
+    """Sample views from data batch, handling both video and control inputs."""
+
+    keep_view_indices = sorted(keep_view_indices)
+    n_keep_views = len(keep_view_indices)
+    n_orig_views = data_batch["sample_n_views"].cpu().item()
+    if n_keep_views == n_orig_views:
+        log.debug("All views are requested to be kept, returning original data batch!")
+        return data_batch
+    num_video_frames_per_view = data_batch["num_video_frames_per_view"].cpu().item()
+
+    select_ids = []
+    for view_id in keep_view_indices:
+        select_ids.extend(list(range(view_id * num_video_frames_per_view, (view_id + 1) * num_video_frames_per_view)))
+    select_ids = torch.tensor(select_ids, device=data_batch["sample_n_views"].device, dtype=torch.int64)
+
+    data_batch["video"] = data_batch["video"][:, :, select_ids]  # (B, C, V * T, H, W)
+    data_batch["view_indices"] = data_batch["view_indices"][:, select_ids]  # (B, V * T)
+    data_batch["view_indices_selection"] = data_batch["view_indices_selection"][:, keep_view_indices]
+    data_batch["sample_n_views"] = 0 * data_batch["sample_n_views"] + n_keep_views
+
+    # process control input keys
+    for key in data_batch.keys():
+        if key.startswith("control_input_"):
+            data_batch[key] = data_batch[key][:, :, select_ids]  # (B, C, V * T, H, W)
+
+    # process list-like keys
+    for sample_captions in data_batch["ai_caption"]:
+        if len(sample_captions) != n_orig_views:
+            raise ValueError(
+                f"Expected {n_orig_views} captions, got {len(sample_captions)} for key {key}. "
+                "If using single caption, view sampling is not currently supported."
+            )
+    data_batch["ai_caption"] = [
+        [caption for i, caption in enumerate(sample_captions) if i in keep_view_indices]
+        for sample_captions in data_batch["ai_caption"]
+    ]
+    data_batch["camera_keys_selection"] = [
+        [camera_key for i, camera_key in enumerate(sample_camera_keys) if i in keep_view_indices]
+        for sample_camera_keys in data_batch["camera_keys_selection"]
+    ]
+
+    return data_batch

cosmos_predict2/_src/predict2_multiview/networks/multiview_cross_dit.py

@@ -0,0 +1,1187 @@
+# 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 os
+from collections import defaultdict, namedtuple
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.amp as amp
+import torch.nn as nn
+from einops import rearrange
+from megatron.core import parallel_state
+from torch.distributed import ProcessGroup, get_process_group_ranks
+from torch.distributed._composable.fsdp import fully_shard
+from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts
+from torchvision import transforms
+
+from cosmos_predict2._src.imaginaire.utils import log
+from cosmos_predict2._src.predict2.conditioner import DataType
+from cosmos_predict2._src.predict2.networks.minimal_v1_lvg_dit import MinimalV1LVGDiT
+from cosmos_predict2._src.predict2.networks.minimal_v4_dit import (
+    Attention,
+    Block,
+    SACConfig,
+)
+from cosmos_predict2._src.predict2_multiview.networks.multiview_dit import (
+    MultiCameraSinCosPosEmbAxis,
+    MultiCameraVideoRopePosition3DEmb,
+)
+
+
+# implementation of MultiViewCrossAttention changes multiview_dit.py
+class MultiViewCrossAttention(Attention):
+    def __init__(self, *args, state_t: int = None, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        assert self.qkv_format == "bshd", "MultiViewCrossAttention only supports qkv_format='bshd'"
+        self.state_t = state_t
+
+    def forward(self, x, context=None, rope_emb=None):
+        assert not self.is_selfattn, "MultiViewCrossAttention does not support self-attention"
+        B, L, D = x.shape
+
+        n_cameras = context.shape[1] // 512
+        x_B_L_D = rearrange(x, "B (V L) D -> (V B) L D", V=n_cameras)
+        context_B_M_D = rearrange(context, "B (V M) D -> (V B) M D", V=n_cameras) if context is not None else None
+        x_B_L_D = super().forward(x_B_L_D, context_B_M_D, rope_emb=rope_emb)
+        x_B_L_D = rearrange(x_B_L_D, "(V B) L D -> B (V L) D", V=n_cameras)
+        return x_B_L_D
+
+
+# ---------------------- Selective Activation Checkpoint Policies -----------------------
+def predict2_2B_crossview_720_context_fn():
+    op_count = defaultdict(int)
+
+    def policy_fn(ctx, func, *args, **kwargs):
+        mode = "recompute" if ctx.is_recompute else "forward"
+
+        if func == torch.ops.aten.mm.default:
+            op_count_key = f"{mode}_mm_count"
+            if op_count[op_count_key] >= 10:
+                result = CheckpointPolicy.MUST_SAVE
+            else:
+                result = CheckpointPolicy.PREFER_RECOMPUTE
+
+            # Update count for next operation
+            op_count[op_count_key] = (op_count[op_count_key] + 1) % 20
+            return result
+
+        if "flash_attn" in str(func):
+            return CheckpointPolicy.MUST_SAVE
+
+        return CheckpointPolicy.PREFER_RECOMPUTE
+
+    return create_selective_checkpoint_contexts(policy_fn)
+
+
+class MultiViewCheckpointMode(str, Enum):
+    """Checkpoint modes for MultiViewCrossDiT architecture."""
+
+    NONE = "none"
+    MM_ONLY = "mm_only"
+    BLOCK_WISE = "block_wise"
+    PREDICT2_2B_CROSSVIEW_720 = "predict2_2b_crossview_720"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+@dataclass
+class MultiViewSACConfig(SACConfig):
+    """Selective Activation Checkpoint Config for MultiViewCrossDiT."""
+
+    def get_context_fn(self):
+        if self.mode == MultiViewCheckpointMode.PREDICT2_2B_CROSSVIEW_720:
+            return predict2_2B_crossview_720_context_fn
+        else:
+            return super().get_context_fn()
+
+
+VideoSize = namedtuple("VideoSize", ["T", "H", "W"])
+
+
+class CrossViewAttention(Attention):
+    def __init__(self, *args, cross_view_attn_map: Dict[int, List[int]], **kwargs):
+        super().__init__(*args, **kwargs)
+        del self.attn_op
+        if self.backend == "transformer_engine":
+            from transformer_engine.pytorch.attention import DotProductAttention
+
+            self.attn_op = DotProductAttention(
+                self.n_heads,
+                self.head_dim,
+                num_gqa_groups=self.n_heads,
+                attention_dropout=0,
+                qkv_format=self.qkv_format,
+                attn_mask_type="padding",  # important
+                attention_type="cross",  # important
+            )
+        else:
+            raise NotImplementedError(f"Backend {self.backend} not supported")
+        self.cross_view_attn_map = cross_view_attn_map
+        self.max_neighbors = max(len(neighbors) for neighbors in cross_view_attn_map.values())
+        self.neighbor_indices = None
+        self.neighbor_mask = None
+
+    def forward(self, x, view_indices_B_V, sv_video_size: VideoSize):
+        """
+        x: (B, V, L, D)
+        view_indices_B_V: (B, V)
+        sv_video_size: VideoSize (T, H, W), where T * H * W = L
+        """
+        assert not self.is_selfattn, "CrossViewAttention does not support self-attention"
+        B, V, L, D = x.shape
+        T, H, W = sv_video_size
+        assert T * H * W == L, f"T * H * W != L: {T * H * W} != {L}"
+
+        # move time dimension to batch dimension
+        x = rearrange(x, "b v (t h w) d -> (b t) v (h w) d", t=T, h=H, w=W)
+        B, V, L, D = x.shape
+
+        view_indices_B_V = view_indices_B_V.repeat_interleave(T, dim=0).long()
+
+        # Create neighbor indices and mask on the fly, only once.
+        if self.neighbor_indices is None or self.neighbor_indices.device != x.device:
+            num_total_views = len(self.cross_view_attn_map)
+            neighbor_indices = torch.zeros((num_total_views, self.max_neighbors), dtype=torch.long, device=x.device)
+            neighbor_mask = torch.zeros((num_total_views, self.max_neighbors), dtype=torch.bool, device=x.device)
+            for i in range(num_total_views):
+                neighbors = self.cross_view_attn_map[i]
+                for j, neighbor_idx in enumerate(neighbors):
+                    neighbor_indices[i, j] = neighbor_idx
+                    neighbor_mask[i, j] = True
+            self.neighbor_indices = neighbor_indices
+            self.neighbor_mask = neighbor_mask
+
+        num_total_views = len(self.cross_view_attn_map)
+        view_indices_to_tensor_pos = torch.full(
+            (B, num_total_views), -1, dtype=torch.long, device=x.device
+        )  # include out of range view index
+        b_indices = torch.arange(B, device=x.device).unsqueeze(1).expand(-1, V).long()
+        view_indices_to_tensor_pos[b_indices, view_indices_B_V] = (
+            torch.arange(V, device=x.device).unsqueeze(0).expand(B, -1)
+        )
+
+        neighbor_view_indices = self.neighbor_indices[view_indices_B_V]  # may include out of range view index
+        gather_tensor_pos = view_indices_to_tensor_pos[
+            b_indices.unsqueeze(2), neighbor_view_indices
+        ]  # [B, V, max_neighbors], out of range view index will be -1
+
+        # Sort to move all -1 to the end, which is convenient for creating attention mask.
+        gather_tensor_pos, sorted_indices = torch.sort(gather_tensor_pos, dim=-1, descending=True)
+
+        b_indices_for_gather = torch.arange(B, device=x.device)[:, None, None]
+        # Clamp to avoid index error. Masked values will be ignored in attention.
+        neighbor_features = x[
+            b_indices_for_gather, torch.clamp(gather_tensor_pos, min=0)
+        ]  # [B, V, max_neighbors, L, C]
+
+        # Prepare for attention
+        query = self.q_proj(rearrange(x, "b v l c -> (b v) l c"))  # [B*V, L, C]
+        context = rearrange(neighbor_features, "b v n l c -> (b v) (n l) c")  # [B*V, max_neighbors*L, C]
+        key = self.k_proj(context)
+        value = self.v_proj(context)
+
+        q, k, v = map(
+            lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
+            (query, key, value),
+        )
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        v = self.v_norm(v)
+
+        # Create attention mask
+        is_neighbor_present = gather_tensor_pos != -1  # [B, V, max_neighbors]
+        mask_for_input_views = self.neighbor_mask[view_indices_B_V]  # [B, V, n]
+
+        # Reorder mask_for_input_views to match the sorted gather_tensor_pos
+        mask_for_input_views = torch.gather(mask_for_input_views, -1, sorted_indices)
+        final_mask = is_neighbor_present & mask_for_input_views
+
+        mask_per_view = rearrange(final_mask, "b v n -> (b v) n")  # [BV, n]
+        mask_kv = mask_per_view.repeat_interleave(L, dim=1)  # [BV, n*L]
+
+        # Reshape mask to [batch_size, 1, 1, max_seqlen_kv] as per official documentation.
+        mask = rearrange(mask_kv, "bv l_kv -> bv 1 1 l_kv")  # [BV, 1, 1, n*L]
+        atten_mask_kv = ~mask  # 0 means keep, 1 means mask
+        atten_mask_q = torch.zeros(query.shape[0], 1, 1, query.shape[1]).to(atten_mask_kv)
+
+        attention_output = self.attn_op(q, k, v, attention_mask=(atten_mask_q, atten_mask_kv))
+        attention_output = attention_output.flatten(2)  # [B*V, L, H*D]
+        output = self.output_dropout(self.output_proj(attention_output))
+        output = rearrange(output, "(b v) l d -> b v l d", v=V)
+        # recover time dimension from batch to seq
+        output = rearrange(output, "(b t) v (h w) d -> b v (t h w) d", t=T, h=H, w=W)
+        return output
+
+    def set_context_parallel_group(self, process_group, ranks, stream):
+        raise NotImplementedError("Cross View Attention doesn't need communication")
+
+
+class MultiViewCrossBlock(Block):
+    """
+    A transformer block that takes n_cameras as input.
+    Self-Attention (Single View) -> Cross-View Attention -> Cross Attention (text and image)
+    """
+
+    def __init__(
+        self,
+        x_dim: int,
+        context_dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        use_adaln_lora: bool = False,
+        adaln_lora_dim: int = 256,
+        cross_view_attn_map: Dict[int, List[int]] = None,
+        state_t: int = None,
+        backend: str = "transformer_engine",
+        image_context_dim: Optional[int] = None,
+        use_wan_fp32_strategy: bool = False,
+        enable_cross_view_attn: bool = False,
+    ):
+        super().__init__(
+            x_dim,
+            context_dim,
+            num_heads,
+            mlp_ratio,
+            use_adaln_lora,
+            adaln_lora_dim,
+            backend,
+            image_context_dim,
+            use_wan_fp32_strategy,
+        )
+        self.state_t = state_t
+        self.cross_view_attn_map = cross_view_attn_map
+        self.enable_cross_view_attn = enable_cross_view_attn
+        if image_context_dim is None:
+            del self.cross_attn
+            # cross attention to text and image condition
+            self.cross_attn = MultiViewCrossAttention(
+                x_dim,
+                context_dim,
+                num_heads,
+                x_dim // num_heads,
+                qkv_format="bshd",
+                state_t=state_t,
+                use_wan_fp32_strategy=use_wan_fp32_strategy,
+            )
+        else:
+            raise NotImplementedError("image_context_dim is not supported for MultiViewBlock")
+
+        if enable_cross_view_attn:
+            self.cross_view_attn = CrossViewAttention(
+                x_dim,
+                x_dim,  # context_dim, can not set to None
+                num_heads,
+                x_dim // num_heads,
+                qkv_format="bshd",
+                use_wan_fp32_strategy=use_wan_fp32_strategy,
+                cross_view_attn_map=cross_view_attn_map,
+            )
+            # no modulation so we set elementwise_affine=True
+            self.layer_norm_cross_view_attn = nn.LayerNorm(x_dim, elementwise_affine=True, eps=1e-6)
+
+    def reset_parameters(self):
+        super().reset_parameters()
+        if self.enable_cross_view_attn:
+            self.layer_norm_cross_view_attn.reset_parameters()
+
+    def init_weights(self):
+        super().init_weights()
+        if self.enable_cross_view_attn:
+            self.cross_view_attn.init_weights()
+
+            # Zero-initialize the output projection
+            torch.nn.init.zeros_(self.cross_view_attn.output_proj.weight)
+            if self.cross_view_attn.output_proj.bias is not None:
+                torch.nn.init.zeros_(self.cross_view_attn.output_proj.bias)
+
+    # most code are copied from parent. insert cross view attn between self attn and cross attn
+    def forward(
+        self,
+        x_B_T_H_W_D: torch.Tensor,
+        view_indices_B_T: torch.Tensor,
+        emb_B_T_D: torch.Tensor,
+        view_embedding_proj_B_V_9D: Optional[torch.Tensor] = None,
+        crossattn_emb: Optional[torch.Tensor] = None,
+        rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
+        adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
+        extra_per_block_pos_emb: Optional[torch.Tensor] = None,
+        block_idx: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        x_B_T_H_W_D: (B, T, H, W, D)
+        view_indices_B_T: (B, T)
+        """
+        num_cameras = torch.unique(view_indices_B_T[0]).shape[0]
+
+        if extra_per_block_pos_emb is not None:
+            x_B_T_H_W_D = x_B_T_H_W_D + extra_per_block_pos_emb
+
+        with amp.autocast("cuda", enabled=self.use_wan_fp32_strategy, dtype=torch.float32):
+            if self.use_adaln_lora:
+                shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = (
+                    self.adaln_modulation_self_attn(emb_B_T_D) + adaln_lora_B_T_3D
+                ).chunk(3, dim=-1)
+                shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = (
+                    self.adaln_modulation_cross_attn(emb_B_T_D) + adaln_lora_B_T_3D
+                ).chunk(3, dim=-1)
+                shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = (
+                    self.adaln_modulation_mlp(emb_B_T_D) + adaln_lora_B_T_3D
+                ).chunk(3, dim=-1)
+            else:
+                shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = self.adaln_modulation_self_attn(
+                    emb_B_T_D
+                ).chunk(3, dim=-1)
+                shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = (
+                    self.adaln_modulation_cross_attn(emb_B_T_D).chunk(3, dim=-1)
+                )
+                shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = self.adaln_modulation_mlp(emb_B_T_D).chunk(3, dim=-1)
+
+        # Reshape tensors from (B, T, D) to (B, T, 1, 1, D) for broadcasting
+        shift_self_attn_B_T_1_1_D = rearrange(shift_self_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        scale_self_attn_B_T_1_1_D = rearrange(scale_self_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        gate_self_attn_B_T_1_1_D = rearrange(gate_self_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+
+        shift_cross_attn_B_T_1_1_D = rearrange(shift_cross_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        scale_cross_attn_B_T_1_1_D = rearrange(scale_cross_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        gate_cross_attn_B_T_1_1_D = rearrange(gate_cross_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+
+        shift_mlp_B_T_1_1_D = rearrange(shift_mlp_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        scale_mlp_B_T_1_1_D = rearrange(scale_mlp_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        gate_mlp_B_T_1_1_D = rearrange(gate_mlp_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+
+        if view_embedding_proj_B_V_9D is not None:
+            (
+                view_shift_self_attn_B_V_D,
+                view_scale_self_attn_B_V_D,
+                view_gate_self_attn_B_V_D,
+                view_shift_cross_attn_B_V_D,
+                view_scale_cross_attn_B_V_D,
+                view_gate_cross_attn_B_V_D,
+                view_shift_mlp_B_V_D,
+                view_scale_mlp_B_V_D,
+                view_gate_mlp_B_V_D,
+            ) = view_embedding_proj_B_V_9D.chunk(9, dim=-1)
+
+            repeat_t = x_B_T_H_W_D.shape[1] // num_cameras
+            assert repeat_t * num_cameras == x_B_T_H_W_D.shape[1]
+
+            def expand_and_rearrange(x_B_V_D):
+                expanded_x_B_V_T_D = x_B_V_D.unsqueeze(2).expand(-1, -1, repeat_t, -1)
+                return rearrange(expanded_x_B_V_T_D, "b v t d -> b (v t) 1 1 d")
+
+            shift_self_attn_B_T_1_1_D = shift_self_attn_B_T_1_1_D + expand_and_rearrange(
+                view_shift_self_attn_B_V_D
+            ).type_as(x_B_T_H_W_D)
+            scale_self_attn_B_T_1_1_D = scale_self_attn_B_T_1_1_D + expand_and_rearrange(
+                view_scale_self_attn_B_V_D
+            ).type_as(x_B_T_H_W_D)
+            gate_self_attn_B_T_1_1_D = gate_self_attn_B_T_1_1_D + expand_and_rearrange(
+                view_gate_self_attn_B_V_D
+            ).type_as(x_B_T_H_W_D)
+            shift_cross_attn_B_T_1_1_D = shift_cross_attn_B_T_1_1_D + expand_and_rearrange(
+                view_shift_cross_attn_B_V_D
+            ).type_as(x_B_T_H_W_D)
+            scale_cross_attn_B_T_1_1_D = scale_cross_attn_B_T_1_1_D + expand_and_rearrange(
+                view_scale_cross_attn_B_V_D
+            ).type_as(x_B_T_H_W_D)
+            gate_cross_attn_B_T_1_1_D = gate_cross_attn_B_T_1_1_D + expand_and_rearrange(
+                view_gate_cross_attn_B_V_D
+            ).type_as(x_B_T_H_W_D)
+            shift_mlp_B_T_1_1_D = shift_mlp_B_T_1_1_D + expand_and_rearrange(view_shift_mlp_B_V_D).type_as(x_B_T_H_W_D)
+            scale_mlp_B_T_1_1_D = scale_mlp_B_T_1_1_D + expand_and_rearrange(view_scale_mlp_B_V_D).type_as(x_B_T_H_W_D)
+            gate_mlp_B_T_1_1_D = gate_mlp_B_T_1_1_D + expand_and_rearrange(view_gate_mlp_B_V_D).type_as(x_B_T_H_W_D)
+
+        B, T, H, W, D = x_B_T_H_W_D.shape
+
+        def _fn(_x_B_T_H_W_D, _norm_layer, _scale_B_T_1_1_D, _shift_B_T_1_1_D):
+            return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D
+
+        normalized_x_B_T_H_W_D = _fn(
+            x_B_T_H_W_D,
+            self.layer_norm_self_attn,
+            scale_self_attn_B_T_1_1_D,
+            shift_self_attn_B_T_1_1_D,
+        )
+
+        rope_emb_L_1_1_D_sv = rearrange(
+            rope_emb_L_1_1_D,
+            "(v m) 1 1 d -> v m 1 1 d",
+            v=num_cameras,
+        )[0]
+
+        result_B_T_H_W_D = rearrange(
+            self.self_attn(
+                rearrange(normalized_x_B_T_H_W_D, "b (v t) h w d -> (b v) (t h w) d", v=num_cameras),
+                None,
+                rope_emb=rope_emb_L_1_1_D_sv,
+            ),
+            "(b v) (t h w) d -> b (v t) h w d",
+            v=num_cameras,
+            h=H,
+            w=W,
+        )
+        x_B_T_H_W_D = x_B_T_H_W_D + gate_self_attn_B_T_1_1_D * result_B_T_H_W_D
+
+        # insert cross view attn here. x_B_T_H_W_D
+        if self.enable_cross_view_attn:
+            num_cameras = torch.unique(view_indices_B_T[0]).shape[0]
+            x_B_V_T_H_W_D = rearrange(x_B_T_H_W_D, "b (v t) h w d -> b v t h w d", v=num_cameras)
+            sv_video_size = VideoSize(T=x_B_V_T_H_W_D.shape[2], H=x_B_V_T_H_W_D.shape[3], W=x_B_V_T_H_W_D.shape[4])
+            x_B_V_L_D = rearrange(x_B_V_T_H_W_D, "b v t h w d -> b v (t h w) d")
+            view_indices_B_V = rearrange(view_indices_B_T, "b (v t) -> b v t", v=num_cameras)[..., 0]
+            result_cross_view_attn_B_T_H_W_D = rearrange(
+                self.cross_view_attn(self.layer_norm_cross_view_attn(x_B_V_L_D), view_indices_B_V, sv_video_size),
+                "b v (t h w) d -> b (v t) h w d",
+                v=num_cameras,
+                t=sv_video_size.T,
+                h=sv_video_size.H,
+                w=sv_video_size.W,
+            )
+            x_B_T_H_W_D = x_B_T_H_W_D + result_cross_view_attn_B_T_H_W_D
+
+        def _x_fn(
+            _x_B_T_H_W_D,
+            layer_norm_cross_attn,
+            _scale_cross_attn_B_T_1_1_D,
+            _shift_cross_attn_B_T_1_1_D,
+            _gate_cross_attn_B_T_1_1_D,
+        ):
+            _normalized_x_B_T_H_W_D = _fn(
+                _x_B_T_H_W_D, layer_norm_cross_attn, _scale_cross_attn_B_T_1_1_D, _shift_cross_attn_B_T_1_1_D
+            )
+            _result_B_T_H_W_D = rearrange(
+                self.cross_attn(
+                    rearrange(_normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d"),
+                    crossattn_emb,
+                    rope_emb=rope_emb_L_1_1_D,
+                ),
+                "b (t h w) d -> b t h w d",
+                t=T,
+                h=H,
+                w=W,
+            )
+            # _x_B_T_H_W_D = _x_B_T_H_W_D + _gate_cross_attn_B_T_1_1_D * _result_B_T_H_W_D
+            return _result_B_T_H_W_D
+
+        result_B_T_H_W_D = _x_fn(
+            x_B_T_H_W_D,
+            self.layer_norm_cross_attn,
+            scale_cross_attn_B_T_1_1_D,
+            shift_cross_attn_B_T_1_1_D,
+            gate_cross_attn_B_T_1_1_D,
+        )
+        x_B_T_H_W_D = result_B_T_H_W_D * gate_cross_attn_B_T_1_1_D + x_B_T_H_W_D
+
+        normalized_x_B_T_H_W_D = _fn(
+            x_B_T_H_W_D,
+            self.layer_norm_mlp,
+            scale_mlp_B_T_1_1_D,
+            shift_mlp_B_T_1_1_D,
+        )
+        result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D)
+        x_B_T_H_W_D = x_B_T_H_W_D + gate_mlp_B_T_1_1_D * result_B_T_H_W_D
+
+        return x_B_T_H_W_D
+
+
+class MultiViewCrossDiT(MinimalV1LVGDiT):
+    def __init__(
+        self,
+        *args,
+        timestep_scale: float = 1.0,
+        crossattn_emb_channels: int = 1024,
+        mlp_ratio: float = 4.0,
+        state_t: int,
+        n_cameras_emb: int,
+        view_condition_dim: int,
+        concat_view_embedding: bool,
+        adaln_view_embedding: bool,
+        layer_mask: Optional[List[bool]] = None,
+        sac_config: MultiViewSACConfig = MultiViewSACConfig(),
+        enable_cross_view_attn: bool = False,
+        cross_view_attn_map_str: Optional[Dict] = None,
+        camera_to_view_id: Optional[Dict] = None,
+        init_cross_view_attn_weight_from: Optional[str] = None,
+        init_cross_view_attn_weight_credentials: Optional[str] = None,
+        **kwargs,
+    ):
+        self.crossattn_emb_channels = crossattn_emb_channels
+        self.mlp_ratio = mlp_ratio
+        self.state_t = state_t
+        self.n_cameras_emb = n_cameras_emb
+        self.view_condition_dim = view_condition_dim
+        self.concat_view_embedding = concat_view_embedding
+        self.adaln_view_embedding = adaln_view_embedding
+        self.enable_cross_view_attn = enable_cross_view_attn
+        self.init_cross_view_attn_weight_from = init_cross_view_attn_weight_from
+        self.init_cross_view_attn_weight_credentials = init_cross_view_attn_weight_credentials
+
+        assert not (self.adaln_view_embedding and self.concat_view_embedding), (
+            "adaln_view_embedding and concat_view_embedding cannot be True at the same time"
+        )
+        assert "in_channels" in kwargs, "in_channels must be provided"
+        kwargs["in_channels"] += (
+            self.view_condition_dim if self.concat_view_embedding else 0
+        )  # this avoids overwritting build_patch_embed which still adds padding_mask channel as appropriate
+        assert layer_mask is None, "layer_mask is not supported for MultiViewDiT"
+        if "n_cameras" in kwargs:
+            del kwargs["n_cameras"]
+        super().__init__(
+            *args,
+            mlp_ratio=mlp_ratio,
+            timestep_scale=timestep_scale,
+            crossattn_emb_channels=crossattn_emb_channels,
+            sac_config=sac_config,
+            **kwargs,
+        )
+
+        cross_view_attn_map = {}
+        for source_view, target_views in cross_view_attn_map_str.items():
+            cross_view_attn_map[int(camera_to_view_id[source_view])] = []
+            for target_view in target_views:
+                cross_view_attn_map[int(camera_to_view_id[source_view])].append(int(camera_to_view_id[target_view]))
+        self.cross_view_attn_map = cross_view_attn_map
+
+        del self.blocks
+        self.blocks = nn.ModuleList(
+            [
+                MultiViewCrossBlock(
+                    x_dim=self.model_channels,
+                    context_dim=self.crossattn_emb_channels,
+                    num_heads=self.num_heads,
+                    mlp_ratio=self.mlp_ratio,
+                    use_adaln_lora=self.use_adaln_lora,
+                    adaln_lora_dim=self.adaln_lora_dim,
+                    backend=self.atten_backend,
+                    image_context_dim=None if self.extra_image_context_dim is None else self.model_channels,
+                    state_t=self.state_t,
+                    use_wan_fp32_strategy=self.use_wan_fp32_strategy,
+                    cross_view_attn_map=self.cross_view_attn_map,
+                    enable_cross_view_attn=self.enable_cross_view_attn,
+                )
+                for _ in range(self.num_blocks)
+            ]
+        )
+
+        if self.concat_view_embedding:
+            self.view_embeddings = nn.Embedding(self.n_cameras_emb, view_condition_dim)
+
+        if self.adaln_view_embedding:
+            self.adaln_view_embedder = nn.Embedding(self.n_cameras_emb, self.model_channels)
+            # cosmos use adaln in self-attn, cross-attn, mlp
+            self.adaln_view_proj = nn.Linear(self.model_channels, self.model_channels * 9)
+
+        self.init_weights()
+        self.enable_selective_checkpoint(sac_config, self.blocks)
+
+    def fully_shard(self, mesh, **fsdp_kwargs):
+        for i, block in enumerate(self.blocks):
+            reshard_after_forward = i < len(self.blocks) - 1
+            fully_shard(block, mesh=mesh, reshard_after_forward=reshard_after_forward, **fsdp_kwargs)
+
+        fully_shard(self.final_layer, mesh=mesh, reshard_after_forward=True, **fsdp_kwargs)
+        if self.extra_per_block_abs_pos_emb:
+            for extra_pos_embedder in self.extra_pos_embedders_options.values():
+                fully_shard(extra_pos_embedder, mesh=mesh, reshard_after_forward=True, **fsdp_kwargs)
+        fully_shard(self.t_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if self.extra_image_context_dim is not None:
+            fully_shard(self.img_context_proj, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+
+        if hasattr(self, "view_embeddings"):
+            fully_shard(self.view_embeddings, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+
+        if hasattr(self, "adaln_view_embedder"):
+            fully_shard(self.adaln_view_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if hasattr(self, "adaln_view_proj"):
+            fully_shard(self.adaln_view_proj, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+
+    def enable_context_parallel(self, process_group: Optional[ProcessGroup] = None):
+        # pos_embedder
+        for pos_embedder in self.pos_embedder_options.values():
+            pos_embedder.enable_context_parallel(process_group=process_group)
+        if self.extra_per_block_abs_pos_emb:
+            for extra_pos_embedder in self.extra_pos_embedders_options.values():
+                extra_pos_embedder.enable_context_parallel(process_group=process_group)
+
+        # attention
+        cp_ranks = get_process_group_ranks(process_group)
+        for block in self.blocks:
+            block.set_context_parallel_group(
+                process_group=process_group,
+                ranks=cp_ranks,
+                stream=torch.cuda.Stream(),
+            )
+
+        self._is_context_parallel_enabled = True
+
+    def disable_context_parallel(self):
+        # pos_embedder
+        for pos_embedder in self.pos_embedder_options.values():
+            pos_embedder.disable_context_parallel()
+        if self.extra_per_block_abs_pos_emb:
+            for extra_pos_embedder in self.extra_pos_embedders_options.values():
+                extra_pos_embedder.disable_context_parallel()
+
+        # attention
+        for block in self.blocks:
+            block.set_context_parallel_group(
+                process_group=None,
+                ranks=None,
+                stream=torch.cuda.Stream(),
+            )
+
+        self._is_context_parallel_enabled = False
+
+    def init_weights(self):
+        self.x_embedder.init_weights()
+        for pos_embedder in self.pos_embedder_options.values():
+            pos_embedder.reset_parameters()
+        if self.extra_per_block_abs_pos_emb:
+            for extra_pos_embedder in self.extra_pos_embedders_options.values():
+                extra_pos_embedder.init_weights()
+
+        self.t_embedder[1].init_weights()
+        for block in self.blocks:
+            block.init_weights()
+
+        self.final_layer.init_weights()
+        self.t_embedding_norm.reset_parameters()
+
+        if self.extra_image_context_dim is not None:
+            self.img_context_proj[0].reset_parameters()
+
+        if hasattr(self, "view_embeddings"):
+            torch.nn.init.normal_(self.view_embeddings.weight, mean=0.0, std=0.02)
+
+        if hasattr(self, "adaln_view_embedder"):
+            torch.nn.init.normal_(self.adaln_view_embedder.weight, mean=0.0, std=0.05)
+
+        if hasattr(self, "adaln_view_proj"):
+            torch.nn.init.zeros_(self.adaln_view_proj.weight)
+            torch.nn.init.zeros_(self.adaln_view_proj.bias)
+
+    def build_pos_embed(self):
+        self.pos_embedder_options = nn.ModuleDict()
+        self.extra_pos_embedders_options = nn.ModuleDict()
+        for n_cameras in range(1, self.n_cameras_emb + 1):
+            pos_embedder, extra_pos_embedder = self.build_pos_embed_for_n_cameras(n_cameras)
+            self.pos_embedder_options[f"n_cameras_{n_cameras}"] = pos_embedder
+            self.extra_pos_embedders_options[f"n_cameras_{n_cameras}"] = extra_pos_embedder
+
+    def build_pos_embed_for_n_cameras(self, n_cameras: int):
+        if self.pos_emb_cls == "rope3d":
+            cls_type = MultiCameraVideoRopePosition3DEmb
+        else:
+            raise ValueError(f"Unknown pos_emb_cls {self.pos_emb_cls}")
+        pos_embedder, extra_pos_embedder = None, None
+        kwargs = dict(
+            model_channels=self.model_channels,
+            len_h=self.max_img_h // self.patch_spatial,
+            len_w=self.max_img_w // self.patch_spatial,
+            len_t=self.max_frames // self.patch_temporal,
+            max_fps=self.max_fps,
+            min_fps=self.min_fps,
+            is_learnable=self.pos_emb_learnable,
+            interpolation=self.pos_emb_interpolation,
+            head_dim=self.model_channels // self.num_heads,
+            h_extrapolation_ratio=self.rope_h_extrapolation_ratio,
+            w_extrapolation_ratio=self.rope_w_extrapolation_ratio,
+            t_extrapolation_ratio=self.rope_t_extrapolation_ratio,
+            enable_fps_modulation=self.rope_enable_fps_modulation,
+            n_cameras=n_cameras,
+        )
+        pos_embedder = cls_type(
+            **kwargs,
+        )
+        assert pos_embedder.enable_fps_modulation == self.rope_enable_fps_modulation, (
+            "enable_fps_modulation must be the same"
+        )
+
+        if self.extra_per_block_abs_pos_emb:
+            raise NotImplementedError("extra_per_block_abs_pos_emb is not tested for multi-view DIT")
+            kwargs["h_extrapolation_ratio"] = self.extra_h_extrapolation_ratio
+            kwargs["w_extrapolation_ratio"] = self.extra_w_extrapolation_ratio
+            kwargs["t_extrapolation_ratio"] = self.extra_t_extrapolation_ratio
+            extra_pos_embedder = MultiCameraSinCosPosEmbAxis(
+                **kwargs,
+            )
+        return pos_embedder, extra_pos_embedder
+
+    def prepare_embedded_sequence(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        view_indices_B_T: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
+        if self.concat_padding_mask:
+            padding_mask = transforms.functional.resize(
+                padding_mask, list(x_B_C_T_H_W.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
+            )
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, padding_mask.unsqueeze(1).repeat(1, 1, x_B_C_T_H_W.shape[2], 1, 1)], dim=1
+            )
+        cp_size = (
+            len(get_process_group_ranks(parallel_state.get_context_parallel_group()))
+            if parallel_state.is_initialized()
+            else 1
+        )
+        n_cameras = (x_B_C_T_H_W.shape[2] * cp_size) // self.state_t
+        pos_embedder = self.pos_embedder_options[f"n_cameras_{n_cameras}"]  # they are all the same if they are rope
+        if self.concat_view_embedding:
+            if view_indices_B_T is None:
+                view_indices = torch.arange(n_cameras).clamp(
+                    max=self.n_cameras_emb - 1
+                )  # View indices [0, 1, ..., V-1]
+                view_indices = view_indices.to(x_B_C_T_H_W.device)
+                view_embedding = self.view_embeddings(view_indices)  # Shape: [V, embedding_dim]
+                view_embedding = rearrange(view_embedding, "V D -> D V")
+                view_embedding = (
+                    view_embedding.unsqueeze(0).unsqueeze(3).unsqueeze(4).unsqueeze(5)
+                )  # Shape: [1, D, V, 1, 1, 1]
+            else:
+                view_indices_B_T = view_indices_B_T.clamp(max=self.n_cameras_emb - 1)
+                view_indices_B_T = view_indices_B_T.to(x_B_C_T_H_W.device).long()
+                view_embedding = self.view_embeddings(view_indices_B_T)  # B, (V T), D
+                view_embedding = rearrange(view_embedding, "B (V T) D -> B D V T", V=n_cameras)
+                view_embedding = view_embedding.unsqueeze(-1).unsqueeze(-1)  # Shape: [B, D, V, T, 1, 1]
+            x_B_C_V_T_H_W = rearrange(x_B_C_T_H_W, "B C (V T) H W -> B C V T H W", V=n_cameras)
+            view_embedding = view_embedding.expand(
+                x_B_C_V_T_H_W.shape[0],
+                view_embedding.shape[1],
+                view_embedding.shape[2],
+                x_B_C_V_T_H_W.shape[3],
+                x_B_C_V_T_H_W.shape[4],
+                x_B_C_V_T_H_W.shape[5],
+            )
+            x_B_C_V_T_H_W = torch.cat([x_B_C_V_T_H_W, view_embedding], dim=1)
+            x_B_C_T_H_W = rearrange(x_B_C_V_T_H_W, " B C V T H W -> B C (V T) H W", V=n_cameras)
+
+        x_B_T_H_W_D = self.x_embedder(x_B_C_T_H_W)
+
+        if self.extra_per_block_abs_pos_emb:
+            extra_pos_embedder = self.extra_pos_embedders_options[str(n_cameras)]
+            extra_pos_emb = extra_pos_embedder(x_B_T_H_W_D, fps=fps)
+        else:
+            extra_pos_emb = None
+
+        if "rope" in self.pos_emb_cls.lower():
+            return x_B_T_H_W_D, pos_embedder(x_B_T_H_W_D, fps=fps), extra_pos_emb
+
+        if "fps_aware" in self.pos_emb_cls:
+            raise NotImplementedError("FPS-aware positional embedding is not supported for multi-view DIT")
+
+        x_B_T_H_W_D = x_B_T_H_W_D + pos_embedder(x_B_T_H_W_D)
+
+        return x_B_T_H_W_D, None, extra_pos_emb
+
+    def forward(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        timesteps_B_T: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        data_type: Optional[DataType] = DataType.VIDEO,
+        view_indices_B_T: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor | List[torch.Tensor] | Tuple[torch.Tensor, List[torch.Tensor]]:
+        # Deletes elements like condition.use_video_condition that are not used in the forward pass
+        del kwargs
+        if data_type == DataType.VIDEO:
+            x_B_C_T_H_W = torch.cat([x_B_C_T_H_W, condition_video_input_mask_B_C_T_H_W.type_as(x_B_C_T_H_W)], dim=1)
+        else:
+            B, _, T, H, W = x_B_C_T_H_W.shape
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, torch.zeros((B, 1, T, H, W), dtype=x_B_C_T_H_W.dtype, device=x_B_C_T_H_W.device)], dim=1
+            )
+
+        assert isinstance(data_type, DataType), (
+            f"Expected DataType, got {type(data_type)}. We need discuss this flag later."
+        )
+        timesteps_B_T = timesteps_B_T * self.timestep_scale
+        x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = self.prepare_embedded_sequence(
+            x_B_C_T_H_W,
+            fps=fps,
+            padding_mask=padding_mask,
+            view_indices_B_T=view_indices_B_T,
+        )
+        if self.use_crossattn_projection:
+            crossattn_emb = self.crossattn_proj(crossattn_emb)
+        with amp.autocast("cuda", enabled=self.use_wan_fp32_strategy, dtype=torch.float32):
+            # (B, 1). input timesteps are (b, 1)
+            if timesteps_B_T.ndim == 1:
+                timesteps_B_T = timesteps_B_T.unsqueeze(1)
+            t_embedding_B_T_D, adaln_lora_B_T_3D = self.t_embedder(timesteps_B_T)
+            t_embedding_B_T_D = self.t_embedding_norm(t_embedding_B_T_D)
+
+        if self.adaln_view_embedding:
+            num_cameras = torch.unique(view_indices_B_T[0]).shape[0]
+            with amp.autocast("cuda", enabled=self.use_wan_fp32_strategy, dtype=torch.float32):
+                view_indices_B_V_T = rearrange(view_indices_B_T, "b (v t) -> b v t", v=num_cameras)
+                view_embedding_B_V = self.adaln_view_embedder(view_indices_B_V_T[..., 0])  # B, V, D
+                view_embedding_proj_B_V_9D = self.adaln_view_proj(view_embedding_B_V)  # B, V, 9D
+        else:
+            view_embedding_proj_B_V_9D = None
+
+        # for logging purpose
+        affline_scale_log_info = {}
+        affline_scale_log_info["t_embedding_B_T_D"] = t_embedding_B_T_D.detach()
+        self.affline_scale_log_info = affline_scale_log_info
+        self.affline_emb = t_embedding_B_T_D
+        self.crossattn_emb = crossattn_emb
+
+        if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
+            assert x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape, (
+                f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}"
+            )
+
+        B, T, H, W, D = x_B_T_H_W_D.shape
+
+        for block_idx, block in enumerate(self.blocks):
+            x_B_T_H_W_D = block(
+                x_B_T_H_W_D,
+                view_indices_B_T,
+                t_embedding_B_T_D,
+                view_embedding_proj_B_V_9D,
+                crossattn_emb,
+                rope_emb_L_1_1_D=rope_emb_L_1_1_D,
+                adaln_lora_B_T_3D=adaln_lora_B_T_3D,
+                extra_per_block_pos_emb=extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
+                block_idx=block_idx,
+            )
+
+        x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
+        x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O)
+
+        return x_B_C_Tt_Hp_Wp
+
+    def init_cross_view_attn_with_self_attn_weights(self, is_ema: bool = False) -> None:
+        """Load self-attention weights from base model checkpoint and initialize cross-view attention."""
+        # Check initialization conditions
+        if self.init_cross_view_attn_weight_from is None:
+            log.info("No checkpoint path provided, skipping cross-view attention initialization")
+            return
+
+        if not self.enable_cross_view_attn:
+            log.info("Cross-view attention not enabled, skipping weight loading")
+            return
+
+        log.critical(
+            f"Loading base model from {self.init_cross_view_attn_weight_from} for cross-view attention initialization"
+        )
+
+        # Import necessary modules
+        import gc
+
+        import torch.distributed.checkpoint as dcp
+        from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
+
+        from cosmos_predict2._src.imaginaire.checkpointer.s3_filesystem import S3StorageReader
+
+        # Prepare checkpoint loading
+        checkpoint_path = os.path.join("s3://bucket/" + self.init_cross_view_attn_weight_from, "model")
+        storage_reader = S3StorageReader(
+            credential_path=self.init_cross_view_attn_weight_credentials,
+            path=checkpoint_path,
+        )
+
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+        # Build minimal state dict - only load required weights
+        log.info("Building minimal state dict (only weights needed for cross-view attention)")
+        minimal_state_dict = self._build_minimal_state_dict(is_ema)
+
+        # Load weights from checkpoint
+        log.info(f"Loading {len(minimal_state_dict)} weight tensors from checkpoint")
+        dcp.load(minimal_state_dict, storage_reader=storage_reader, planner=DefaultLoadPlanner(allow_partial_load=True))
+
+        # Verify that loaded weights maintain correct sharding (if FSDP is used)
+        if torch.distributed.is_initialized():
+            self._verify_loaded_checkpoint_sharding(minimal_state_dict, is_ema)
+
+        # Initialize cross-view attention weights
+        log.info("Starting to initialize cross-view attention from self-attention weights")
+        initialized_count, weight_stats = self._copy_weights_to_cross_view_attn(minimal_state_dict, is_ema)
+
+        log.info(f"Successfully initialized cross-view attention for {initialized_count}/{len(self.blocks)} layers")
+
+        # Release memory
+        log.info("Releasing checkpoint memory")
+        del minimal_state_dict
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+        # Verify weight loading
+        self.weight_stats_before, self.weight_stats_after = weight_stats
+        self._verify_cross_view_attn_weights_loaded()
+
+    def _build_minimal_state_dict(self, is_ema: bool) -> dict[str, torch.Tensor]:
+        """Build minimal state dict containing only required weights.
+
+        Note: When using FSDP, torch.empty_like preserves DTensor sharding metadata,
+        allowing dcp.load to correctly load only the local shard for each rank.
+        """
+        minimal_state_dict = {}
+        num_layers = len(self.blocks)
+
+        # Determine parameter names to load
+        param_names = ["q_proj.weight", "k_proj.weight", "v_proj.weight"]
+        if hasattr(self.blocks[0], "cross_view_attn"):
+            cross_view_attn = self.blocks[0].cross_view_attn
+            if hasattr(cross_view_attn, "q_norm") and hasattr(cross_view_attn.q_norm, "weight"):
+                param_names.extend(["q_norm.weight", "k_norm.weight"])
+
+        # Create placeholder for each parameter in each layer
+        prefix = f"net{'_ema' if is_ema else ''}"
+        for layer_idx in range(num_layers):
+            for param_name in param_names:
+                key = f"{prefix}.blocks.{layer_idx}.self_attn.{param_name}"
+
+                # Get target parameter to determine shape and dtype
+                # IMPORTANT: This preserves DTensor sharding metadata if model is FSDP-wrapped
+                target_param = self._get_nested_attr(self.blocks[layer_idx].cross_view_attn, param_name.split("."))
+
+                if target_param is not None:
+                    # empty_like preserves DTensor sharding spec, which tells dcp.load
+                    # which shard to load for the current rank
+                    placeholder = torch.empty_like(target_param)
+                    minimal_state_dict[key] = placeholder
+
+                    # Log sharding info for first layer to verify FSDP setup
+                    if layer_idx == 0:
+                        from torch.distributed._tensor.api import DTensor
+
+                        if isinstance(placeholder, DTensor):
+                            log.info(f"Parameter {param_name} is DTensor with placement: {placeholder.placements}")
+                        else:
+                            log.info(f"Parameter {param_name} is regular tensor (not sharded)")
+                else:
+                    log.warning(f"Parameter {param_name} does not exist in layer {layer_idx} cross_view_attn")
+
+        return minimal_state_dict
+
+    def _get_nested_attr(self, obj, attr_path: list[str]):
+        """Recursively get nested attribute."""
+        try:
+            for attr in attr_path:
+                obj = getattr(obj, attr)
+            return obj
+        except AttributeError:
+            return None
+
+    def _verify_loaded_checkpoint_sharding(self, state_dict: dict[str, torch.Tensor], is_ema: bool) -> None:
+        """Verify that checkpoint was loaded with correct FSDP sharding.
+
+        This checks that:
+        1. If local model uses DTensor, loaded weights are also DTensor with matching sharding
+        2. The loaded shard size matches what we expect for the current rank
+        """
+        from torch.distributed._tensor.api import DTensor
+
+        rank = torch.distributed.get_rank()
+        world_size = torch.distributed.get_world_size()
+
+        # Check a sample parameter from the first layer
+        prefix = f"net{'_ema' if is_ema else ''}"
+        sample_key = f"{prefix}.blocks.0.self_attn.q_proj.weight"
+
+        if sample_key in state_dict:
+            loaded_param = state_dict[sample_key]
+            target_param = self.blocks[0].cross_view_attn.q_proj.weight
+
+            loaded_is_dtensor = isinstance(loaded_param, DTensor)
+            target_is_dtensor = isinstance(target_param, DTensor)
+
+            if target_is_dtensor and not loaded_is_dtensor:
+                log.warning(
+                    f"[Rank {rank}] Mismatch: Local model uses DTensor (FSDP sharded), "
+                    f"but checkpoint loaded as regular tensor. This may cause OOM or incorrect behavior."
+                )
+            elif not target_is_dtensor and loaded_is_dtensor:
+                log.warning(
+                    f"[Rank {rank}] Mismatch: Local model uses regular tensor, "
+                    f"but checkpoint loaded as DTensor. This is unusual."
+                )
+            elif target_is_dtensor and loaded_is_dtensor:
+                # Both are DTensor - verify sharding matches
+                if loaded_param.placements != target_param.placements:
+                    log.warning(
+                        f"[Rank {rank}] DTensor placement mismatch:\n"
+                        f"  Loaded: {loaded_param.placements}\n"
+                        f"  Target: {target_param.placements}\n"
+                        f"This may cause incorrect weight copying."
+                    )
+                else:
+                    log.info(
+                        f"[Rank {rank}] ✅ Checkpoint sharding verified: "
+                        f"DTensor with placements {loaded_param.placements}"
+                    )
+
+                # Check local shard size
+                loaded_local = loaded_param.to_local()
+                target_local = target_param.to_local()
+                log.info(
+                    f"[Rank {rank}] Local shard shape - Loaded: {loaded_local.shape}, Target: {target_local.shape}"
+                )
+            else:
+                # Both are regular tensors
+                log.info(f"[Rank {rank}] Both checkpoint and model use regular tensors (no FSDP sharding)")
+
+    def _copy_weights_to_cross_view_attn(
+        self, state_dict: dict[str, torch.Tensor], is_ema: bool
+    ) -> tuple[int, tuple[dict, dict]]:
+        """Copy loaded weights to cross-view attention layers and record statistics."""
+        from torch.distributed._tensor.api import DTensor
+
+        num_layers = len(self.blocks)
+        initialized_count = 0
+        weight_stats_before = {}
+        weight_stats_after = {}
+        prefix = f"net{'_ema' if is_ema else ''}"
+
+        for layer_idx in range(num_layers):
+            block = self.blocks[layer_idx]
+            cross_view_attn = block.cross_view_attn
+
+            # Determine parameters to copy for current layer
+            param_names = ["q_proj.weight", "k_proj.weight", "v_proj.weight"]
+            if hasattr(cross_view_attn, "q_norm") and hasattr(cross_view_attn.q_norm, "weight"):
+                param_names.extend(["q_norm.weight", "k_norm.weight"])
+
+            copied_params = []
+            should_record_stats = layer_idx % 5 == 0  # Sample every 5 layers to record statistics
+
+            for param_name in param_names:
+                key = f"{prefix}.blocks.{layer_idx}.self_attn.{param_name}"
+
+                if key not in state_dict:
+                    log.warning(f"Key {key} not found in checkpoint, skipping")
+                    continue
+
+                # Get target parameter
+                target_param = self._get_nested_attr(cross_view_attn, param_name.split("."))
+
+                if not isinstance(target_param, torch.nn.Parameter):
+                    log.warning(f"Parameter {param_name} in block {layer_idx} is not a Parameter object")
+                    continue
+
+                # Record statistics before copy (sampled)
+                if should_record_stats and param_name == "q_proj.weight":
+                    before_local = target_param.to_local() if isinstance(target_param, DTensor) else target_param
+                    weight_stats_before[layer_idx] = {
+                        "mean": before_local.mean().item(),
+                        "std": before_local.std().item(),
+                        "abs_max": before_local.abs().max().item(),
+                        "dtype": str(before_local.dtype),
+                    }
+
+                # Copy weights (ensure dtype consistency)
+                source_weight = state_dict[key]
+                if source_weight.dtype != target_param.dtype:
+                    log.warning(
+                        f"Layer {layer_idx} {param_name}: dtype mismatch ({source_weight.dtype} -> {target_param.dtype}), converting"
+                    )
+                    source_weight = source_weight.to(dtype=target_param.dtype)
+
+                with torch.no_grad():
+                    target_param.copy_(source_weight)
+
+                # Record statistics after copy (sampled)
+                if should_record_stats and param_name == "q_proj.weight":
+                    after_local = target_param.to_local() if isinstance(target_param, DTensor) else target_param
+                    weight_stats_after[layer_idx] = {
+                        "mean": after_local.mean().item(),
+                        "std": after_local.std().item(),
+                        "abs_max": after_local.abs().max().item(),
+                        "dtype": str(after_local.dtype),
+                    }
+
+                copied_params.append(param_name)
+
+            if copied_params:
+                initialized_count += 1
+                # Print detailed info only every 10 layers to reduce log noise
+                if layer_idx % 10 == 0:
+                    log.info(f"Initialized cross_view_attn for block {layer_idx}, copied parameters: {copied_params}")
+
+        return initialized_count, (weight_stats_before, weight_stats_after)
+
+    def _verify_cross_view_attn_weights_loaded(self) -> None:
+        """Verify cross-view attention weights are correctly loaded by comparing statistics before and after copy."""
+        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+
+        num_updated = 0
+        num_failed = 0
+        dtype_mismatches = []
+
+        # Check weight changes and dtype consistency for sampled layers
+        for layer_idx in sorted(self.weight_stats_before.keys()):
+            if layer_idx not in self.weight_stats_after:
+                num_failed += 1
+                continue
+
+            before = self.weight_stats_before[layer_idx]
+            after = self.weight_stats_after[layer_idx]
+
+            # Verify dtype consistency
+            if before["dtype"] != after["dtype"]:
+                dtype_mismatches.append(f"Layer {layer_idx}: {before['dtype']} -> {after['dtype']}")
+                num_failed += 1
+                continue
+
+            # Calculate weight changes (threshold: 1e-4)
+            mean_change = abs(after["mean"] - before["mean"])
+            std_change = abs(after["std"] - before["std"])
+            max_change = abs(after["abs_max"] - before["abs_max"])
+
+            if mean_change > 1e-4 or std_change > 1e-4 or max_change > 1e-4:
+                num_updated += 1
+            else:
+                num_failed += 1
+
+        # Output verification results
+        total_sampled = len(self.weight_stats_before)
+        if num_failed == 0:
+            log.info(
+                f"[Rank {rank}] ✅ Weight verification passed: {num_updated}/{total_sampled} sampled layers successfully updated",
+                rank0_only=False,
+            )
+            if total_sampled > 0:
+                first_layer_dtype = self.weight_stats_after[list(self.weight_stats_after.keys())[0]]["dtype"]
+                log.info(
+                    f"[Rank {rank}] ✅ All weights maintain consistent dtype: {first_layer_dtype}", rank0_only=False
+                )
+        else:
+            log.warning(
+                f"[Rank {rank}] ⚠️ Weight verification: {num_updated}/{total_sampled} layers updated, {num_failed} layers may have failed",
+                rank0_only=False,
+            )
+            if dtype_mismatches:
+                log.error(
+                    f"[Rank {rank}] ❌ Detected dtype mismatches:\n" + "\n".join(dtype_mismatches), rank0_only=False
+                )
+
+        # Clean up temporary attributes
+        delattr(self, "weight_stats_before")
+        delattr(self, "weight_stats_after")

cosmos_predict2/_src/predict2_multiview/networks/multiview_dit.py

@@ -0,0 +1,618 @@
+# 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.
+
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.amp as amp
+import torch.nn as nn
+from einops import rearrange, repeat
+from megatron.core import parallel_state
+from torch.distributed import ProcessGroup, get_process_group_ranks
+from torch.distributed._composable.fsdp import fully_shard
+from torchvision import transforms
+
+from cosmos_predict2._src.imaginaire.utils.context_parallel import split_inputs_cp
+from cosmos_predict2._src.predict2.conditioner import DataType
+from cosmos_predict2._src.predict2.networks.minimal_v1_lvg_dit import MinimalV1LVGDiT
+from cosmos_predict2._src.predict2.networks.minimal_v4_dit import (
+    Attention,
+    Block,
+    SACConfig,
+    VideoPositionEmb,
+    VideoRopePosition3DEmb,
+)
+
+
+class MultiViewCrossAttention(Attention):
+    def __init__(self, *args, state_t: int = None, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        assert self.qkv_format == "bshd", "MultiViewCrossAttention only supports qkv_format='bshd'"
+        self.state_t = state_t
+
+    def forward(self, x, context=None, rope_emb=None):
+        assert not self.is_selfattn, "MultiViewCrossAttention does not support self-attention"
+        B, L, D = x.shape
+
+        n_cameras = context.shape[1] // 512
+        x_B_L_D = rearrange(x, "B (V L) D -> (V B) L D", V=n_cameras)
+        context_B_M_D = rearrange(context, "B (V M) D -> (V B) M D", V=n_cameras) if context is not None else None
+        x_B_L_D = super().forward(x_B_L_D, context_B_M_D, rope_emb=rope_emb)
+        x_B_L_D = rearrange(x_B_L_D, "(V B) L D -> B (V L) D", V=n_cameras)
+        return x_B_L_D
+
+
+class MultiViewBlock(Block):
+    """
+    A transformer block that takes n_cameras as input. This block
+    """
+
+    def __init__(
+        self,
+        x_dim: int,
+        context_dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        use_adaln_lora: bool = False,
+        adaln_lora_dim: int = 256,
+        state_t: int = None,
+        backend: str = "transformer_engine",
+        image_context_dim: Optional[int] = None,
+        use_wan_fp32_strategy: bool = False,
+    ):
+        super().__init__(
+            x_dim,
+            context_dim,
+            num_heads,
+            mlp_ratio,
+            use_adaln_lora,
+            adaln_lora_dim,
+            backend,
+            image_context_dim,
+            use_wan_fp32_strategy,
+        )
+        self.state_t = state_t
+        if image_context_dim is None:
+            del self.cross_attn
+            self.cross_attn = MultiViewCrossAttention(
+                x_dim,
+                context_dim,
+                num_heads,
+                x_dim // num_heads,
+                qkv_format="bshd",
+                state_t=state_t,
+                use_wan_fp32_strategy=use_wan_fp32_strategy,
+            )
+        else:
+            raise NotImplementedError("image_context_dim is not supported for MultiViewBlock")
+
+
+class MultiCameraVideoRopePosition3DEmb(VideoRopePosition3DEmb):
+    def __init__(self, *args, n_cameras: int = 1, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.n_cameras = n_cameras
+
+    def generate_embeddings(
+        self,
+        B_T_H_W_C: torch.Size,
+        fps: Optional[torch.Tensor] = None,
+        h_ntk_factor: Optional[float] = None,
+        w_ntk_factor: Optional[float] = None,
+        t_ntk_factor: Optional[float] = None,
+    ):
+        B, T, H, W, C = B_T_H_W_C
+        single_camera_B_T_H_W_C = (B, T // self.n_cameras, H, W, C)
+        em_T_H_W_D = []
+        for _ in range(self.n_cameras):
+            em_L_1_1_D = super().generate_embeddings(
+                single_camera_B_T_H_W_C,
+                fps=fps,
+                h_ntk_factor=h_ntk_factor,
+                w_ntk_factor=w_ntk_factor,
+                t_ntk_factor=t_ntk_factor,
+            )
+            em_T_H_W_D.append(rearrange(em_L_1_1_D, "(t h w) 1 1 d -> t h w d", t=T // self.n_cameras, h=H, w=W))
+        em_T_H_W_D = torch.cat(em_T_H_W_D, dim=0)
+        return em_T_H_W_D.float()
+
+    def generate_embeddings_with_refs(
+        self,
+        B_Te_H_W_C: torch.Size,
+        num_ref: int,
+        ref_positions: List[int],
+        fps: Optional[torch.Tensor] = None,
+        h_ntk_factor: Optional[float] = None,
+        w_ntk_factor: Optional[float] = None,
+        t_ntk_factor: Optional[float] = None,
+    ):
+        """RoPE for a per-view grid of (state_t real + num_ref reference) frames.
+
+        Per view the temporal positions are ``[0 .. state_t-1]`` for the real frames followed by the fixed
+        ``ref_positions`` for the appended reference frames; views are concatenated on the temporal axis
+        (view-major), matching the token layout ``[view0 real, view0 ref, view1 real, view1 ref, ...]``.
+        Returns ``(V*(state_t+num_ref), H, W, D)`` (un-split; caller applies ``_split_for_context_parallel``).
+        """
+        B, Te, H, W, C = B_Te_H_W_C
+        frames_per_view = Te // self.n_cameras
+        state_t = frames_per_view - num_ref
+        assert state_t >= 0 and len(ref_positions) == num_ref, (state_t, num_ref, ref_positions)
+        temporal_positions = torch.tensor(
+            list(range(state_t)) + list(ref_positions), dtype=torch.float
+        )  # (frames_per_view,)
+        single_camera_B_T_H_W_C = (B, frames_per_view, H, W, C)
+        em_T_H_W_D = []
+        for _ in range(self.n_cameras):
+            em_L_1_1_D = super().generate_embeddings(
+                single_camera_B_T_H_W_C,
+                fps=fps,
+                h_ntk_factor=h_ntk_factor,
+                w_ntk_factor=w_ntk_factor,
+                t_ntk_factor=t_ntk_factor,
+                temporal_positions=temporal_positions,
+            )
+            em_T_H_W_D.append(
+                rearrange(em_L_1_1_D, "(t h w) 1 1 d -> t h w d", t=frames_per_view, h=H, w=W)
+            )
+        em_T_H_W_D = torch.cat(em_T_H_W_D, dim=0)
+        return em_T_H_W_D.float()
+
+    @property
+    def seq_dim(self):
+        return 1
+
+    def _split_for_context_parallel(self, embeddings):
+        if self._cp_group is not None:
+            embeddings = rearrange(embeddings, "(V T) H W D -> V (T H W) 1 1 D", V=self.n_cameras)
+            embeddings = split_inputs_cp(x=embeddings, seq_dim=self.seq_dim, cp_group=self._cp_group)
+            embeddings = rearrange(embeddings, "V T 1 1 D -> (V T) 1 1 D", V=self.n_cameras)
+        else:
+            embeddings = rearrange(embeddings, "t h w d -> (t h w) 1 1 d")
+        return embeddings
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+class MultiCameraSinCosPosEmbAxis(VideoPositionEmb):
+    def __init__(
+        self,
+        *,  # enforce keyword arguments
+        interpolation: str,
+        model_channels: int,
+        len_h: int,
+        len_w: int,
+        len_t: int,
+        h_extrapolation_ratio: float = 1.0,
+        w_extrapolation_ratio: float = 1.0,
+        t_extrapolation_ratio: float = 1.0,
+        n_cameras: int = 4,
+        **kwargs,
+    ):
+        """
+        Args:
+            interpolation (str): we curretly only support "crop", ideally when we need extrapolation capacity, we should adjust frequency or other more advanced methods. they are not implemented yet.
+        """
+        del kwargs  # unused
+        self.n_cameras = n_cameras
+        super().__init__()
+        self.interpolation = interpolation
+        assert self.interpolation in ["crop"], f"Unknown interpolation method {self.interpolation}"
+        self.model_channels = model_channels
+        self.len_h = len_h
+        self.len_w = len_w
+        self.len_t = len_t
+        self.h_extrapolation_ratio = h_extrapolation_ratio
+        self.w_extrapolation_ratio = w_extrapolation_ratio
+        self.t_extrapolation_ratio = t_extrapolation_ratio
+
+        emb_h, emb_w, emb_t = self.compute_1d_pos_embeddings()
+
+        self.register_buffer("pos_emb_h", emb_h, persistent=False)
+        self.register_buffer("pos_emb_w", emb_w, persistent=False)
+        self.register_buffer("pos_emb_t", emb_t, persistent=False)
+
+    def compute_1d_pos_embeddings(self):
+        """
+        Compute 1d pos embed for each axis
+        """
+        dim = self.model_channels
+        dim_h = dim // 6 * 2
+        dim_w = dim_h
+        dim_t = dim - 2 * dim_h
+        assert dim == dim_h + dim_w + dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}"
+
+        emb_h = get_1d_sincos_pos_embed_from_grid(dim_h, pos=np.arange(self.len_h) * 1.0 / self.h_extrapolation_ratio)
+        emb_w = get_1d_sincos_pos_embed_from_grid(dim_w, pos=np.arange(self.len_w) * 1.0 / self.w_extrapolation_ratio)
+        emb_t = get_1d_sincos_pos_embed_from_grid(dim_t, pos=np.arange(self.len_t) * 1.0 / self.t_extrapolation_ratio)
+
+        emb_h = torch.from_numpy(emb_h).float()
+        emb_w = torch.from_numpy(emb_w).float()
+        emb_t = torch.from_numpy(emb_t).float()
+
+        return emb_h, emb_w, emb_t
+
+    def reset_parameters(self):
+        emb_h, emb_w, emb_t = self.compute_1d_pos_embeddings()
+        self.pos_emb_h = emb_h
+        self.pos_emb_w = emb_w
+        self.pos_emb_t = emb_t
+
+    def generate_embeddings(self, B_T_H_W_C: torch.Size, fps=Optional[torch.Tensor]) -> torch.Tensor:
+        B, T, H, W, C = B_T_H_W_C
+
+        single_camera_T = T // self.n_cameras
+
+        if self.interpolation == "crop":
+            emb_h_H = self.pos_emb_h[:H]
+            emb_w_W = self.pos_emb_w[:W]
+            emb_t_T = self.pos_emb_t[:single_camera_T]
+            emb = torch.cat(
+                [
+                    torch.cat(
+                        [
+                            repeat(emb_t_T, "t d-> b t h w d", b=B, h=H, w=W),
+                            repeat(emb_h_H, "h d-> b t h w d", b=B, t=single_camera_T, w=W),
+                            repeat(emb_w_W, "w d-> b t h w d", b=B, t=single_camera_T, h=H),
+                        ],
+                        dim=-1,
+                    )
+                    for _ in range(self.n_cameras)
+                ],
+                1,
+            )
+            assert list(emb.shape)[:4] == [B, T, H, W], f"bad shape: {list(emb.shape)[:4]} != {B, T, H, W}"
+            return emb
+
+    @property
+    def seq_dim(self):
+        return 1
+
+    def _split_for_context_parallel(self, embeddings):
+        if self._cp_group is not None:
+            embeddings = rearrange(embeddings, "B (V T) H W C -> (B V) T H W C", V=self.n_cameras)
+            embeddings = split_inputs_cp(x=embeddings, seq_dim=self.seq_dim, cp_group=self._cp_group)
+            embeddings = rearrange(embeddings, "(B V) T H W C -> B (V T) H W C", V=self.n_cameras)
+        return embeddings
+
+
+class MultiViewDiT(MinimalV1LVGDiT):
+    def __init__(
+        self,
+        *args,
+        timestep_scale: float = 1.0,
+        crossattn_emb_channels: int = 1024,
+        mlp_ratio: float = 4.0,
+        state_t: int,
+        n_cameras_emb: int,
+        view_condition_dim: int,
+        concat_view_embedding: bool,
+        layer_mask: Optional[List[bool]] = None,
+        sac_config: SACConfig = SACConfig(),
+        **kwargs,
+    ):
+        self.state_t = state_t
+        self.n_cameras_emb = n_cameras_emb
+        self.view_condition_dim = view_condition_dim
+        self.concat_view_embedding = concat_view_embedding
+        assert "in_channels" in kwargs, "in_channels must be provided"
+        kwargs["in_channels"] += (
+            self.view_condition_dim if self.concat_view_embedding else 0
+        )  # this avoids overwritting build_patch_embed which still adds padding_mask channel as appropriate
+        assert layer_mask is None, "layer_mask is not supported for MultiViewDiT"
+        if "n_cameras" in kwargs:
+            del kwargs["n_cameras"]
+        super().__init__(
+            *args,
+            mlp_ratio=mlp_ratio,
+            timestep_scale=timestep_scale,
+            crossattn_emb_channels=crossattn_emb_channels,
+            sac_config=sac_config,
+            **kwargs,
+        )
+        del self.blocks
+        self.blocks = nn.ModuleList(
+            [
+                MultiViewBlock(
+                    x_dim=self.model_channels,
+                    context_dim=crossattn_emb_channels,
+                    num_heads=self.num_heads,
+                    mlp_ratio=mlp_ratio,
+                    use_adaln_lora=self.use_adaln_lora,
+                    adaln_lora_dim=self.adaln_lora_dim,
+                    backend=self.atten_backend,
+                    image_context_dim=None if self.extra_image_context_dim is None else self.model_channels,
+                    state_t=self.state_t,
+                    use_wan_fp32_strategy=self.use_wan_fp32_strategy,
+                )
+                for _ in range(self.num_blocks)
+            ]
+        )
+
+        if self.concat_view_embedding:
+            self.view_embeddings = nn.Embedding(self.n_cameras_emb, view_condition_dim)
+
+        self.init_weights()
+        self.enable_selective_checkpoint(sac_config, self.blocks)
+
+    def fully_shard(self, mesh, **fsdp_kwargs):
+        for i, block in enumerate(self.blocks):
+            reshard_after_forward = i < len(self.blocks) - 1
+            fully_shard(block, mesh=mesh, reshard_after_forward=reshard_after_forward, **fsdp_kwargs)
+
+        fully_shard(self.final_layer, mesh=mesh, reshard_after_forward=True, **fsdp_kwargs)
+        if self.extra_per_block_abs_pos_emb:
+            for extra_pos_embedder in self.extra_pos_embedders_options.values():
+                fully_shard(extra_pos_embedder, mesh=mesh, reshard_after_forward=True, **fsdp_kwargs)
+        fully_shard(self.t_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if self.extra_image_context_dim is not None:
+            fully_shard(self.img_context_proj, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+
+    def enable_context_parallel(self, process_group: Optional[ProcessGroup] = None):
+        # pos_embedder
+        for pos_embedder in self.pos_embedder_options.values():
+            pos_embedder.enable_context_parallel(process_group=process_group)
+        if self.extra_per_block_abs_pos_emb:
+            for extra_pos_embedder in self.extra_pos_embedders_options.values():
+                extra_pos_embedder.enable_context_parallel(process_group=process_group)
+
+        # attention
+        cp_ranks = get_process_group_ranks(process_group)
+        for block in self.blocks:
+            block.set_context_parallel_group(
+                process_group=process_group,
+                ranks=cp_ranks,
+                stream=torch.cuda.Stream(),
+            )
+
+        self._is_context_parallel_enabled = True
+
+    def disable_context_parallel(self):
+        # pos_embedder
+        for pos_embedder in self.pos_embedder_options.values():
+            pos_embedder.disable_context_parallel()
+        if self.extra_per_block_abs_pos_emb:
+            for extra_pos_embedder in self.extra_pos_embedders_options.values():
+                extra_pos_embedder.disable_context_parallel()
+
+        # attention
+        for block in self.blocks:
+            block.set_context_parallel_group(
+                process_group=None,
+                ranks=None,
+                stream=torch.cuda.Stream(),
+            )
+
+        self._is_context_parallel_enabled = False
+
+    def init_weights(self):
+        self.x_embedder.init_weights()
+        for pos_embedder in self.pos_embedder_options.values():
+            pos_embedder.reset_parameters()
+        if self.extra_per_block_abs_pos_emb:
+            for extra_pos_embedder in self.extra_pos_embedders_options.values():
+                extra_pos_embedder.init_weights()
+
+        self.t_embedder[1].init_weights()
+        for block in self.blocks:
+            block.init_weights()
+
+        self.final_layer.init_weights()
+        self.t_embedding_norm.reset_parameters()
+
+        if self.extra_image_context_dim is not None:
+            self.img_context_proj[0].reset_parameters()
+
+    def build_pos_embed(self):
+        self.pos_embedder_options = nn.ModuleDict()
+        self.extra_pos_embedders_options = nn.ModuleDict()
+        for n_cameras in range(1, self.n_cameras_emb + 1):
+            pos_embedder, extra_pos_embedder = self.build_pos_embed_for_n_cameras(n_cameras)
+            self.pos_embedder_options[f"n_cameras_{n_cameras}"] = pos_embedder
+            self.extra_pos_embedders_options[f"n_cameras_{n_cameras}"] = extra_pos_embedder
+
+    def build_pos_embed_for_n_cameras(self, n_cameras: int):
+        if self.pos_emb_cls == "rope3d":
+            cls_type = MultiCameraVideoRopePosition3DEmb
+        else:
+            raise ValueError(f"Unknown pos_emb_cls {self.pos_emb_cls}")
+        pos_embedder, extra_pos_embedder = None, None
+        kwargs = dict(
+            model_channels=self.model_channels,
+            len_h=self.max_img_h // self.patch_spatial,
+            len_w=self.max_img_w // self.patch_spatial,
+            len_t=self.max_frames // self.patch_temporal,
+            max_fps=self.max_fps,
+            min_fps=self.min_fps,
+            is_learnable=self.pos_emb_learnable,
+            interpolation=self.pos_emb_interpolation,
+            head_dim=self.model_channels // self.num_heads,
+            h_extrapolation_ratio=self.rope_h_extrapolation_ratio,
+            w_extrapolation_ratio=self.rope_w_extrapolation_ratio,
+            t_extrapolation_ratio=self.rope_t_extrapolation_ratio,
+            enable_fps_modulation=self.rope_enable_fps_modulation,
+            n_cameras=n_cameras,
+        )
+        pos_embedder = cls_type(
+            **kwargs,
+        )
+        assert pos_embedder.enable_fps_modulation == self.rope_enable_fps_modulation, (
+            "enable_fps_modulation must be the same"
+        )
+
+        if self.extra_per_block_abs_pos_emb:
+            raise NotImplementedError("extra_per_block_abs_pos_emb is not tested for multi-view DIT")
+            kwargs["h_extrapolation_ratio"] = self.extra_h_extrapolation_ratio
+            kwargs["w_extrapolation_ratio"] = self.extra_w_extrapolation_ratio
+            kwargs["t_extrapolation_ratio"] = self.extra_t_extrapolation_ratio
+            extra_pos_embedder = MultiCameraSinCosPosEmbAxis(
+                **kwargs,
+            )
+        return pos_embedder, extra_pos_embedder
+
+    def prepare_embedded_sequence(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        view_indices_B_T: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
+        if self.concat_padding_mask:
+            padding_mask = transforms.functional.resize(
+                padding_mask, list(x_B_C_T_H_W.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
+            )
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, padding_mask.unsqueeze(1).repeat(1, 1, x_B_C_T_H_W.shape[2], 1, 1)], dim=1
+            )
+        cp_size = (
+            len(get_process_group_ranks(parallel_state.get_context_parallel_group()))
+            if parallel_state.is_initialized()
+            else 1
+        )
+        n_cameras = (x_B_C_T_H_W.shape[2] * cp_size) // self.state_t
+        pos_embedder = self.pos_embedder_options[f"n_cameras_{n_cameras}"]
+        if self.concat_view_embedding:
+            if view_indices_B_T is None:
+                view_indices = torch.arange(n_cameras).clamp(
+                    max=self.n_cameras_emb - 1
+                )  # View indices [0, 1, ..., V-1]
+                view_indices = view_indices.to(x_B_C_T_H_W.device)
+                view_embedding = self.view_embeddings(view_indices)  # Shape: [V, embedding_dim]
+                view_embedding = rearrange(view_embedding, "V D -> D V")
+                view_embedding = (
+                    view_embedding.unsqueeze(0).unsqueeze(3).unsqueeze(4).unsqueeze(5)
+                )  # Shape: [1, D, V, 1, 1, 1]
+            else:
+                view_indices_B_T = view_indices_B_T.clamp(max=self.n_cameras_emb - 1)
+                view_indices_B_T = view_indices_B_T.to(x_B_C_T_H_W.device).long()
+                view_embedding = self.view_embeddings(view_indices_B_T)  # B, (V T), D
+                view_embedding = rearrange(view_embedding, "B (V T) D -> B D V T", V=n_cameras)
+                view_embedding = view_embedding.unsqueeze(-1).unsqueeze(-1)  # Shape: [B, D, V, T, 1, 1]
+            x_B_C_V_T_H_W = rearrange(x_B_C_T_H_W, "B C (V T) H W -> B C V T H W", V=n_cameras)
+            view_embedding = view_embedding.expand(
+                x_B_C_V_T_H_W.shape[0],
+                view_embedding.shape[1],
+                view_embedding.shape[2],
+                x_B_C_V_T_H_W.shape[3],
+                x_B_C_V_T_H_W.shape[4],
+                x_B_C_V_T_H_W.shape[5],
+            )
+            x_B_C_V_T_H_W = torch.cat([x_B_C_V_T_H_W, view_embedding], dim=1)
+            x_B_C_T_H_W = rearrange(x_B_C_V_T_H_W, " B C V T H W -> B C (V T) H W", V=n_cameras)
+
+        x_B_T_H_W_D = self.x_embedder(x_B_C_T_H_W)
+
+        if self.extra_per_block_abs_pos_emb:
+            extra_pos_embedder = self.extra_pos_embedders_options[str(n_cameras)]
+            extra_pos_emb = extra_pos_embedder(x_B_T_H_W_D, fps=fps)
+        else:
+            extra_pos_emb = None
+
+        if "rope" in self.pos_emb_cls.lower():
+            return x_B_T_H_W_D, pos_embedder(x_B_T_H_W_D, fps=fps), extra_pos_emb
+
+        if "fps_aware" in self.pos_emb_cls:
+            raise NotImplementedError("FPS-aware positional embedding is not supported for multi-view DIT")
+
+        x_B_T_H_W_D = x_B_T_H_W_D + pos_embedder(x_B_T_H_W_D)
+
+        return x_B_T_H_W_D, None, extra_pos_emb
+
+    def forward(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        timesteps_B_T: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        data_type: Optional[DataType] = DataType.VIDEO,
+        view_indices_B_T: Optional[torch.Tensor] = None,
+        intermediate_feature_ids: Optional[List[int]] = None,
+        **kwargs,
+    ) -> torch.Tensor | List[torch.Tensor] | Tuple[torch.Tensor, List[torch.Tensor]]:
+        # Deletes elements like condition.use_video_condition that are not used in the forward pass
+        del kwargs
+        if data_type == DataType.VIDEO:
+            x_B_C_T_H_W = torch.cat([x_B_C_T_H_W, condition_video_input_mask_B_C_T_H_W.type_as(x_B_C_T_H_W)], dim=1)
+        else:
+            B, _, T, H, W = x_B_C_T_H_W.shape
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, torch.zeros((B, 1, T, H, W), dtype=x_B_C_T_H_W.dtype, device=x_B_C_T_H_W.device)], dim=1
+            )
+
+        assert isinstance(data_type, DataType), (
+            f"Expected DataType, got {type(data_type)}. We need discuss this flag later."
+        )
+        timesteps_B_T = timesteps_B_T * self.timestep_scale
+        x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = self.prepare_embedded_sequence(
+            x_B_C_T_H_W,
+            fps=fps,
+            padding_mask=padding_mask,
+            view_indices_B_T=view_indices_B_T,
+        )
+        if self.use_crossattn_projection:
+            crossattn_emb = self.crossattn_proj(crossattn_emb)
+        with amp.autocast("cuda", enabled=self.use_wan_fp32_strategy, dtype=torch.float32):
+            if timesteps_B_T.ndim == 1:
+                timesteps_B_T = timesteps_B_T.unsqueeze(1)
+            t_embedding_B_T_D, adaln_lora_B_T_3D = self.t_embedder(timesteps_B_T)
+            t_embedding_B_T_D = self.t_embedding_norm(t_embedding_B_T_D)
+
+        # for logging purpose
+        affline_scale_log_info = {}
+        affline_scale_log_info["t_embedding_B_T_D"] = t_embedding_B_T_D.detach()
+        self.affline_scale_log_info = affline_scale_log_info
+        self.affline_emb = t_embedding_B_T_D
+        self.crossattn_emb = crossattn_emb
+
+        if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
+            assert x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape, (
+                f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}"
+            )
+
+        B, T, H, W, D = x_B_T_H_W_D.shape
+
+        for block in self.blocks:
+            x_B_T_H_W_D = block(
+                x_B_T_H_W_D,
+                t_embedding_B_T_D,
+                crossattn_emb,
+                rope_emb_L_1_1_D=rope_emb_L_1_1_D,
+                adaln_lora_B_T_3D=adaln_lora_B_T_3D,
+                extra_per_block_pos_emb=extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
+            )
+
+        x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
+        x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O)
+        return x_B_C_Tt_Hp_Wp

cosmos_predict2/_src/predict2_multiview/networks/multiview_pose_dit.py

@@ -0,0 +1,526 @@
+# 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.
+
+"""Multi-view DiT with a pose encoder + warped/visibility conditioning for 2-actor joint generation.
+
+Design (additive, zero-initialized -- base checkpoint loads unchanged):
+  * ``pose_encoder``  : a 6x Conv3d (UniAnimate-DiT style) stack that turns a per-actor projected 2D pose RGB
+                        map into a patch-grid feature that is ADDED to the patch tokens.
+  * ``cond_embedder`` : a PatchEmbed over ``cat([warped_latent(16), visibility_mask(1)])`` whose output is
+                        ADDED to the patch tokens. This is mathematically identical to concatenating the
+                        warped latent + visibility mask onto the input channels of ``x_embedder`` (since
+                        ``Linear([x; w; v]) == W_x x + W_w w + W_v v``), but it leaves ``x_embedder``
+                        unchanged so the pretrained 2B checkpoint loads with no shape surgery.
+
+Both extra paths have a zero-initialized final projection, so at step 0 the network output is identical to
+the warm-started base ``MultiViewDiT`` (useful sanity check + stable fine-tuning start).
+
+Actor 1 / Actor 2 are treated as the two "views" of the existing multi-view machinery: tensors are laid out
+``(B, C, V*T, H, W)`` with ``V=2``; self-attention is shared across all ``V*T`` tokens (already the case in
+``MultiViewBlock``), so the two actors are jointly denoised.
+"""
+
+from typing import List, Optional, Tuple
+
+import torch
+import torch.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.distributed._composable.fsdp import fully_shard
+
+from cosmos_predict2._src.predict2.conditioner import DataType
+from cosmos_predict2._src.predict2.networks.minimal_v4_dit import PatchEmbed
+from cosmos_predict2._src.predict2_multiview.networks.multiview_dit import MultiViewDiT
+
+
+class PoseEncoder(nn.Module):
+    """3D-conv driving-pose encoder, mirroring UniAnimate-DiT (arXiv:2504.11289).
+
+    Input : ``(N, 3, T_pixel, H_pixel, W_pixel)`` projected-pose RGB in uint8/float [0,255].
+    Output: ``(N, out_channels, T_latent, H_pixel/16, W_pixel/16)`` where ``T_latent`` matches the VAE's
+            temporal compression (/4). The final projection conv is zero-initialized.
+    """
+
+    def __init__(self, out_channels: int, hidden_dim: int = 16, in_channels: int = 3, proj_init_std: float = 0.0):
+        super().__init__()
+        self.proj_init_std = proj_init_std  # 0.0 -> zero-init proj (stable); >0 -> small-random (faster pose learning)
+        h = hidden_dim
+        self.conv = nn.Sequential(
+            nn.Conv3d(in_channels, h, (3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1)),
+            nn.SiLU(),
+            nn.Conv3d(h, h, (3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1)),
+            nn.SiLU(),
+            nn.Conv3d(h, h, (3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1)),
+            nn.SiLU(),
+            nn.Conv3d(h, h, (3, 3, 3), stride=(1, 2, 2), padding=(1, 1, 1)),  # spatial /2
+            nn.SiLU(),
+            nn.Conv3d(h, h, (3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1)),  # temporal/2, spatial/2
+            nn.SiLU(),
+            nn.Conv3d(h, h, (3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1)),  # temporal/2, spatial/2
+            nn.SiLU(),
+        )
+        # final projection to model dim, last spatial /2  -> total spatial /16, temporal /4
+        self.proj = nn.Conv3d(h, out_channels, (1, 2, 2), stride=(1, 2, 2), padding=0)
+
+    def reset_parameters(self) -> None:
+        for m in self.conv.modules():
+            if isinstance(m, nn.Conv3d):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        # final projection: zero-init (proj_init_std=0) keeps day-0 output == base but starves the conv stack
+        # of gradient until proj grows; a small-random init lets the conv feature extractor learn from step 0.
+        if self.proj_init_std > 0:
+            nn.init.normal_(self.proj.weight, mean=0.0, std=self.proj_init_std)
+        else:
+            nn.init.zeros_(self.proj.weight)
+        if self.proj.bias is not None:
+            nn.init.zeros_(self.proj.bias)
+
+    def forward(self, pose_N_C_T_H_W: torch.Tensor) -> torch.Tensor:
+        # normalize pose RGB [0,255] -> [0,1] in fp32, then match the conv weight dtype (bf16 during training)
+        w_dtype = self.conv[0].weight.dtype
+        x = pose_N_C_T_H_W.to(torch.float32) / 255.0
+        x = x.to(w_dtype)
+        x = self.conv(x)
+        x = self.proj(x)
+        return x
+
+
+class MultiViewPoseDiT(MultiViewDiT):
+    def __init__(
+        self,
+        *args,
+        warped_latent_channels: int = 16,
+        visibility_channels: int = 1,
+        pose_in_channels: int = 3,
+        pose_hidden_dim: int = 16,
+        pose_proj_init_std: float = 0.0,  # 0 -> zero-init pose proj; >0 -> small-random (faster pose learning, e.g. stage 1)
+        freeze_view_embedding: bool = False,  # stage-1 single-view: view embedding is meaningless -> freeze it
+        # pose routing mode:
+        #   "encoder"     : Conv3d PoseEncoder over the pixel pose RGB, ADDED to tokens (default, original).
+        #   "vae_concat"  : frozen-VAE pose latent (16ch) CHANNEL-CONCATENATED with warped+vis into cond_embedder.
+        #   "vae_mlp_add" : frozen-VAE pose latent -> trainable PatchEmbed+MLP -> ADDED to tokens.
+        pose_mode: str = "encoder",
+        pose_latent_channels: int = 16,  # VAE latent ch for the pose map (== warped_latent_channels)
+        # in-context reference-frame appearance conditioning: append R clean reference frames per view as
+        # extra tokens (via the shared x_embedder) that participate in self-attention, at fixed non-contiguous
+        # temporal RoPE positions, then strip them before the output. Zero new params except a zero-init gate.
+        enable_reference_frames: bool = False,
+        num_reference_frames: int = 0,       # R reference frames per view-slot (total appended = V*R)
+        ref_rope_offset: int = 50,           # temporal RoPE position of the first reference frame
+        ref_rope_stride: int = 5,            # spacing between consecutive reference-frame positions
+        # SHARED reference mode: the V*R appended refs are ONE greedy-selected shared set (split across the
+        # per-view slots), not per-view pools. With full cross-view self-attention both views attend all of
+        # them, so the view-slot placement is arbitrary and the per-view view-embedding on refs is dropped
+        # (refs carry no view identity; their only geometric id is the posed-reference Plücker). Real frames
+        # are unaffected.
+        shared_reference: bool = False,
+        # per-pixel Plücker ray conditioning: 6-ch ray map (dir+moment) -> zero-init PatchEmbed -> ADD to tokens
+        # (same additive pattern as warped/pose; mathematically == channel-concat + wider x_embedder, but keeps
+        # x_embedder unchanged so the base checkpoint warm-starts cleanly). Grounds both views in a shared frame.
+        enable_plucker: bool = False,
+        plucker_channels: int = 6,
+        enable_reference_plucker: bool = False,  # also add Plücker to the reference tokens (posed references)
+        # also add a POSE latent to the reference tokens: VAE-encoded per-person skeleton render of each ref
+        # frame (16ch) -> zero-init PatchEmbed -> ADD to the reference tokens. Gives the in-context references
+        # the same pose signal the main frames get from control_input_pose. Same additive pattern as ref-Plücker.
+        enable_reference_pose: bool = False,
+        # composite DEPTH conditioning (warped scene depth + human mesh depth, RGB-encoded -> frozen VAE ->
+        # depth_latent 16ch) -> zero-init PatchEmbed -> ADD to tokens. Same additive pattern as warped/pose/plucker
+        # (mathematically == VAE-channel-concat into cond_embedder, but a separate zero-init branch keeps the base
+        # checkpoint's cond_embedder unchanged so it warm-starts cleanly).
+        enable_depth: bool = False,
+        depth_latent_channels: int = 16,
+        **kwargs,
+    ):
+        self.enable_depth = enable_depth
+        self.depth_latent_channels = depth_latent_channels
+        self.warped_latent_channels = warped_latent_channels
+        self.visibility_channels = visibility_channels
+        self.pose_in_channels = pose_in_channels
+        self.pose_hidden_dim = pose_hidden_dim
+        self.pose_proj_init_std = pose_proj_init_std
+        self.freeze_view_embedding = freeze_view_embedding
+        assert pose_mode in ("encoder", "vae_concat", "vae_mlp_add"), pose_mode
+        self.pose_mode = pose_mode
+        self.pose_latent_channels = pose_latent_channels
+        self.enable_reference_frames = enable_reference_frames
+        self.shared_reference = shared_reference
+        self.num_reference_frames = int(num_reference_frames)
+        self.ref_rope_offset = int(ref_rope_offset)
+        self.ref_rope_stride = int(ref_rope_stride)
+        self.ref_positions = [self.ref_rope_offset + self.ref_rope_stride * i for i in range(self.num_reference_frames)]
+        self.enable_plucker = enable_plucker
+        self.plucker_channels = plucker_channels
+        self.enable_reference_plucker = enable_reference_plucker
+        self.enable_reference_pose = enable_reference_pose
+        # NOTE: we intentionally do NOT widen x_embedder. The warped+visibility "channel concat" is realized
+        # by a separate zero-init PatchEmbed added to the tokens (see module docstring).
+        super().__init__(*args, **kwargs)
+        if self.enable_reference_frames and self.num_reference_frames > 0:
+            # zero-init per-CHANNEL gate on the reference-token contribution -> day-0 output ~= warm-started
+            # base. Shape (model_channels,) (not a scalar): FSDP2 shards params per-parameter on dim 0, and a
+            # size-1 param yields a degenerate empty shard on some ranks that intermittently breaks the DCP
+            # save collective (gather_object NCCL error). A (D,) vector shards evenly and is more expressive.
+            self.ref_gate = nn.Parameter(torch.zeros(self.model_channels))
+
+        # additive conditioning path for warped latent + visibility mask (concatenated on channels). In
+        # "vae_concat" mode the VAE pose latent is concatenated into the SAME embedder (channel concat).
+        cond_in = self.warped_latent_channels + self.visibility_channels
+        if self.pose_mode == "vae_concat":
+            cond_in += self.pose_latent_channels
+        self.cond_embedder = PatchEmbed(
+            spatial_patch_size=self.patch_spatial,
+            temporal_patch_size=self.patch_temporal,
+            in_channels=cond_in,
+            out_channels=self.model_channels,
+        )
+        # Plücker ray conditioning: PatchEmbed over the 6-ch ray map, ADDED to tokens (zero-init proj)
+        if self.enable_plucker:
+            self.plucker_embedder = PatchEmbed(
+                spatial_patch_size=self.patch_spatial,
+                temporal_patch_size=self.patch_temporal,
+                in_channels=self.plucker_channels,
+                out_channels=self.model_channels,
+            )
+        # composite-depth conditioning: PatchEmbed over the 16-ch depth VAE latent, ADDED to tokens (zero-init proj)
+        if self.enable_depth:
+            self.depth_embedder = PatchEmbed(
+                spatial_patch_size=self.patch_spatial,
+                temporal_patch_size=self.patch_temporal,
+                in_channels=self.depth_latent_channels,
+                out_channels=self.model_channels,
+            )
+        # reference-POSE conditioning: PatchEmbed over the 16-ch VAE latent of each reference frame's skeleton
+        # render, ADDED to the reference tokens (zero-init). Mirrors the reference-Plücker additive path.
+        if self.enable_reference_pose:
+            self.reference_pose_embedder = PatchEmbed(
+                spatial_patch_size=self.patch_spatial,
+                temporal_patch_size=self.patch_temporal,
+                in_channels=self.pose_latent_channels,
+                out_channels=self.model_channels,
+            )
+        # pose conditioning path (varies by mode)
+        if self.pose_mode == "encoder":
+            self.pose_encoder = PoseEncoder(
+                out_channels=self.model_channels,
+                hidden_dim=self.pose_hidden_dim,
+                in_channels=self.pose_in_channels,
+                proj_init_std=self.pose_proj_init_std,
+            )
+        elif self.pose_mode == "vae_mlp_add":
+            # frozen-VAE pose latent -> patch-embed to tokens -> trainable MLP -> ADD. MLP last layer
+            # zero-initialized so day-0 output == warm-started base.
+            self.pose_latent_embedder = PatchEmbed(
+                spatial_patch_size=self.patch_spatial,
+                temporal_patch_size=self.patch_temporal,
+                in_channels=self.pose_latent_channels,
+                out_channels=self.model_channels,
+            )
+            self.pose_mlp = nn.Sequential(
+                nn.Linear(self.model_channels, self.model_channels),
+                nn.SiLU(),
+                nn.Linear(self.model_channels, self.model_channels),
+            )
+
+        # re-run init now that the new submodules exist (zero-inits their final projections)
+        self.init_weights()
+
+        # stage-1 (single-view): the view embedding can't distinguish actors with V=1 -> exclude it from
+        # training (architecture/x_embedder unchanged, so weights still transfer cleanly to the joint stage 2).
+        if self.freeze_view_embedding and hasattr(self, "view_embeddings"):
+            self.view_embeddings.requires_grad_(False)
+
+    # ------------------------------------------------------------------ weights / sharding
+    def init_weights(self):
+        super().init_weights()
+        if hasattr(self, "cond_embedder"):
+            # zero-init the warped/visibility(+pose, in vae_concat) projection -> day-0 contribution is 0 (== base)
+            nn.init.zeros_(self.cond_embedder.proj[1].weight)
+        if hasattr(self, "plucker_embedder"):
+            nn.init.zeros_(self.plucker_embedder.proj[1].weight)  # day-0 Plücker contribution is 0 (== base)
+        if hasattr(self, "depth_embedder"):
+            nn.init.zeros_(self.depth_embedder.proj[1].weight)  # day-0 depth contribution is 0 (== base)
+        if hasattr(self, "reference_pose_embedder"):
+            nn.init.zeros_(self.reference_pose_embedder.proj[1].weight)  # day-0 ref-pose contribution is 0 (== base)
+        if hasattr(self, "pose_encoder"):
+            self.pose_encoder.reset_parameters()
+        if hasattr(self, "pose_mlp"):
+            # zero-init the MLP's last layer -> the (trainable) VAE-pose path contributes 0 at step 0
+            nn.init.zeros_(self.pose_mlp[-1].weight)
+            if self.pose_mlp[-1].bias is not None:
+                nn.init.zeros_(self.pose_mlp[-1].bias)
+
+    def fully_shard(self, mesh, **fsdp_kwargs):
+        super().fully_shard(mesh, **fsdp_kwargs)
+        if hasattr(self, "cond_embedder"):
+            fully_shard(self.cond_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if hasattr(self, "plucker_embedder"):
+            fully_shard(self.plucker_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if hasattr(self, "depth_embedder"):
+            fully_shard(self.depth_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if hasattr(self, "reference_pose_embedder"):
+            fully_shard(self.reference_pose_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if hasattr(self, "pose_encoder"):
+            fully_shard(self.pose_encoder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if hasattr(self, "pose_latent_embedder"):
+            fully_shard(self.pose_latent_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if hasattr(self, "pose_mlp"):
+            fully_shard(self.pose_mlp, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+
+    # ------------------------------------------------------------------ conditioning helpers
+    def _embed_warped_visibility(
+        self,
+        warped_latent_B_C_T_H_W: torch.Tensor,
+        visibility_mask_B_C_T_H_W: torch.Tensor,
+        pose_latent_B_C_T_H_W: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """PatchEmbed over cat([warped_latent, visibility(, pose_latent)]) -> token grid (B, T, H, W, D).
+
+        In "vae_concat" mode the frozen-VAE pose latent is concatenated on the channel dim alongside the
+        warped latent + visibility mask, so a single linear patch projection mixes all conditioning latents.
+        """
+        ref_dtype = self.cond_embedder.proj[1].weight.dtype
+        parts = [warped_latent_B_C_T_H_W.to(ref_dtype), visibility_mask_B_C_T_H_W.to(ref_dtype)]
+        if pose_latent_B_C_T_H_W is not None:
+            parts.append(pose_latent_B_C_T_H_W.to(ref_dtype))
+        cond_B_C_T_H_W = torch.cat(parts, dim=1)
+        return self.cond_embedder(cond_B_C_T_H_W)
+
+    def _embed_pose_latent_mlp(self, pose_latent_B_C_T_H_W: torch.Tensor) -> torch.Tensor:
+        """VAE pose latent -> trainable PatchEmbed -> MLP -> token grid (B, T, H, W, D). Last MLP layer
+        zero-init so day-0 contribution is 0."""
+        ref_dtype = self.pose_latent_embedder.proj[1].weight.dtype
+        feat = self.pose_latent_embedder(pose_latent_B_C_T_H_W.to(ref_dtype))  # (B, T, H, W, D)
+        return self.pose_mlp(feat)
+
+    def _embed_pose(
+        self, pose_map_B_C_T_H_W: torch.Tensor, n_cameras: int, grid_thw: Tuple[int, int, int]
+    ) -> torch.Tensor:
+        """Run the pose encoder per-view (no cross-actor temporal leak) -> token grid (B, V*T, H, W, D)."""
+        T_tok, H_tok, W_tok = grid_thw
+        # split actors so the temporal convs never span the actor boundary
+        pose_NV = rearrange(pose_map_B_C_T_H_W, "B C (V T) H W -> (B V) C T H W", V=n_cameras)
+        feat = self.pose_encoder(pose_NV)  # (B*V, D, t, h, w)
+        feat = feat.to(self.cond_embedder.proj[1].weight.dtype)
+        # be robust to off-by-one between conv-downsampled grid and the latent patch grid
+        if tuple(feat.shape[-3:]) != (T_tok // n_cameras, H_tok, W_tok):
+            feat = F.interpolate(
+                feat, size=(T_tok // n_cameras, H_tok, W_tok), mode="trilinear", align_corners=False
+            )
+        feat = rearrange(feat, "(B V) D t h w -> B (V t) h w D", V=n_cameras)
+        return feat
+
+    def _embed_reference_tokens(
+        self, reference_latent_B_C_VR_H_W: torch.Tensor, n_cameras: int, reference_plucker_map=None,
+        reference_pose_latent=None,
+    ) -> torch.Tensor:
+        """Assemble the SAME input channels as prepare_embedded_sequence for the reference latents and run the
+        SHARED x_embedder -> reference token grid (B, V*R, H_tok, W_tok, D). Reference frames are marked
+        'given/clean' via an all-ones condition mask (reusing the existing mask channel; no x_embedder widening).
+        Each view's R refs get that view's view-embedding (broadcast), so refs are tagged to their stream."""
+        ref_dtype = self.x_embedder.proj[1].weight.dtype
+        x = reference_latent_B_C_VR_H_W.to(ref_dtype)
+        B, C, VR, H, W = x.shape
+        R = VR // n_cameras
+        # (1) condition-mask channel = ones (refs are given/clean) -- mirrors forward()'s cond-mask concat
+        x = torch.cat([x, torch.ones((B, 1, VR, H, W), dtype=ref_dtype, device=x.device)], dim=1)
+        # (2) padding-mask channel = zeros (refs have no padding) -- mirrors prepare_embedded_sequence
+        if self.concat_padding_mask:
+            x = torch.cat([x, torch.zeros((B, 1, VR, H, W), dtype=ref_dtype, device=x.device)], dim=1)
+        # (3) view-embedding channels. The SHARED x_embedder expects these channels (real frames always carry
+        # them), so the width must match either way -- "removing the view embedding" for shared refs means
+        # feeding a NEUTRAL (zero) view signal, not dropping the channels:
+        #   - per-view mode: each view's refs get that view's learned embedding (broadcast), tagging them.
+        #   - shared_reference mode: all refs get an all-ZERO view channel -> no view identity (refs are one
+        #     shared set with arbitrary view-slot placement); their only geometric id is the posed Plücker.
+        if self.concat_view_embedding:
+            if self.shared_reference:
+                ve = torch.zeros((B, self.view_condition_dim, VR, H, W), dtype=ref_dtype, device=x.device)
+                x = torch.cat([x, ve], dim=1)
+            else:
+                view_idx = torch.arange(n_cameras, device=x.device).clamp(max=self.n_cameras_emb - 1)
+                ve = self.view_embeddings(view_idx).to(ref_dtype)  # (V, Dv)
+                ve = rearrange(ve, "V D -> D V")[None, :, :, None, None, None]  # (1, Dv, V, 1, 1, 1)
+                xv = rearrange(x, "B C (V R) H W -> B C V R H W", V=n_cameras)
+                ve = ve.expand(B, ve.shape[1], n_cameras, R, H, W)
+                x = rearrange(torch.cat([xv, ve], dim=1), "B C V R H W -> B C (V R) H W")
+        ref_tokens = self.x_embedder(x)  # (B, V*R, H_tok, W_tok, D)
+        # posed-reference Plücker: ground the reference tokens geometrically (reuse the zero-init plucker embedder)
+        if self.enable_reference_plucker and hasattr(self, "plucker_embedder") and reference_plucker_map is not None:
+            plk = self.plucker_embedder(reference_plucker_map.to(ref_dtype))  # (B, V*R, H_tok, W_tok, D)
+            ref_tokens = ref_tokens + plk.to(ref_tokens.dtype)
+        # reference-POSE: add the per-person skeleton latent of each reference frame (zero-init embedder)
+        if self.enable_reference_pose and hasattr(self, "reference_pose_embedder") and reference_pose_latent is not None:
+            rp = self.reference_pose_embedder(reference_pose_latent.to(ref_dtype))  # (B, V*R, H_tok, W_tok, D)
+            ref_tokens = ref_tokens + rp.to(ref_tokens.dtype)
+        return ref_tokens
+
+    # ------------------------------------------------------------------ forward
+    def forward(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        timesteps_B_T: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        data_type: Optional[DataType] = DataType.VIDEO,
+        view_indices_B_T: Optional[torch.Tensor] = None,
+        intermediate_feature_ids: Optional[List[int]] = None,
+        # NEW conditioning inputs (declared explicitly so they survive `del kwargs`)
+        pose_map_B_C_T_H_W: Optional[torch.Tensor] = None,
+        pose_latent_B_C_T_H_W: Optional[torch.Tensor] = None,
+        warped_latent_B_C_T_H_W: Optional[torch.Tensor] = None,
+        visibility_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
+        reference_latent_B_C_R_H_W: Optional[torch.Tensor] = None,
+        plucker_map_B_C_T_H_W: Optional[torch.Tensor] = None,
+        reference_plucker_map_B_C_R_H_W: Optional[torch.Tensor] = None,
+        reference_pose_latent_B_C_R_H_W: Optional[torch.Tensor] = None,
+        depth_latent_B_C_T_H_W: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del kwargs
+        if data_type == DataType.VIDEO:
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, condition_video_input_mask_B_C_T_H_W.type_as(x_B_C_T_H_W)], dim=1
+            )
+        else:
+            B, _, T, H, W = x_B_C_T_H_W.shape
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, torch.zeros((B, 1, T, H, W), dtype=x_B_C_T_H_W.dtype, device=x_B_C_T_H_W.device)],
+                dim=1,
+            )
+
+        assert isinstance(data_type, DataType)
+        timesteps_B_T = timesteps_B_T * self.timestep_scale
+        x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb = self.prepare_embedded_sequence(
+            x_B_C_T_H_W,
+            fps=fps,
+            padding_mask=padding_mask,
+            view_indices_B_T=view_indices_B_T,
+        )
+
+        B, T_tok, H_tok, W_tok, D = x_B_T_H_W_D.shape
+        n_cameras = T_tok // self.state_t
+
+        # additive warped-latent + visibility conditioning (zero-init -> 0 at step 0). In "vae_concat" mode
+        # the VAE pose latent is channel-concatenated into the same cond_embedder.
+        # cast the additive term back to x's dtype so we never silently upcast the token stream.
+        if warped_latent_B_C_T_H_W is not None and visibility_mask_B_C_T_H_W is not None:
+            pose_for_concat = pose_latent_B_C_T_H_W if self.pose_mode == "vae_concat" else None
+            cond_emb = self._embed_warped_visibility(
+                warped_latent_B_C_T_H_W, visibility_mask_B_C_T_H_W, pose_for_concat
+            )
+            x_B_T_H_W_D = x_B_T_H_W_D + cond_emb.to(x_B_T_H_W_D.dtype)
+
+        # additive pose conditioning (zero-init -> 0 at step 0)
+        if self.pose_mode == "encoder" and pose_map_B_C_T_H_W is not None:
+            pose_emb = self._embed_pose(pose_map_B_C_T_H_W, n_cameras, (T_tok, H_tok, W_tok))
+            x_B_T_H_W_D = x_B_T_H_W_D + pose_emb.to(x_B_T_H_W_D.dtype)
+        elif self.pose_mode == "vae_mlp_add" and pose_latent_B_C_T_H_W is not None:
+            pose_emb = self._embed_pose_latent_mlp(pose_latent_B_C_T_H_W)
+            x_B_T_H_W_D = x_B_T_H_W_D + pose_emb.to(x_B_T_H_W_D.dtype)
+
+        # additive Plücker ray conditioning (zero-init -> 0 at step 0). Per-pixel camera rays (canonical frame)
+        # patch-embedded and added to the real-frame tokens -> cross-view shared-space grounding.
+        if self.enable_plucker and plucker_map_B_C_T_H_W is not None:
+            plk_dtype = self.plucker_embedder.proj[1].weight.dtype
+            plucker_emb = self.plucker_embedder(plucker_map_B_C_T_H_W.to(plk_dtype))  # (B, V*state_t, h, w, D)
+            x_B_T_H_W_D = x_B_T_H_W_D + plucker_emb.to(x_B_T_H_W_D.dtype)
+
+        # additive composite-DEPTH conditioning (zero-init -> 0 at step 0). Depth VAE latent (warped scene depth +
+        # human mesh depth) patch-embedded and added to the real-frame tokens (== VAE-channel-concat, warm-clean).
+        if self.enable_depth and depth_latent_B_C_T_H_W is not None:
+            d_dtype = self.depth_embedder.proj[1].weight.dtype
+            depth_emb = self.depth_embedder(depth_latent_B_C_T_H_W.to(d_dtype))
+            x_B_T_H_W_D = x_B_T_H_W_D + depth_emb.to(x_B_T_H_W_D.dtype)
+
+        # in-context reference-frame appearance conditioning: append R reference frames per view onto the
+        # temporal axis (view-major: [view0 real, view0 ref, view1 real, view1 ref, ...]) with fixed
+        # non-contiguous temporal RoPE, gated by a zero-init scalar. Stripped before the output layer.
+        use_refs = (
+            self.enable_reference_frames
+            and self.num_reference_frames > 0
+            and reference_latent_B_C_R_H_W is not None
+        )
+        if use_refs:
+            R = self.num_reference_frames
+            ref_tokens = self._embed_reference_tokens(
+                reference_latent_B_C_R_H_W, n_cameras, reference_plucker_map_B_C_R_H_W,
+                reference_pose_latent_B_C_R_H_W,
+            )  # (B, V*R, H, W, D)
+            gate = self.ref_gate.to(x_B_T_H_W_D.dtype)
+            x_real = rearrange(x_B_T_H_W_D, "B (V t) H W D -> B V t H W D", V=n_cameras)
+            x_ref = rearrange(gate * ref_tokens.to(x_B_T_H_W_D.dtype), "B (V r) H W D -> B V r H W D", V=n_cameras)
+            x_B_T_H_W_D = rearrange(torch.cat([x_real, x_ref], dim=2), "B V te H W D -> B (V te) H W D")
+            # regenerate RoPE for the augmented per-view [real(state_t) ; ref(R)] grid (real positions
+            # [0..state_t-1], ref positions [offset, offset+stride, ...]); order matches the token layout above.
+            Bx, Te, Hx, Wx, Dx = x_B_T_H_W_D.shape
+            pos_embedder = self.pos_embedder_options[f"n_cameras_{n_cameras}"]
+            ref_rope = pos_embedder.generate_embeddings_with_refs(
+                (Bx, Te, Hx, Wx, Dx), num_ref=R, ref_positions=self.ref_positions, fps=fps
+            )
+            rope_emb_L_1_1_D = pos_embedder._split_for_context_parallel(ref_rope)
+
+        if self.use_crossattn_projection:
+            crossattn_emb = self.crossattn_proj(crossattn_emb)
+
+        with amp.autocast("cuda", enabled=self.use_wan_fp32_strategy, dtype=torch.float32):
+            if timesteps_B_T.ndim == 1:
+                timesteps_B_T = timesteps_B_T.unsqueeze(1)
+            t_embedding_B_T_D, adaln_lora_B_T_3D = self.t_embedder(timesteps_B_T)
+            t_embedding_B_T_D = self.t_embedding_norm(t_embedding_B_T_D)
+
+        affline_scale_log_info = {}
+        affline_scale_log_info["t_embedding_B_T_D"] = t_embedding_B_T_D.detach()
+        self.affline_scale_log_info = affline_scale_log_info
+        self.affline_emb = t_embedding_B_T_D
+        self.crossattn_emb = crossattn_emb
+
+        # when reference frames are appended, the per-frame adaLN modulation must span (state_t+R) frames/view.
+        # Expand the (per-frame) timestep embedding + adaLN-LoRA by padding each view's R ref slots with its
+        # last real-frame value (refs are stripped before the loss, so this padded value never reaches output).
+        te_blocks, lora_blocks = t_embedding_B_T_D, adaln_lora_B_T_3D
+        if use_refs:
+            def _pad_ref_slots(emb):
+                if emb is None or emb.shape[1] != n_cameras * self.state_t:
+                    return emb  # T==1 (broadcast) or already augmented -> leave as-is
+                e = rearrange(emb, "B (V t) X -> B V t X", V=n_cameras)
+                e = torch.cat([e, e[:, :, -1:, :].expand(-1, -1, self.num_reference_frames, -1)], dim=2)
+                return rearrange(e, "B V te X -> B (V te) X")
+            te_blocks = _pad_ref_slots(t_embedding_B_T_D)
+            lora_blocks = _pad_ref_slots(adaln_lora_B_T_3D)
+
+        for block in self.blocks:
+            x_B_T_H_W_D = block(
+                x_B_T_H_W_D,
+                te_blocks,
+                crossattn_emb,
+                rope_emb_L_1_1_D=rope_emb_L_1_1_D,
+                adaln_lora_B_T_3D=lora_blocks,
+                extra_per_block_pos_emb=extra_pos_emb,
+            )
+
+        # strip the appended reference frames per view -> back to (B, V*state_t, H, W, D) for the output layer
+        if use_refs:
+            xg = rearrange(x_B_T_H_W_D, "B (V te) H W D -> B V te H W D", V=n_cameras)
+            x_B_T_H_W_D = rearrange(xg[:, :, : self.state_t], "B V t H W D -> B (V t) H W D")
+
+        x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
+        x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O)
+        return x_B_C_Tt_Hp_Wp

cosmos_predict2/_src/predict2_multiview/scripts/inference.py

@@ -0,0 +1,330 @@
+# 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.
+"""
+This script is based on projects/cosmos/diffusion/v2/inference/vid2vid.py
+
+To run inference on the training data (as visualization/debugging), use:
+```bash
+EXP=buttercup_predict2p5_2b_7views_res720p_fps30_t8_joint_alpamayo1capviewprefix_allcapsviewprefix_29frames_nofps_uniform_dropoutt0
+ckpt_path=s3://bucket/cosmos_predict2_multiview/cosmos2_mv/buttercup_predict2p5_2b_7views_res720p_fps30_t8_joint_alpamayo1capviewprefix_allcapsviewprefix_29frames_nofps_uniform_dropoutt0-0/checkpoints/iter_000012000/
+PYTHONPATH=. torchrun --nproc_per_node=8 --master_port=12341 -m cosmos_predict2._src.predict2_multiview.scripts.inference --experiment ${EXP} --ckpt_path ${ckpt_path} --context_parallel_size 8 --input_is_train_data --max_samples 1 --num_conditional_frames 0 --guidance 3 --save_root results/predict2_multiview/
+
+EXP=predict2p5_2b_mv_7train7_res480p_fps15_t24_alpamayo_only_allcaption_uniform_nofps
+ckpt_path=s3://bucket/cosmos_predict2_multiview/cosmos2p5_mv/predict2p5_2b_mv_7train7_res480p_fps15_t24_alpamayo_only_allcaption_uniform_nofps-0/checkpoints/iter_000020000/
+PYTHONPATH=. torchrun --nproc_per_node=1 --master_port=12341 -m cosmos_predict2._src.predict2_multiview.scripts.inference --experiment ${EXP} --ckpt_path ${ckpt_path} --context_parallel_size 1 --input_is_train_data --max_samples 1 --num_conditional_frames 0 --guidance 3 --save_root results/predict2_multiview_480p_20k/
+```
+"""
+
+import argparse
+import os
+
+import torch as th
+from einops import rearrange
+from megatron.core import parallel_state
+
+from cosmos_predict2._src.imaginaire.lazy_config import instantiate
+from cosmos_predict2._src.imaginaire.utils import distributed, log
+from cosmos_predict2._src.imaginaire.visualize.video import save_img_or_video
+from cosmos_predict2._src.predict2.utils.model_loader import load_model_from_checkpoint
+from cosmos_predict2._src.predict2_multiview.scripts.mv_visualize_helper import arrange_video_visualization
+
+NUM_CONDITIONAL_FRAMES_KEY = "num_conditional_frames"
+
+
+def to_model_input(data_batch, model):
+    """
+    Similar to misc.to, but avoid converting uint8 "video" to float
+    """
+    for k, v in data_batch.items():
+        _v = v
+        if isinstance(v, th.Tensor):
+            _v = _v.cuda()
+            if th.is_floating_point(v):
+                _v = _v.to(**model.tensor_kwargs)
+        data_batch[k] = _v
+    return data_batch
+
+
+class Vid2VidInference:
+    """
+    Handles the Vid2Vid inference process, including model loading, data preparation,
+    and video generation from an image/video and text prompt. Now supports context parallelism.
+    """
+
+    def __init__(
+        self,
+        experiment_name: str,
+        ckpt_path: str,
+        s3_credential_path: str = "",
+        context_parallel_size: int = 1,
+        experiment_opts: list[str] = [],
+    ):
+        """
+        Initializes the Vid2VidInference 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.
+        """
+        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.experiment_opts = experiment_opts
+
+        if "RANK" in os.environ:
+            self._init_distributed()
+
+        # Load the model and config
+        model, config = load_model_from_checkpoint(
+            experiment_name=self.experiment_name,
+            s3_checkpoint_dir=self.ckpt_path,
+            config_file="cosmos_predict2/_src/predict2_multiview/configs/vid2vid/config.py",
+            load_ema_to_reg=True,
+            experiment_opts=self.experiment_opts,
+        )
+
+        # Enable context parallel on the model if using context parallelism
+        self.rank0 = True
+        if self.context_parallel_size > 1:
+            model.net.enable_context_parallel(self.process_group)
+            self.rank0 = distributed.get_rank() == 0
+
+        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()
+
+        log.info(f"Initialized context parallel with size {self.context_parallel_size}")
+        log.info(f"Current rank: {distributed.get_rank()}, World size: {distributed.get_world_size()}")
+
+    def generate_from_batch(
+        self,
+        data_batch,
+        guidance: int = 7,
+        seed: int = 1,
+        num_steps: int = 35,
+        stack_mode: str = "time",
+        use_negative_prompt: bool = True,
+    ):
+        """Generate video tensor from batch.
+
+        Returns:
+            Tensor with values in the range [0, 1]
+            If stack mode is "time", the tensor is of shape (1, 3, v * t, h, w)
+            If stack mode is "height", the tensor is of shape (1, 3, t, v * h, w)
+            If stack mode is "width", the tensor is of shape (1, 3, t, h, v * w)
+            If stack mode is "grid", the tensor is of shape (1, 3, t, 3 * h, 3 * w)
+        """
+        data_batch = to_model_input(data_batch, self.model)
+        if self.model.config.text_encoder_config is not None and self.model.config.text_encoder_config.compute_online:
+            self.model.inplace_compute_text_embeddings_online(data_batch)
+        raw_data, x0, condition = self.model.get_data_and_condition(data_batch)
+        sample = self.model.generate_samples_from_batch(
+            data_batch,
+            guidance=guidance,
+            # make sure no mismatch and also works for cp
+            state_shape=x0.shape[1:],
+            n_sample=x0.shape[0],
+            seed=seed,  # Fixed seed for reproducibility
+            num_steps=num_steps,
+            is_negative_prompt=use_negative_prompt,
+        )
+        # (bsz = 1, c = 3, t = n_camera * t, h, w)
+        video = ((self.model.decode(sample) + 1.0) / 2.0).clamp(0, 1)
+
+        # Arrange video according to stack_mode
+        video = arrange_video_visualization(video, data_batch, method=stack_mode)
+        return video
+
+    def generate_from_batch_autoregressive(
+        self,
+        data_batch,
+        num_chunks=2,
+        chunk_overlap=2,
+        guidance: int = 7,
+        seed: int = 1,
+        num_steps: int = 35,
+        stack_mode: str = "time",
+        use_negative_prompt: bool = True,
+    ):
+        """Generate video tensor from batch, with autoregressive mode enabled
+        num_chunks: total number of single generation
+        chunk_overlap: overlap the
+        """
+        data_batch = to_model_input(data_batch, self.model)
+        if self.model.config.text_encoder_config is not None and self.model.config.text_encoder_config.compute_online:
+            self.model.inplace_compute_text_embeddings_online(data_batch)
+
+        n_views = len(data_batch["camera_keys_selection"][0])
+        num_video_frames_per_view = data_batch["num_video_frames_per_view"][0]
+
+        generated_chunks = []
+
+        for i in range(num_chunks):
+            log.info(f"start generate chunk {i + 1} / {num_chunks}")
+            _, x0, _ = self.model.get_data_and_condition(data_batch)
+            sample = self.model.generate_samples_from_batch(
+                data_batch,
+                guidance=guidance,
+                # make sure no mismatch and also works for cp
+                state_shape=x0.shape[1:],
+                n_sample=x0.shape[0],
+                seed=seed,  # Fixed seed for reproducibility
+                num_steps=num_steps,
+                is_negative_prompt=use_negative_prompt,
+            )
+            # (bsz = 1, c = 3, t = n_camera * t, h, w)
+            decoded = self.model.decode(sample)
+            chunk_video = ((decoded + 1.0) / 2.0).clamp(0, 1)[0]
+            chunk_video = rearrange(chunk_video, "C (V T) H W -> V C T H W", V=n_views)
+            if i == 0:
+                generated_chunks.append(chunk_video)
+            else:
+                generated_chunks.append(chunk_video[:, :, chunk_overlap:])
+            data_batch["num_conditional_frames"] = chunk_overlap
+            data_batch["video"].zero_()
+            for v in range(n_views):
+                start_idx = num_video_frames_per_view * v
+                overlaps = (
+                    chunk_video[v, :, num_video_frames_per_view - chunk_overlap : num_video_frames_per_view] * 255
+                )
+                overlaps = overlaps.to(th.uint8).clamp(0, 255)
+                data_batch["video"][:, :, start_idx : start_idx + chunk_overlap] = overlaps
+
+        video = th.cat(generated_chunks, dim=2)
+        video = rearrange(video, "V C T H W -> C (V T) H W", V=n_views).unsqueeze(0)
+
+        # Arrange video according to stack_mode
+        video = arrange_video_visualization(video, data_batch, method=stack_mode)
+        return video
+
+    def cleanup(self):
+        """Clean up distributed resources."""
+        if "RANK" in os.environ:
+            import torch.distributed as dist
+            from megatron.core import parallel_state
+
+            if parallel_state.is_initialized():
+                parallel_state.destroy_model_parallel()
+            dist.destroy_process_group()
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parses command-line arguments for the Vid2Vid inference script."""
+    parser = argparse.ArgumentParser(description="Image2World/Video2World inference script")
+    parser.add_argument("--experiment", type=str, required=True, help="Experiment config")
+    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(
+        "--context_parallel_size",
+        type=int,
+        default=1,
+        help="Context parallel size (number of GPUs to split context over). Set to 8 for 8 GPUs",
+    )
+    # generation
+    parser.add_argument("--guidance", type=int, default=7, help="Guidance value")
+    parser.add_argument("--fps", type=int, default=10, help="Frames per second")
+    parser.add_argument("--seed", type=int, default=1, help="Guidance value")
+    parser.add_argument("--num_conditional_frames", type=int, default=1, help="Number of conditional frames")
+    # input
+    parser.add_argument(
+        "--input_is_train_data",
+        action="store_true",
+        help="Inference on the training data, the input_root will be ignored if this is set",
+    )
+    parser.add_argument("--run_mads_verification", action="store_true", help="Run MADS verification")
+    parser.add_argument(
+        "--mads_verification_prompt",
+        type=str,
+        default='The video opens with a view from inside a vehicle, positioned at an intersection under a clear blue sky. The camera angle is from the dashboard, offering a first-person perspective of the road ahead. The intersection is marked by multiple traffic lights and street signs, including one that reads "E Garden Blvd." A white van with "TM Stuckateur" branding is seen driving through the intersection, heading in the same direction as the viewer\'s vehicle. Other cars are also present, moving smoothly along the multi-lane road. As the vehicle starts to move forward, the camera pans slightly to the right, revealing more of the surroundings. The road is lined with trees on both sides, providing a natural canopy that filters the sunlight. The trees are lush and green, indicating it might be spring or summer. On the left side of the road, there is a large building with a sign that reads "GROCERY OUTLET," suggesting the presence of a retail store nearby. Further down the road, additional buildings and residential structures can be seen, hinting at a suburban or semi-urban area. The sun is bright and high in the sky, casting long shadows across the road. The light creates a warm, inviting atmosphere, enhancing the clarity of the scene. The road itself is well-maintained, with clear lane markings and directional arrows painted on the asphalt. Overhead, power lines run parallel to the road, supported by poles that also hold traffic lights and street lamps. As the vehicle continues its journey, the camera maintains a steady focus on the road ahead, capturing the smooth flow of traffic and the serene environment. The absence of heavy traffic or congestion adds to the tranquil mood of the scene. The overall ambiance is one of calm and order, with the interplay of natural and man-made elements creating a harmonious urban landscape. The gentle curve of the road and the soft glow of the setting sun add a sense of peacefulness to the drive, making the viewer feel as though they are part of this quiet, picturesque neighborhood.',
+    )
+    parser.add_argument(
+        "--stack_mode",
+        type=str,
+        default="time",
+        choices=["height", "width", "time", "grid"],
+        help="Video stacking mode for visualization. grid will create a 3x3 grid of views.",
+    )
+    parser.add_argument("--input_root", type=str, default="assets/image2world", help="Input root")
+    parser.add_argument("--save_root", type=str, default="results/image2world", help="Save root")
+    parser.add_argument("--max_samples", type=int, default=20, help="Maximum number of samples to generate")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    os.environ["NVTE_FUSED_ATTN"] = "0"
+    th.backends.cudnn.benchmark = False
+    th.backends.cudnn.deterministic = True
+    th.enable_grad(False)
+    args = parse_arguments()
+    # Initialize the inference handler with context parallel support
+    vid2vid_cli = Vid2VidInference(
+        args.experiment, args.ckpt_path, args.s3_cred, context_parallel_size=args.context_parallel_size
+    )
+    mem_bytes = th.cuda.memory_allocated(device=th.device("cuda" if th.cuda.is_available() else "cpu"))
+    log.info(f"GPU memory usage after model dcp.load: {mem_bytes / (1024**3):.2f} GB")
+
+    # Only process files on rank 0 if using distributed processing
+    rank0 = True
+    if args.context_parallel_size > 1:
+        rank0 = distributed.get_rank() == 0
+
+    os.makedirs(args.save_root, exist_ok=True)
+    if args.input_is_train_data:
+        dataloader = instantiate(vid2vid_cli.config.dataloader_train)
+        for i, batch in enumerate(dataloader):
+            if i >= args.max_samples:
+                break
+            if args.run_mads_verification:
+                assert args.num_conditional_frames == 0, "MADS verification only supports 0 conditional frame"
+                log.warning(f"Running MADS verification with prompt: {args.mads_verification_prompt[0:100]}...")
+                batch["ai_caption"] = [args.mads_verification_prompt]
+            batch[NUM_CONDITIONAL_FRAMES_KEY] = args.num_conditional_frames
+            video = vid2vid_cli.generate_from_batch(
+                batch, guidance=args.guidance, seed=args.seed, stack_mode=args.stack_mode
+            )
+            if rank0:
+                save_name = f"mads_verification_{i}" if args.run_mads_verification else f"infer_from_train_{i}"
+                save_img_or_video(video[0], f"{args.save_root}/{save_name}", fps=args.fps)
+            if args.run_mads_verification:
+                break
+    else:
+        raise NotImplementedError("Not implemented")

cosmos_predict2/_src/predict2_multiview/scripts/inference_cli.py

@@ -0,0 +1,582 @@
+# 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.
+"""
+Inference script for constructing data_batch from videos and captions, then running predict2_multiview model.
+
+Expected directory structure:
+```
+input_root/
+├── videos/                                    # Video folder
+│   ├── ftheta_camera_front_wide_120fov/      # or camera_front_wide_120fov/
+│   │   ├── video_id_1.mp4
+│   │   ├── video_id_2.mp4
+│   │   └── ...
+│   ├── ftheta_camera_cross_right_120fov/     # or camera_cross_right_120fov/
+│   │   ├── video_id_1.mp4
+│   │   ├── video_id_2.mp4
+│   │   └── ...
+│   ├── ftheta_camera_rear_right_70fov/       # or camera_rear_right_70fov/
+│   ├── ftheta_camera_rear_tele_30fov/        # or camera_rear_tele_30fov/
+│   ├── ftheta_camera_rear_left_70fov/        # or camera_rear_left_70fov/
+│   ├── ftheta_camera_cross_left_120fov/      # or camera_cross_left_120fov/
+│   └── ftheta_camera_front_tele_30fov/       # or camera_front_tele_30fov/
+│
+└── captions/                                  # Caption folder (optional, uses default prompt if not present)
+    ├── ftheta_camera_front_wide_120fov/      # or camera_front_wide_120fov/
+    │   ├── video_id_1.txt
+    │   ├── video_id_2.txt
+    │   └── ...
+    ├── ftheta_camera_cross_right_120fov/     # or camera_cross_right_120fov/
+    │   ├── video_id_1.txt
+    │   ├── video_id_2.txt
+    │   └── ...
+    ├── ftheta_camera_rear_right_70fov/       # or camera_rear_right_70fov/
+    ├── ftheta_camera_rear_tele_30fov/        # or camera_rear_tele_30fov/
+    ├── ftheta_camera_rear_left_70fov/        # or camera_rear_left_70fov/
+    ├── ftheta_camera_cross_left_120fov/      # or camera_cross_left_120fov/
+    └── ftheta_camera_front_tele_30fov/       # or camera_front_tele_30fov/
+
+Notes:
+- The videos/ folder is required (unless num_conditional_frames=0, which uses dummy all-zero videos)
+- The captions/ folder is optional; if not present, a preset default driving scene description is used
+- Each camera's subfolder name supports two formats: "ftheta_{camera_name}" or "{camera_name}"
+- video_id must be consistent across all camera folders
+- All 7 camera views must have corresponding subfolders and files
+
+Camera view to View Index mapping:
+- camera_front_wide_120fov: 0
+- camera_cross_right_120fov: 1
+- camera_rear_right_70fov: 2
+- camera_rear_tele_30fov: 3
+- camera_rear_left_70fov: 4
+- camera_cross_left_120fov: 5
+- camera_front_tele_30fov: 6
+```
+
+Usage:
+```bash
+EXP=predict2p5_2b_mv_7train7_res720p_fps10_t24_frombase2p5avfinetune_alpamayo_only_allcaption_uniform_nofps
+ckpt_path=s3://bucket/cosmos_predict2_multiview/cosmos2_mv2/predict2p5_2b_mv_7train7_res720p_fps10_t24_frombase2p5avfinetune_alpamayo_only_allcaption_uniform_nofps_resume1-0/checkpoints/iter_000024500/
+
+PYTHONPATH=. torchrun --nproc_per_node=8 --master_port=12341 -m cosmos_predict2._src.predict2_multiview.scripts.inference_cli \
+    --experiment ${EXP} \
+    --ckpt_path ${ckpt_path} \
+    --context_parallel_size 8 \
+    --input_root /project/cosmos/yiflu/project_official_i4/condition_assets/multiview-inference-assets-1203 \
+    --num_conditional_frames 0 \
+    --guidance 5 \
+    --fps 10 \
+    --save_root results/predict2_multiview_av_grid/ \
+    --max_samples 5 --stack_mode grid
+```
+
+```bash
+EXP=predict2p5_2b_mv_7train7_res480p_fps15_t24_alpamayo_only_allcaption_uniform_nofps
+ckpt_path=s3://bucket/cosmos_predict2_multiview/cosmos2p5_mv/predict2p5_2b_mv_7train7_res480p_fps15_t24_alpamayo_only_allcaption_uniform_nofps-0/checkpoints/iter_000020000/
+
+PYTHONPATH=. torchrun --nproc_per_node=8 --master_port=12341 -m cosmos_predict2._src.predict2_multiview.scripts.inference_cli \
+    --experiment ${EXP} \
+    --ckpt_path ${ckpt_path} \
+    --context_parallel_size 8 \
+    --input_root /project/cosmos/yiflu/project_official_i4/condition_assets/multiview-inference-assets-1203 \
+    --num_conditional_frames 1 \
+    --guidance 5 \
+    --fps 15 \
+    --save_root results/predict2_multiview_av_480p_i2v_grid/ \
+    --max_samples 5 --stack_mode grid \
+    --target_height 480 --target_width 832 \
+    model.config.net.init_cross_view_attn_weight_from=null
+```
+"""
+
+import argparse
+import os
+from pathlib import Path
+
+import torch as th
+import torchvision
+
+from cosmos_predict2._src.imaginaire.utils import distributed, log
+from cosmos_predict2._src.imaginaire.utils.easy_io import easy_io
+from cosmos_predict2._src.imaginaire.visualize.video import save_img_or_video
+from cosmos_predict2._src.predict2_multiview.scripts.inference import Vid2VidInference
+
+NUM_CONDITIONAL_FRAMES_KEY = "num_conditional_frames"
+
+# Camera name to view index mapping
+CAMERA_TO_VIEW_INDEX = {
+    "camera_front_wide_120fov": 0,
+    "camera_cross_right_120fov": 1,
+    "camera_rear_right_70fov": 2,
+    "camera_rear_tele_30fov": 3,
+    "camera_rear_left_70fov": 4,
+    "camera_cross_left_120fov": 5,
+    "camera_front_tele_30fov": 6,
+}
+
+DEFAULT_CAMERA_ORDER = list(CAMERA_TO_VIEW_INDEX.keys())
+
+# Camera-specific caption prefixes describing camera position and orientation
+CAMERA_TO_CAPTION_PREFIX = {
+    "camera_front_wide_120fov": "The video is captured from a camera mounted on a car. The camera is facing forward.",
+    "camera_cross_right_120fov": "The video is captured from a camera mounted on a car. The camera is facing to the right.",
+    "camera_rear_right_70fov": "The video is captured from a camera mounted on a car. The camera is facing the rear right side.",
+    "camera_rear_tele_30fov": "The video is captured from a camera mounted on a car. The camera is facing backwards.",
+    "camera_rear_left_70fov": "The video is captured from a camera mounted on a car. The camera is facing the rear left side.",
+    "camera_cross_left_120fov": "The video is captured from a camera mounted on a car. The camera is facing to the left.",
+    "camera_front_tele_30fov": "The video is captured from a telephoto camera mounted on a car. The camera is facing forward.",
+}
+
+DEFAULT_DRIVING_SCENE_PROMPT = """
+A clear daytime driving scene on an open road. The weather is sunny with bright natural lighting and good visibility. 
+The sky is partly cloudy with scattered white clouds. The road surface is dry and well-maintained. 
+The overall atmosphere is calm and peaceful with moderate traffic conditions. The lighting creates clear 
+shadows and provides excellent contrast for safe navigation."""
+
+
+def load_video(video_path: str, target_frames: int = 93, target_size: tuple[int, int] = (720, 1280)) -> th.Tensor:
+    """
+    Load video and process it to target size and frame count.
+
+    Args:
+        video_path: Path to video file
+        target_frames: Target number of frames
+        target_size: Target resolution (H, W)
+
+    Returns:
+        Video tensor with shape (C, T, H, W), dtype uint8
+    """
+    try:
+        # Load video using easy_io
+        video_frames, video_metadata = easy_io.load(video_path)  # Returns (T, H, W, C) numpy array
+    except Exception as e:
+        raise ValueError(f"Failed to load video {video_path}: {e}")
+
+    # Convert to tensor: (T, H, W, C) -> (C, T, H, W)
+    video_tensor = th.from_numpy(video_frames).float() / 255.0
+    video_tensor = video_tensor.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
+
+    C, T, H, W = video_tensor.shape
+
+    # Adjust frame count: if video is too long, take first target_frames; if too short, pad with last frame
+    if T > target_frames:
+        video_tensor = video_tensor[:, :target_frames, :, :]
+    elif T < target_frames:
+        # Pad with last frame
+        last_frame = video_tensor[:, -1:, :, :]
+        padding_frames = target_frames - T
+        last_frame_repeated = last_frame.repeat(1, padding_frames, 1, 1)
+        video_tensor = th.cat([video_tensor, last_frame_repeated], dim=1)
+
+    # Convert to uint8: (C, T, H, W) -> (T, C, H, W)
+    video_tensor = video_tensor.permute(1, 0, 2, 3)
+    video_tensor = (video_tensor * 255.0).to(th.uint8)
+
+    # Adjust resolution
+    target_h, target_w = target_size
+    if H != target_h or W != target_w:
+        # Use resize and center crop
+        video_tensor = resize_and_crop(video_tensor, target_size)
+
+    # Convert back to (C, T, H, W)
+    video_tensor = video_tensor.permute(1, 0, 2, 3)
+
+    return video_tensor
+
+
+def resize_and_crop(video: th.Tensor, target_size: tuple[int, int]) -> th.Tensor:
+    """
+    Resize video and center crop.
+
+    Args:
+        video: Input video with shape (T, C, H, W)
+        target_size: Target resolution (H, W)
+
+    Returns:
+        Resized video with shape (T, C, target_H, target_W)
+    """
+    orig_h, orig_w = video.shape[2], video.shape[3]
+    target_h, target_w = target_size
+
+    # Calculate scaling ratio to match the smaller dimension to target
+    scaling_ratio = max((target_w / orig_w), (target_h / orig_h))
+    resizing_shape = (int(scaling_ratio * orig_h), int(scaling_ratio * orig_w))
+
+    video_resized = torchvision.transforms.functional.resize(video, resizing_shape)
+    video_cropped = torchvision.transforms.functional.center_crop(video_resized, target_size)
+
+    return video_cropped
+
+
+def load_multiview_videos(
+    input_root: Path,
+    video_id: str,
+    camera_order: list[str],
+    target_frames: int = 93,
+    target_size: tuple[int, int] = (720, 1280),
+) -> th.Tensor:
+    """
+    Load multi-view videos.
+
+    Args:
+        input_root: Input root directory
+        video_id: Video ID (filename without extension)
+        camera_order: List of camera names in order
+        target_frames: Target number of frames per view
+        target_size: Target resolution (H, W)
+
+    Returns:
+        Multi-view video tensor with shape (C, V*T, H, W)
+    """
+    videos_dir = input_root / "videos"
+    video_tensors = []
+
+    for camera in camera_order:
+        if (videos_dir / f"ftheta_{camera}").exists():
+            sub_dir = f"ftheta_{camera}"
+        elif (videos_dir / camera).exists():
+            sub_dir = camera
+        else:
+            raise FileNotFoundError(f"Folder not found: {videos_dir / f'ftheta_{camera}'} or {videos_dir / camera}")
+
+        video_path = videos_dir / sub_dir / f"{video_id}.mp4"
+
+        if not video_path.exists():
+            raise FileNotFoundError(f"Video not found: {video_path}")
+
+        # Load single view video: (C, T, H, W)
+        video_tensor = load_video(str(video_path), target_frames, target_size)
+        video_tensors.append(video_tensor)
+
+    # Concatenate all views: (C, V*T, H, W)
+    multiview_video = th.cat(video_tensors, dim=1)
+
+    return multiview_video
+
+
+def load_multiview_captions(
+    input_root: Path, video_id: str, camera_order: list[str], add_camera_prefix: bool = True
+) -> list[str]:
+    """
+    Load multi-view captions. Uses default prompt if captions directory does not exist.
+
+    Args:
+        input_root: Input root directory
+        video_id: Video ID (filename without extension)
+        camera_order: List of camera names in order
+        add_camera_prefix: Whether to add camera-specific prefix to captions
+
+    Returns:
+        List of captions, one per view
+    """
+    captions_dir = input_root / "captions"
+
+    # If captions directory does not exist, use default prompt
+    if not captions_dir.exists():
+        log.warning(
+            f"Captions directory not found: {captions_dir}. Using default driving scene prompt for all cameras."
+        )
+        return [DEFAULT_DRIVING_SCENE_PROMPT] * len(camera_order)
+
+    captions = []
+
+    for camera in camera_order:
+        if (captions_dir / f"ftheta_{camera}").exists():
+            sub_dir = f"ftheta_{camera}"
+        elif (captions_dir / camera).exists():
+            sub_dir = camera
+        else:
+            raise FileNotFoundError(f"Folder not found: {captions_dir / f'ftheta_{camera}'} or {captions_dir / camera}")
+
+        caption_filename = f"{sub_dir}/{video_id}.txt"
+        caption_path = captions_dir / caption_filename
+
+        if not caption_path.exists():
+            raise FileNotFoundError(f"Caption file not found: {caption_path}")
+
+        with open(caption_path, "r", encoding="utf-8") as f:
+            caption = f.read().strip()
+
+        # Add camera-specific prefix if enabled
+        if add_camera_prefix and camera in CAMERA_TO_CAPTION_PREFIX:
+            caption = f"{CAMERA_TO_CAPTION_PREFIX[camera]} {caption}"
+
+        captions.append(caption)
+
+    return captions
+
+
+def construct_data_batch(
+    multiview_video: th.Tensor,
+    captions: list[str],
+    camera_order: list[str],
+    num_conditional_frames: int = 0,
+    fps: float = 15.0,
+    target_frames_per_view: int = 93,
+) -> dict:
+    """
+    Construct data_batch for model inference.
+
+    Args:
+        multiview_video: Multi-view video tensor with shape (C, V*T, H, W)
+        captions: List of captions
+        camera_order: List of camera names in order
+        num_conditional_frames: Number of conditional frames
+        fps: Frames per second
+        target_frames_per_view: Number of frames per view
+
+    Returns:
+        data_batch dictionary
+    """
+    C, VT, H, W = multiview_video.shape
+    n_views = len(camera_order)
+    T = VT // n_views
+
+    # Add batch dimension: (C, V*T, H, W) -> (1, C, V*T, H, W)
+    multiview_video = multiview_video.unsqueeze(0)
+
+    # Construct correct view_indices based on camera order
+    # Each view's T frames all use that view's corresponding view index
+    view_indices_list = []
+    for camera in camera_order:
+        view_idx = CAMERA_TO_VIEW_INDEX[camera]
+        view_indices_list.extend([view_idx] * T)
+    view_indices = th.tensor(view_indices_list, dtype=th.int64).unsqueeze(0)  # (1, V*T)
+
+    # Construct view_indices_selection: view indices of cameras in camera_order
+    view_indices_selection = th.tensor(
+        [CAMERA_TO_VIEW_INDEX[camera] for camera in camera_order], dtype=th.int64
+    ).unsqueeze(0)  # (1, n_views)
+
+    # Find position of front_wide_120fov in camera_order as ref_cam_view_idx_sample_position
+    ref_cam_position = (
+        camera_order.index("camera_front_wide_120fov") if "camera_front_wide_120fov" in camera_order else 0
+    )
+
+    # Construct data_batch
+    data_batch = {
+        "video": multiview_video,
+        "ai_caption": [captions],
+        "view_indices": view_indices,  # (1, V*T), using correct view index
+        "fps": th.tensor([fps], dtype=th.float64),
+        "chunk_index": th.tensor([0], dtype=th.int64),
+        "frame_indices": th.arange(target_frames_per_view).unsqueeze(0),  # (1, T)
+        "num_video_frames_per_view": th.tensor([target_frames_per_view], dtype=th.int64),
+        "view_indices_selection": view_indices_selection,  # (1, n_views), using correct view index
+        "camera_keys_selection": [camera_order],
+        "sample_n_views": th.tensor([n_views], dtype=th.int64),
+        "padding_mask": th.zeros(1, 1, H, W, dtype=th.float32),
+        "ref_cam_view_idx_sample_position": th.tensor([ref_cam_position], dtype=th.int64),
+        "front_cam_view_idx_sample_position": [None],
+        "original_hw": th.tensor([[[H, W]] * n_views], dtype=th.int64),  # (1, n_views, 2)
+        NUM_CONDITIONAL_FRAMES_KEY: num_conditional_frames,
+    }
+
+    return data_batch
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description="Predict2 Multiview inference from videos and captions")
+    parser.add_argument("--experiment", type=str, required=True, help="Experiment config")
+    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(
+        "--context_parallel_size",
+        type=int,
+        default=1,
+        help="Context parallel size (number of GPUs to split context over). Set to 8 for 8 GPUs",
+    )
+    # Generation parameters
+    parser.add_argument("--guidance", type=int, default=5, help="Guidance value")
+    parser.add_argument("--fps", type=int, default=15, help="Output video FPS")
+    parser.add_argument("--seed", type=int, default=1, help="Random seed")
+    parser.add_argument("--num_conditional_frames", type=int, default=0, help="Number of conditional frames")
+    parser.add_argument("--num_steps", type=int, default=35, help="Number of diffusion steps")
+    # Input/output
+    parser.add_argument(
+        "--input_root",
+        type=str,
+        required=True,
+        help="Input root directory containing videos/ and captions/ subdirectories",
+    )
+    parser.add_argument("--save_root", type=str, default="results/predict2_multiview_av/", help="Save root")
+    parser.add_argument("--max_samples", type=int, default=5, help="Maximum number of samples to generate")
+    parser.add_argument(
+        "--stack_mode",
+        type=str,
+        default="time",
+        choices=["height", "width", "time", "grid"],
+        help="Video stacking mode for visualization. grid will create a 3x3 grid of views.",
+    )
+    # Video parameters
+    parser.add_argument("--target_frames", type=int, default=93, help="Target number of frames per view")
+    parser.add_argument("--target_height", type=int, default=720, help="Target video height")
+    parser.add_argument("--target_width", type=int, default=1280, help="Target video width")
+    # Caption parameters
+    parser.add_argument(
+        "--add_camera_prefix",
+        action="store_true",
+        default=True,
+        help="Add camera-specific position/orientation prefix to captions",
+    )
+    parser.add_argument(
+        "--no_camera_prefix",
+        action="store_false",
+        dest="add_camera_prefix",
+        help="Do not add camera-specific prefix to captions",
+    )
+    # Experiment options
+    parser.add_argument(
+        "opts",
+        help="""
+        Modify config options at the end of the command. For Yacs configs, use
+        space-separated "PATH.KEY VALUE" pairs.
+        For python-based LazyConfig, use "path.key=value".
+        """.strip(),
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    os.environ["NVTE_FUSED_ATTN"] = "0"
+    th.backends.cudnn.benchmark = False
+    th.backends.cudnn.deterministic = True
+    th.enable_grad(False)
+
+    args = parse_arguments()
+
+    # Initialize inference handler
+    vid2vid_cli = Vid2VidInference(
+        args.experiment,
+        args.ckpt_path,
+        args.s3_cred,
+        context_parallel_size=args.context_parallel_size,
+        experiment_opts=args.opts,
+    )
+    mem_bytes = th.cuda.memory_allocated(device=th.device("cuda" if th.cuda.is_available() else "cpu"))
+    log.info(f"GPU memory usage after model dcp.load: {mem_bytes / (1024**3):.2f} GB")
+
+    # Only process files on rank 0
+    rank0 = True
+    if args.context_parallel_size > 1:
+        rank0 = distributed.get_rank() == 0
+
+    # Create output directory
+    os.makedirs(args.save_root, exist_ok=True)
+
+    input_root = Path(args.input_root)
+    videos_dir = input_root / "videos"
+
+    # Get all video IDs (from first camera directory)
+    if (videos_dir / f"ftheta_{DEFAULT_CAMERA_ORDER[0]}").exists():
+        first_camera_dir = videos_dir / f"ftheta_{DEFAULT_CAMERA_ORDER[0]}"
+    else:
+        first_camera_dir = videos_dir / DEFAULT_CAMERA_ORDER[0]
+
+    video_files = sorted(first_camera_dir.glob("*.mp4"))
+    video_ids = [f.stem for f in video_files[: args.max_samples]]
+
+    log.info(f"Found {len(video_ids)} video IDs, processing {min(len(video_ids), args.max_samples)} samples")
+
+    for i, video_id in enumerate(video_ids):
+        if rank0:
+            log.info(f"Processing sample {i + 1}/{len(video_ids)}: {video_id}")
+
+        try:
+            # Load multi-view captions
+            captions = load_multiview_captions(
+                input_root, video_id, DEFAULT_CAMERA_ORDER, add_camera_prefix=args.add_camera_prefix
+            )
+
+            # Decide whether to load real videos based on num_conditional_frames
+            if args.num_conditional_frames == 0:
+                # Pure text-to-video generation, use dummy all-zero video
+                n_views = len(DEFAULT_CAMERA_ORDER)
+                multiview_video = th.zeros(
+                    3,  # C
+                    n_views * args.target_frames,  # V*T
+                    args.target_height,  # H
+                    args.target_width,  # W
+                    dtype=th.uint8,
+                )
+                if rank0:
+                    log.info(f"Using dummy video (all zeros) for text-to-video generation: {multiview_video.shape}")
+            else:
+                # Need video conditioning, load real videos
+                multiview_video = load_multiview_videos(
+                    input_root,
+                    video_id,
+                    DEFAULT_CAMERA_ORDER,
+                    target_frames=args.target_frames,
+                    target_size=(args.target_height, args.target_width),
+                )
+                if rank0:
+                    log.info(f"Loaded multiview video: {multiview_video.shape}")
+
+            if rank0:
+                log.info(f"Loaded {len(captions)} captions")
+                log.info(f"First caption preview: {captions[0][:100]}...")
+
+            # Construct data_batch
+            data_batch = construct_data_batch(
+                multiview_video,
+                captions,
+                DEFAULT_CAMERA_ORDER,
+                num_conditional_frames=args.num_conditional_frames,
+                fps=args.fps,
+                target_frames_per_view=args.target_frames,
+            )
+
+            # Run inference, already arranged the video according to stack_mode
+            video = vid2vid_cli.generate_from_batch(
+                data_batch,
+                guidance=args.guidance,
+                seed=args.seed + i,
+                num_steps=args.num_steps,
+                stack_mode=args.stack_mode,
+            )
+
+            # Save results
+            if rank0:
+                save_name = f"inference_av_{video_id}"
+                save_img_or_video(video[0], f"{args.save_root}/{save_name}", fps=args.fps)
+                log.info(f"Saved video to {args.save_root}/{save_name}")
+
+        except Exception as e:
+            log.error(f"Error processing {video_id}: {e}")
+            import traceback
+
+            traceback.print_exc()
+            continue
+
+    # Synchronize all processes
+    if args.context_parallel_size > 1:
+        th.distributed.barrier()
+
+    # Cleanup distributed resources
+    vid2vid_cli.cleanup()
+
+
+if __name__ == "__main__":
+    main()

cosmos_predict2/_src/predict2_multiview/scripts/mv_visualize_helper.py

@@ -0,0 +1,164 @@
+# 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 os
+
+import torch as th
+
+from cosmos_predict2._src.imaginaire.utils import log
+from cosmos_predict2._src.imaginaire.visualize.video import save_img_or_video
+
+# Visualization layouts for multi-view video arrangement
+VISUALIZE_LAYOUTS_MADS = {
+    "width": [
+        [
+            "camera_rear_left_70fov",
+            "camera_cross_left_120fov",
+            "camera_front_wide_120fov",
+            "camera_cross_right_120fov",
+            "camera_rear_right_70fov",
+            "camera_rear_tele_30fov",
+            "camera_front_tele_30fov",
+        ]
+    ],
+    "height": [
+        ["camera_rear_left_70fov"],
+        ["camera_cross_left_120fov"],
+        ["camera_front_wide_120fov"],
+        ["camera_cross_right_120fov"],
+        ["camera_rear_right_70fov"],
+        ["camera_rear_tele_30fov"],
+        ["camera_front_tele_30fov"],
+    ],
+    "grid": [
+        [None, "camera_front_tele_30fov", None],
+        ["camera_cross_left_120fov", "camera_front_wide_120fov", "camera_cross_right_120fov"],
+        ["camera_rear_left_70fov", "camera_rear_tele_30fov", "camera_rear_right_70fov"],
+    ],
+}
+
+VISUALIZE_LAYOUTS_AGIBOT = {
+    "width": [
+        ["head_color", "hand_left", "hand_right"],
+    ],
+}
+
+VISUALIZE_LAYOUTS = {
+    "mads": VISUALIZE_LAYOUTS_MADS,
+    "agibot": VISUALIZE_LAYOUTS_AGIBOT,
+}
+
+
+def arrange_video_visualization(mv_video, data_batch, method="width", dataset="mads"):
+    """
+    Rearrange multi-view video based on specified layout method.
+
+    Args:
+        mv_video: (B, C, V * T, H, W) - Multi-view video tensor
+        data_batch: Batch containing camera order information
+        method: Method to arrange video visualization. Can be "width", "height", "grid", or "time".
+                - "width": Arrange all 7 views in a single horizontal row
+                - "height": Arrange all 7 views in a single vertical column
+                - "grid": Arrange views in a 3x3 grid with None values for empty positions
+                - "time": Keep original format (V*T in time dimension, no spatial rearrangement)
+    Returns:
+        Video tensor arranged according to the layout:
+        - For "width": (B, C, T, H, V*W) where V=7
+        - For "height": (B, C, T, V*H, W) where V=7
+        - For "grid": (B, C, T, 3*H, 3*W) with black padding for None positions
+        - For "time": (B, C, V*T, H, W) (unchanged)
+    """
+    # Handle "time" mode - return video unchanged
+    if method == "time":
+        return mv_video
+
+    if method not in VISUALIZE_LAYOUTS[dataset]:
+        raise ValueError(
+            f"Unsupported visualization method: {method}. Choose from {list(VISUALIZE_LAYOUTS[dataset].keys()) + ['time']}"
+        )
+
+    current_view_order = data_batch["camera_keys_selection"][0]
+    n_views = len(current_view_order)
+    B, C, VT, H, W = mv_video.shape
+    T = VT // n_views
+
+    # Reshape to separate view and time dimensions: B C (V T) H W -> B C V T H W
+    video = mv_video.view(B, C, n_views, T, H, W)
+
+    # Create mapping from view name to tensor index
+    view_name_to_video_tensor_idx = {view_name: idx for idx, view_name in enumerate(current_view_order)}
+
+    # Create black view for None positions (used in grid layout)
+    black_view = th.zeros(B, C, T, H, W, dtype=video.dtype, device=video.device)
+
+    # Get layout definition
+    layout_definition = VISUALIZE_LAYOUTS[dataset][method]
+
+    # Arrange video according to layout
+    grid_rows = []
+    for row_of_view_names in layout_definition:
+        row_tensors = []
+        for view_name in row_of_view_names:
+            if view_name is not None and view_name in view_name_to_video_tensor_idx:
+                tensor_idx = view_name_to_video_tensor_idx[view_name]
+                # video is B C V T H W. Get tensor for view: B C T H W
+                row_tensors.append(video[:, :, tensor_idx])
+            else:
+                # Use black view for None positions or missing views
+                row_tensors.append(black_view)
+        grid_rows.append(th.cat(row_tensors, dim=-1))  # Concat on W dimension
+
+    # Concatenate rows on H dimension
+    video = th.cat(grid_rows, dim=-2)  # Concat on H dimension
+
+    return video
+
+
+def save_each_view_separately(
+    mv_video: th.Tensor,
+    data_batch: dict,
+    save_dir: str,
+    fps: float = 10.0,
+    prefix: str = "",
+) -> None:
+    """
+    Save each camera view as a separate video file.
+
+    Args:
+        mv_video: Multi-view video tensor with shape (C, V*T, H, W) where V is number of views
+        data_batch: Data batch containing camera_keys_selection with actual camera order
+        save_dir: Directory to save individual view videos
+        fps: Frames per second for saved videos
+        prefix: Optional prefix for saved filenames (e.g., "video_" or "control_")
+    """
+    os.makedirs(save_dir, exist_ok=True)
+
+    # Extract actual camera order from data_batch
+    camera_order = data_batch["camera_keys_selection"][0]
+
+    C, VT, H, W = mv_video.shape
+    n_views = len(camera_order)
+    T = VT // n_views
+
+    # Reshape to separate views: (C, V*T, H, W) -> (C, V, T, H, W)
+    video_views = mv_video.view(C, n_views, T, H, W)
+
+    # Save each view
+    for view_idx, camera_name in enumerate(camera_order):
+        view_video = video_views[:, view_idx, :, :, :]  # (C, T, H, W)
+        filename = f"{prefix}{camera_name}" if prefix else camera_name
+        view_save_path = os.path.join(save_dir, filename)
+        save_img_or_video(view_video, view_save_path, fps=fps)
+        log.info(f"Saved view {camera_name} to {view_save_path}")

cosmos_predict2/_src/predict2_multiview/utils/optim_instantiate.py

@@ -0,0 +1,234 @@
+# 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 re
+
+import hydra
+import torch
+from omegaconf import ListConfig
+from torch import nn
+
+from cosmos_predict2._src.imaginaire.utils import log
+from cosmos_predict2._src.predict2.utils.fused_adam_dtensor import FusedAdam
+
+
+def get_regular_param_group(net: nn.Module):
+    """
+    seperate the parameters of the network into two groups: decay and no_decay.
+    based on nano_gpt codebase.
+    """
+    param_dict = {pn: p for pn, p in net.named_parameters()}
+    param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+
+    decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+    nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+    return decay_params, nodecay_params
+
+
+def get_base_optimizer(
+    model: nn.Module,
+    lr: float,
+    weight_decay: float,
+    optim_type: str = "adamw",
+    **kwargs,
+) -> torch.optim.Optimizer:
+    net_decay_param, net_nodecay_param = get_regular_param_group(model)
+
+    num_decay_params = sum(p.numel() for p in net_decay_param)
+    num_nodecay_params = sum(p.numel() for p in net_nodecay_param)
+    net_param_total = num_decay_params + num_nodecay_params
+    log.critical(f"total num parameters : {net_param_total:,}")
+
+    param_group = [
+        {
+            "params": net_decay_param + net_nodecay_param,
+            "lr": lr,
+            "weight_decay": weight_decay,
+        },
+    ]
+
+    if optim_type == "adamw":
+        opt_cls = torch.optim.AdamW
+    elif optim_type == "fusedadam":
+        opt_cls = FusedAdam
+    else:
+        raise ValueError(f"Unknown optimizer type: {optim_type}")
+
+    for k, v in kwargs.items():
+        if isinstance(v, ListConfig):
+            kwargs[k] = list(v)
+
+    return opt_cls(param_group, **kwargs)
+
+
+def get_base_scheduler(
+    optimizer: torch.optim.Optimizer,
+    model: nn.Module,
+    scheduler_config: dict,
+):
+    net_scheduler = hydra.utils.instantiate(scheduler_config)
+    net_scheduler.model = model
+
+    return torch.optim.lr_scheduler.LambdaLR(
+        optimizer,
+        lr_lambda=[
+            net_scheduler.schedule,
+        ],
+    )
+
+
+def get_multiple_optimizer(
+    model: nn.Module,
+    lr: float,
+    weight_decay: float,
+    optim_type: str = "adamw",
+    lr_overrides: list[dict] = None,
+    **kwargs,
+) -> torch.optim.Optimizer:
+    """
+    Get an optimizer with multiple learning rates for different parts of the model,
+    allowing pattern matching for parameter names.
+
+    The logic is:
+    1. All parameters are initially considered for the default learning rate.
+    2. We iterate through lr_overrides. If a parameter's name matches a pattern,
+       it's moved to a group with the specified learning rate. A parameter is only
+       assigned to the *first* pattern it matches.
+
+    Args:
+        model (nn.Module): The model to optimize.
+        lr (float): The default learning rate.
+        weight_decay (float): The default weight decay.
+        optim_type (str): The type of optimizer to use ('adamw' or 'fusedadam').
+        lr_overrides (list[dict], optional): A list of dicts with keys:
+              - 'pattern' (str): The pattern to match (required)
+              - 'lr' (float): The learning rate for matching params (required)
+              - 'match_type' (str): 'regex', 'contains', 'startswith', 'endswith' (default: 'contains')
+
+            Example:
+            [
+                {'pattern': 'cross_view_attn', 'lr': 2e-4, 'match_type': 'contains'},
+                {'pattern': 'text_encoder', 'lr': 1e-5, 'match_type': 'contains'},
+            ]
+
+            This is Hydra-friendly and can be overridden from command line like:
+            optimizer.lr_overrides.0.lr=1e-4
+            optimizer.lr_overrides.1.pattern=vision_encoder
+        **kwargs: Additional arguments for the optimizer.
+
+    Returns:
+        torch.optim.Optimizer: The configured optimizer.
+    """
+    param_dict = {pn: p for pn, p in model.named_parameters() if p.requires_grad}
+
+    # Initialize groups for parameters with overridden LR
+    override_groups = {}  # (lr, has_decay) -> [params]
+    override_groups_name = {}  # (lr, has_decay) -> [name]
+
+    # Initialize lists for parameters with default LR
+    default_decay_params = []
+    default_nodecay_params = []
+
+    # Temporarily hold all params to check against overrides
+    unassigned_params = list(param_dict.items())
+
+    # First, assign params that match an override pattern
+    if lr_overrides:
+        override_list = lr_overrides
+
+        for name, p in list(unassigned_params):
+            assigned = False
+            for override_item in override_list:
+                pattern = override_item["pattern"]
+                special_lr = override_item["lr"]
+                match_type = override_item.get("match_type", "contains")
+
+                # Determine if the parameter name matches
+                matched = False
+                if match_type == "regex":
+                    matched = re.match(pattern, name) is not None
+                elif match_type == "contains":
+                    matched = pattern in name
+                elif match_type == "startswith":
+                    matched = name.startswith(pattern)
+                elif match_type == "endswith":
+                    matched = name.endswith(pattern)
+                else:
+                    raise ValueError(
+                        f"Unknown match_type: {match_type}. Must be one of: regex, contains, startswith, endswith"
+                    )
+
+                if matched:
+                    has_decay = p.dim() >= 2
+                    group_key = (special_lr, has_decay)
+                    if group_key not in override_groups:
+                        override_groups[group_key] = []
+                    if group_key not in override_groups_name:
+                        override_groups_name[group_key] = []
+                    override_groups[group_key].append(p)
+                    override_groups_name[group_key].append(name)
+                    assigned = True
+                    break  # Assign to first matching pattern
+            if assigned:
+                # Remove from unassigned list; this is a bit inefficient but clear
+                unassigned_params = [(n, param) for n, param in unassigned_params if n != name]
+
+    # Assign all remaining params to default groups
+    for name, p in unassigned_params:
+        if p.dim() >= 2:
+            default_decay_params.append(p)
+        else:
+            default_nodecay_params.append(p)
+
+    # Build final param_groups list for the optimizer
+    final_param_groups = []
+    if default_decay_params:
+        final_param_groups.append({"params": default_decay_params, "lr": lr, "weight_decay": weight_decay})
+    if default_nodecay_params:
+        final_param_groups.append({"params": default_nodecay_params, "lr": lr, "weight_decay": 0.0})
+
+    for (special_lr, has_decay), params in override_groups.items():
+        final_param_groups.append(
+            {"params": params, "lr": special_lr, "weight_decay": weight_decay if has_decay else 0.0}
+        )
+
+    # print the parameter names in each group
+    for (special_lr, has_decay), params in override_groups.items():
+        log.critical(f"special_lr {special_lr}: {override_groups_name[(special_lr, has_decay)]}")
+
+    # Log parameter group information
+    total_params = 0
+    log.critical("Optimizer parameter groups:")
+    for i, group in enumerate(final_param_groups):
+        group_params = sum(p.numel() for p in group["params"])
+        total_params += group_params
+        log.critical(
+            f"  Group {i}: num_params={group_params:,}, lr={group['lr']:.1e}, weight_decay={group['weight_decay']}"
+        )
+    log.critical(f"Total trainable parameters: {total_params:,}")
+
+    if optim_type == "adamw":
+        opt_cls = torch.optim.AdamW
+    elif optim_type == "fusedadam":
+        opt_cls = FusedAdam
+    else:
+        raise ValueError(f"Unknown optimizer type: {optim_type}")
+
+    for k, v in kwargs.items():
+        if isinstance(v, ListConfig):
+            kwargs[k] = list(v)
+
+    return opt_cls(final_param_groups, **kwargs)

sh/train_nymeria_longer.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Launch 2-actor joint pose-conditioned full fine-tune on the LONGER Nymeria dataset (200+77 frames),
+# 4-GPU FSDP. Warm-starts the 2B multiview backbone from the HF checkpoint; VAE = local Wan2.1.
+#
+# Usage:
+#   bash sh/train_nymeria_longer.sh            # full run (max_iter from the experiment)
+#   NPROC=4 bash sh/train_nymeria_longer.sh    # override #GPUs
+#   bash sh/train_nymeria_longer.sh trainer.max_iter=5   # pass extra hydra overrides
+#
+# Outputs (checkpoints every save_iter + validation videos) -> /data/model_output/train/<job.name>/
+# W&B: entity=$WANDB_ENTITY, project=$WANDB_PROJECT, run name == <job.name> == the output leaf folder.
+set -euo pipefail
+
+REPO="/home/ubuntu/dahyun/cosmos-predict2.5"
+PY="/home/ubuntu/miniforge3/envs/ego_dh/bin"
+EXP="${EXP:-nymeria_pose_2actor_2b_longer}"
+NPROC="${NPROC:-4}"
+PORT="${PORT:-29581}"
+
+# secrets (HF_TOKEN, WANDB_API_KEY) — kept outside the repo
+source /home/ubuntu/.ego_dh_secrets.sh
+
+export WANDB_ENTITY="${WANDB_ENTITY:-dhyunchung}"          # W&B entity
+export WANDB_PROJECT="${WANDB_PROJECT:-multi-ego}"          # W&B project (decoupled from the output folder)
+export COSMOS_QWEN_TOKENIZER_DIR="${COSMOS_QWEN_TOKENIZER_DIR:-/data/cosmos_reason1_7b}"  # local Qwen tokenizer (offline text encoder)
+export IMAGINAIRE_OUTPUT_ROOT="${IMAGINAIRE_OUTPUT_ROOT:-/data/model_output}"  # -> /data/model_output/train/<name>
+export PYTHONPATH="$REPO:${PYTHONPATH:-}"
+
+cd "$REPO"
+# stdout/stderr log (in addition to the framework's own console.log under the output dir)
+LOGDIR="${LOGDIR:-/data3/nymeria_train_logs}"
+mkdir -p "$LOGDIR"
+LOGFILE="$LOGDIR/${EXP}_$(date +%Y%m%d_%H%M%S).log"
+echo "logging to $LOGFILE"
+"$PY/torchrun" --nproc_per_node="$NPROC" --master_port="$PORT" \
+  -m cosmos_oss.scripts.train \
+  --config=cosmos_predict2/_src/predict2_multiview/configs/vid2vid/config.py -- \
+  experiment="$EXP" "$@" 2>&1 | tee "$LOGFILE"