minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] optimize: optimize Cosmos3 i2v latent prep (#27084)

Browse Source
This commit is contained in:
Chi McIsaac
2026-06-03 10:23:37 +08:00
committed by GitHub
parent 3715a07e87
commit 3d54056391
1 changed files with 2 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py
+2 -5
View File
@@ -321,11 +321,8 @@ class Cosmos3LatentPreparationStage(PipelineStage):
if is_i2v: if is_i2v:
vae_dtype = next(self.vae.parameters()).dtype vae_dtype = next(self.vae.parameters()).dtype
pixel_video = ( pixel_video = batch.preprocessed_image.unsqueeze(2).to(
batch.preprocessed_image.unsqueeze(2) device=device, dtype=vae_dtype
.expand(-1, -1, batch.num_frames, -1, -1)
.contiguous()
.to(device=device, dtype=vae_dtype)
) )
with torch.no_grad(): with torch.no_grad():
cond_latent = self._vae_encode(pixel_video).to(dtype) cond_latent = self._vae_encode(pixel_video).to(dtype)