nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[misc]: simplify cosmos3 audio (fold sound split into _split_flat_latent)

Browse Source 
Collapse the byte-identical _split_sound_latent into the generic
_split_flat_latent (duck-typed on numel/shape). Behavior-preserving;
130 tests green.
This commit is contained in:
SolitaryThinker
2026-07-13 03:23:45 -07:00
parent 41e2d3ee8b
commit c72fc2de2c
1 changed files with 7 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py
+7 -13
View File
@@ -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: