From 2447740308d58e3071e08c4a0e5022ef6839338b Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Mon, 22 Jun 2026 17:51:47 +0800 Subject: [PATCH] [Perf]perf: skip unused Cosmos3 I2V conditioning latents (#4614) Signed-off-by: bjf-frz --- .../diffusion/diffusion_benchmark_serving.py | 8 +++ .../models/cosmos3/test_cosmos3_pipeline.py | 29 +++++--- .../models/cosmos3/pipeline_cosmos3.py | 66 +++++++------------ 3 files changed, 53 insertions(+), 50 deletions(-) diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index bf08b8c2b..7e0bd4ed1 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/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 diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 73e1e9851..493929f4a 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/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( diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 724d87319..618511c1d 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/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]: