[feat] cosmos3 PR2: native AVAE sound decoder + bit-exact parity

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

@@ -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

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)

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 <worktree> && <fv-cosmos3 python> -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)