From 9de2ab750541f5217ac85585ffcd408cfb0bfb08 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 7 Jun 2026 09:50:51 -0700 Subject: [PATCH] [feat] cosmos3 PR2: native AVAE sound decoder + bit-exact parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First audio component: the Cosmos3 sound tokenizer's decode path (t2vs needs only DECODE — generate sound latents, decode to waveform). Finding: the shipped `sound_tokenizer` checkpoint is decoder-only (`decoder.*`, the SpectrogramConvNeXt encoder is not exported) in diffusers AutoencoderOobleck naming, but with SnakeBeta activations (alpha+beta, logscale) and weight_g/weight_v weight-norm — i.e. exactly FastVideo's existing native `OobleckVAE` decoder. So no from-scratch port: reuse it. - `OobleckDecoderBlock`: add `output_padding=stride % 2` to the transpose conv. A provable no-op for even strides (Stable Audio: [2,4,4,8,8]); required for the odd stride in Cosmos3 (`[2,4,5,6,8]`), where the framework's `output_padding=stride%2` otherwise makes the decode 1 sample longer per odd-stride block. Without it, parity diverged (60 vs 59 samples). - `Cosmos3SoundVAE` (`fastvideo/models/audio/cosmos3_avae.py`): decoder-only wrapper — `decode` (Oobleck decoder + clamp to [-1,1], matching AVAEModel), `get_latent_num_samples` (N // hop_size=1920), `from_pretrained` (reads the checkpoint config, strict-loads `decoder.*`). 48 kHz stereo, 473M params. - Parity: `test_cosmos3_avae_parity.py` maps the framework OobleckDecoder (`avae_utils.models`, Sequential naming, the oracle) into the FastVideo decoder and asserts bit-exact decode across stride patterns incl. odd (5) — max abs diff 0.0. Verified the real 1.9GB checkpoint strict-loads and decodes [1,64,25] -> [1,2,48000] (1s @ 48kHz stereo). Full cosmos3 suite: 121 passed, 0 skipped. --- fastvideo/models/audio/cosmos3_avae.py | 133 ++++++++++++++++++ fastvideo/models/vaes/oobleck.py | 3 + .../cosmos3/test_cosmos3_avae_parity.py | 131 +++++++++++++++++ 3 files changed, 267 insertions(+) create mode 100644 fastvideo/models/audio/cosmos3_avae.py create mode 100644 tests/local_tests/cosmos3/test_cosmos3_avae_parity.py diff --git a/fastvideo/models/audio/cosmos3_avae.py b/fastvideo/models/audio/cosmos3_avae.py new file mode 100644 index 000000000..6db721975 --- /dev/null +++ b/fastvideo/models/audio/cosmos3_avae.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 sound tokenizer (AVAE) — decode path. + +The Cosmos3 ``sound_tokenizer`` is an AVAE (audio VAE). Its shipped diffusers +checkpoint is **decoder-only** (``decoder.*``) in ``AutoencoderOobleck`` naming +with SnakeBeta activations and ``weight_g``/``weight_v`` weight-norm — exactly +FastVideo's native :class:`~fastvideo.models.vaes.oobleck.OobleckDecoder` +(verified bit-exact vs the framework in ``test_cosmos3_avae_parity``). Text-to- +video+sound (t2vs) only needs DECODE: the DiT generates the sound latent and +this module decodes it to a waveform, so only the decoder is ported (the +SpectrogramConvNeXt encoder is not exported in the checkpoint). + +Mirrors the framework ``AVAEModel.decode``: run the Oobleck decoder, then clamp +to [-1, 1]. The VAE bottleneck's decode is the identity (the DiT already emits +the post-bottleneck latent), so there is no bottleneck step here. +""" +from __future__ import annotations + +import json +import os +from dataclasses import dataclass, field + +import numpy as np +import torch +import torch.nn as nn + +from fastvideo.logger import init_logger +from fastvideo.models.vaes.oobleck import OobleckDecoder + +logger = init_logger(__name__) + + +@dataclass +class Cosmos3SoundVAEArchConfig: + """Cosmos3 AVAE decoder constants (from ``sound_tokenizer/config.json``).""" + + dec_dim: int = 320 # decoder base channels + vocoder_input_dim: int = 64 # latent channels in + dec_c_mults: list[int] = field(default_factory=lambda: [1, 2, 4, 8, 16]) + dec_strides: list[int] = field(default_factory=lambda: [2, 4, 5, 6, 8]) + audio_channels: int = 2 # stereo + sampling_rate: int = 48000 + + @property + def hop_size(self) -> int: + return int(np.prod(self.dec_strides)) # 1920 + + +class Cosmos3SoundVAE(nn.Module): + """Decoder-only Cosmos3 AVAE: latent ``[B, z, T]`` -> waveform ``[B, C, N]``.""" + + def __init__(self, arch: Cosmos3SoundVAEArchConfig | None = None) -> None: + super().__init__() + self.arch = arch or Cosmos3SoundVAEArchConfig() + self.decoder = OobleckDecoder( + channels=self.arch.dec_dim, + input_channels=self.arch.vocoder_input_dim, + audio_channels=self.arch.audio_channels, + # The framework builds decoder blocks from ``reversed(dec_strides)`` + # (deepest first), so block strides are e.g. [8,6,5,4,2]. + upsampling_ratios=list(reversed(self.arch.dec_strides)), + channel_multiples=list(self.arch.dec_c_mults), + ) + + @property + def sample_rate(self) -> int: + return self.arch.sampling_rate + + @property + def audio_channels(self) -> int: + return self.arch.audio_channels + + @property + def hop_size(self) -> int: + return self.arch.hop_size + + def get_latent_num_samples(self, num_audio_samples: int) -> int: + """Latent length for a given audio length (``AVAEInterface``: ``N // hop``).""" + return int(num_audio_samples) // self.arch.hop_size + + @torch.no_grad() + def decode(self, latent: torch.Tensor) -> torch.Tensor: + """Decode normalized latent ``[B, z, T]`` to waveform ``[B, C, N]`` in [-1, 1]. + + Matches ``AVAEModel.decode``: Oobleck decoder then clamp to [-1, 1] (the + VAE bottleneck decode is identity). + """ + audio = self.decoder(latent) # [B, C, N] + return audio.clamp(-1.0, 1.0) + + @classmethod + def from_pretrained( + cls, + model_path: str, + *, + torch_dtype: torch.dtype | None = None, + ) -> "Cosmos3SoundVAE": + """Build + load the decoder from a ``sound_tokenizer`` directory. + + Reads ``config.json`` (``dec_dim`` / ``vocoder_input_dim`` / + ``dec_c_mults`` / ``dec_strides`` / ``sampling_rate`` / ``stereo``) and + loads the ``decoder.*`` weights (the checkpoint is decoder-only). + """ + from safetensors.torch import load_file + + cfg_path = os.path.join(model_path, "config.json") + with open(cfg_path) as f: + cfg = json.load(f) + arch = Cosmos3SoundVAEArchConfig( + dec_dim=int(cfg["dec_dim"]), + vocoder_input_dim=int(cfg["vocoder_input_dim"]), + dec_c_mults=list(cfg["dec_c_mults"]), + dec_strides=list(cfg["dec_strides"]), + audio_channels=2 if cfg.get("stereo", True) else 1, + sampling_rate=int(cfg.get("sampling_rate", 48000)), + ) + model = cls(arch) + + weights_path = os.path.join(model_path, "diffusion_pytorch_model.safetensors") + state = load_file(weights_path) + # Decoder-only checkpoint: strip the ``decoder.`` prefix. + dec_state = {k[len("decoder."):]: v for k, v in state.items() if k.startswith("decoder.")} + model.decoder.load_state_dict(dec_state, strict=True) + logger.info("Loaded Cosmos3 sound AVAE decoder (%d params) from %s", + sum(p.numel() for p in model.parameters()), model_path) + + if torch_dtype is not None: + model = model.to(dtype=torch_dtype) + model.eval() + return model + + +EntryClass = Cosmos3SoundVAE diff --git a/fastvideo/models/vaes/oobleck.py b/fastvideo/models/vaes/oobleck.py index 8ba2fd7c8..1c5c6ff7b 100644 --- a/fastvideo/models/vaes/oobleck.py +++ b/fastvideo/models/vaes/oobleck.py @@ -94,6 +94,9 @@ class OobleckDecoderBlock(nn.Module): input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2), + # Clean L*stride upsample for both parities; a no-op (0) for even + # strides (Stable Audio), needed for odd strides (Cosmos3: 5). + output_padding=stride % 2, )) self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1) self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3) diff --git a/tests/local_tests/cosmos3/test_cosmos3_avae_parity.py b/tests/local_tests/cosmos3/test_cosmos3_avae_parity.py new file mode 100644 index 000000000..02669e1bf --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_avae_parity.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo Cosmos3 sound decoder vs the framework AVAE. + +The Cosmos3 ``sound_tokenizer`` is an AVAE (audio VAE). Its shipped diffusers +checkpoint is **decoder-only** (``decoder.*``; the SpectrogramConvNeXt encoder is +not exported) in diffusers ``AutoencoderOobleck`` naming, but with **SnakeBeta** +activations (alpha+beta, logscale) and ``weight_g``/``weight_v`` weight-norm — +i.e. exactly FastVideo's existing native ``OobleckVAE`` decoder +(``fastvideo/models/vaes/oobleck.py``). t2vs only needs DECODE (generate sound +latents -> waveform), so this pins the decoder. + +The framework decoder +(``cosmos_framework.model.vfm.tokenizers.audio.avae_utils.models.OobleckDecoder``, +``nn.Sequential`` naming, ``output_padding=stride%2`` on the transpose convs) is +the parity ORACLE. We build a tiny framework decoder, map its weights into the +FastVideo decoder (Sequential -> conv1/block.N/res_unitM/snake1/conv2), and +assert bit-exact decode. Strides include an ODD value (5, as in the real config +``[2,4,5,6,8]``) to exercise the ``output_padding`` path that diverged before. + +CPU / float32. Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_avae_parity.py -q +""" +from __future__ import annotations + +import pytest +import torch + +# The official framework provides the parity oracle. +_fw_models = pytest.importorskip( + "cosmos_framework.model.vfm.tokenizers.audio.avae_utils.models", + reason="cosmos_framework not installed; run in fv-cosmos3 env.", +) +from cosmos_framework.model.vfm.tokenizers.audio.avae_utils.env import ( # noqa: E402 + AttrDict, +) + +from fastvideo.models.vaes.oobleck import OobleckDecoder as FvOobleckDecoder # noqa: E402 + +pytestmark = [pytest.mark.local] + +FwOobleckDecoder = _fw_models.OobleckDecoder + + +def _framework_decoder(dec_dim, vocoder_input_dim, dec_c_mults, dec_strides): + """Framework OobleckDecoder (the parity oracle), non-causal / no-antialias.""" + h = AttrDict({ + "vocoder_input_dim": vocoder_input_dim, + "input_channels": 1, + "stereo": True, # 2 audio channels + "dec_dim": dec_dim, + "dec_c_mults": dec_c_mults, + "dec_strides": dec_strides, + "dec_use_snake": True, + "dec_use_nearest_upsample": False, + "dec_anti_aliasing": False, + "causal": False, + "dec_use_tanh_at_final": False, + "padding_mode": "zeros", + }) + return FwOobleckDecoder(h).eval() + + +def _framework_to_fastvideo_decoder_state(fw_decoder, num_blocks): + """Map framework Sequential decoder weights -> FastVideo decoder names. + + framework: layers.0=first conv; layers.{1..K}=OobleckDecoderBlock + (.layers.0 snake, .1 conv_t, .{2,3,4} ResidualUnit{.layers.0 snake, + .1 conv, .2 snake, .3 conv}); layers.{1+K}=final snake; layers.{2+K}=final conv. + FastVideo: conv1; block.{b}.{snake1,conv_t1,res_unit{1,2,3}.{snake1,conv1,snake2,conv2}}; + snake1; conv2. Snake alpha/beta: framework [C] -> FastVideo [1,C,1]. + """ + out = {} + for k, v in fw_decoder.state_dict().items(): + p = k.split(".") + li = int(p[1]) + if li == 0: + nk = "conv1." + ".".join(p[2:]) + elif li == 1 + num_blocks: + nk = "snake1." + ".".join(p[2:]) + elif li == 2 + num_blocks: + nk = "conv2." + ".".join(p[2:]) + else: + b = li - 1 + sub = int(p[3]) + if sub == 0: + nk = f"block.{b}.snake1." + ".".join(p[4:]) + elif sub == 1: + nk = f"block.{b}.conv_t1." + ".".join(p[4:]) + else: + r = sub - 2 # ResidualUnit index 0..2 + m = {0: "snake1", 1: "conv1", 2: "snake2", 3: "conv2"}[int(p[5])] + nk = f"block.{b}.res_unit{r + 1}.{m}." + ".".join(p[6:]) + if nk.endswith(".alpha") or nk.endswith(".beta"): + v = v.reshape(1, -1, 1) + out[nk] = v + return out + + +# (dec_dim, vocoder_input_dim, dec_c_mults, dec_strides) — tiny; strides incl odd. +_CASES = [ + pytest.param(4, 8, [1, 2], [5, 2], id="odd_stride5"), + pytest.param(6, 8, [1, 2, 4], [2, 5, 6], id="real_stride_pattern_tiny"), + pytest.param(4, 4, [1, 2], [4, 8], id="even_strides"), +] + + +class TestCosmos3AVAEParity: + + @pytest.mark.parametrize(("dec_dim", "vin", "cmults", "strides"), _CASES) + def test_decode_matches_framework(self, dec_dim, vin, cmults, strides): + torch.manual_seed(0) + fw = _framework_decoder(dec_dim, vin, cmults, strides) + fv = FvOobleckDecoder( + channels=dec_dim, + input_channels=vin, + audio_channels=2, + upsampling_ratios=list(reversed(strides)), # framework reverses dec_strides + channel_multiples=cmults, + ).eval() + state = _framework_to_fastvideo_decoder_state(fw, num_blocks=len(strides)) + fv.load_state_dict(state, strict=True) # exact name + shape match + + z = torch.randn(1, vin, 5) + with torch.no_grad(): + a = fw(z) + b = fv(z) + assert a.shape == b.shape, f"shape: fw={a.shape} fv={b.shape}" + max_abs = (a - b).abs().max().item() + print(f"\n[avae_decode dim={dec_dim} strides={strides}] max abs diff = {max_abs:.3e}") + torch.testing.assert_close(b, a, atol=1e-6, rtol=1e-5)