From 3d540563916394115f9612f30601ae3fa97ada7c Mon Sep 17 00:00:00 2001 From: Chi McIsaac <153383231+qimcis@users.noreply.github.com> Date: Tue, 2 Jun 2026 19:23:37 -0700 Subject: [PATCH] [diffusion] optimize: optimize Cosmos3 i2v latent prep (#27084) --- .../pipelines_core/stages/model_specific_stages/cosmos3.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 09f6d8685..84987622a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -321,11 +321,8 @@ class Cosmos3LatentPreparationStage(PipelineStage): if is_i2v: vae_dtype = next(self.vae.parameters()).dtype - pixel_video = ( - batch.preprocessed_image.unsqueeze(2) - .expand(-1, -1, batch.num_frames, -1, -1) - .contiguous() - .to(device=device, dtype=vae_dtype) + pixel_video = batch.preprocessed_image.unsqueeze(2).to( + device=device, dtype=vae_dtype ) with torch.no_grad(): cond_latent = self._vae_encode(pixel_video).to(dtype)