From f0d7fe18f2c16dcbcaf4032009a13e2060db112b Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Mon, 22 Jun 2026 09:59:03 +0800 Subject: [PATCH] [Perf][Cosmos3] Optimize I2V conditioning transfer (8.63% faster x-inference) (#4467) Signed-off-by: david6666666 <530634352@qq.com> Co-authored-by: Hongsheng Liu --- .../models/cosmos3/test_cosmos3_pipeline.py | 32 ++++++++++++ .../models/cosmos3/pipeline_cosmos3.py | 52 ++++++++++++------- 2 files changed, 65 insertions(+), 19 deletions(-) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index ddca947cd..73e1e9851 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -783,6 +783,38 @@ 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: + pipeline = make_cosmos3_pipeline() + calls: list[tuple[str, tuple[int, ...]]] = [] + + def record_to_vae_device(tensor: torch.Tensor, *, pin_cpu: bool = False) -> torch.Tensor: + assert pin_cpu is True + calls.append(("to_vae_device", tuple(tensor.shape))) + return tensor + + 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) + + latent = pipeline._encode_conditioning_video( + torch.zeros(1, 3, 16, 24), + num_frames=5, + height=16, + width=24, + ) + + assert calls == [ + ("to_vae_device", (1, 3, 16, 24)), + ("encode", (1, 3, 5, 16, 24)), + ] + assert latent.shape == (1, 2, 2, 2, 3) + + 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) diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index b3dace893..724d87319 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1644,18 +1644,36 @@ class Cosmos3OmniDiffusersPipeline( # -- VAE decode ---------------------------------------------------------- + def _get_latents_mean_std(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + cached = getattr(self, "_latents_mean_std", None) + if cached is not None: + latents_mean, latents_std = cached + if latents_mean.device == device and latents_mean.dtype == dtype: + return latents_mean, latents_std + + latents_mean = torch.as_tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1) + latents_std = torch.as_tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(1, -1, 1, 1, 1) + self._latents_mean_std = (latents_mean, latents_std) + return latents_mean, latents_std + + def _to_vae_device(self, tensor: torch.Tensor, *, pin_cpu: bool = False) -> torch.Tensor: + if tensor.device == self.device and tensor.dtype == self.vae.dtype: + return tensor + + non_blocking = False + if tensor.device.type == "cpu" and self.device.type == "cuda": + if pin_cpu and not tensor.is_pinned(): + tensor = tensor.pin_memory() + non_blocking = tensor.is_pinned() + + return tensor.to(device=self.device, dtype=self.vae.dtype, non_blocking=non_blocking) + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: latents = latents.to(self.vae.dtype) if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): - if not hasattr(self, "_latents_mean"): - self._latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(self.device, self.vae.dtype) - ) - self._latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(self.device, self.vae.dtype) - ) - latents = (latents * self._latents_std) + self._latents_mean + latents_mean, latents_std = self._get_latents_mean_std(latents.device, latents.dtype) + latents = (latents * latents_std) + latents_mean else: scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) latents = latents / scaling_factor @@ -1769,18 +1787,17 @@ class Cosmos3OmniDiffusersPipeline( The caller keeps only the conditioned latent frame(s) and replaces the rest with noise. """ - # image_tensor: [1, 3, H, W] -> [1, 3, num_frames, H, W] + # 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() - video = video.to(device=self.device, dtype=self.vae.dtype) latent = self.vae.encode(video).latent_dist.mode() # Normalize (inverse of _decode_latents denormalization) if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) - ) - latents_std = torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) + 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) @@ -1818,14 +1835,11 @@ class Cosmos3OmniDiffusersPipeline( if 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)}.") - video = video_tensor.to(device=self.device, dtype=self.vae.dtype) + 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 = ( - torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) - ) - latents_std = torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) + 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)