[Model] Cosmos3: Skip unused action conditioning frames (#4789)

Signed-off-by: Sean Gillen <sgillen@nvidia.com>

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -916,6 +916,49 @@ def test_prepare_latents_i2v_encodes_only_conditioning_frame(make_cosmos3_pipeli
     assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]]
 
 
+@pytest.mark.parametrize("mode", ["policy", "forward_dynamics"])
+def test_prepare_action_video_latents_encodes_only_conditioning_frame(make_cosmos3_pipeline, mode: str) -> None:
+    pipeline = make_cosmos3_pipeline()
+    video = torch.zeros(1, 3, 9, 24, 32)
+
+    latents, velocity_mask, condition_latents = pipeline._prepare_latents_action_video(
+        video,
+        mode,
+        24,
+        32,
+        9,
+        torch.Generator(device="cpu").manual_seed(0),
+        image_size=torch.tensor([1, 3, 16, 24]),
+    )
+
+    assert pipeline.vae.encode_input_shapes == [(1, 3, 1, 24, 32)]
+    assert latents.shape == (1, 2, 3, 2, 3)
+    assert condition_latents.shape == latents.shape
+    torch.testing.assert_close(condition_latents[:, :, 0], torch.ones(1, 2, 2, 3))
+    assert torch.count_nonzero(condition_latents[:, :, 1:]) == 0
+    torch.testing.assert_close(latents[:, :, 0:1], condition_latents[:, :, 0:1])
+    assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]]
+
+
+def test_prepare_inverse_dynamics_latents_encodes_full_video(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    video = torch.zeros(1, 3, 9, 16, 24)
+
+    latents, velocity_mask, condition_latents = pipeline._prepare_latents_action_video(
+        video,
+        "inverse_dynamics",
+        16,
+        24,
+        9,
+        torch.Generator(device="cpu").manual_seed(0),
+    )
+
+    assert pipeline.vae.encode_input_shapes == [(1, 3, 9, 16, 24)]
+    assert condition_latents.shape == (1, 2, 3, 2, 3)
+    torch.testing.assert_close(latents, condition_latents)
+    assert torch.count_nonzero(velocity_mask) == 0
+
+
 def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     latents = torch.zeros(1, 2, 1, 1, 1)

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -77,6 +77,7 @@ from .action import (
     normalize_action_mode,
     pad_action_to_dim,
     resolve_domain_id,
+    vision_condition_indexes,
 )
 from .transfer import (
     Cosmos3TransferConfig,
@@ -2009,8 +2010,21 @@ class Cosmos3OmniDiffusersPipeline(
         generator: torch.Generator,
         image_size: Any | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Prepare video latents for action modes with mode-specific conditioning."""
+        """Prepare video latents for action modes with mode-specific conditioning.
+
+        Policy and forward-dynamics modes condition only latent frame zero.
+        The Wan VAE is temporally causal, so encoding only the first pixel
+        frame preserves that latent while avoiding unused future-frame work.
+        Inverse dynamics conditions every latent and keeps the full encode.
+        """
         del height, width
+        if video_tensor.ndim == 4:
+            video_tensor = video_tensor.unsqueeze(0)
+        if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
+            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
+        if video_tensor.shape[2] < 1:
+            raise ValueError("Cosmos3 action video tensor must contain at least one frame.")
+
         C = self.transformer.latent_channel_size
         T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
         latent_hw = self._latent_hw_from_image_size(image_size)
@@ -2026,12 +2040,24 @@ class Cosmos3OmniDiffusersPipeline(
             device=self.device,
             dtype=self.dtype,
         )
-        cond_latent = self._encode_video_tensor(video_tensor, image_size=image_size)
-        if cond_latent.shape[2:] != noise.shape[2:]:
+        condition_indexes = vision_condition_indexes(mode, num_frames, self.vae_scale_factor_temporal)
+        condition_video = video_tensor[:, :, :1] if condition_indexes == [0] else video_tensor
+        cond_prefix_latent = self._encode_video_tensor(condition_video, image_size=image_size)
+        expected_prefix = (1, C, max(condition_indexes) + 1, H_lat, W_lat)
+        if (
+            cond_prefix_latent.shape[0] != expected_prefix[0]
+            or cond_prefix_latent.shape[1] != expected_prefix[1]
+            or cond_prefix_latent.shape[2] < expected_prefix[2]
+            or cond_prefix_latent.shape[3:] != expected_prefix[3:]
+        ):
             raise ValueError(
                 "Cosmos3 action video latent shape mismatch: "
-                f"encoded={tuple(cond_latent.shape)}, expected={tuple(noise.shape)}."
+                f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}."
             )
+
+        condition_latents = torch.zeros_like(noise)
+        for index in condition_indexes:
+            condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1]
         condition_mask = build_vision_condition_mask(
             mode,
             num_frames,
@@ -2039,9 +2065,9 @@ class Cosmos3OmniDiffusersPipeline(
             device=self.device,
             dtype=self.dtype,
         )
-        latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise
+        latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
         velocity_mask = 1.0 - condition_mask
-        return latents, velocity_mask, cond_latent
+        return latents, velocity_mask, condition_latents
 
     def _prepare_action_latents(
         self,