[diffusion] feat: enable parallel decode for cosmos3(#27037)

This commit is contained in:
Mick
2026-06-02 18:18:18 +08:00
committed by GitHub
parent 84e1108312
commit a777672939
@@ -56,11 +56,10 @@ class Cosmos3Config(PipelineConfig):
# Encoder is needed for I2V; T2V/T2I never invoke it. # Encoder is needed for I2V; T2V/T2I never invoke it.
self.vae_config.load_encoder = True self.vae_config.load_encoder = True
self.vae_config.load_decoder = True self.vae_config.load_decoder = True
# WanVAE defaults use_parallel_encode/decode to True, which silently # keep WanVAE encode replicated because parallel encode changes I2V
# activates an SP-sharded VAE path when sp_world_size > 1 and produces # conditioning latents when sp_world_size > 1
# garbled pixels for cosmos3's latent shape.
self.vae_config.use_parallel_encode = False self.vae_config.use_parallel_encode = False
self.vae_config.use_parallel_decode = False self.vae_config.use_parallel_decode = True
def adjust_num_frames(self, num_frames: int) -> int: def adjust_num_frames(self, num_frames: int) -> int:
"""Round ``num_frames`` so ``(n - 1) % 4 == 0`` for the VAE. """Round ``num_frames`` so ``(n - 1) % 4 == 0`` for the VAE.