From c72fc2de2cef95076d38b81050832110a5f23662 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 7 Jun 2026 16:37:55 -0700 Subject: [PATCH] [misc]: simplify cosmos3 audio (fold sound split into _split_flat_latent) Collapse the byte-identical _split_sound_latent into the generic _split_flat_latent (duck-typed on numel/shape). Behavior-preserving; 130 tests green. --- .../basic/cosmos3/cosmos3_pipeline.py | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py index 8348531f0..4882fe222 100644 --- a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py @@ -182,8 +182,12 @@ class Cosmos3VisionSpec: # =========================================================================== # Pure denoise/CFG math (parity oracle target) # =========================================================================== -def _split_flat_latent(flat: torch.Tensor, specs: list[Cosmos3VisionSpec]) -> list[torch.Tensor]: - """Split a flat latent vector into per-vision-item ``[C, T, H, W]`` tensors.""" +def _split_flat_latent(flat: torch.Tensor, specs: list[Any]) -> list[torch.Tensor]: + """Split a flat vector into per-item tensors via each spec's ``numel``/``shape``. + + Shared by vision (``[C, T, H, W]``) and sound (``[C, T]``) specs — every spec + exposes ``numel`` and ``shape``. + """ out: list[torch.Tensor] = [] offset = 0 for spec in specs: @@ -211,16 +215,6 @@ class Cosmos3SoundSpec: return int(math.prod(self.shape)) -def _split_sound_latent(flat: torch.Tensor, specs: list[Cosmos3SoundSpec]) -> list[torch.Tensor]: - """Split a flat sound vector into per-item ``[C, T]`` tensors.""" - out: list[torch.Tensor] = [] - offset = 0 - for spec in specs: - out.append(flat[offset:offset + spec.numel].reshape(spec.shape)) - offset += spec.numel - return out - - def cosmos3_get_cfg_velocity( *, transformer: Any, @@ -267,7 +261,7 @@ def cosmos3_get_cfg_velocity( # framework per-sample concat order ([vision_i | sound_i]); single sample here. vision_total = sum(spec.numel for spec in specs) noise_x_vision = _split_flat_latent(flat_latent[:vision_total], specs) - noise_x_sound = (_split_sound_latent(flat_latent[vision_total:], sound_specs) if sound_specs else None) + noise_x_sound = (_split_flat_latent(flat_latent[vision_total:], sound_specs) if sound_specs else None) device = next(transformer.parameters()).device def _run(token_ids: list[int]) -> torch.Tensor: