[feat] cosmos3 PR2: t2vs pipeline (joint denoise + AVAE decode + AV mux)

Wire text-to-video+sound end-to-end, completing the audio PR.

- CFG velocity (`cosmos3_get_cfg_velocity`) + `Cosmos3DenoiseEngine` take an
  optional `sound_specs`: the denoise latent becomes a combined [vision | sound]
  flat vector; each step packs vision+sound, forwards the DiT, zeros the
  prediction on conditioning frames for both, and returns the concatenated
  velocity. The UniPC scheduler steps the joint vector. `Cosmos3SoundSpec` +
  `_split_sound_latent` added.
- `Cosmos3DenoisingStage` t2vs path (gated on `COSMOS3_T2VS`): size the sound
  latent from the video duration (framework `create_placeholder_audio` +
  `get_latent_num_samples`), append sound noise, joint-denoise, split, AVAE-decode
  the sound latent, and set `batch.extra["audio"]` / `["audio_sample_rate"]` so
  the generator muxes a stereo 48 kHz AAC track into the mp4. Sound AVAE is
  lazy-loaded from `<model_path>/sound_tokenizer` (`_get_sound_vae`).
- Parity (`test_cosmos3_sound_parity::test_t2vs_cfg_velocity_matches_framework`):
  the combined [vision|sound] sequential-CFG velocity vs a framework-DiT oracle —
  max abs diff = 0.0, mean abs diff = 0.0 across t2vs/i2vs cases.
- Example `basic_cosmos3_t2vs_new_api.py`.

Verified on B200 (real weights, 1280x704, 29f, 35 steps): coherent ocean/sunset
video WITH a real stereo 48 kHz audio track (mean -10.2 dB, peak -0.4 dB; not
silence). Full cosmos3 suite: 130 passed, 0 skipped.

examples/inference/basic/basic_cosmos3_t2vs_new_api.py

@@ -0,0 +1,67 @@
+import os
+import time
+
+# t2vs (text -> video + sound). The Cosmos3 denoise stage generates a joint
+# [vision | sound] latent and AVAE-decodes the sound to a waveform muxed into the
+# mp4. The joint-sound path is gated on COSMOS3_T2VS (set here for the example).
+os.environ.setdefault("COSMOS3_T2VS", "1")
+
+from fastvideo import VideoGenerator  # noqa: E402
+from fastvideo.api import (  # noqa: E402
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+OUTPUT_PATH = "video_samples_cosmos3_t2vs"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(text_encoder=True, pin_cpu_memory=True, dit=False, vae=False),
+        ),
+    )
+
+    load_start = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start
+
+    prompt = (
+        "Ocean waves crash against a rocky shore at sunset, white foam spraying "
+        "into the air as seagulls wheel overhead. Golden light, cinematic wide "
+        "shot, the rhythmic roar of the surf."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        sampling=SamplingConfig(
+            num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
+            height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(output_path=OUTPUT_PATH, save_video=True, return_frames=False),
+    )
+
+    start = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video+sound: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py

@@ -38,7 +38,7 @@ third-party, both explicitly allowed.
 from __future__ import annotations
 
 import math
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import torch
@@ -49,6 +49,7 @@ from fastvideo.models.schedulers.scheduling_unipc_multistep import (
     UniPCMultistepScheduler, )
 from fastvideo.pipelines.basic.cosmos3.sequence_packing import (
     Cosmos3SampleInputs,
+    Cosmos3SoundItem,
     Cosmos3VisionItem,
     pack_cosmos3_video_sequence,
 )
@@ -191,6 +192,35 @@ def _split_flat_latent(flat: torch.Tensor, specs: list[Cosmos3VisionSpec]) -> li
     return out
 
 
+@dataclass
+class Cosmos3SoundSpec:
+    """Geometry + conditioning for one sound item in a denoise run.
+
+    Args:
+        shape: ``(C, T)`` of the sound latent (channels, temporal frames).
+        condition_frame_indexes: Latent-frame indices kept clean (``[]`` for t2vs).
+        fps: Sound latent FPS (``sound_latent_fps``); used iff fps modulation is on.
+    """
+
+    shape: tuple[int, int]
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    fps: float | None = None
+
+    @property
+    def numel(self) -> int:
+        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,
@@ -210,6 +240,8 @@ def cosmos3_get_cfg_velocity(
     include_end_of_generation_token: bool = False,
     fps_per_item: list[float] | None = None,
     normalize_cfg: bool = False,
+    sound_specs: list[Cosmos3SoundSpec] | None = None,
+    sound_fps_per_item: list[float] | None = None,
 ) -> torch.Tensor:
     """Sequential-CFG velocity for one denoise step (framework math).
 
@@ -231,10 +263,23 @@ def cosmos3_get_cfg_velocity(
     assert timestep.numel() == 1, "timestep must be a scalar"
     timestep_value = float(timestep.reshape(()).item())
 
-    noise_x_vision = _split_flat_latent(flat_latent, specs)
+    # Combined flat layout (t2vs): [all vision | all sound], matching the
+    # 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)
     device = next(transformer.parameters()).device
 
     def _run(token_ids: list[int]) -> torch.Tensor:
+        sound_items: list[Cosmos3SoundItem] = []
+        if sound_specs is not None and noise_x_sound is not None:
+            sound_items = [
+                Cosmos3SoundItem(
+                    latent=noise_x_sound[i],
+                    condition_frame_indexes=list(ss.condition_frame_indexes),
+                    fps=(sound_fps_per_item[i] if sound_fps_per_item is not None else None),
+                ) for i, ss in enumerate(sound_specs)
+            ]
         samples = [
             Cosmos3SampleInputs(
                 text_ids=list(token_ids),
@@ -243,6 +288,7 @@ def cosmos3_get_cfg_velocity(
                     condition_frame_indexes=list(spec.condition_frame_indexes),
                     fps=(fps_per_item[i] if fps_per_item is not None else None),
                 ),
+                sound=(sound_items[i] if i < len(sound_items) else None),
                 timestep=timestep_value,
             ) for i, (latent, spec) in enumerate(zip(noise_x_vision, specs, strict=False))
         ]
@@ -258,20 +304,35 @@ def cosmos3_get_cfg_velocity(
             temporal_compression_factor=temporal_compression_factor,
         )
         out = transformer(**packed.to_dit_kwargs(device=device))
+
+        # Vision velocity: zero on conditioning frames, per item, flattened.
+        vision_vel = torch.zeros(vision_total, device=flat_latent.device, dtype=flat_latent.dtype)
         preds = out.get("preds_vision")
-        if preds is None:
-            # Fully-conditioned clip: no noisy frames, zero velocity.
-            return torch.zeros_like(flat_latent)
-        # Zero the prediction on conditioning frames, per item.
-        velocity_items: list[torch.Tensor] = []
-        for pred, cond_mask in zip(preds, packed.vision_condition_mask, strict=False):
-            pred = pred.squeeze(0) if pred.dim() == 5 else pred  # [C, T, H, W]
-            keep = (1.0 - cond_mask).to(dtype=pred.dtype, device=pred.device)  # [T,1,1]
-            if keep.sum() > 0:
-                velocity_items.append(pred * keep)
-            else:
-                velocity_items.append(torch.zeros_like(pred))
-        return torch.cat([v.reshape(-1) for v in velocity_items])
+        if preds is not None:
+            items: list[torch.Tensor] = []
+            for pred, cond_mask in zip(preds, packed.vision_condition_mask, strict=False):
+                pred = pred.squeeze(0) if pred.dim() == 5 else pred  # [C, T, H, W]
+                keep = (1.0 - cond_mask).to(dtype=pred.dtype, device=pred.device)  # [T,1,1]
+                items.append(pred * keep if keep.sum() > 0 else torch.zeros_like(pred))
+            vision_vel = torch.cat([v.reshape(-1) for v in items]).to(flat_latent.dtype)
+
+        if not sound_specs:
+            return vision_vel
+
+        # Sound velocity: preds_sound are per-item [C, T], already zero on clean
+        # frames (unpack fills only noisy frames); zero on cond frames defensively.
+        sound_total = sum(spec.numel for spec in sound_specs)
+        sound_vel = torch.zeros(sound_total, device=flat_latent.device, dtype=flat_latent.dtype)
+        preds_s = out.get("preds_sound")
+        if preds_s is not None:
+            s_items: list[torch.Tensor] = []
+            for pred, cond_mask in zip(preds_s, packed.sound_condition_mask, strict=False):
+                pred = pred.squeeze(0) if pred.dim() == 3 else pred  # [C, T]
+                keep = (1.0 - cond_mask).reshape(1, -1).to(dtype=pred.dtype, device=pred.device)  # [1, T]
+                s_items.append(pred * keep)
+            sound_vel = torch.cat([v.reshape(-1) for v in s_items]).to(flat_latent.dtype)
+
+        return torch.cat([vision_vel, sound_vel])
 
     cond_v = _run(cond_token_ids)
     uncond_v = _run(uncond_token_ids)
@@ -325,6 +386,8 @@ class Cosmos3DenoiseEngine:
         cond_token_ids: list[int],
         uncond_token_ids: list[int],
         fps_per_item: list[float] | None = None,
+        sound_specs: list[Cosmos3SoundSpec] | None = None,
+        sound_fps_per_item: list[float] | None = None,
     ) -> torch.Tensor:
         return cosmos3_get_cfg_velocity(
             transformer=self.transformer,
@@ -343,6 +406,8 @@ class Cosmos3DenoiseEngine:
             temporal_compression_factor=self.temporal_compression_factor,
             include_end_of_generation_token=self.include_end_of_generation_token,
             fps_per_item=fps_per_item,
+            sound_specs=sound_specs,
+            sound_fps_per_item=sound_fps_per_item,
         )
 
     def denoise(
@@ -356,12 +421,16 @@ class Cosmos3DenoiseEngine:
         uncond_token_ids: list[int],
         fps_per_item: list[float] | None = None,
         progress_bar: Any | None = None,
+        sound_specs: list[Cosmos3SoundSpec] | None = None,
+        sound_fps_per_item: list[float] | None = None,
     ) -> torch.Tensor:
         """Run the full UniPC denoise loop, returning the final flat latent.
 
         For each timestep: compute the sequential-CFG velocity, then
         ``scheduler.step(model_output=v, timestep, sample=latent.unsqueeze(0))``
         (the framework steps with a leading batch axis), squeezing back to flat.
+        For t2vs the flat latent is ``[vision | sound]`` and the velocity covers
+        both; the scheduler steps the combined vector jointly.
         """
         latent = flat_latent
         iterator = progress_bar(timesteps) if progress_bar is not None else timesteps
@@ -374,6 +443,8 @@ class Cosmos3DenoiseEngine:
                 cond_token_ids=cond_token_ids,
                 uncond_token_ids=uncond_token_ids,
                 fps_per_item=fps_per_item,
+                sound_specs=sound_specs,
+                sound_fps_per_item=sound_fps_per_item,
             )
             stepped = self.scheduler.step(
                 model_output=v_pred,

fastvideo/pipelines/stages/cosmos3_stages.py

@@ -20,6 +20,7 @@ This mirrors the framework ``Cosmos3OmniDiffusersPipeline.__call__``.
 """
 from __future__ import annotations
 
+import os
 import weakref
 from typing import Any
 
@@ -31,6 +32,7 @@ from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
     Cosmos3DenoiseEngine,
+    Cosmos3SoundSpec,
     Cosmos3VisionSpec,
     _VaeNorm,
     cosmos3_special_tokens,
@@ -186,6 +188,29 @@ class Cosmos3DenoisingStage(PipelineStage):
 
         flat_latent = init_latent.reshape(-1)
         fps_per_item = [fps] if bool(arch.enable_fps_modulation) else None
+
+        # ---- t2vs: jointly generate sound (combined [vision | sound] latent) ----
+        # Mirrors the framework: a placeholder audio sized to the video duration
+        # sets the sound latent length; sound shares the denoise/CFG with vision.
+        with_audio = is_video and os.environ.get("COSMOS3_T2VS", "") not in ("", "0")
+        sound_specs = None
+        sound_fps_per_item = None
+        sound_vae = None
+        sound_shape: tuple[int, int] | None = None
+        if with_audio:
+            sound_vae = self._get_sound_vae(pipe, device, dtype)
+            sound_dim = int(arch.sound_dim)
+            sound_latent_fps = float(arch.sound_latent_fps)
+            # Framework ``create_placeholder_audio`` + ``get_latent_num_samples``.
+            num_audio_samples = int(num_frames / fps * sound_vae.sample_rate)
+            sound_latent_t = max(1, sound_vae.get_latent_num_samples(num_audio_samples))
+            sound_shape = (sound_dim, sound_latent_t)
+            sound_noise = randn_tensor((sound_dim, sound_latent_t), generator=generator, device=device,
+                                       dtype=dtype).float()
+            flat_latent = torch.cat([flat_latent, sound_noise.reshape(-1)])
+            sound_specs = [Cosmos3SoundSpec(shape=sound_shape, condition_frame_indexes=[], fps=sound_latent_fps)]
+            sound_fps_per_item = [sound_latent_fps] if bool(arch.enable_fps_modulation) else None
+
         final_flat = engine.denoise(
             flat_latent=flat_latent,
             timesteps=timesteps,
@@ -195,15 +220,25 @@ class Cosmos3DenoisingStage(PipelineStage):
             uncond_token_ids=uncond_ids,
             fps_per_item=fps_per_item,
             progress_bar=lambda it: tqdm(it, desc="Cosmos3 denoising"),
+            sound_specs=sound_specs,
+            sound_fps_per_item=sound_fps_per_item,
         )
 
-        # ---- Decode: [C, T, H, W] -> pixels [B, 3, T, H, W] in [0, 1] ----
-        result_latent = final_flat.reshape(latent_shape).unsqueeze(0).to(device=device, dtype=dtype)
+        # ---- Decode vision: [C, T, H, W] -> pixels [B, 3, T, H, W] in [0, 1] ----
+        vision_flat = final_flat[:spec.numel]
+        result_latent = vision_flat.reshape(latent_shape).unsqueeze(0).to(device=device, dtype=dtype)
         decoded = cosmos3_vae_decode(self.vae, result_latent, norm)  # [B, 3, T, H, W] in [-1, 1]
         video = ((1.0 + decoded) / 2.0).clamp(0.0, 1.0)
 
         batch.latents = result_latent
         batch.output = video
+
+        # ---- Decode sound: AVAE latent [C, T] -> waveform [C, N] in [-1, 1] ----
+        if with_audio and sound_vae is not None and sound_shape is not None:
+            sound_latent = final_flat[spec.numel:].reshape(sound_shape).unsqueeze(0).to(device=device, dtype=dtype)
+            waveform = sound_vae.decode(sound_latent)  # [1, C_audio, N]
+            batch.extra["audio"] = waveform[0].detach().float().cpu()  # [C_audio, N]
+            batch.extra["audio_sample_rate"] = int(sound_vae.sample_rate)
         return batch
 
     # ------------------------------------------------------------------
@@ -223,6 +258,24 @@ class Cosmos3DenoisingStage(PipelineStage):
         img = TF.resize(img, [resize_h, resize_w])
         return TF.center_crop(img, [target_h, target_w])
 
+    @staticmethod
+    def _get_sound_vae(pipe: Any, device: torch.device, dtype: torch.dtype) -> Any:
+        """Lazily load + cache the Cosmos3 sound AVAE decoder from the checkpoint.
+
+        The video path does not load ``sound_tokenizer``; t2vs needs only its
+        decoder, so we load it on first use from ``<model_path>/sound_tokenizer``.
+        """
+        cached = getattr(pipe, "_sound_vae", None) if pipe is not None else None
+        if cached is not None:
+            return cached
+        from fastvideo.models.audio.cosmos3_avae import Cosmos3SoundVAE
+        model_path = pipe.model_path
+        sound_dir = os.path.join(model_path, "sound_tokenizer")
+        sound_vae = Cosmos3SoundVAE.from_pretrained(sound_dir, torch_dtype=dtype).to(device)
+        if pipe is not None:
+            pipe._sound_vae = sound_vae
+        return sound_vae
+
     @classmethod
     def _image_to_video_tensor(
         cls,

tests/local_tests/cosmos3/test_cosmos3_sound_parity.py

@@ -233,3 +233,56 @@ class TestCosmos3SoundParity:
         torch.testing.assert_close(fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3)
         torch.testing.assert_close(fv_out["preds_sound"][0], fw_out["preds_sound"][0], atol=1e-4, rtol=1e-3)
         torch.testing.assert_close(fv_on_fw["preds_sound"][0], fw_out["preds_sound"][0], atol=1e-4, rtol=1e-3)
+
+    @pytest.mark.parametrize(("grid_t", "lh", "lw", "snd_t", "n_text", "cond_v", "cond_s"), _CASES)
+    def test_t2vs_cfg_velocity_matches_framework(self, grid_t, lh, lw, snd_t, n_text, cond_v, cond_s):
+        """The combined [vision|sound] sequential-CFG velocity (the t2vs denoise
+        step's pipeline glue) matches a framework-DiT oracle."""
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
+            Cosmos3SoundSpec,
+            Cosmos3VisionSpec,
+            cosmos3_get_cfg_velocity,
+        )
+
+        vfm, dit = self._build()
+        vision_shape = (_LATENT_CHANNEL, grid_t, lh // _LATENT_PATCH_SIZE, lw // _LATENT_PATCH_SIZE)
+        # _make_inputs builds the vision LATENT [1,C,T,H,W]; here we drive the
+        # combined flat latent directly, so use the un-patchified latent shape.
+        vlat_shape = (_LATENT_CHANNEL, grid_t, lh, lw)
+        sound_shape = (_SOUND_DIM, snd_t)
+        torch.manual_seed(3)
+        cond_ids = torch.randint(0, 60, (n_text,)).tolist()
+        uncond_ids = torch.randint(0, 60, (max(1, n_text - 1),)).tolist()
+        vis_numel = int(torch.tensor(vlat_shape).prod())
+        snd_numel = int(torch.tensor(sound_shape).prod())
+        flat = torch.randn(vis_numel + snd_numel)
+        guidance, ts = 6.0, 500.0
+
+        def _fw_velocity(ids):
+            vision = flat[:vis_numel].reshape(vlat_shape).unsqueeze(0)  # [1,C,T,H,W]
+            sound = flat[vis_numel:].reshape(sound_shape)  # [C,T]
+            ps = _framework_pack_sound(text_ids=ids, vision=vision, sound=sound,
+                                       cond_vision=cond_v, cond_sound=cond_s,
+                                       timestep=ts, is_image_batch=(grid_t == 1))
+            with torch.no_grad():
+                out = vfm(packed_seq=ps)
+            pv = out["preds_vision"][0].squeeze(0)  # [C,T,H,W] (zero on clean)
+            psd = out["preds_sound"][0]  # [C,T] (zero on clean)
+            return torch.cat([pv.reshape(-1), psd.reshape(-1)])
+
+        fw_cond, fw_uncond = _fw_velocity(cond_ids), _fw_velocity(uncond_ids)
+        fw_v = fw_uncond + guidance * (fw_cond - fw_uncond)
+
+        fv_v = cosmos3_get_cfg_velocity(
+            transformer=dit, flat_latent=flat, timestep=torch.tensor([ts]), guidance=guidance,
+            specs=[Cosmos3VisionSpec(shape=vlat_shape, condition_frame_indexes=list(cond_v))],
+            sound_specs=[Cosmos3SoundSpec(shape=sound_shape, condition_frame_indexes=list(cond_s))],
+            cond_token_ids=cond_ids, uncond_token_ids=uncond_ids,
+            special_tokens=_SPECIAL_TOKENS, latent_patch_size=_LATENT_PATCH_SIZE,
+            temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, reset_spatial_ids=_RESET_SPATIAL_IDS,
+            enable_fps_modulation=False, base_fps=24.0, temporal_compression_factor=_TCF,
+        )
+        assert fv_v.shape == fw_v.shape, f"shape fv={fv_v.shape} fw={fw_v.shape}"
+        mx, mn = _diffs(fv_v, fw_v)
+        print(f"\n[t2vs_cfg_velocity {grid_t}x{lh}x{lw} snd={snd_t}] max={mx:.3e} mean={mn:.3e}")
+        torch.testing.assert_close(fv_v, fw_v, atol=1e-4, rtol=1e-3)