[Perf]perf: skip unused Cosmos3 I2V conditioning latents (#4614)

Signed-off-by: bjf-frz <frz123db@gmail.com>

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -1031,6 +1031,14 @@ def _save_generated_outputs(
         if not out.success or not out.response_body:
             continue
 
+        if isinstance(out.response_body, bytes):
+            fname = f"req_{idx:04d}.mp4"
+            fpath = os.path.join(save_dir, fname)
+            with open(fpath, "wb") as f:
+                f.write(out.response_body)
+            saved += 1
+            continue
+
         media_urls: list[str] = []
 
         # Chat-completions style: choices[*].message.content[*].image_url.url

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -54,8 +54,10 @@ class StubCosmos3VAE:
             latents_mean=[0.0] * z_dim,
             latents_std=[1.0] * z_dim,
         )
+        self.encode_input_shapes: list[tuple[int, ...]] = []
 
     def encode(self, video: torch.Tensor):
+        self.encode_input_shapes.append(tuple(video.shape))
         latent_frames = (video.shape[2] - 1) // self.config.scale_factor_temporal + 1
         latent_height = video.shape[-2] // self.config.scale_factor_spatial
         latent_width = video.shape[-1] // self.config.scale_factor_spatial
@@ -733,7 +735,7 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
     latents = pipeline._prepare_latents(16, 24, 5, torch.Generator(device="cpu").manual_seed(0))
     assert latents.shape == (1, 2, 2, 2, 3)
 
-    pipeline._encode_conditioning_video = lambda *args, **kwargs: torch.full((1, 2, 2, 2, 3), 5.0)
+    pipeline._encode_conditioning_image_latent = lambda *args, **kwargs: torch.full((1, 2, 1, 2, 3), 5.0)
     i2v_latents, velocity_mask, image_latent = pipeline._prepare_latents_i2v(
         torch.zeros(1, 3, 16, 24), 16, 24, 5, torch.Generator(device="cpu").manual_seed(0)
     )
@@ -783,7 +785,7 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
     torch.testing.assert_close(action, clean)
 
 
-def test_i2v_conditioning_moves_single_frame_before_video_expand(make_cosmos3_pipeline) -> None:
+def test_prepare_latents_i2v_encodes_only_conditioning_frame(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     calls: list[tuple[str, tuple[int, ...]]] = []
 
@@ -795,24 +797,30 @@ def test_i2v_conditioning_moves_single_frame_before_video_expand(make_cosmos3_pi
     class RecordingVAE(StubCosmos3VAE):
         def encode(self, video: torch.Tensor):
             calls.append(("encode", tuple(video.shape)))
-            assert video.is_contiguous()
             return super().encode(video)
 
     pipeline._to_vae_device = record_to_vae_device
     pipeline.vae = RecordingVAE(z_dim=2)
+    generator = torch.Generator(device="cpu").manual_seed(0)
 
-    latent = pipeline._encode_conditioning_video(
+    latents, velocity_mask, image_latent = pipeline._prepare_latents_i2v(
         torch.zeros(1, 3, 16, 24),
-        num_frames=5,
-        height=16,
-        width=24,
+        16,
+        24,
+        9,
+        generator,
     )
 
     assert calls == [
         ("to_vae_device", (1, 3, 16, 24)),
-        ("encode", (1, 3, 5, 16, 24)),
+        ("encode", (1, 3, 1, 16, 24)),
     ]
-    assert latent.shape == (1, 2, 2, 2, 3)
+    assert pipeline.vae.encode_input_shapes[-1] == (1, 3, 1, 16, 24)
+    assert latents.shape == (1, 2, 3, 2, 3)
+    assert image_latent.shape == (1, 2, 1, 2, 3)
+    torch.testing.assert_close(latents[:, :, 0:1], image_latent)
+    assert not torch.allclose(latents[:, :, 1:], image_latent.expand(-1, -1, 2, -1, -1))
+    assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]]
 
 
 def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
@@ -846,6 +854,9 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
         image_latent=torch.full((1, 2, 1, 1, 1), 7.0),
     )
     torch.testing.assert_close(i2v[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0))
+    i2v_noise = pipeline.scheduler.step_calls[-1][0]
+    torch.testing.assert_close(i2v_noise[:, :, 0:1], torch.zeros(1, 2, 1, 1, 1))
+    torch.testing.assert_close(i2v_noise[:, :, 1:2], torch.full((1, 2, 1, 1, 1), 2.0))
 
     pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
     video_result, action_result = pipeline.diffuse(

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -1771,31 +1771,7 @@ class Cosmos3OmniDiffusersPipeline(
 
     # -- I2V latent preparation ---------------------------------------------
 
-    def _encode_conditioning_video(
-        self,
-        image_tensor: torch.Tensor,
-        num_frames: int,
-        height: int,
-        width: int,
-    ) -> torch.Tensor:
-        """VAE-encode a conditioning image as a full-length video.
-
-        The WAN VAE has temporal compression (factor 4), so encoding a single
-        frame produces degenerate temporal features.  We fill the entire
-        pixel-space video with the conditioning image (repeating it across all
-        frames) so the temporal encoder sees plausible content everywhere.
-        The caller keeps only the conditioned latent frame(s) and replaces
-        the rest with noise.
-        """
-        # Move the single conditioning frame first, then materialize the
-        # repeated video on GPU.  Materializing [1, 3, T, H, W] on CPU creates
-        # a large pageable H2D copy on I2V requests.
-        image_tensor = self._to_vae_device(image_tensor, pin_cpu=True)
-        video = image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous()
-
-        latent = self.vae.encode(video).latent_dist.mode()
-
-        # Normalize (inverse of _decode_latents denormalization)
+    def _normalize_vae_latent(self, latent: torch.Tensor) -> torch.Tensor:
         if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"):
             latents_mean, latents_std = self._get_latents_mean_std(latent.device, latent.dtype)
             latent = (latent - latents_mean) / latents_std
@@ -1803,7 +1779,20 @@ class Cosmos3OmniDiffusersPipeline(
             scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
             latent = latent * scaling_factor
 
-        return latent.to(self.dtype)
+        return latent
+
+    def _encode_conditioning_image_latent(self, image_tensor: torch.Tensor) -> torch.Tensor:
+        """VAE-encode the first I2V conditioning frame.
+
+        I2V only consumes the first conditioning latent frame.  Encoding the
+        input image as a one-frame video keeps that latent while avoiding VAE
+        work for repeated frames that are later replaced by noise.
+        """
+        image_tensor = self._to_vae_device(image_tensor, pin_cpu=True)
+        video = image_tensor.unsqueeze(2)
+        latent = self.vae.encode(video).latent_dist.mode()
+        latent = self._normalize_vae_latent(latent)
+        return latent[:, :, 0:1, :, :].to(self.dtype)
 
     def _latent_hw_from_image_size(self, image_size: Any | None) -> tuple[int, int] | None:
         if image_size is None:
@@ -1837,14 +1826,7 @@ class Cosmos3OmniDiffusersPipeline(
 
         video = self._to_vae_device(video_tensor)
         latent = self.vae.encode(video).latent_dist.mode()
-
-        if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"):
-            latents_mean, latents_std = self._get_latents_mean_std(latent.device, latent.dtype)
-            latent = (latent - latents_mean) / latents_std
-        else:
-            scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
-            latent = latent * scaling_factor
-
+        latent = self._normalize_vae_latent(latent)
         latent = self._crop_latent_to_image_size(latent, image_size)
         return latent.to(self.dtype)
 
@@ -1875,13 +1857,12 @@ class Cosmos3OmniDiffusersPipeline(
             dtype=self.dtype,
         )
 
-        cond_latent = self._encode_conditioning_video(image_tensor, num_frames, height, width)
-        image_latent = cond_latent[:, :, 0:1, :, :]
+        image_latent = self._encode_conditioning_image_latent(image_tensor)
+        latents = noise
+        latents[:, :, 0:1, :, :] = image_latent
 
-        condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype)
-        condition_mask[:, :, 0, :, :] = 1.0
-        latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise
-        velocity_mask = 1.0 - condition_mask
+        velocity_mask = torch.ones(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype)
+        velocity_mask[:, :, 0, :, :] = 0.0
         return latents, velocity_mask, image_latent
 
     def _prepare_latents_v2v(
@@ -2191,7 +2172,10 @@ class Cosmos3OmniDiffusersPipeline(
         ) -> torch.Tensor | tuple[torch.Tensor, ...]:
             video_pred, action_pred, sound_pred = _split_noise_pred(noise_pred)
             if velocity_mask is not None:
-                video_pred = video_pred * velocity_mask
+                if image_latent is not None and condition_latents is None:
+                    video_pred[:, :, 0:1, :, :] = 0
+                else:
+                    video_pred = video_pred * velocity_mask
             if action_pred is not None and action_velocity_mask is not None:
                 action_pred = action_pred * action_velocity_mask
                 if raw_action_dim is not None and 0 < raw_action_dim < action_pred.shape[-1]: