From 8c40e1c6d89d99474950961f31f86ae26ba30c82 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 22:30:22 -0700 Subject: [PATCH 01/25] Add design spec for audio_image2video (A2V) inference Conditions video generation on a real input audio clip + input image using the nano_diffusers_sound_encoder checkpoint. Reuses the existing ts2v sound condition plan plus preserved image first-frame conditioning. Co-Authored-By: Claude Opus 4.8 (1M context) --- ...6-11-a2v-sound-encoder-inference-design.md | 231 ++++++++++++++++++ 1 file changed, 231 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md diff --git a/docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md b/docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md new file mode 100644 index 00000000..58480aed --- /dev/null +++ b/docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md @@ -0,0 +1,231 @@ +# Audio+Image→Video (A2V) Inference with the Sound-Encoder Checkpoint + +**Date:** 2026-06-11 +**Status:** Approved design — ready for implementation plan +**Branch:** `a2v-sound-encoder-inference` + +## Summary + +Add a new inference mode, `audio_image2video`, that conditions video generation on a +**real input audio clip** plus an **input image (first frame)**, using the +`nano_diffusers_sound_encoder` checkpoint +(`nvidia/Cosmos3-Experimental`, subfolder `nano_diffusers_sound_encoder`). + +Today the only audio path in inference is *generation*: `enable_sound: true` injects a +zero-filled placeholder waveform and the model **generates** sound (internally +`mode="t2vs"`). There is no way to feed a real audio clip as a **condition**. This spec +closes that gap on the inference side only — the model and data layers already support +sound-as-condition. + +## Background / Current State + +- **Checkpoint format.** The default `Cosmos3-Nano` checkpoint is diffusers-format + (`model_index.json`, `transformer/`, `vae/`, `sound_tokenizer/`, `vision_encoder/`) + and already ships a `sound_tokenizer` (AVAE). The framework's loader already handles + this format. The new `nano_diffusers_sound_encoder` checkpoint has the same layout + with different `transformer` + `sound_tokenizer` weights; its `config.json` maps to the + same `OmniMoTModelConfig` (`sound_gen=true`, `sound_dim=64`, AVAE + `sample_rate=48000`, `audio_channels=2`). +- **Sound generation today.** `cosmos_framework/inference/inference.py:get_sample_data` + (≈ line 610) calls `create_placeholder_audio` → `inject_sound_into_batch`, which + hardcodes `mode="t2vs"` in `cosmos_framework/inference/sound.py:105`. All sound tokens + are generated; the placeholder only establishes the target length. +- **Sound conditioning already exists at the data/model layer.** + `cosmos_framework/data/vfm/sound_data_utils.py` defines `ts2v` (Text+Sound→Video, + sound conditioned) and `ti2sv`. The MOT network handles clean vs. noisy sound tokens + via the condition mask (`_encode_sound`/`_decode_sound`). These plans are **not + reachable** from any `model_mode`. +- **Vision-condition preservation.** `inject_sound_into_batch` captures and re-applies + the existing `condition_frame_indexes_vision` (sound.py lines ~90–94, 110–111). So the + image's first-frame condition survives whichever sound mode is chosen. +- **Output path is generic.** `inference.py` (≈ lines 1489–1544) decodes any `"sound"` in + the model outputs (`model.decode_sound`) and muxes it into the `.mp4` + (`mux_audio_into_video`). No change is needed for conditioned audio. + +### Current `ModelMode` values + +`text2image`, `text2video`, `image2image`, `image2video`, `video2video`, +`forward_dynamics`, `inverse_dynamics`, `policy`, `reasoner`. None takes audio as input. + +## Goals + +1. Register the `nano_diffusers_sound_encoder` checkpoint as a named checkpoint. +2. Add `audio_image2video`: image (first frame) + real audio clip → video, with the audio + used as a **clean condition** and muxed into the output video. +3. Full deliverable: input loading, defaults, example input, colocated tests, docs. + +## Non-Goals + +- Audio-only conditioning (`ts2v` with no image) as a separate user-facing mode (YAGNI). +- Any change to the model architecture, sequence packing, AVAE tokenizer, training, or + the output/save path. +- A new `tis2v` sequence plan: the audio+image combination falls out of `ts2v` (conditions + sound) plus the preserved image vision-condition. + +## Design + +### Key leverage point + +"Audio+image→video" needs no new sequence-plan combination. `inject_sound_into_batch` +already preserves the image's `condition_frame_indexes_vision=[0]`. Selecting the existing +`ts2v` plan (all sound latents conditioned) and letting that preservation re-apply `[0]` +yields exactly: image first-frame conditioned + audio conditioned + remaining video +generated. + +### Changes by component + +**1. Checkpoint registry** — `cosmos_framework/inference/args.py` (`_CHECKPOINTS`, line 1051) + +Add: + +```python +"Cosmos3-Nano-SoundEncoder": CheckpointConfig( + model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["8B"], + config_file=str(CONFIG_DIR / "model/Cosmos3-Nano.yaml"), + s3_uri=..., # match the Nano entry's pattern; unused for HF-backed load + hf=CheckpointDirHf( + repository="nvidia/Cosmos3-Experimental", + revision="main", + subdirectory="nano_diffusers_sound_encoder", + ), +) +``` + +Reuses `Cosmos3-Nano.yaml` (same `OmniMoTModelConfig`). Weight compatibility is verified +on load and end-to-end by the slurm run. + +**2. New model mode** — `cosmos_framework/inference/args.py` (`ModelMode`, line 157) + +- Add `AUDIO_IMAGE2VIDEO = "audio_image2video"`. +- Add a `_SOUND_CONDITION_MODES: frozenset[ModelMode] = {ModelMode.AUDIO_IMAGE2VIDEO}` and + an `is_sound_condition` property mirroring `is_action`/`is_reasoner`. +- `condition_vision_mode` already resolves to `image` from an image `vision_path` + (lines 355–368), giving `condition_frame_indexes_vision=[0]` via existing defaults. + +**3. Audio input field** — `cosmos_framework/inference/args.py` + +- `SoundDataArgs` (line 514): add `sound_path: ResolvedFilePath | None = None`. +- `SoundDataOverrides` (line 518): add `sound_path: ResolvedFilePathOrUrl | None = None` + with a docstring; add a `download()` override that calls + `self.sound_path = download_file(self.sound_path, output_dir, "sound")`. +- `_build_sound_data` (line 524): for `AUDIO_IMAGE2VIDEO`, require `sound_path` set and + force `enable_sound = True`; keep the existing `sound_gen` validation. + +**4. Audio decode helper** — `cosmos_framework/inference/sound.py` + +```python +def load_conditioning_audio( + path: Path, + *, + sample_rate: int, + audio_channels: int, + num_samples: int, +) -> torch.Tensor: + """Decode an audio file to a [1, C, N] waveform aligned to the video duration. + + Reads via soundfile, resamples to ``sample_rate``, conforms channel count to + ``audio_channels`` (mono->stereo duplicate / stereo->mono mean), and trims or + zero-pads to ``num_samples`` so the audio and video latent streams align temporally. + """ +``` + +`num_samples` = `int(num_frames / fps * sample_rate)` (matches `create_placeholder_audio`). + +**5. Wire conditioning** + +- `cosmos_framework/inference/sound.py`: add `condition_sound: bool = False` to + `inject_sound_into_batch`. When `True`, build the plan with `mode="ts2v"` (instead of + `"t2vs"`); the existing vision-cond preservation keeps the image's `[0]`. +- `cosmos_framework/inference/inference.py` `get_sample_data` (≈ line 610): when + `sample_args.sound_path` is set, call `load_conditioning_audio(...)` and + `inject_sound_into_batch(out, audio, model, condition_sound=True)`. Otherwise keep the + existing placeholder-generation branch unchanged. + +**6. Defaults + example** + +- `cosmos_framework/inference/defaults/audio_image2video/sample_args.json`: copy of + `image2video/sample_args.json` with `"enable_sound": true`. +- `inputs/omni/a2v.json`: `model_mode="audio_image2video"`, image `vision_path`, + `sound_path`, and a prompt. + +**7. Output** — no code change. `inference.py:1489` decodes and muxes the (clean, +conditioned) sound into `vision.mp4`. + +### Data flow + +``` +a2v.json (vision_path=image, sound_path=audio, model_mode=audio_image2video) + -> download() resolves both paths + -> get_sample_data: + condition_vision_mode=image -> load_conditioning_image -> build_conditioned_video_batch + (condition_frame_indexes_vision=[0], sequence_plan with image cond) + sound_path set -> load_conditioning_audio -> [1,C,N] real waveform + inject_sound_into_batch(condition_sound=True): + mode="ts2v" -> condition_frame_indexes_sound = all + preserve condition_frame_indexes_vision=[0] + -> model encodes audio (AVAE) as clean sound tokens; image as clean first frame + -> diffusion generates remaining video frames conditioned on both + -> outputs["sound"] decoded + muxed into vision.mp4 +``` + +## Testing + +Colocated unit tests (matching the `colocated-tests-ci` convention): + +- ➕ `cosmos_framework/inference/sound_test.py` + - `load_conditioning_audio`: resample rate change, mono↔stereo conformance, + trim/pad to `num_samples`, returns `[1, C, num_samples]`. + - `inject_sound_into_batch(condition_sound=True)`: produces a plan with all sound + latents conditioned **and** preserves a pre-set `condition_frame_indexes_vision=[0]`; + `condition_sound=False` keeps the existing `t2vs` behavior. +- ✏️ `cosmos_framework/inference/args_test.py` + - `audio_image2video` resolves `condition_vision_mode=image` and + `condition_frame_indexes_vision=[0]`. + - `_build_sound_data` requires `sound_path` and sets `enable_sound=True`; validation + fails on a model with `sound_gen=False`. +- ➕ `cosmos_framework/data/vfm/sound_data_utils_test.py` (if absent): `ts2v` plan sets + `condition_frame_indexes_sound = range(sound_latent_length)` and empty vision cond. + +CPU-only tests use small synthetic tensors / a tiny generated `.wav`; no GPU or checkpoint +download required. + +## End-to-end verification (slurm) + +Run at repo root inside the i4 container (per the `slurm-node` skill): + +```shell +python -m cosmos_framework.scripts.inference \ + --parallelism-preset=latency \ + -i "inputs/omni/a2v.json" \ + -o outputs/a2v \ + --checkpoint-path Cosmos3-Nano-SoundEncoder \ + --seed=0 +``` + +Success = `outputs/a2v//vision.mp4` exists, plays, and contains an audio track +matching the input clip; no shape/dtype errors during sound encode/condition. + +## Example inputs (sourced by implementer) + +- **Image:** the robot image already referenced by `inputs/omni/i2v.json`. +- **Audio:** extract the audio track from a Cosmos3-Nano example sound output + (e.g. `assets/example_t2vs_output.mp4` in the `nvidia/Cosmos3-Nano` HF repo) into a + short `.wav`, referenced by URL or local path. + +## Risks + +1. **Conditioned sound tokens must stay frozen during sampling.** The model uses the same + forward path as `ts2v` training, so this should hold; the slurm run is the gate. If the + sampler does not freeze clean sound tokens, a small inference-loop fix may be needed + (out of scope until observed). +2. **Checkpoint/config compatibility.** `nano_diffusers_sound_encoder/config.json` matches + `Cosmos3-Nano.yaml`'s `OmniMoTModelConfig`; confirmed by inspection, verified on load. +3. **Audio/video temporal alignment.** Handled by trimming/padding audio to the video + duration in `load_conditioning_audio`. + +## Files touched (summary) + +Modify: `inference/args.py`, `inference/sound.py`, `inference/inference.py`, +`inference/args_test.py`, `docs/inference.md`. +Create: `inference/sound_test.py`, `inference/defaults/audio_image2video/sample_args.json`, +`inputs/omni/a2v.json`, and `data/vfm/sound_data_utils_test.py` (if absent). From 10ff2bd63fd0cea24441d667fcc132a7919d1510 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 22:49:58 -0700 Subject: [PATCH 02/25] Add implementation plan for audio_image2video (A2V) inference Co-Authored-By: Claude Opus 4.8 (1M context) --- .../2026-06-11-a2v-sound-encoder-inference.md | 770 ++++++++++++++++++ 1 file changed, 770 insertions(+) create mode 100644 docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md diff --git a/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md b/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md new file mode 100644 index 00000000..55c593b9 --- /dev/null +++ b/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md @@ -0,0 +1,770 @@ +# Audio+Image→Video (A2V) Inference Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add an `audio_image2video` inference mode that conditions video generation on a real input audio clip + an input image (first frame), using the `nano_diffusers_sound_encoder` checkpoint. + +**Architecture:** Reuse the existing `ts2v` (sound-conditioned) sequence plan plus the existing image first-frame vision conditioning, which `inject_sound_into_batch` already preserves. Add a real-audio loader and a `model_mode` + `sound_path` input field; no model/network/tokenizer changes. + +**Tech Stack:** Python, PyTorch, pydantic args, `soundfile`/`torchaudio` for audio I/O, `pytest` (colocated `*_test.py`), the Cosmos3 OmniMoT diffusers-format checkpoint loader. + +**Spec:** `docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md` + +**Branch:** `a2v-sound-encoder-inference` + +--- + +## Background the implementer must know + +- **No model changes.** The MOT model + AVAE already support sound-as-condition (`ts2v`). The only gap is the inference pipeline: it never loads real audio and never selects a conditioning plan. +- **`inject_sound_into_batch` preserves vision conditioning** (`cosmos_framework/inference/sound.py` ~lines 90–94, 110–111). So image-first-frame + audio-conditioned falls out of `mode="ts2v"` automatically. +- **Output decode+mux is generic** (`cosmos_framework/inference/inference.py` ~lines 1489–1544): any `"sound"` in model outputs is decoded and muxed into the `.mp4`. No change needed. +- **Sample-arg machinery:** `OmniSampleOverrides.build_sample(model_config=...)` runs `_build_vision_data` then `_build_sound_data` (`args.py` ~line 1004). `download()` methods cascade via `super().download()` through the MRO (see `VisionDataOverrides.download` at `args.py:447`). +- **Run tests** with the repo's pytest. Audio unit tests are CPU-only (synthetic tensors / tiny generated WAV). Run via the `slurm-node` skill if pytest needs the i4 container; otherwise locally. + +--- + +## File Structure + +- `cosmos_framework/inference/args.py` — checkpoint registry entry, `ModelMode.AUDIO_IMAGE2VIDEO`, `sound_path` field + validation/download. +- `cosmos_framework/inference/sound.py` — `load_conditioning_audio` helper; `condition_sound` param on `inject_sound_into_batch`. +- `cosmos_framework/inference/inference.py` — `get_sample_data` branch that loads real audio and conditions on it. +- `cosmos_framework/inference/sound_test.py` (new) — unit tests for the two `sound.py` additions. +- `cosmos_framework/inference/args_test.py` — unit test for the new mode + `sound_path` validation. +- `cosmos_framework/inference/defaults/audio_image2video/sample_args.json` (new) — per-mode defaults. +- `inputs/omni/a2v.json` (new) + `inputs/omni/assets/a2v_audio.wav` (new) — example input. +- `docs/inference.md` — Modes table row + checkpoint note. + +--- + +## Task 1: Register the sound-encoder checkpoint + +**Files:** +- Modify: `cosmos_framework/inference/args.py` (`_CHECKPOINTS`, line ~1051) +- Test: `cosmos_framework/inference/args_test.py::test_checkpoints` (existing — auto-covers the new entry) + +- [ ] **Step 1: Add the registry entry** + +In `_CHECKPOINTS`, after the `"Cosmos3-Nano"` entry, add: + +```python + # Diffusers HF checkpoint whose transformer is trained to condition on + # (encode) input sound, enabling audio_image2video (A2V). Reuses the + # Cosmos3-Nano architecture (OmniMoTModelConfig, sound_gen=True). + "Cosmos3-Nano-SoundEncoder": CheckpointConfig( + model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["8B"], + config_file=str(CONFIG_DIR / "model/Cosmos3-Nano.yaml"), + s3_uri="", # unused for HF-backed checkpoints + hf=CheckpointDirHf( + repository="nvidia/Cosmos3-Experimental", + revision="main", + subdirectory="nano_diffusers_sound_encoder", + ), + ), +``` + +- [ ] **Step 2: Verify the registry resolves and downloads `checkpoint.json`** + +Run: `pytest cosmos_framework/inference/args_test.py::test_checkpoints -v` +Expected: PASS. (Requires `HF_TOKEN` with access to `nvidia/Cosmos3-Experimental`; run inside the dev/slurm env where the token is set. The test downloads `nano_diffusers_sound_encoder/checkpoint.json` via the `subdirectory` filter and parses it.) + +- [ ] **Step 3: Commit** + +```bash +git add cosmos_framework/inference/args.py +git commit -m "Register Cosmos3-Nano-SoundEncoder checkpoint" +``` + +--- + +## Task 2: Add the `audio_image2video` model mode + +**Files:** +- Modify: `cosmos_framework/inference/args.py` (`ModelMode` enum line ~157; frozensets line ~181–188) +- Test: `cosmos_framework/inference/args_test.py` (new test added in Task 3, Step 5) + +- [ ] **Step 1: Add the enum member** + +In `class ModelMode(StrEnum)` (after `VIDEO2VIDEO = "video2video"`): + +```python + AUDIO_IMAGE2VIDEO = "audio_image2video" +``` + +- [ ] **Step 2: Add the mode-group frozenset and property** + +After `REASONER_MODEL_MODES` (line ~188) add: + +```python +# Modes that condition generation on a real input audio clip (require a model +# with ``sound_gen=True`` and a ``sound_path``). +SOUND_CONDITION_MODEL_MODES: frozenset[ModelMode] = frozenset({ModelMode.AUDIO_IMAGE2VIDEO}) +``` + +In `class ModelMode`, alongside `is_action` / `is_reasoner`: + +```python + @property + def is_sound_condition(self) -> bool: + return self in SOUND_CONDITION_MODEL_MODES +``` + +- [ ] **Step 3: Verify import + enum value** + +Run: `python -c "from cosmos_framework.inference.args import ModelMode; print(ModelMode.AUDIO_IMAGE2VIDEO.value, ModelMode.AUDIO_IMAGE2VIDEO.is_sound_condition)"` +Expected: `audio_image2video True` + +- [ ] **Step 4: Commit** + +```bash +git add cosmos_framework/inference/args.py +git commit -m "Add audio_image2video model mode" +``` + +--- + +## Task 3: Add the `sound_path` input field + validation/download + +**Files:** +- Modify: `cosmos_framework/inference/args.py` (`SoundDataArgs` line ~514; `SoundDataOverrides` line ~518) +- Test: `cosmos_framework/inference/args_test.py` + +- [ ] **Step 1: Write the failing test** + +Add to `cosmos_framework/inference/args_test.py` (note the new imports `ModelMode` is already imported; add `SoundDataOverrides`): + +```python +import types + +from cosmos_framework.inference.args import SoundDataOverrides + + +def test_build_sound_data_requires_sound_path_for_a2v(): + model_config = types.SimpleNamespace(sound_gen=True) + sample_meta = types.SimpleNamespace(model_mode=ModelMode.AUDIO_IMAGE2VIDEO) + + # Missing sound_path -> error + overrides = SoundDataOverrides(sound_path=None) + with pytest.raises(ValueError, match="sound_path"): + overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) + + # sound_path set -> enable_sound forced True + overrides = SoundDataOverrides(sound_path="clip.wav") + overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) + assert overrides.enable_sound is True + + +def test_build_sound_data_rejects_model_without_sound_gen(): + model_config = types.SimpleNamespace(sound_gen=False) + sample_meta = types.SimpleNamespace(model_mode=ModelMode.AUDIO_IMAGE2VIDEO) + overrides = SoundDataOverrides(sound_path="clip.wav") + with pytest.raises(ValueError, match="sound tokenizer"): + overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `pytest cosmos_framework/inference/args_test.py::test_build_sound_data_requires_sound_path_for_a2v -v` +Expected: FAIL — `SoundDataOverrides` has no `sound_path` (TypeError/validation error). + +- [ ] **Step 3: Add the field + download + validation** + +In `class SoundDataArgs(ArgsBase)` (line ~514): + +```python +class SoundDataArgs(ArgsBase): + enable_sound: bool = False + sound_path: ResolvedFilePath | None = None +``` + +In `class SoundDataOverrides(OverridesBase)` (line ~518) add the field and a `download()` override, and extend `_build_sound_data`: + +```python +class SoundDataOverrides(OverridesBase): + """Sound data overrides.""" + + enable_sound: Training[bool | None] = None + """Enable joint video+sound generation (t2vs mode). Requires a checkpoint with sound modules.""" + sound_path: Training[ResolvedFilePathOrUrl | None] = None + """Path or URL to a conditioning audio clip (e.g. .wav/.mp3/.flac). Required for + audio_image2video; the clip is encoded by the AVAE and used as a clean condition.""" + + @override + def download(self, output_dir: Path): + super().download(output_dir) + self.sound_path = download_file(self.sound_path, output_dir, "sound") + + def _build_sound_data(self, model_config: "OmniMoTModelConfig", sample_meta: SampleMeta): + if sample_meta.model_mode.is_sound_condition: + if self.sound_path is None: + raise ValueError( + f"model_mode={sample_meta.model_mode.value} requires a `sound_path` " + "(a conditioning audio clip)" + ) + self.enable_sound = True + if self.enable_sound is None: + self.enable_sound = False + if self.enable_sound and not model_config.sound_gen: + raise ValueError( + "enable_sound=True requires a model with a sound tokenizer " + "(model.config.sound_gen=True), but the loaded checkpoint has no sound tokenizer" + ) +``` + +(`@override`, `Path`, `download_file`, `ResolvedFilePathOrUrl`, `Training` are already imported in `args.py`.) + +- [ ] **Step 4: Run test to verify it passes** + +Run: `pytest cosmos_framework/inference/args_test.py::test_build_sound_data_requires_sound_path_for_a2v cosmos_framework/inference/args_test.py::test_build_sound_data_rejects_model_without_sound_gen -v` +Expected: PASS (both). + +- [ ] **Step 5: Add the mode-resolution test** + +Add to `args_test.py` (reuses `model_dict.config` pattern from `test_sample_args`; place it as its own test so it can build a sample for the new mode): + +```python +def test_audio_image2video_conditions_image_and_sound(tmp_path: Path): + import omegaconf + from cosmos_framework.inference.common.config import structure_config + + setup_args = OmniSetupOverrides( + checkpoint_path=DEFAULT_CHECKPOINT_NAME, + output_dir=tmp_path / "outputs", + ).build_setup() + model_dict = structure_config(setup_args.load_model_config_dict(), omegaconf.DictConfig) + + args = OmniSampleOverrides( + name="a2v", + output_dir=tmp_path / "a2v", + model_mode=ModelMode.AUDIO_IMAGE2VIDEO, + vision_path="robot.jpg", # image extension -> first-frame condition + sound_path="clip.wav", + ).build_sample(model_config=model_dict.config) + + assert args.condition_vision_mode.value == "image" + assert args.condition_frame_indexes_vision == [0] + assert args.enable_sound is True + assert args.sound_path == "clip.wav" +``` + +- [ ] **Step 6: Run it** + +Run: `pytest cosmos_framework/inference/args_test.py::test_audio_image2video_conditions_image_and_sound -v` +Expected: PASS. (Downloads the default Nano model config; run in the dev/slurm env.) + +- [ ] **Step 7: Commit** + +```bash +git add cosmos_framework/inference/args.py cosmos_framework/inference/args_test.py +git commit -m "Add sound_path input + audio_image2video arg validation" +``` + +--- + +## Task 4: `load_conditioning_audio` helper + +**Files:** +- Modify: `cosmos_framework/inference/sound.py` +- Test: `cosmos_framework/inference/sound_test.py` (new) + +- [ ] **Step 1: Write the failing test** + +Create `cosmos_framework/inference/sound_test.py`: + +```python +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from pathlib import Path + +import soundfile as sf +import torch + +from cosmos_framework.inference.sound import load_conditioning_audio + + +def _write_wav(path: Path, sample_rate: int, channels: int, num_samples: int) -> None: + data = torch.zeros(num_samples, channels).numpy() if channels > 1 else torch.zeros(num_samples).numpy() + sf.write(str(path), data, sample_rate) + + +def test_load_conditioning_audio_resamples_and_pads(tmp_path: Path): + src = tmp_path / "in.wav" + _write_wav(src, sample_rate=44100, channels=1, num_samples=44100) # 1.0s mono @44.1k + + out = load_conditioning_audio(src, sample_rate=48000, audio_channels=2, num_samples=96000) + + assert out.shape == (1, 2, 96000) # [1, C, N], stereo, exactly num_samples (2.0s @48k -> pad) + assert out.dtype == torch.float32 + + +def test_load_conditioning_audio_trims(tmp_path: Path): + src = tmp_path / "in.wav" + _write_wav(src, sample_rate=48000, channels=2, num_samples=48000 * 4) # 4s stereo @48k + + out = load_conditioning_audio(src, sample_rate=48000, audio_channels=2, num_samples=48000 * 2) + + assert out.shape == (1, 2, 48000 * 2) # trimmed to 2s +``` + +- [ ] **Step 2: Run to verify it fails** + +Run: `pytest cosmos_framework/inference/sound_test.py -v` +Expected: FAIL — `ImportError: cannot import name 'load_conditioning_audio'`. + +- [ ] **Step 3: Implement the helper** + +Add to `cosmos_framework/inference/sound.py` (after `create_placeholder_audio`): + +```python +def load_conditioning_audio( + path: Path, + *, + sample_rate: int, + audio_channels: int, + num_samples: int, +) -> torch.Tensor: + """Decode an audio file into a conditioning waveform aligned to the video. + + Reads ``path`` with soundfile, resamples to ``sample_rate``, conforms the + channel count to ``audio_channels`` (mono->stereo duplicate, stereo->mono + mean), and trims or zero-pads to exactly ``num_samples`` so the audio and + video latent streams cover the same duration. + + Returns: + Audio tensor of shape (1, C, N) where C == audio_channels and + N == num_samples, dtype float32. + """ + import soundfile as sf # type: ignore[import-not-found] + + data, src_sr = sf.read(str(path), dtype="float32", always_2d=True) # [N, C] + waveform = torch.from_numpy(data).transpose(0, 1).contiguous() # [C, N] + + # Resample to the tokenizer's rate. + if src_sr != sample_rate: + import torchaudio + + waveform = torchaudio.functional.resample(waveform, orig_freq=src_sr, new_freq=sample_rate) + + # Conform channels. + cur_channels = waveform.shape[0] + if cur_channels != audio_channels: + if cur_channels == 1 and audio_channels == 2: + waveform = waveform.repeat(2, 1) + elif cur_channels == 2 and audio_channels == 1: + waveform = waveform.mean(dim=0, keepdim=True) + else: + raise ValueError( + f"Cannot convert {cur_channels}-channel audio to {audio_channels} channels" + ) + + # Trim or zero-pad to num_samples. + n = waveform.shape[-1] + if n > num_samples: + waveform = waveform[:, :num_samples] + elif n < num_samples: + waveform = torch.nn.functional.pad(waveform, (0, num_samples - n)) + + return waveform.unsqueeze(0).to(dtype=torch.float32) # [1, C, N] +``` + +- [ ] **Step 4: Run to verify it passes** + +Run: `pytest cosmos_framework/inference/sound_test.py -v` +Expected: PASS (both tests). + +- [ ] **Step 5: Commit** + +```bash +git add cosmos_framework/inference/sound.py cosmos_framework/inference/sound_test.py +git commit -m "Add load_conditioning_audio helper" +``` + +--- + +## Task 5: `condition_sound` mode in `inject_sound_into_batch` + +**Files:** +- Modify: `cosmos_framework/inference/sound.py` (`inject_sound_into_batch`, line ~62) +- Test: `cosmos_framework/inference/sound_test.py` + +- [ ] **Step 1: Write the failing test** + +Append to `cosmos_framework/inference/sound_test.py`: + +```python +import types + +from cosmos_framework.data.vfm.sequence_packing import SequencePlan + + +def _fake_model(sound_latent_t: int, temporal_cf: int = 4): + sound_tok = types.SimpleNamespace( + get_latent_num_samples=lambda n: sound_latent_t, + audio_channels=2, + ) + vision_tok = types.SimpleNamespace(temporal_compression_factor=temporal_cf) + return types.SimpleNamespace(tokenizer_sound_gen=sound_tok, tokenizer_vision_gen=vision_tok) + + +def test_inject_sound_conditions_sound_and_preserves_image(tmp_path: Path): + from cosmos_framework.inference.sound import inject_sound_into_batch + + model = _fake_model(sound_latent_t=50) + # Video tensor [1,3,T,H,W] with T=48 -> 12 video latents at cf=4. + video = torch.zeros(1, 3, 48, 16, 16) + audio = torch.zeros(1, 2, 96000) + batch = { + "video": [video], + # Pre-existing image first-frame condition (as set by build_conditioned_video_batch). + "sequence_plan": [SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=[0])], + } + + inject_sound_into_batch(batch, audio, model, condition_sound=True) + + plan = batch["sequence_plan"][0] + assert plan.has_sound is True + assert plan.condition_frame_indexes_sound == list(range(50)) # all sound conditioned (ts2v) + assert plan.condition_frame_indexes_vision == [0] # image cond preserved + + +def test_inject_sound_default_generates_sound(tmp_path: Path): + from cosmos_framework.inference.sound import inject_sound_into_batch + + model = _fake_model(sound_latent_t=50) + video = torch.zeros(1, 3, 48, 16, 16) + audio = torch.zeros(1, 2, 96000) + batch = {"video": [video], "sequence_plan": [SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=[])]} + + inject_sound_into_batch(batch, audio, model) # default condition_sound=False + + plan = batch["sequence_plan"][0] + assert plan.condition_frame_indexes_sound == [] # t2vs: sound generated +``` + +- [ ] **Step 2: Run to verify it fails** + +Run: `pytest cosmos_framework/inference/sound_test.py::test_inject_sound_conditions_sound_and_preserves_image -v` +Expected: FAIL — `inject_sound_into_batch` got an unexpected keyword argument `condition_sound`. + +- [ ] **Step 3: Add the `condition_sound` parameter** + +In `cosmos_framework/inference/sound.py`, change the signature and the mode selection: + +```python +def inject_sound_into_batch( + data_batch: dict[str, Any], + audio_tensor: torch.Tensor | None, + model: Any, + *, + condition_sound: bool = False, +) -> dict[str, Any]: +``` + +Update the docstring's Args to add: + +``` + condition_sound: When True, the provided audio is used as a clean + condition (mode "ts2v") and the video is generated from it. When + False (default), sound is generated jointly (mode "t2vs"). +``` + +In the `if has_sound:` block, replace the hardcoded `mode="t2vs"`: + +```python + sequence_plan = build_sequence_plan_for_sound( + mode="ts2v" if condition_sound else "t2vs", + video_latent_length=video_latent_t, + sound_latent_length=sound_latent_t, + ) +``` + +(The existing `existing_vision_cond` preservation below it already re-applies the image's `condition_frame_indexes_vision`.) + +- [ ] **Step 4: Run to verify both pass** + +Run: `pytest cosmos_framework/inference/sound_test.py -v` +Expected: PASS (all four tests). + +- [ ] **Step 5: Commit** + +```bash +git add cosmos_framework/inference/sound.py cosmos_framework/inference/sound_test.py +git commit -m "Support sound conditioning (ts2v) in inject_sound_into_batch" +``` + +--- + +## Task 6: Wire real audio into `get_sample_data` + +**Files:** +- Modify: `cosmos_framework/inference/inference.py` (`get_sample_data`, line ~610) + +- [ ] **Step 1: Replace the sound-injection block** + +In `get_sample_data`, replace the existing `if sample_args.enable_sound:` block (lines ~610–625) with: + +```python + if sample_args.enable_sound: + from cosmos_framework.inference.sound import ( + create_placeholder_audio, + get_audio_tokenizer_info, + inject_sound_into_batch, + load_conditioning_audio, + ) + + audio_info = get_audio_tokenizer_info(model) + if not audio_info.has_sound: + raise ValueError("enable_sound=True but model has no sound tokenizer") + + condition_sound = sample_args.sound_path is not None + if condition_sound: + num_samples = int(sample_args.num_frames / sample_args.fps * audio_info.sample_rate) + audio = load_conditioning_audio( + Path(sample_args.sound_path), + sample_rate=audio_info.sample_rate, + audio_channels=getattr(audio_info.tokenizer, "audio_channels", 2), + num_samples=num_samples, + ) + else: + audio = create_placeholder_audio( + num_frames=sample_args.num_frames, + conditioning_fps=sample_args.fps, + audio_info=audio_info, + ) + inject_sound_into_batch(out, audio, model, condition_sound=condition_sound) +``` + +(`Path` is already imported in `inference.py`.) + +- [ ] **Step 2: Sanity-check the module imports** + +Run: `python -c "import cosmos_framework.inference.inference"` +Expected: no error (module imports cleanly). + +- [ ] **Step 3: Commit** + +```bash +git add cosmos_framework/inference/inference.py +git commit -m "Load and condition on real input audio in get_sample_data" +``` + +--- + +## Task 7: Defaults + example input + conditioning audio asset + +**Files:** +- Create: `cosmos_framework/inference/defaults/audio_image2video/sample_args.json` +- Create: `inputs/omni/assets/a2v_audio.wav` +- Create: `inputs/omni/a2v.json` + +- [ ] **Step 1: Create the per-mode defaults** + +Create `cosmos_framework/inference/defaults/audio_image2video/sample_args.json` (image2video defaults + `enable_sound: true`): + +```json +{ + "num_steps": 35, + "guidance": 6.0, + "shift": 10.0, + "sigma_max": 80.0, + "normalize_cfg": false, + "autoregressive": false, + "negative_prompt": null, + "negative_prompt_file": "neg_prompts.json", + "duration_template": "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS.", + "resolution_template": "This video is of {height}x{width} resolution.", + "negative_metadata_mode": "none", + "inverse_duration_template": "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS.", + "inverse_resolution_template": "This video is not of {height}x{width} resolution.", + "negative_prompt_keep_metadata": true, + "aspect_ratio": "16,9", + "fps": 24, + "num_frames": 189, + "video_save_quality": 10, + "image_save_quality": 95, + "enable_sound": true +} +``` + +- [ ] **Step 2: Extract a conditioning audio clip from a published example** + +Run inside the i4 container (via the `slurm-node` skill), from the repo root: + +```bash +mkdir -p inputs/omni/assets +curl -sS -H "Authorization: Bearer $HF_TOKEN" \ + "https://huggingface.co/nvidia/Cosmos3-Nano/resolve/main/assets/example_t2vs_output.mp4" \ + -o /tmp/example_t2vs_output.mp4 +python - <<'PY' +import av, numpy as np, soundfile as sf +container = av.open("/tmp/example_t2vs_output.mp4") +astream = container.streams.audio[0] +sr = astream.codec_context.sample_rate +frames = [f.to_ndarray() for f in container.decode(astream)] # each [C, n] fltp or [n*C] packed +# Normalize to [N, C] float32. +import numpy as np +chunks = [] +for arr in frames: + if arr.ndim == 2: # planar [C, n] + chunks.append(arr.T) + else: # packed + chunks.append(arr.reshape(-1, astream.channels)) +audio = np.concatenate(chunks, axis=0).astype("float32") +audio = audio[: sr * 4] # keep first 4 seconds to bound file size +sf.write("inputs/omni/assets/a2v_audio.wav", audio, sr) +print("wrote", audio.shape, "@", sr) +PY +``` + +Expected: prints the shape and writes `inputs/omni/assets/a2v_audio.wav` (~1–1.5 MB). + +- [ ] **Step 3: Create the example input file** + +Create `inputs/omni/a2v.json` (image from the existing i2v example; audio is the extracted clip, referenced relative to the input file): + +```json +{ + "model_mode": "audio_image2video", + "name": "a2v", + "prompt": "{\"temporal_caption\": \"A silver robotic arm in a clean lab pours water from a glass jar into a white ceramic cup; soft mechanical whirring and gentle water trickling are audible.\", \"audio_description\": \"Gentle splashing and trickling of water with a faint mechanical whir from servo motors; no speech or music.\", \"resolution\": {\"H\": 720, \"W\": 1280}, \"aspect_ratio\": \"16,9\", \"fps\": 24}", + "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/robot_153.jpg", + "sound_path": "assets/a2v_audio.wav" +} +``` + +- [ ] **Step 4: Validate the example parses into sample args** + +Run inside the dev/slurm env: + +```bash +python -c " +import json, pathlib +d = json.loads(pathlib.Path('inputs/omni/a2v.json').read_text()) +assert d['model_mode'] == 'audio_image2video' +assert d['sound_path'].endswith('.wav') +print('ok') +" +``` +Expected: `ok`. (Full arg-building is exercised by the slurm run in Task 9.) + +- [ ] **Step 5: Commit** + +```bash +git add cosmos_framework/inference/defaults/audio_image2video/sample_args.json inputs/omni/a2v.json inputs/omni/assets/a2v_audio.wav +git commit -m "Add audio_image2video defaults and a2v example input" +``` + +--- + +## Task 8: Documentation + +**Files:** +- Modify: `docs/inference.md` + +- [ ] **Step 1: Add the Modes table row** + +In the Modes table (after the `video2video` row), add: + +```markdown +| `audio_image2video` | text prompt + image + audio | `vision.mp4` (with conditioning audio muxed in) | `prompt`, `vision_path`, `sound_path` | [`inputs/omni/a2v.json`](../inputs/omni/a2v.json) | +``` + +- [ ] **Step 2: Note the checkpoint under the sound sentence** + +Replace the existing sentence (line ~148): + +```markdown +Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To run every example in one batch, use `-i "inputs/omni/*.json"`. +``` + +with: + +```markdown +Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To instead **condition** generation on a real audio clip (audio+image → video), use `model_mode: audio_image2video` with a `sound_path` and the sound-encoder checkpoint `--checkpoint-path Cosmos3-Nano-SoundEncoder` (see [`inputs/omni/a2v.json`](../inputs/omni/a2v.json)). To run every example in one batch, use `-i "inputs/omni/*.json"`. +``` + +- [ ] **Step 3: Add the checkpoint to the Models table** + +In the Models table, add a row: + +```markdown +| Cosmos3-Nano-SoundEncoder | `--checkpoint-path=Cosmos3-Nano-SoundEncoder` | `audio_image2video` (audio+image → video), plus all Nano modes | +``` + +- [ ] **Step 4: Commit** + +```bash +git add docs/inference.md +git commit -m "Document audio_image2video mode and sound-encoder checkpoint" +``` + +--- + +## Task 9: End-to-end A2V run on slurm (verification gate) + +**Files:** none (verification only) + +- [ ] **Step 1: Run A2V inference** via the `slurm-node` skill, from the repo root inside the i4 container: + +```bash +python -m cosmos_framework.scripts.inference \ + --parallelism-preset=latency \ + -i "inputs/omni/a2v.json" \ + -o outputs/a2v \ + --checkpoint-path Cosmos3-Nano-SoundEncoder \ + --seed=0 +``` + +Expected: completes without shape/dtype errors during sound encode/condition; writes `outputs/a2v/a2v/vision.mp4` and `sample_args.json`. + +- [ ] **Step 2: Verify the output** + +```bash +python - <<'PY' +import av +c = av.open("outputs/a2v/a2v/vision.mp4") +assert len(c.streams.video) == 1, "no video stream" +assert len(c.streams.audio) >= 1, "no audio track muxed" +print("video frames:", c.streams.video[0].frames, "| audio streams:", len(c.streams.audio)) +PY +``` +Expected: prints a positive video frame count and ≥1 audio stream. + +- [ ] **Step 3: Inspect `sample_args.json`** to confirm conditioning was applied: + +```bash +python -c " +import json +a = json.load(open('outputs/a2v/a2v/sample_args.json')) +assert a['model_mode'] == 'audio_image2video' +assert a['enable_sound'] is True +assert a['condition_frame_indexes_vision'] == [0] +print('conditioning ok') +" +``` +Expected: `conditioning ok`. + +- [ ] **Step 4: If sound tokens are not frozen during sampling** (risk #1 in the spec — e.g. output audio does not match the input clip), STOP and use `superpowers:systematic-debugging` to trace whether the diffusion loop respects `condition_frame_indexes_sound`. This is the one place a model-side fix might be needed; do not paper over it. + +--- + +## Self-Review + +**Spec coverage:** +- Checkpoint registration → Task 1. ✅ +- `audio_image2video` mode → Task 2. ✅ +- `sound_path` field + validation/download → Task 3. ✅ +- `load_conditioning_audio` → Task 4. ✅ +- `condition_sound` wiring (`ts2v` + preserved image cond) → Tasks 5–6. ✅ +- Defaults + example + audio asset → Task 7. ✅ +- Output decode+mux (no change) → confirmed in Task 9 verification. ✅ +- Tests (colocated) → Tasks 3, 4, 5. ✅ (`sound_data_utils` `ts2v` is exercised indirectly by Task 5's plan assertions; the spec's optional `sound_data_utils_test.py` is dropped as redundant — Task 5 already asserts the resulting `condition_frame_indexes_sound`.) +- Docs → Task 8. ✅ +- E2E slurm verification → Task 9. ✅ + +**Placeholder scan:** No TBD/TODO; every code step shows complete code. ✅ + +**Type consistency:** `load_conditioning_audio(path, *, sample_rate, audio_channels, num_samples) -> Tensor[1,C,N]` is defined in Task 4 and called identically in Task 6. `inject_sound_into_batch(..., *, condition_sound=False)` defined in Task 5, called with `condition_sound=condition_sound` in Task 6. `ModelMode.AUDIO_IMAGE2VIDEO` / `is_sound_condition` / `SOUND_CONDITION_MODEL_MODES` consistent across Tasks 2–3. ✅ + +**Note on `sound_data_utils_test.py`:** The spec listed it as "if absent." Dropped to avoid duplicate coverage; Task 5 asserts the `ts2v` plan's effect end-to-end through `inject_sound_into_batch`. If a direct unit test is preferred, add one asserting `build_sequence_plan_for_sound("ts2v", v, s).condition_frame_indexes_sound == list(range(s))`. From 24262925b6ceca98abaddd57f852a299cd52cdbd Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:00:18 -0700 Subject: [PATCH 03/25] Plan: resample via scipy not torchaudio (absent from inference container) Co-Authored-By: Claude Opus 4.8 (1M context) --- .../2026-06-11-a2v-sound-encoder-inference.md | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md b/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md index 55c593b9..015f4101 100644 --- a/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md +++ b/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md @@ -6,7 +6,7 @@ **Architecture:** Reuse the existing `ts2v` (sound-conditioned) sequence plan plus the existing image first-frame vision conditioning, which `inject_sound_into_batch` already preserves. Add a real-audio loader and a `model_mode` + `sound_path` input field; no model/network/tokenizer changes. -**Tech Stack:** Python, PyTorch, pydantic args, `soundfile`/`torchaudio` for audio I/O, `pytest` (colocated `*_test.py`), the Cosmos3 OmniMoT diffusers-format checkpoint loader. +**Tech Stack:** Python, PyTorch, pydantic args, `soundfile` (read) + `scipy.signal` (resample) for audio I/O, `pytest` (colocated `*_test.py`), the Cosmos3 OmniMoT diffusers-format checkpoint loader. Note: `torchaudio` is NOT a project dependency and is absent from the inference container — do not use it. **Spec:** `docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md` @@ -340,11 +340,18 @@ def load_conditioning_audio( data, src_sr = sf.read(str(path), dtype="float32", always_2d=True) # [N, C] waveform = torch.from_numpy(data).transpose(0, 1).contiguous() # [C, N] - # Resample to the tokenizer's rate. + # Resample to the tokenizer's rate. Uses scipy (a declared dependency); + # torchaudio is intentionally avoided as it is not a project dependency + # and is absent from the inference container. if src_sr != sample_rate: - import torchaudio + from math import gcd - waveform = torchaudio.functional.resample(waveform, orig_freq=src_sr, new_freq=sample_rate) + import scipy.signal + + g = gcd(int(src_sr), int(sample_rate)) + up, down = int(sample_rate) // g, int(src_sr) // g + resampled = scipy.signal.resample_poly(waveform.numpy(), up, down, axis=-1) # [C, N'] + waveform = torch.from_numpy(resampled.astype("float32")).contiguous() # Conform channels. cur_channels = waveform.shape[0] From 74679dde4cbdcbdb435a67e45ae8be07dae1f7bd Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:03:42 -0700 Subject: [PATCH 04/25] Register Cosmos3-Nano-SoundEncoder checkpoint Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/args.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 3c234f74..0ab3c307 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -1058,6 +1058,19 @@ def build_sample(self, *, model_config: Any) -> OmniSampleArgs: revision="main", ), ), + # Diffusers HF checkpoint whose transformer is trained to condition on + # (encode) input sound, enabling audio_image2video (A2V). Reuses the + # Cosmos3-Nano architecture (OmniMoTModelConfig, sound_gen=True). + "Cosmos3-Nano-SoundEncoder": CheckpointConfig( + model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["8B"], + config_file=str(CONFIG_DIR / "model/Cosmos3-Nano.yaml"), + s3_uri="", # unused for HF-backed checkpoints + hf=CheckpointDirHf( + repository="nvidia/Cosmos3-Experimental", + revision="main", + subdirectory="nano_diffusers_sound_encoder", + ), + ), "Cosmos3-Super": CheckpointConfig( model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["32B"], config_file=str(CONFIG_DIR / "model/Cosmos3-Super.yaml"), From c2e181320642a381f28a6b4a911d8e4455ab2881 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:06:52 -0700 Subject: [PATCH 05/25] Add audio_image2video model mode Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/args.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 0ab3c307..45d4e2eb 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -160,6 +160,7 @@ class ModelMode(StrEnum): IMAGE2IMAGE = "image2image" IMAGE2VIDEO = "image2video" VIDEO2VIDEO = "video2video" + AUDIO_IMAGE2VIDEO = "audio_image2video" # Action FORWARD_DYNAMICS = "forward_dynamics" @@ -176,6 +177,10 @@ def is_action(self) -> bool: def is_reasoner(self) -> bool: return self in REASONER_MODEL_MODES + @property + def is_sound_condition(self) -> bool: + return self in SOUND_CONDITION_MODEL_MODES + # Image-output modes: ``num_frames`` defaults to 1 and the output is saved as a still image. _IMAGE_OUTPUT_MODES: frozenset[ModelMode] = frozenset({ModelMode.TEXT2IMAGE, ModelMode.IMAGE2IMAGE}) @@ -187,6 +192,10 @@ def is_reasoner(self) -> bool: REASONER_MODEL_MODES: frozenset[ModelMode] = frozenset({ModelMode.REASONER}) +# Modes that condition generation on a real input audio clip (require a model +# with ``sound_gen=True`` and a ``sound_path``). +SOUND_CONDITION_MODEL_MODES: frozenset[ModelMode] = frozenset({ModelMode.AUDIO_IMAGE2VIDEO}) + class VisionMode(StrEnum): IMAGE = "image" From 162a8c2c64c62064b40fce35bfd280eec277ef63 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:15:49 -0700 Subject: [PATCH 06/25] Add sound_path input + audio_image2video arg validation Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/args.py | 16 +++++++++ cosmos_framework/inference/args_test.py | 47 +++++++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 45d4e2eb..a2307442 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -522,6 +522,7 @@ def _build_vision_data(self, model_config: "OmniMoTModelConfig", sample_meta: Sa class SoundDataArgs(ArgsBase): enable_sound: bool = False + sound_path: ResolvedFilePath | None = None class SoundDataOverrides(OverridesBase): @@ -529,8 +530,23 @@ class SoundDataOverrides(OverridesBase): enable_sound: Training[bool | None] = None """Enable joint video+sound generation (t2vs mode). Requires a checkpoint with sound modules.""" + sound_path: str | None = None + """Path or URL to a conditioning audio clip (e.g. .wav/.mp3/.flac). Required for + audio_image2video; the clip is encoded by the AVAE and used as a clean condition.""" + + @override + def download(self, output_dir: Path): + super().download(output_dir) + self.sound_path = download_file(self.sound_path, output_dir, "sound") def _build_sound_data(self, model_config: "OmniMoTModelConfig", sample_meta: SampleMeta): + if sample_meta.model_mode.is_sound_condition: + if self.sound_path is None: + raise ValueError( + f"model_mode={sample_meta.model_mode.value} requires a `sound_path` " + "(a conditioning audio clip)" + ) + self.enable_sound = True if self.enable_sound is None: self.enable_sound = False if self.enable_sound and not model_config.sound_gen: diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index 3bf37032..a14ef2b7 100644 --- a/cosmos_framework/inference/args_test.py +++ b/cosmos_framework/inference/args_test.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: OpenMDW-1.1 import json +import types from pathlib import Path import omegaconf @@ -15,6 +16,7 @@ ModelMode, OmniSampleOverrides, OmniSetupOverrides, + SoundDataOverrides, ) from cosmos_framework.inference.common.config import structure_config @@ -156,3 +158,48 @@ def test_sample_args(tmp_path: Path): assert text2image_args.num_steps == 50 assert text2image_args.guidance == 4.0 assert text2image_args.shift == 3.0 + + +def test_build_sound_data_requires_sound_path_for_a2v(): + model_config = types.SimpleNamespace(sound_gen=True) + sample_meta = types.SimpleNamespace(model_mode=ModelMode.AUDIO_IMAGE2VIDEO) + + overrides = SoundDataOverrides(sound_path=None) + with pytest.raises(ValueError, match="sound_path"): + overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) + + overrides = SoundDataOverrides(sound_path="clip.wav") + overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) + assert overrides.enable_sound is True + + +def test_build_sound_data_rejects_model_without_sound_gen(): + model_config = types.SimpleNamespace(sound_gen=False) + sample_meta = types.SimpleNamespace(model_mode=ModelMode.AUDIO_IMAGE2VIDEO) + overrides = SoundDataOverrides(sound_path="clip.wav") + with pytest.raises(ValueError, match="sound tokenizer"): + overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) + + +def test_audio_image2video_conditions_image_and_sound(tmp_path: Path): + import omegaconf + from cosmos_framework.inference.common.config import structure_config + + setup_args = OmniSetupOverrides( + checkpoint_path=DEFAULT_CHECKPOINT_NAME, + output_dir=tmp_path / "outputs", + ).build_setup() + model_dict = structure_config(setup_args.load_model_config_dict(), omegaconf.DictConfig) + + args = OmniSampleOverrides( + name="a2v", + output_dir=tmp_path / "a2v", + model_mode=ModelMode.AUDIO_IMAGE2VIDEO, + vision_path="robot.jpg", + sound_path="clip.wav", + ).build_sample(model_config=model_dict.config) + + assert args.condition_vision_mode.value == "image" + assert args.condition_frame_indexes_vision == [0] + assert args.enable_sound is True + assert args.sound_path == "clip.wav" From 99203d4c2d340588ff3bc26db7ebcad87824eb9d Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:24:15 -0700 Subject: [PATCH 07/25] Fix sound_path override type and test fixtures Use ResolvedFilePathOrUrl (matching vision_path) for the SoundDataOverrides sound_path field; fix unit tests to use a URL / real tmp_path files. Add missing defaults/audio_image2video/sample_args.json (copied from image2video with enable_sound=true) required by build_sample(). Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/args.py | 2 +- cosmos_framework/inference/args_test.py | 15 ++++++++----- .../audio_image2video/sample_args.json | 22 +++++++++++++++++++ 3 files changed, 33 insertions(+), 6 deletions(-) create mode 100644 cosmos_framework/inference/defaults/audio_image2video/sample_args.json diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index a2307442..6b787bc1 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -530,7 +530,7 @@ class SoundDataOverrides(OverridesBase): enable_sound: Training[bool | None] = None """Enable joint video+sound generation (t2vs mode). Requires a checkpoint with sound modules.""" - sound_path: str | None = None + sound_path: ResolvedFilePathOrUrl | None = None """Path or URL to a conditioning audio clip (e.g. .wav/.mp3/.flac). Required for audio_image2video; the clip is encoded by the AVAE and used as a clean condition.""" diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index a14ef2b7..631f87ef 100644 --- a/cosmos_framework/inference/args_test.py +++ b/cosmos_framework/inference/args_test.py @@ -168,7 +168,7 @@ def test_build_sound_data_requires_sound_path_for_a2v(): with pytest.raises(ValueError, match="sound_path"): overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) - overrides = SoundDataOverrides(sound_path="clip.wav") + overrides = SoundDataOverrides(sound_path="https://example.com/clip.wav") overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) assert overrides.enable_sound is True @@ -176,7 +176,7 @@ def test_build_sound_data_requires_sound_path_for_a2v(): def test_build_sound_data_rejects_model_without_sound_gen(): model_config = types.SimpleNamespace(sound_gen=False) sample_meta = types.SimpleNamespace(model_mode=ModelMode.AUDIO_IMAGE2VIDEO) - overrides = SoundDataOverrides(sound_path="clip.wav") + overrides = SoundDataOverrides(sound_path="https://example.com/clip.wav") with pytest.raises(ValueError, match="sound tokenizer"): overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) @@ -191,15 +191,20 @@ def test_audio_image2video_conditions_image_and_sound(tmp_path: Path): ).build_setup() model_dict = structure_config(setup_args.load_model_config_dict(), omegaconf.DictConfig) + img = tmp_path / "robot.jpg" + img.write_bytes(b"\xff\xd8\xff\xe0") # minimal non-empty file; not actually decoded here + clip = tmp_path / "clip.wav" + clip.write_bytes(b"RIFF") + args = OmniSampleOverrides( name="a2v", output_dir=tmp_path / "a2v", model_mode=ModelMode.AUDIO_IMAGE2VIDEO, - vision_path="robot.jpg", - sound_path="clip.wav", + vision_path=str(img), + sound_path=str(clip), ).build_sample(model_config=model_dict.config) assert args.condition_vision_mode.value == "image" assert args.condition_frame_indexes_vision == [0] assert args.enable_sound is True - assert args.sound_path == "clip.wav" + assert Path(args.sound_path).name == "clip.wav" diff --git a/cosmos_framework/inference/defaults/audio_image2video/sample_args.json b/cosmos_framework/inference/defaults/audio_image2video/sample_args.json new file mode 100644 index 00000000..1fe7360b --- /dev/null +++ b/cosmos_framework/inference/defaults/audio_image2video/sample_args.json @@ -0,0 +1,22 @@ +{ + "num_steps": 35, + "guidance": 6.0, + "shift": 10.0, + "sigma_max": 80.0, + "normalize_cfg": false, + "autoregressive": false, + "negative_prompt": null, + "negative_prompt_file": "neg_prompts.json", + "duration_template": "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS.", + "resolution_template": "This video is of {height}x{width} resolution.", + "negative_metadata_mode": "none", + "inverse_duration_template": "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS.", + "inverse_resolution_template": "This video is not of {height}x{width} resolution.", + "negative_prompt_keep_metadata": true, + "aspect_ratio": "16,9", + "fps": 24, + "num_frames": 189, + "video_save_quality": 10, + "image_save_quality": 95, + "enable_sound": true +} From c0d89fabd217f4f4ed95f96416e1488ae41c8743 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:28:31 -0700 Subject: [PATCH 08/25] Add load_conditioning_audio helper Reads an arbitrary audio file via soundfile, resamples with scipy.signal.resample_poly (avoiding torchaudio which is absent from the inference container), conforms channel count, and trim/pads to an exact sample count. Covered by two new pytest tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/sound.py | 58 ++++++++++++++++++++++++ cosmos_framework/inference/sound_test.py | 36 +++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 cosmos_framework/inference/sound_test.py diff --git a/cosmos_framework/inference/sound.py b/cosmos_framework/inference/sound.py index 1ecf3fbe..076a565d 100644 --- a/cosmos_framework/inference/sound.py +++ b/cosmos_framework/inference/sound.py @@ -59,6 +59,64 @@ def create_placeholder_audio( return torch.zeros(1, sound_channels, sound_num_samples) # [1,C_audio,N_samples] +def load_conditioning_audio( + path: Path, + *, + sample_rate: int, + audio_channels: int, + num_samples: int, +) -> torch.Tensor: + """Decode an audio file into a conditioning waveform aligned to the video. + + Reads ``path`` with soundfile, resamples to ``sample_rate``, conforms the + channel count to ``audio_channels`` (mono->stereo duplicate, stereo->mono + mean), and trims or zero-pads to exactly ``num_samples`` so the audio and + video latent streams cover the same duration. + + Returns: + Audio tensor of shape (1, C, N) where C == audio_channels and + N == num_samples, dtype float32. + """ + import soundfile as sf # type: ignore[import-not-found] + + data, src_sr = sf.read(str(path), dtype="float32", always_2d=True) # [N, C] + waveform = torch.from_numpy(data).transpose(0, 1).contiguous() # [C, N] + + # Resample to the tokenizer's rate. Uses scipy (a declared dependency); + # torchaudio is intentionally avoided as it is not a project dependency + # and is absent from the inference container. + if src_sr != sample_rate: + from math import gcd + + import scipy.signal + + g = gcd(int(src_sr), int(sample_rate)) + up, down = int(sample_rate) // g, int(src_sr) // g + resampled = scipy.signal.resample_poly(waveform.numpy(), up, down, axis=-1) # [C, N'] + waveform = torch.from_numpy(resampled.astype("float32")).contiguous() + + # Conform channels. + cur_channels = waveform.shape[0] + if cur_channels != audio_channels: + if cur_channels == 1 and audio_channels == 2: + waveform = waveform.repeat(2, 1) + elif cur_channels == 2 and audio_channels == 1: + waveform = waveform.mean(dim=0, keepdim=True) + else: + raise ValueError( + f"Cannot convert {cur_channels}-channel audio to {audio_channels} channels" + ) + + # Trim or zero-pad to num_samples. + n = waveform.shape[-1] + if n > num_samples: + waveform = waveform[:, :num_samples] + elif n < num_samples: + waveform = torch.nn.functional.pad(waveform, (0, num_samples - n)) + + return waveform.unsqueeze(0).to(dtype=torch.float32) # [1, C, N] + + def inject_sound_into_batch( data_batch: dict[str, Any], audio_tensor: torch.Tensor | None, diff --git a/cosmos_framework/inference/sound_test.py b/cosmos_framework/inference/sound_test.py new file mode 100644 index 00000000..691306bf --- /dev/null +++ b/cosmos_framework/inference/sound_test.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from pathlib import Path + +import soundfile as sf +import torch + +from cosmos_framework.inference.sound import load_conditioning_audio + + +def _write_wav(path: Path, sample_rate: int, channels: int, num_samples: int) -> None: + if channels > 1: + data = torch.zeros(num_samples, channels).numpy() + else: + data = torch.zeros(num_samples).numpy() + sf.write(str(path), data, sample_rate) + + +def test_load_conditioning_audio_resamples_and_pads(tmp_path: Path): + src = tmp_path / "in.wav" + _write_wav(src, sample_rate=44100, channels=1, num_samples=44100) # 1.0s mono @44.1k + + out = load_conditioning_audio(src, sample_rate=48000, audio_channels=2, num_samples=96000) + + assert out.shape == (1, 2, 96000) # [1, C, N]; stereo, padded to 2.0s @48k + assert out.dtype == torch.float32 + + +def test_load_conditioning_audio_trims(tmp_path: Path): + src = tmp_path / "in.wav" + _write_wav(src, sample_rate=48000, channels=2, num_samples=48000 * 4) # 4s stereo @48k + + out = load_conditioning_audio(src, sample_rate=48000, audio_channels=2, num_samples=48000 * 2) + + assert out.shape == (1, 2, 48000 * 2) # trimmed to 2s From 5efceb2a2e182d6588526003bb01f806b88f786f Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:31:48 -0700 Subject: [PATCH 09/25] Support sound conditioning (ts2v) in inject_sound_into_batch Add keyword-only parameter `condition_sound: bool = False` that selects mode "ts2v" (all sound conditioned, video generated) when True, preserving the existing "t2vs" default (joint generation). Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/sound.py | 7 +++- cosmos_framework/inference/sound_test.py | 51 ++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/cosmos_framework/inference/sound.py b/cosmos_framework/inference/sound.py index 076a565d..261bea41 100644 --- a/cosmos_framework/inference/sound.py +++ b/cosmos_framework/inference/sound.py @@ -121,6 +121,8 @@ def inject_sound_into_batch( data_batch: dict[str, Any], audio_tensor: torch.Tensor | None, model: Any, + *, + condition_sound: bool = False, ) -> dict[str, Any]: """Add sound data and upgrade the SequencePlan in an existing data batch. @@ -131,6 +133,9 @@ def inject_sound_into_batch( data_batch: Existing data batch (from get_video_sample_batch or build_conditioned_video_batch). audio_tensor: Audio waveform tensor (1, C, N) or None. model: The OmniMoTModel instance. + condition_sound: When True, the provided audio is used as a clean + condition (mode "ts2v") and the video is generated from it. When + False (default), sound is generated jointly (mode "t2vs"). Returns: The same data_batch dict, mutated in-place with sound fields added. @@ -161,7 +166,7 @@ def inject_sound_into_batch( # existing vision conditioning is preserved in the sequence plan for i2v and v2v modes sequence_plan = build_sequence_plan_for_sound( - mode="t2vs", + mode="ts2v" if condition_sound else "t2vs", video_latent_length=video_latent_t, sound_latent_length=sound_latent_t, ) diff --git a/cosmos_framework/inference/sound_test.py b/cosmos_framework/inference/sound_test.py index 691306bf..06664294 100644 --- a/cosmos_framework/inference/sound_test.py +++ b/cosmos_framework/inference/sound_test.py @@ -34,3 +34,54 @@ def test_load_conditioning_audio_trims(tmp_path: Path): out = load_conditioning_audio(src, sample_rate=48000, audio_channels=2, num_samples=48000 * 2) assert out.shape == (1, 2, 48000 * 2) # trimmed to 2s + + +import types + +from cosmos_framework.data.vfm.sequence_packing import SequencePlan +from cosmos_framework.inference.sound import inject_sound_into_batch + + +def _fake_model(sound_latent_t: int, temporal_cf: int = 4): + sound_tok = types.SimpleNamespace( + get_latent_num_samples=lambda n: sound_latent_t, + audio_channels=2, + ) + vision_tok = types.SimpleNamespace(temporal_compression_factor=temporal_cf) + return types.SimpleNamespace(tokenizer_sound_gen=sound_tok, tokenizer_vision_gen=vision_tok) + + +def test_inject_sound_conditions_sound_and_preserves_image(): + model = _fake_model(sound_latent_t=50) + video = torch.zeros(1, 3, 48, 16, 16) # [1,3,T,H,W], T=48 -> 12 video latents @cf=4 + audio = torch.zeros(1, 2, 96000) + batch = { + "video": [video], + "sequence_plan": [ + SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=[0]) + ], + } + + inject_sound_into_batch(batch, audio, model, condition_sound=True) + + plan = batch["sequence_plan"][0] + assert plan.has_sound is True + assert plan.condition_frame_indexes_sound == list(range(50)) # all sound conditioned (ts2v) + assert plan.condition_frame_indexes_vision == [0] # image cond preserved + + +def test_inject_sound_default_generates_sound(): + model = _fake_model(sound_latent_t=50) + video = torch.zeros(1, 3, 48, 16, 16) + audio = torch.zeros(1, 2, 96000) + batch = { + "video": [video], + "sequence_plan": [ + SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=[]) + ], + } + + inject_sound_into_batch(batch, audio, model) # default condition_sound=False + + plan = batch["sequence_plan"][0] + assert plan.condition_frame_indexes_sound == [] # t2vs: sound generated From fb9401880be17d0b68ba5c694aa964640d100a4a Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:33:50 -0700 Subject: [PATCH 10/25] Load and condition on real input audio in get_sample_data Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/inference.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index 83ff6655..abc99311 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -612,17 +612,29 @@ def get_sample_data( create_placeholder_audio, get_audio_tokenizer_info, inject_sound_into_batch, + load_conditioning_audio, ) audio_info = get_audio_tokenizer_info(model) if not audio_info.has_sound: raise ValueError("enable_sound=True but model has no sound tokenizer") - audio_placeholder = create_placeholder_audio( - num_frames=sample_args.num_frames, - conditioning_fps=sample_args.fps, - audio_info=audio_info, - ) - inject_sound_into_batch(out, audio_placeholder, model) + + condition_sound = sample_args.sound_path is not None + if condition_sound: + num_samples = int(sample_args.num_frames / sample_args.fps * audio_info.sample_rate) + audio = load_conditioning_audio( + Path(sample_args.sound_path), + sample_rate=audio_info.sample_rate, + audio_channels=getattr(audio_info.tokenizer, "audio_channels", 2), + num_samples=num_samples, + ) + else: + audio = create_placeholder_audio( + num_frames=sample_args.num_frames, + conditioning_fps=sample_args.fps, + audio_info=audio_info, + ) + inject_sound_into_batch(out, audio, model, condition_sound=condition_sound) return out From 0ffc6c8fac585346f4a402f7bc7d645d8bdf1cec Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:36:33 -0700 Subject: [PATCH 11/25] Add audio_image2video (a2v) example input and conditioning audio Co-Authored-By: Claude Opus 4.8 (1M context) --- inputs/omni/a2v.json | 7 +++++++ inputs/omni/assets/a2v_audio.wav | 3 +++ 2 files changed, 10 insertions(+) create mode 100644 inputs/omni/a2v.json create mode 100644 inputs/omni/assets/a2v_audio.wav diff --git a/inputs/omni/a2v.json b/inputs/omni/a2v.json new file mode 100644 index 00000000..45d52542 --- /dev/null +++ b/inputs/omni/a2v.json @@ -0,0 +1,7 @@ +{ + "model_mode": "audio_image2video", + "name": "a2v", + "prompt": "{\"temporal_caption\": \"A silver robotic arm in a clean lab pours water from a glass jar into a white ceramic cup; soft mechanical whirring and gentle water trickling are audible.\", \"audio_description\": \"Gentle splashing and trickling of water with a faint mechanical whir from servo motors; no speech or music.\", \"resolution\": {\"H\": 720, \"W\": 1280}, \"aspect_ratio\": \"16,9\", \"fps\": 24}", + "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/robot_153.jpg", + "sound_path": "assets/a2v_audio.wav" +} diff --git a/inputs/omni/assets/a2v_audio.wav b/inputs/omni/assets/a2v_audio.wav new file mode 100644 index 00000000..645472b3 --- /dev/null +++ b/inputs/omni/assets/a2v_audio.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:87fb26f71721290f8f1251abefd07e9086896532a0db1bbb3ef901d2a3d78a52 +size 768044 From 753b715e08681889eb10b7c8c93c59b9785e8501 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 11 Jun 2026 23:55:36 -0700 Subject: [PATCH 12/25] Document audio_image2video mode and sound-encoder checkpoint Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/inference.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/inference.md b/docs/inference.md index 15801476..8f9230ce 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -130,6 +130,7 @@ The four `--{dp,cp,cfgp}-*-size` flags override the auto-selected values from `- | ------------- | --------------------------------- | ---------------------------------------------- | | Cosmos3-Nano | `--checkpoint-path=Cosmos3-Nano` | All | | Cosmos3-Super | `--checkpoint-path=Cosmos3-Super` | `text2image`, `text2video`, `image2video` | +| Cosmos3-Nano-SoundEncoder | `--checkpoint-path=Cosmos3-Nano-SoundEncoder` | `audio_image2video` (audio+image → video), plus all Nano modes | ## Modes @@ -141,11 +142,12 @@ The four `--{dp,cp,cfgp}-*-size` flags override the auto-selected values from `- | `text2video` | text prompt | `vision.mp4` | `prompt` | [`inputs/omni/t2v.json`](../inputs/omni/t2v.json) | | `image2video` | text prompt + image | `vision.mp4` | `prompt`, `vision_path` | [`inputs/omni/i2v.json`](../inputs/omni/i2v.json) | | `video2video` | text prompt + video | `vision.mp4` | `prompt`, `vision_path` | [`inputs/omni/v2v.json`](../inputs/omni/v2v.json) | +| `audio_image2video` | text prompt + image + audio | `vision.mp4` (with the conditioning audio muxed in) | `prompt`, `vision_path`, `sound_path` | [`inputs/omni/a2v.json`](../inputs/omni/a2v.json) | | `forward_dynamics` | observation image/video + prompt + actions | future visual rollout in `vision.mp4` | `domain_name`, `vision_path`, `action_path` | [`inputs/omni/action_forward_dynamics_av.json`](../inputs/omni/action_forward_dynamics_av.json), [`inputs/omni/action_forward_dynamics_camera.json`](../inputs/omni/action_forward_dynamics_camera.json), [`inputs/omni/action_forward_dynamics_robot.json`](../inputs/omni/action_forward_dynamics_robot.json), [`inputs/omni/action_forward_dynamics_batch.jsonl`](../inputs/omni/action_forward_dynamics_batch.jsonl) | | `inverse_dynamics` | observation video + prompt | predicted action sequence in `sample_outputs.json` | `domain_name`, `vision_path` | [`inputs/omni/action_inverse_dynamics_av.json`](../inputs/omni/action_inverse_dynamics_av.json), [`inputs/omni/action_inverse_dynamics_robot.json`](../inputs/omni/action_inverse_dynamics_robot.json), [`inputs/omni/action_inverse_dynamics_batch.jsonl`](../inputs/omni/action_inverse_dynamics_batch.jsonl) | | `policy` | observation image/video + prompt | predicted action sequence in `sample_outputs.json` + future visual rollout in `vision.mp4` | `domain_name`, `vision_path` | [`inputs/omni/action_policy_av.json`](../inputs/omni/action_policy_av.json), [`inputs/omni/action_policy_robot.json`](../inputs/omni/action_policy_robot.json), [`inputs/omni/action_policy_batch.jsonl`](../inputs/omni/action_policy_batch.jsonl) | -Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To run every example in one batch, use `-i "inputs/omni/*.json"`. +Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To instead **condition** generation on a real audio clip (audio+image → video), use `model_mode: audio_image2video` with a `sound_path` and the sound-encoder checkpoint `--checkpoint-path Cosmos3-Nano-SoundEncoder` (see [`inputs/omni/a2v.json`](../inputs/omni/a2v.json)). To run every example in one batch, use `-i "inputs/omni/*.json"`. ## Parallelism Arguments From c713bafc67020d210b5c3b7bd54f32491ceac607 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 01:58:01 -0700 Subject: [PATCH 13/25] Load full AVAE (encoder+decoder) for sound-conditioned A2V inference The "AVAE" checkpoint registration pointed at nvidia/Cosmos3-Nano/sound_tokenizer, which is decoder-only (182 decoder keys, 0 encoder keys). With strict=False the 67 encoder params silently stayed at random init, so encode(input_audio) produced noise-dominated latents (encode-twice absdiff 0.79; decode(encode(sine)) corr 0.002) and audio_image2video output bore no relation to the conditioning clip. Sound generation (t2vs) was unaffected because it only decodes diffusion latents. Point the AVAE source at nvidia/Cosmos3-Experimental/nano_diffusers_sound_encoder/ sound_tokenizer, which ships the full AVAE: diffusers OobleckDecoder (decoder.block.*) plus the native SpecConvNeXt encoder (encoder.layers.*). _materialize_avae_ckpt already remaps decoder keys and passes the native encoder keys through unchanged. Verified: AVAE round-trip sine corr 1.000, real-audio envelope corr 0.998 (25/25 taps), encode determinism absdiff 0.009; end-to-end A2V output audio envelope corr 0.940 vs the conditioning clip. Also update the i2vs a2v example (hammer image + metallic-tapping audio). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../inference/common/checkpoints.py | 20 ++++++++++++------- inputs/omni/a2v.json | 4 ++-- inputs/omni/assets/a2v_audio.wav | 4 ++-- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/cosmos_framework/inference/common/checkpoints.py b/cosmos_framework/inference/common/checkpoints.py index 4041cb93..d4d51ed4 100644 --- a/cosmos_framework/inference/common/checkpoints.py +++ b/cosmos_framework/inference/common/checkpoints.py @@ -76,8 +76,10 @@ def _materialize_avae_ckpt(local_dir: str) -> None: ``[C]`` and loads via ``load_state_dict(strict=False)`` — so without remapping the keys, none match and every decoder weight is silently left at init (noise). We invert the forward conversion (key remap + snake reshape) and wrap the result - under ``state_dict``. Decoder-only is sufficient: generation only decodes sound - latents to a waveform. Idempotent. + under ``state_dict``. The encoder ships in the already-native ``encoder.layers.*`` + layout, so it passes through ``_avae_block_key_to_legacy`` unchanged and loads + directly (required for encoding input audio in sound-conditioned modes). + Idempotent. """ import torch from safetensors.torch import load_file @@ -245,13 +247,17 @@ def register_checkpoints(): uri="s3://bucket/pretrained/tokenizers/audio/avae", ), hf=CheckpointDirHf( - repository="nvidia/Cosmos3-Nano", + repository="nvidia/Cosmos3-Experimental", revision="main", - subdirectory="sound_tokenizer", + subdirectory="nano_diffusers_sound_encoder/sound_tokenizer", ), - # The sound_tokenizer/ safetensors are decoder-only and use the diffusers - # OobleckDecoder key layout; _materialize_avae_ckpt remaps them back to the - # legacy decoder.layers.* layout the native AVAE loader expects. + # The nano_diffusers_sound_encoder/sound_tokenizer/ safetensors ship the FULL + # AVAE: the diffusers OobleckDecoder decoder (decoder.block.*) plus the native + # SpecConvNeXt encoder (encoder.layers.*). _materialize_avae_ckpt remaps the + # decoder keys back to the legacy decoder.layers.* layout and passes the + # already-native encoder keys through unchanged. The encoder is required to + # ENCODE input audio for sound-conditioned modes (audio_image2video); the base + # nvidia/Cosmos3-Nano/sound_tokenizer is decoder-only and cannot encode. post_download=_materialize_avae_ckpt, ), ) diff --git a/inputs/omni/a2v.json b/inputs/omni/a2v.json index 45d52542..d0da628c 100644 --- a/inputs/omni/a2v.json +++ b/inputs/omni/a2v.json @@ -1,7 +1,7 @@ { "model_mode": "audio_image2video", "name": "a2v", - "prompt": "{\"temporal_caption\": \"A silver robotic arm in a clean lab pours water from a glass jar into a white ceramic cup; soft mechanical whirring and gentle water trickling are audible.\", \"audio_description\": \"Gentle splashing and trickling of water with a faint mechanical whir from servo motors; no speech or music.\", \"resolution\": {\"H\": 720, \"W\": 1280}, \"aspect_ratio\": \"16,9\", \"fps\": 24}", - "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/robot_153.jpg", + "prompt": "{\"subjects\": [{\"description\": \"A claw hammer with a wooden handle and a polished steel head, positioned above a nail embedded in a wooden plank\", \"appearance_details\": \"The hammer has a warm-toned, slightly worn wooden handle with visible wood grain and a smooth, silvery steel head showing minor surface wear and patina\", \"relationship\": \"The hammer is directly above and aligned with the nail, poised to strike it\", \"location\": \"Upper center of frame\", \"relative_size\": \"Large within frame\", \"orientation\": \"Horizontal, head facing downward toward the nail\", \"pose\": \"Suspended above the nail in striking position\", \"action\": \"Repeatedly striking the nail into the wood\", \"state_changes\": \"The hammer rises and falls in a rhythmic motion, making contact with the nail head on each downstroke\", \"clothing\": \"\", \"expression\": \"\", \"gender\": \"\", \"age\": \"\", \"skin_tone_and_texture\": \"\", \"facial_features\": \"\", \"number_of_subjects\": 1, \"number_of_arms\": 0, \"number_of_legs\": 0}, {\"description\": \"A small steel nail partially driven into a wooden plank, with a flat round head and a thin shank showing ring-shank grooves\", \"appearance_details\": \"Silver-colored steel nail with a flat circular head, approximately 2 inches of the shank still visible above the wood surface, with fine ring grooves along the shaft\", \"relationship\": \"The nail is the target of the hammer's strikes, embedded vertically in the wooden surface below\", \"location\": \"Lower center of frame\", \"relative_size\": \"Small within frame\", \"orientation\": \"Vertical, pointing straight up\", \"pose\": \"Standing upright, partially embedded in wood\", \"action\": \"Being driven deeper into the wood with each hammer strike\", \"state_changes\": \"The nail progressively sinks deeper into the wood plank with each successive hammer blow\", \"clothing\": \"\", \"expression\": \"\", \"gender\": \"\", \"age\": \"\", \"skin_tone_and_texture\": \"\", \"facial_features\": \"\", \"number_of_subjects\": 1, \"number_of_arms\": 0, \"number_of_legs\": 0}], \"background_setting\": \"A woodworking workshop or garage workspace. The background is softly blurred, revealing what appears to be other tools, wooden materials, and a window letting in natural light from the right side. The work surface is a thick plank of light-colored natural wood with visible grain patterns.\", \"lighting\": {\"conditions\": \"Warm natural daylight mixed with ambient indoor lighting, creating a soft, diffused illumination\", \"direction\": \"Side-lit from the right, with secondary ambient light from the left\", \"shadows\": \"Soft shadows cast by the hammer onto the wood surface, with gentle shadow beneath the nail; the hammer head casts a subtle shadow on the plank below\", \"illumination_effect\": \"Warm, golden tone that enhances the natural wood textures and gives the metal surfaces a gentle gleam\"}, \"aesthetics\": {\"composition\": \"Centered composition with the hammer and nail aligned vertically in the middle of the frame, creating a strong focal point with the action happening at the intersection\", \"color_scheme\": \"Warm earth tones dominated by honey-brown wood, silver-gray steel, with soft muted background tones of beige and blue-gray\", \"mood_atmosphere\": \"Industrious, focused, craftsmanship, satisfying manual labor\", \"patterns\": \"Horizontal wood grain lines across the plank surface\"}, \"cinematography\": {\"camera_motion\": \"Static with very subtle vibration from hammer impacts\", \"framing\": \"Close-up shot focused tightly on the hammer head and nail\", \"camera_angle\": \"Eye-level, slightly below the hammer looking straight at the point of contact\", \"depth_of_field\": \"Shallow\", \"focus\": \"Sharp focus on the nail and hammer head, with the background and edges of the wood plank softly blurred\", \"lens_focal_length\": \"Macro or short telephoto, approximately 85-100mm equivalent\"}, \"style_medium\": \"Live-action video\", \"artistic_style\": \"Cinematic close-up, documentary-style craftsmanship footage\", \"context\": \"A woodworking or carpentry scene capturing the satisfying act of hammering a nail into wood, suitable for craftsmanship content, ASMR, or instructional material\", \"actions\": [{\"time\": \"0:00-0:01\", \"description\": \"The hammer is poised above the nail, then lifts slightly in preparation for the first strike\"}, {\"time\": \"0:01-0:02\", \"description\": \"The hammer swings down and strikes the nail with a sharp metallic tap, the nail sinks slightly into the wood\"}, {\"time\": \"0:02-0:03\", \"description\": \"The hammer rebounds upward and swings down again for a second strike, driving the nail further\"}, {\"time\": \"0:03-0:04\", \"description\": \"A third strike lands on the nail head, producing another sharp tap; the nail visibly shorter above the wood surface\"}, {\"time\": \"0:04-0:05\", \"description\": \"The hammer rises and delivers a fourth firm blow, the nail now about halfway driven in\"}, {\"time\": \"0:05-0:06\", \"description\": \"A fifth strike drives the nail deeper, with small wood fibers slightly displaced around the entry point\"}, {\"time\": \"0:06-0:07\", \"description\": \"A final strong blow drives the nail nearly flush with the wood surface, the hammer rests momentarily above\"}], \"text_and_signage_elements\": [], \"segments\": [{\"segment_index\": 0, \"time_range\": \"0:00-0:03\", \"description\": \"The hammer begins its rhythmic striking pattern, delivering the first three blows to the nail with measured force\", \"key_changes\": \"The nail progressively sinks from its initial position into the wood; each strike produces a visible impact moment\", \"camera\": \"Static close-up, maintaining focus on the point of contact between hammer and nail\"}, {\"segment_index\": 1, \"time_range\": \"0:03-0:07\", \"description\": \"The hammer continues with stronger, more confident strikes as the nail is driven deeper into the plank\", \"key_changes\": \"The nail becomes significantly shorter above the wood surface; the rhythm of strikes may slightly intensify; small wood displacement visible around the nail\", \"camera\": \"Static close-up, same framing maintained throughout with slight vibration on each impact\"}], \"transitions\": [], \"temporal_caption\": \"The video opens with a close-up of a steel hammer poised just above a nail standing upright in a wooden plank. Within the first second, the hammer lifts slightly and then swings down to deliver the first sharp strike against the nail head. The metallic tap resonates as the nail sinks a small amount into the wood. The hammer rebounds and continues in a steady rhythm, striking the nail approximately once per second. With each successive blow, the nail is driven progressively deeper into the honey-colored wood. By the midpoint of the video, the nail is roughly halfway embedded. The strikes continue with consistent force, and by the final second, the nail is nearly flush with the wood surface. The camera remains static throughout, capturing the satisfying repetitive motion in intimate close-up detail.\", \"audio_description\": \"Sharp, crisp metallic tapping sounds as the steel hammer head contacts the nail with each strike. Each tap is followed by a brief thud as the vibration transfers into the plank. Between strikes, there is a soft whoosh of the hammer moving through air. Ambient workshop sounds are faintly present in the background \\u2014 perhaps a distant hum or subtle room tone. No music or speech is present.\", \"resolution\": {\"H\": 720, \"W\": 1280}, \"aspect_ratio\": \"16,9\", \"duration\": \"7s\", \"fps\": 24}", + "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/01_hammer_nail.png", "sound_path": "assets/a2v_audio.wav" } diff --git a/inputs/omni/assets/a2v_audio.wav b/inputs/omni/assets/a2v_audio.wav index 645472b3..e1efd066 100644 --- a/inputs/omni/assets/a2v_audio.wav +++ b/inputs/omni/assets/a2v_audio.wav @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:87fb26f71721290f8f1251abefd07e9086896532a0db1bbb3ef901d2a3d78a52 -size 768044 +oid sha256:61317f4ce3b3e1e2e23109702d8e175f8742f4b9e818bd8a55d3cdb2789861be +size 1515564 From 0b00fa8454dcde4ed71012b4ec25db8885cdd2a5 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 02:27:32 -0700 Subject: [PATCH 14/25] Use Cosmos3-Nano (not the experimental repo) for sound-conditioned A2V Revert checkpoint/doc references from nvidia/Cosmos3-Experimental (nano_diffusers_sound_encoder) back to the standard Cosmos3-Nano: drop the Cosmos3-Nano-SoundEncoder registry entry, point the AVAE sound_tokenizer source at nvidia/Cosmos3-Nano/sound_tokenizer, and drop the SoundEncoder rows from docs. All A2V logic is unchanged (audio_image2video mode, sound_path input, load_conditioning_audio, condition_sound wiring, encoder-key passthrough in _materialize_avae_ckpt). The published Cosmos3-Nano sound_tokenizer will be updated to ship the full AVAE (encoder + decoder), at which point A2V works on the default checkpoint with no further changes. Also remove the brainstorming design spec. Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/args.py | 13 - .../inference/common/checkpoints.py | 18 +- docs/inference.md | 3 +- ...6-11-a2v-sound-encoder-inference-design.md | 231 ------------------ 4 files changed, 10 insertions(+), 255 deletions(-) delete mode 100644 docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 6b787bc1..62c3c899 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -1083,19 +1083,6 @@ def build_sample(self, *, model_config: Any) -> OmniSampleArgs: revision="main", ), ), - # Diffusers HF checkpoint whose transformer is trained to condition on - # (encode) input sound, enabling audio_image2video (A2V). Reuses the - # Cosmos3-Nano architecture (OmniMoTModelConfig, sound_gen=True). - "Cosmos3-Nano-SoundEncoder": CheckpointConfig( - model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["8B"], - config_file=str(CONFIG_DIR / "model/Cosmos3-Nano.yaml"), - s3_uri="", # unused for HF-backed checkpoints - hf=CheckpointDirHf( - repository="nvidia/Cosmos3-Experimental", - revision="main", - subdirectory="nano_diffusers_sound_encoder", - ), - ), "Cosmos3-Super": CheckpointConfig( model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["32B"], config_file=str(CONFIG_DIR / "model/Cosmos3-Super.yaml"), diff --git a/cosmos_framework/inference/common/checkpoints.py b/cosmos_framework/inference/common/checkpoints.py index d4d51ed4..86877ce2 100644 --- a/cosmos_framework/inference/common/checkpoints.py +++ b/cosmos_framework/inference/common/checkpoints.py @@ -247,17 +247,17 @@ def register_checkpoints(): uri="s3://bucket/pretrained/tokenizers/audio/avae", ), hf=CheckpointDirHf( - repository="nvidia/Cosmos3-Experimental", + repository="nvidia/Cosmos3-Nano", revision="main", - subdirectory="nano_diffusers_sound_encoder/sound_tokenizer", + subdirectory="sound_tokenizer", ), - # The nano_diffusers_sound_encoder/sound_tokenizer/ safetensors ship the FULL - # AVAE: the diffusers OobleckDecoder decoder (decoder.block.*) plus the native - # SpecConvNeXt encoder (encoder.layers.*). _materialize_avae_ckpt remaps the - # decoder keys back to the legacy decoder.layers.* layout and passes the - # already-native encoder keys through unchanged. The encoder is required to - # ENCODE input audio for sound-conditioned modes (audio_image2video); the base - # nvidia/Cosmos3-Nano/sound_tokenizer is decoder-only and cannot encode. + # The sound_tokenizer/ safetensors ship the full AVAE: the diffusers + # OobleckDecoder decoder (decoder.block.*) plus the native SpecConvNeXt + # encoder (encoder.layers.*). _materialize_avae_ckpt remaps the decoder keys + # back to the legacy decoder.layers.* layout and passes the already-native + # encoder keys through unchanged. The encoder is required to ENCODE input + # audio for sound-conditioned modes (audio_image2video); the decoder alone + # only supports sound generation. post_download=_materialize_avae_ckpt, ), ) diff --git a/docs/inference.md b/docs/inference.md index 8f9230ce..23c41b4d 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -130,7 +130,6 @@ The four `--{dp,cp,cfgp}-*-size` flags override the auto-selected values from `- | ------------- | --------------------------------- | ---------------------------------------------- | | Cosmos3-Nano | `--checkpoint-path=Cosmos3-Nano` | All | | Cosmos3-Super | `--checkpoint-path=Cosmos3-Super` | `text2image`, `text2video`, `image2video` | -| Cosmos3-Nano-SoundEncoder | `--checkpoint-path=Cosmos3-Nano-SoundEncoder` | `audio_image2video` (audio+image → video), plus all Nano modes | ## Modes @@ -147,7 +146,7 @@ The four `--{dp,cp,cfgp}-*-size` flags override the auto-selected values from `- | `inverse_dynamics` | observation video + prompt | predicted action sequence in `sample_outputs.json` | `domain_name`, `vision_path` | [`inputs/omni/action_inverse_dynamics_av.json`](../inputs/omni/action_inverse_dynamics_av.json), [`inputs/omni/action_inverse_dynamics_robot.json`](../inputs/omni/action_inverse_dynamics_robot.json), [`inputs/omni/action_inverse_dynamics_batch.jsonl`](../inputs/omni/action_inverse_dynamics_batch.jsonl) | | `policy` | observation image/video + prompt | predicted action sequence in `sample_outputs.json` + future visual rollout in `vision.mp4` | `domain_name`, `vision_path` | [`inputs/omni/action_policy_av.json`](../inputs/omni/action_policy_av.json), [`inputs/omni/action_policy_robot.json`](../inputs/omni/action_policy_robot.json), [`inputs/omni/action_policy_batch.jsonl`](../inputs/omni/action_policy_batch.jsonl) | -Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To instead **condition** generation on a real audio clip (audio+image → video), use `model_mode: audio_image2video` with a `sound_path` and the sound-encoder checkpoint `--checkpoint-path Cosmos3-Nano-SoundEncoder` (see [`inputs/omni/a2v.json`](../inputs/omni/a2v.json)). To run every example in one batch, use `-i "inputs/omni/*.json"`. +Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To instead **condition** generation on a real audio clip (audio+image → video), use `model_mode: audio_image2video` with a `sound_path` (see [`inputs/omni/a2v.json`](../inputs/omni/a2v.json)). To run every example in one batch, use `-i "inputs/omni/*.json"`. ## Parallelism Arguments diff --git a/docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md b/docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md deleted file mode 100644 index 58480aed..00000000 --- a/docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md +++ /dev/null @@ -1,231 +0,0 @@ -# Audio+Image→Video (A2V) Inference with the Sound-Encoder Checkpoint - -**Date:** 2026-06-11 -**Status:** Approved design — ready for implementation plan -**Branch:** `a2v-sound-encoder-inference` - -## Summary - -Add a new inference mode, `audio_image2video`, that conditions video generation on a -**real input audio clip** plus an **input image (first frame)**, using the -`nano_diffusers_sound_encoder` checkpoint -(`nvidia/Cosmos3-Experimental`, subfolder `nano_diffusers_sound_encoder`). - -Today the only audio path in inference is *generation*: `enable_sound: true` injects a -zero-filled placeholder waveform and the model **generates** sound (internally -`mode="t2vs"`). There is no way to feed a real audio clip as a **condition**. This spec -closes that gap on the inference side only — the model and data layers already support -sound-as-condition. - -## Background / Current State - -- **Checkpoint format.** The default `Cosmos3-Nano` checkpoint is diffusers-format - (`model_index.json`, `transformer/`, `vae/`, `sound_tokenizer/`, `vision_encoder/`) - and already ships a `sound_tokenizer` (AVAE). The framework's loader already handles - this format. The new `nano_diffusers_sound_encoder` checkpoint has the same layout - with different `transformer` + `sound_tokenizer` weights; its `config.json` maps to the - same `OmniMoTModelConfig` (`sound_gen=true`, `sound_dim=64`, AVAE - `sample_rate=48000`, `audio_channels=2`). -- **Sound generation today.** `cosmos_framework/inference/inference.py:get_sample_data` - (≈ line 610) calls `create_placeholder_audio` → `inject_sound_into_batch`, which - hardcodes `mode="t2vs"` in `cosmos_framework/inference/sound.py:105`. All sound tokens - are generated; the placeholder only establishes the target length. -- **Sound conditioning already exists at the data/model layer.** - `cosmos_framework/data/vfm/sound_data_utils.py` defines `ts2v` (Text+Sound→Video, - sound conditioned) and `ti2sv`. The MOT network handles clean vs. noisy sound tokens - via the condition mask (`_encode_sound`/`_decode_sound`). These plans are **not - reachable** from any `model_mode`. -- **Vision-condition preservation.** `inject_sound_into_batch` captures and re-applies - the existing `condition_frame_indexes_vision` (sound.py lines ~90–94, 110–111). So the - image's first-frame condition survives whichever sound mode is chosen. -- **Output path is generic.** `inference.py` (≈ lines 1489–1544) decodes any `"sound"` in - the model outputs (`model.decode_sound`) and muxes it into the `.mp4` - (`mux_audio_into_video`). No change is needed for conditioned audio. - -### Current `ModelMode` values - -`text2image`, `text2video`, `image2image`, `image2video`, `video2video`, -`forward_dynamics`, `inverse_dynamics`, `policy`, `reasoner`. None takes audio as input. - -## Goals - -1. Register the `nano_diffusers_sound_encoder` checkpoint as a named checkpoint. -2. Add `audio_image2video`: image (first frame) + real audio clip → video, with the audio - used as a **clean condition** and muxed into the output video. -3. Full deliverable: input loading, defaults, example input, colocated tests, docs. - -## Non-Goals - -- Audio-only conditioning (`ts2v` with no image) as a separate user-facing mode (YAGNI). -- Any change to the model architecture, sequence packing, AVAE tokenizer, training, or - the output/save path. -- A new `tis2v` sequence plan: the audio+image combination falls out of `ts2v` (conditions - sound) plus the preserved image vision-condition. - -## Design - -### Key leverage point - -"Audio+image→video" needs no new sequence-plan combination. `inject_sound_into_batch` -already preserves the image's `condition_frame_indexes_vision=[0]`. Selecting the existing -`ts2v` plan (all sound latents conditioned) and letting that preservation re-apply `[0]` -yields exactly: image first-frame conditioned + audio conditioned + remaining video -generated. - -### Changes by component - -**1. Checkpoint registry** — `cosmos_framework/inference/args.py` (`_CHECKPOINTS`, line 1051) - -Add: - -```python -"Cosmos3-Nano-SoundEncoder": CheckpointConfig( - model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["8B"], - config_file=str(CONFIG_DIR / "model/Cosmos3-Nano.yaml"), - s3_uri=..., # match the Nano entry's pattern; unused for HF-backed load - hf=CheckpointDirHf( - repository="nvidia/Cosmos3-Experimental", - revision="main", - subdirectory="nano_diffusers_sound_encoder", - ), -) -``` - -Reuses `Cosmos3-Nano.yaml` (same `OmniMoTModelConfig`). Weight compatibility is verified -on load and end-to-end by the slurm run. - -**2. New model mode** — `cosmos_framework/inference/args.py` (`ModelMode`, line 157) - -- Add `AUDIO_IMAGE2VIDEO = "audio_image2video"`. -- Add a `_SOUND_CONDITION_MODES: frozenset[ModelMode] = {ModelMode.AUDIO_IMAGE2VIDEO}` and - an `is_sound_condition` property mirroring `is_action`/`is_reasoner`. -- `condition_vision_mode` already resolves to `image` from an image `vision_path` - (lines 355–368), giving `condition_frame_indexes_vision=[0]` via existing defaults. - -**3. Audio input field** — `cosmos_framework/inference/args.py` - -- `SoundDataArgs` (line 514): add `sound_path: ResolvedFilePath | None = None`. -- `SoundDataOverrides` (line 518): add `sound_path: ResolvedFilePathOrUrl | None = None` - with a docstring; add a `download()` override that calls - `self.sound_path = download_file(self.sound_path, output_dir, "sound")`. -- `_build_sound_data` (line 524): for `AUDIO_IMAGE2VIDEO`, require `sound_path` set and - force `enable_sound = True`; keep the existing `sound_gen` validation. - -**4. Audio decode helper** — `cosmos_framework/inference/sound.py` - -```python -def load_conditioning_audio( - path: Path, - *, - sample_rate: int, - audio_channels: int, - num_samples: int, -) -> torch.Tensor: - """Decode an audio file to a [1, C, N] waveform aligned to the video duration. - - Reads via soundfile, resamples to ``sample_rate``, conforms channel count to - ``audio_channels`` (mono->stereo duplicate / stereo->mono mean), and trims or - zero-pads to ``num_samples`` so the audio and video latent streams align temporally. - """ -``` - -`num_samples` = `int(num_frames / fps * sample_rate)` (matches `create_placeholder_audio`). - -**5. Wire conditioning** - -- `cosmos_framework/inference/sound.py`: add `condition_sound: bool = False` to - `inject_sound_into_batch`. When `True`, build the plan with `mode="ts2v"` (instead of - `"t2vs"`); the existing vision-cond preservation keeps the image's `[0]`. -- `cosmos_framework/inference/inference.py` `get_sample_data` (≈ line 610): when - `sample_args.sound_path` is set, call `load_conditioning_audio(...)` and - `inject_sound_into_batch(out, audio, model, condition_sound=True)`. Otherwise keep the - existing placeholder-generation branch unchanged. - -**6. Defaults + example** - -- `cosmos_framework/inference/defaults/audio_image2video/sample_args.json`: copy of - `image2video/sample_args.json` with `"enable_sound": true`. -- `inputs/omni/a2v.json`: `model_mode="audio_image2video"`, image `vision_path`, - `sound_path`, and a prompt. - -**7. Output** — no code change. `inference.py:1489` decodes and muxes the (clean, -conditioned) sound into `vision.mp4`. - -### Data flow - -``` -a2v.json (vision_path=image, sound_path=audio, model_mode=audio_image2video) - -> download() resolves both paths - -> get_sample_data: - condition_vision_mode=image -> load_conditioning_image -> build_conditioned_video_batch - (condition_frame_indexes_vision=[0], sequence_plan with image cond) - sound_path set -> load_conditioning_audio -> [1,C,N] real waveform - inject_sound_into_batch(condition_sound=True): - mode="ts2v" -> condition_frame_indexes_sound = all - preserve condition_frame_indexes_vision=[0] - -> model encodes audio (AVAE) as clean sound tokens; image as clean first frame - -> diffusion generates remaining video frames conditioned on both - -> outputs["sound"] decoded + muxed into vision.mp4 -``` - -## Testing - -Colocated unit tests (matching the `colocated-tests-ci` convention): - -- ➕ `cosmos_framework/inference/sound_test.py` - - `load_conditioning_audio`: resample rate change, mono↔stereo conformance, - trim/pad to `num_samples`, returns `[1, C, num_samples]`. - - `inject_sound_into_batch(condition_sound=True)`: produces a plan with all sound - latents conditioned **and** preserves a pre-set `condition_frame_indexes_vision=[0]`; - `condition_sound=False` keeps the existing `t2vs` behavior. -- ✏️ `cosmos_framework/inference/args_test.py` - - `audio_image2video` resolves `condition_vision_mode=image` and - `condition_frame_indexes_vision=[0]`. - - `_build_sound_data` requires `sound_path` and sets `enable_sound=True`; validation - fails on a model with `sound_gen=False`. -- ➕ `cosmos_framework/data/vfm/sound_data_utils_test.py` (if absent): `ts2v` plan sets - `condition_frame_indexes_sound = range(sound_latent_length)` and empty vision cond. - -CPU-only tests use small synthetic tensors / a tiny generated `.wav`; no GPU or checkpoint -download required. - -## End-to-end verification (slurm) - -Run at repo root inside the i4 container (per the `slurm-node` skill): - -```shell -python -m cosmos_framework.scripts.inference \ - --parallelism-preset=latency \ - -i "inputs/omni/a2v.json" \ - -o outputs/a2v \ - --checkpoint-path Cosmos3-Nano-SoundEncoder \ - --seed=0 -``` - -Success = `outputs/a2v//vision.mp4` exists, plays, and contains an audio track -matching the input clip; no shape/dtype errors during sound encode/condition. - -## Example inputs (sourced by implementer) - -- **Image:** the robot image already referenced by `inputs/omni/i2v.json`. -- **Audio:** extract the audio track from a Cosmos3-Nano example sound output - (e.g. `assets/example_t2vs_output.mp4` in the `nvidia/Cosmos3-Nano` HF repo) into a - short `.wav`, referenced by URL or local path. - -## Risks - -1. **Conditioned sound tokens must stay frozen during sampling.** The model uses the same - forward path as `ts2v` training, so this should hold; the slurm run is the gate. If the - sampler does not freeze clean sound tokens, a small inference-loop fix may be needed - (out of scope until observed). -2. **Checkpoint/config compatibility.** `nano_diffusers_sound_encoder/config.json` matches - `Cosmos3-Nano.yaml`'s `OmniMoTModelConfig`; confirmed by inspection, verified on load. -3. **Audio/video temporal alignment.** Handled by trimming/padding audio to the video - duration in `load_conditioning_audio`. - -## Files touched (summary) - -Modify: `inference/args.py`, `inference/sound.py`, `inference/inference.py`, -`inference/args_test.py`, `docs/inference.md`. -Create: `inference/sound_test.py`, `inference/defaults/audio_image2video/sample_args.json`, -`inputs/omni/a2v.json`, and `data/vfm/sound_data_utils_test.py` (if absent). From 0bee46d2cb4e2d64c867c9ebf2fb77902dbf1804 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 02:28:50 -0700 Subject: [PATCH 15/25] Remove A2V implementation plan doc Co-Authored-By: Claude Opus 4.8 (1M context) --- .../2026-06-11-a2v-sound-encoder-inference.md | 777 ------------------ 1 file changed, 777 deletions(-) delete mode 100644 docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md diff --git a/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md b/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md deleted file mode 100644 index 015f4101..00000000 --- a/docs/superpowers/plans/2026-06-11-a2v-sound-encoder-inference.md +++ /dev/null @@ -1,777 +0,0 @@ -# Audio+Image→Video (A2V) Inference Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Add an `audio_image2video` inference mode that conditions video generation on a real input audio clip + an input image (first frame), using the `nano_diffusers_sound_encoder` checkpoint. - -**Architecture:** Reuse the existing `ts2v` (sound-conditioned) sequence plan plus the existing image first-frame vision conditioning, which `inject_sound_into_batch` already preserves. Add a real-audio loader and a `model_mode` + `sound_path` input field; no model/network/tokenizer changes. - -**Tech Stack:** Python, PyTorch, pydantic args, `soundfile` (read) + `scipy.signal` (resample) for audio I/O, `pytest` (colocated `*_test.py`), the Cosmos3 OmniMoT diffusers-format checkpoint loader. Note: `torchaudio` is NOT a project dependency and is absent from the inference container — do not use it. - -**Spec:** `docs/superpowers/specs/2026-06-11-a2v-sound-encoder-inference-design.md` - -**Branch:** `a2v-sound-encoder-inference` - ---- - -## Background the implementer must know - -- **No model changes.** The MOT model + AVAE already support sound-as-condition (`ts2v`). The only gap is the inference pipeline: it never loads real audio and never selects a conditioning plan. -- **`inject_sound_into_batch` preserves vision conditioning** (`cosmos_framework/inference/sound.py` ~lines 90–94, 110–111). So image-first-frame + audio-conditioned falls out of `mode="ts2v"` automatically. -- **Output decode+mux is generic** (`cosmos_framework/inference/inference.py` ~lines 1489–1544): any `"sound"` in model outputs is decoded and muxed into the `.mp4`. No change needed. -- **Sample-arg machinery:** `OmniSampleOverrides.build_sample(model_config=...)` runs `_build_vision_data` then `_build_sound_data` (`args.py` ~line 1004). `download()` methods cascade via `super().download()` through the MRO (see `VisionDataOverrides.download` at `args.py:447`). -- **Run tests** with the repo's pytest. Audio unit tests are CPU-only (synthetic tensors / tiny generated WAV). Run via the `slurm-node` skill if pytest needs the i4 container; otherwise locally. - ---- - -## File Structure - -- `cosmos_framework/inference/args.py` — checkpoint registry entry, `ModelMode.AUDIO_IMAGE2VIDEO`, `sound_path` field + validation/download. -- `cosmos_framework/inference/sound.py` — `load_conditioning_audio` helper; `condition_sound` param on `inject_sound_into_batch`. -- `cosmos_framework/inference/inference.py` — `get_sample_data` branch that loads real audio and conditions on it. -- `cosmos_framework/inference/sound_test.py` (new) — unit tests for the two `sound.py` additions. -- `cosmos_framework/inference/args_test.py` — unit test for the new mode + `sound_path` validation. -- `cosmos_framework/inference/defaults/audio_image2video/sample_args.json` (new) — per-mode defaults. -- `inputs/omni/a2v.json` (new) + `inputs/omni/assets/a2v_audio.wav` (new) — example input. -- `docs/inference.md` — Modes table row + checkpoint note. - ---- - -## Task 1: Register the sound-encoder checkpoint - -**Files:** -- Modify: `cosmos_framework/inference/args.py` (`_CHECKPOINTS`, line ~1051) -- Test: `cosmos_framework/inference/args_test.py::test_checkpoints` (existing — auto-covers the new entry) - -- [ ] **Step 1: Add the registry entry** - -In `_CHECKPOINTS`, after the `"Cosmos3-Nano"` entry, add: - -```python - # Diffusers HF checkpoint whose transformer is trained to condition on - # (encode) input sound, enabling audio_image2video (A2V). Reuses the - # Cosmos3-Nano architecture (OmniMoTModelConfig, sound_gen=True). - "Cosmos3-Nano-SoundEncoder": CheckpointConfig( - model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["8B"], - config_file=str(CONFIG_DIR / "model/Cosmos3-Nano.yaml"), - s3_uri="", # unused for HF-backed checkpoints - hf=CheckpointDirHf( - repository="nvidia/Cosmos3-Experimental", - revision="main", - subdirectory="nano_diffusers_sound_encoder", - ), - ), -``` - -- [ ] **Step 2: Verify the registry resolves and downloads `checkpoint.json`** - -Run: `pytest cosmos_framework/inference/args_test.py::test_checkpoints -v` -Expected: PASS. (Requires `HF_TOKEN` with access to `nvidia/Cosmos3-Experimental`; run inside the dev/slurm env where the token is set. The test downloads `nano_diffusers_sound_encoder/checkpoint.json` via the `subdirectory` filter and parses it.) - -- [ ] **Step 3: Commit** - -```bash -git add cosmos_framework/inference/args.py -git commit -m "Register Cosmos3-Nano-SoundEncoder checkpoint" -``` - ---- - -## Task 2: Add the `audio_image2video` model mode - -**Files:** -- Modify: `cosmos_framework/inference/args.py` (`ModelMode` enum line ~157; frozensets line ~181–188) -- Test: `cosmos_framework/inference/args_test.py` (new test added in Task 3, Step 5) - -- [ ] **Step 1: Add the enum member** - -In `class ModelMode(StrEnum)` (after `VIDEO2VIDEO = "video2video"`): - -```python - AUDIO_IMAGE2VIDEO = "audio_image2video" -``` - -- [ ] **Step 2: Add the mode-group frozenset and property** - -After `REASONER_MODEL_MODES` (line ~188) add: - -```python -# Modes that condition generation on a real input audio clip (require a model -# with ``sound_gen=True`` and a ``sound_path``). -SOUND_CONDITION_MODEL_MODES: frozenset[ModelMode] = frozenset({ModelMode.AUDIO_IMAGE2VIDEO}) -``` - -In `class ModelMode`, alongside `is_action` / `is_reasoner`: - -```python - @property - def is_sound_condition(self) -> bool: - return self in SOUND_CONDITION_MODEL_MODES -``` - -- [ ] **Step 3: Verify import + enum value** - -Run: `python -c "from cosmos_framework.inference.args import ModelMode; print(ModelMode.AUDIO_IMAGE2VIDEO.value, ModelMode.AUDIO_IMAGE2VIDEO.is_sound_condition)"` -Expected: `audio_image2video True` - -- [ ] **Step 4: Commit** - -```bash -git add cosmos_framework/inference/args.py -git commit -m "Add audio_image2video model mode" -``` - ---- - -## Task 3: Add the `sound_path` input field + validation/download - -**Files:** -- Modify: `cosmos_framework/inference/args.py` (`SoundDataArgs` line ~514; `SoundDataOverrides` line ~518) -- Test: `cosmos_framework/inference/args_test.py` - -- [ ] **Step 1: Write the failing test** - -Add to `cosmos_framework/inference/args_test.py` (note the new imports `ModelMode` is already imported; add `SoundDataOverrides`): - -```python -import types - -from cosmos_framework.inference.args import SoundDataOverrides - - -def test_build_sound_data_requires_sound_path_for_a2v(): - model_config = types.SimpleNamespace(sound_gen=True) - sample_meta = types.SimpleNamespace(model_mode=ModelMode.AUDIO_IMAGE2VIDEO) - - # Missing sound_path -> error - overrides = SoundDataOverrides(sound_path=None) - with pytest.raises(ValueError, match="sound_path"): - overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) - - # sound_path set -> enable_sound forced True - overrides = SoundDataOverrides(sound_path="clip.wav") - overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) - assert overrides.enable_sound is True - - -def test_build_sound_data_rejects_model_without_sound_gen(): - model_config = types.SimpleNamespace(sound_gen=False) - sample_meta = types.SimpleNamespace(model_mode=ModelMode.AUDIO_IMAGE2VIDEO) - overrides = SoundDataOverrides(sound_path="clip.wav") - with pytest.raises(ValueError, match="sound tokenizer"): - overrides._build_sound_data(model_config=model_config, sample_meta=sample_meta) -``` - -- [ ] **Step 2: Run test to verify it fails** - -Run: `pytest cosmos_framework/inference/args_test.py::test_build_sound_data_requires_sound_path_for_a2v -v` -Expected: FAIL — `SoundDataOverrides` has no `sound_path` (TypeError/validation error). - -- [ ] **Step 3: Add the field + download + validation** - -In `class SoundDataArgs(ArgsBase)` (line ~514): - -```python -class SoundDataArgs(ArgsBase): - enable_sound: bool = False - sound_path: ResolvedFilePath | None = None -``` - -In `class SoundDataOverrides(OverridesBase)` (line ~518) add the field and a `download()` override, and extend `_build_sound_data`: - -```python -class SoundDataOverrides(OverridesBase): - """Sound data overrides.""" - - enable_sound: Training[bool | None] = None - """Enable joint video+sound generation (t2vs mode). Requires a checkpoint with sound modules.""" - sound_path: Training[ResolvedFilePathOrUrl | None] = None - """Path or URL to a conditioning audio clip (e.g. .wav/.mp3/.flac). Required for - audio_image2video; the clip is encoded by the AVAE and used as a clean condition.""" - - @override - def download(self, output_dir: Path): - super().download(output_dir) - self.sound_path = download_file(self.sound_path, output_dir, "sound") - - def _build_sound_data(self, model_config: "OmniMoTModelConfig", sample_meta: SampleMeta): - if sample_meta.model_mode.is_sound_condition: - if self.sound_path is None: - raise ValueError( - f"model_mode={sample_meta.model_mode.value} requires a `sound_path` " - "(a conditioning audio clip)" - ) - self.enable_sound = True - if self.enable_sound is None: - self.enable_sound = False - if self.enable_sound and not model_config.sound_gen: - raise ValueError( - "enable_sound=True requires a model with a sound tokenizer " - "(model.config.sound_gen=True), but the loaded checkpoint has no sound tokenizer" - ) -``` - -(`@override`, `Path`, `download_file`, `ResolvedFilePathOrUrl`, `Training` are already imported in `args.py`.) - -- [ ] **Step 4: Run test to verify it passes** - -Run: `pytest cosmos_framework/inference/args_test.py::test_build_sound_data_requires_sound_path_for_a2v cosmos_framework/inference/args_test.py::test_build_sound_data_rejects_model_without_sound_gen -v` -Expected: PASS (both). - -- [ ] **Step 5: Add the mode-resolution test** - -Add to `args_test.py` (reuses `model_dict.config` pattern from `test_sample_args`; place it as its own test so it can build a sample for the new mode): - -```python -def test_audio_image2video_conditions_image_and_sound(tmp_path: Path): - import omegaconf - from cosmos_framework.inference.common.config import structure_config - - setup_args = OmniSetupOverrides( - checkpoint_path=DEFAULT_CHECKPOINT_NAME, - output_dir=tmp_path / "outputs", - ).build_setup() - model_dict = structure_config(setup_args.load_model_config_dict(), omegaconf.DictConfig) - - args = OmniSampleOverrides( - name="a2v", - output_dir=tmp_path / "a2v", - model_mode=ModelMode.AUDIO_IMAGE2VIDEO, - vision_path="robot.jpg", # image extension -> first-frame condition - sound_path="clip.wav", - ).build_sample(model_config=model_dict.config) - - assert args.condition_vision_mode.value == "image" - assert args.condition_frame_indexes_vision == [0] - assert args.enable_sound is True - assert args.sound_path == "clip.wav" -``` - -- [ ] **Step 6: Run it** - -Run: `pytest cosmos_framework/inference/args_test.py::test_audio_image2video_conditions_image_and_sound -v` -Expected: PASS. (Downloads the default Nano model config; run in the dev/slurm env.) - -- [ ] **Step 7: Commit** - -```bash -git add cosmos_framework/inference/args.py cosmos_framework/inference/args_test.py -git commit -m "Add sound_path input + audio_image2video arg validation" -``` - ---- - -## Task 4: `load_conditioning_audio` helper - -**Files:** -- Modify: `cosmos_framework/inference/sound.py` -- Test: `cosmos_framework/inference/sound_test.py` (new) - -- [ ] **Step 1: Write the failing test** - -Create `cosmos_framework/inference/sound_test.py`: - -```python -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -from pathlib import Path - -import soundfile as sf -import torch - -from cosmos_framework.inference.sound import load_conditioning_audio - - -def _write_wav(path: Path, sample_rate: int, channels: int, num_samples: int) -> None: - data = torch.zeros(num_samples, channels).numpy() if channels > 1 else torch.zeros(num_samples).numpy() - sf.write(str(path), data, sample_rate) - - -def test_load_conditioning_audio_resamples_and_pads(tmp_path: Path): - src = tmp_path / "in.wav" - _write_wav(src, sample_rate=44100, channels=1, num_samples=44100) # 1.0s mono @44.1k - - out = load_conditioning_audio(src, sample_rate=48000, audio_channels=2, num_samples=96000) - - assert out.shape == (1, 2, 96000) # [1, C, N], stereo, exactly num_samples (2.0s @48k -> pad) - assert out.dtype == torch.float32 - - -def test_load_conditioning_audio_trims(tmp_path: Path): - src = tmp_path / "in.wav" - _write_wav(src, sample_rate=48000, channels=2, num_samples=48000 * 4) # 4s stereo @48k - - out = load_conditioning_audio(src, sample_rate=48000, audio_channels=2, num_samples=48000 * 2) - - assert out.shape == (1, 2, 48000 * 2) # trimmed to 2s -``` - -- [ ] **Step 2: Run to verify it fails** - -Run: `pytest cosmos_framework/inference/sound_test.py -v` -Expected: FAIL — `ImportError: cannot import name 'load_conditioning_audio'`. - -- [ ] **Step 3: Implement the helper** - -Add to `cosmos_framework/inference/sound.py` (after `create_placeholder_audio`): - -```python -def load_conditioning_audio( - path: Path, - *, - sample_rate: int, - audio_channels: int, - num_samples: int, -) -> torch.Tensor: - """Decode an audio file into a conditioning waveform aligned to the video. - - Reads ``path`` with soundfile, resamples to ``sample_rate``, conforms the - channel count to ``audio_channels`` (mono->stereo duplicate, stereo->mono - mean), and trims or zero-pads to exactly ``num_samples`` so the audio and - video latent streams cover the same duration. - - Returns: - Audio tensor of shape (1, C, N) where C == audio_channels and - N == num_samples, dtype float32. - """ - import soundfile as sf # type: ignore[import-not-found] - - data, src_sr = sf.read(str(path), dtype="float32", always_2d=True) # [N, C] - waveform = torch.from_numpy(data).transpose(0, 1).contiguous() # [C, N] - - # Resample to the tokenizer's rate. Uses scipy (a declared dependency); - # torchaudio is intentionally avoided as it is not a project dependency - # and is absent from the inference container. - if src_sr != sample_rate: - from math import gcd - - import scipy.signal - - g = gcd(int(src_sr), int(sample_rate)) - up, down = int(sample_rate) // g, int(src_sr) // g - resampled = scipy.signal.resample_poly(waveform.numpy(), up, down, axis=-1) # [C, N'] - waveform = torch.from_numpy(resampled.astype("float32")).contiguous() - - # Conform channels. - cur_channels = waveform.shape[0] - if cur_channels != audio_channels: - if cur_channels == 1 and audio_channels == 2: - waveform = waveform.repeat(2, 1) - elif cur_channels == 2 and audio_channels == 1: - waveform = waveform.mean(dim=0, keepdim=True) - else: - raise ValueError( - f"Cannot convert {cur_channels}-channel audio to {audio_channels} channels" - ) - - # Trim or zero-pad to num_samples. - n = waveform.shape[-1] - if n > num_samples: - waveform = waveform[:, :num_samples] - elif n < num_samples: - waveform = torch.nn.functional.pad(waveform, (0, num_samples - n)) - - return waveform.unsqueeze(0).to(dtype=torch.float32) # [1, C, N] -``` - -- [ ] **Step 4: Run to verify it passes** - -Run: `pytest cosmos_framework/inference/sound_test.py -v` -Expected: PASS (both tests). - -- [ ] **Step 5: Commit** - -```bash -git add cosmos_framework/inference/sound.py cosmos_framework/inference/sound_test.py -git commit -m "Add load_conditioning_audio helper" -``` - ---- - -## Task 5: `condition_sound` mode in `inject_sound_into_batch` - -**Files:** -- Modify: `cosmos_framework/inference/sound.py` (`inject_sound_into_batch`, line ~62) -- Test: `cosmos_framework/inference/sound_test.py` - -- [ ] **Step 1: Write the failing test** - -Append to `cosmos_framework/inference/sound_test.py`: - -```python -import types - -from cosmos_framework.data.vfm.sequence_packing import SequencePlan - - -def _fake_model(sound_latent_t: int, temporal_cf: int = 4): - sound_tok = types.SimpleNamespace( - get_latent_num_samples=lambda n: sound_latent_t, - audio_channels=2, - ) - vision_tok = types.SimpleNamespace(temporal_compression_factor=temporal_cf) - return types.SimpleNamespace(tokenizer_sound_gen=sound_tok, tokenizer_vision_gen=vision_tok) - - -def test_inject_sound_conditions_sound_and_preserves_image(tmp_path: Path): - from cosmos_framework.inference.sound import inject_sound_into_batch - - model = _fake_model(sound_latent_t=50) - # Video tensor [1,3,T,H,W] with T=48 -> 12 video latents at cf=4. - video = torch.zeros(1, 3, 48, 16, 16) - audio = torch.zeros(1, 2, 96000) - batch = { - "video": [video], - # Pre-existing image first-frame condition (as set by build_conditioned_video_batch). - "sequence_plan": [SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=[0])], - } - - inject_sound_into_batch(batch, audio, model, condition_sound=True) - - plan = batch["sequence_plan"][0] - assert plan.has_sound is True - assert plan.condition_frame_indexes_sound == list(range(50)) # all sound conditioned (ts2v) - assert plan.condition_frame_indexes_vision == [0] # image cond preserved - - -def test_inject_sound_default_generates_sound(tmp_path: Path): - from cosmos_framework.inference.sound import inject_sound_into_batch - - model = _fake_model(sound_latent_t=50) - video = torch.zeros(1, 3, 48, 16, 16) - audio = torch.zeros(1, 2, 96000) - batch = {"video": [video], "sequence_plan": [SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=[])]} - - inject_sound_into_batch(batch, audio, model) # default condition_sound=False - - plan = batch["sequence_plan"][0] - assert plan.condition_frame_indexes_sound == [] # t2vs: sound generated -``` - -- [ ] **Step 2: Run to verify it fails** - -Run: `pytest cosmos_framework/inference/sound_test.py::test_inject_sound_conditions_sound_and_preserves_image -v` -Expected: FAIL — `inject_sound_into_batch` got an unexpected keyword argument `condition_sound`. - -- [ ] **Step 3: Add the `condition_sound` parameter** - -In `cosmos_framework/inference/sound.py`, change the signature and the mode selection: - -```python -def inject_sound_into_batch( - data_batch: dict[str, Any], - audio_tensor: torch.Tensor | None, - model: Any, - *, - condition_sound: bool = False, -) -> dict[str, Any]: -``` - -Update the docstring's Args to add: - -``` - condition_sound: When True, the provided audio is used as a clean - condition (mode "ts2v") and the video is generated from it. When - False (default), sound is generated jointly (mode "t2vs"). -``` - -In the `if has_sound:` block, replace the hardcoded `mode="t2vs"`: - -```python - sequence_plan = build_sequence_plan_for_sound( - mode="ts2v" if condition_sound else "t2vs", - video_latent_length=video_latent_t, - sound_latent_length=sound_latent_t, - ) -``` - -(The existing `existing_vision_cond` preservation below it already re-applies the image's `condition_frame_indexes_vision`.) - -- [ ] **Step 4: Run to verify both pass** - -Run: `pytest cosmos_framework/inference/sound_test.py -v` -Expected: PASS (all four tests). - -- [ ] **Step 5: Commit** - -```bash -git add cosmos_framework/inference/sound.py cosmos_framework/inference/sound_test.py -git commit -m "Support sound conditioning (ts2v) in inject_sound_into_batch" -``` - ---- - -## Task 6: Wire real audio into `get_sample_data` - -**Files:** -- Modify: `cosmos_framework/inference/inference.py` (`get_sample_data`, line ~610) - -- [ ] **Step 1: Replace the sound-injection block** - -In `get_sample_data`, replace the existing `if sample_args.enable_sound:` block (lines ~610–625) with: - -```python - if sample_args.enable_sound: - from cosmos_framework.inference.sound import ( - create_placeholder_audio, - get_audio_tokenizer_info, - inject_sound_into_batch, - load_conditioning_audio, - ) - - audio_info = get_audio_tokenizer_info(model) - if not audio_info.has_sound: - raise ValueError("enable_sound=True but model has no sound tokenizer") - - condition_sound = sample_args.sound_path is not None - if condition_sound: - num_samples = int(sample_args.num_frames / sample_args.fps * audio_info.sample_rate) - audio = load_conditioning_audio( - Path(sample_args.sound_path), - sample_rate=audio_info.sample_rate, - audio_channels=getattr(audio_info.tokenizer, "audio_channels", 2), - num_samples=num_samples, - ) - else: - audio = create_placeholder_audio( - num_frames=sample_args.num_frames, - conditioning_fps=sample_args.fps, - audio_info=audio_info, - ) - inject_sound_into_batch(out, audio, model, condition_sound=condition_sound) -``` - -(`Path` is already imported in `inference.py`.) - -- [ ] **Step 2: Sanity-check the module imports** - -Run: `python -c "import cosmos_framework.inference.inference"` -Expected: no error (module imports cleanly). - -- [ ] **Step 3: Commit** - -```bash -git add cosmos_framework/inference/inference.py -git commit -m "Load and condition on real input audio in get_sample_data" -``` - ---- - -## Task 7: Defaults + example input + conditioning audio asset - -**Files:** -- Create: `cosmos_framework/inference/defaults/audio_image2video/sample_args.json` -- Create: `inputs/omni/assets/a2v_audio.wav` -- Create: `inputs/omni/a2v.json` - -- [ ] **Step 1: Create the per-mode defaults** - -Create `cosmos_framework/inference/defaults/audio_image2video/sample_args.json` (image2video defaults + `enable_sound: true`): - -```json -{ - "num_steps": 35, - "guidance": 6.0, - "shift": 10.0, - "sigma_max": 80.0, - "normalize_cfg": false, - "autoregressive": false, - "negative_prompt": null, - "negative_prompt_file": "neg_prompts.json", - "duration_template": "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS.", - "resolution_template": "This video is of {height}x{width} resolution.", - "negative_metadata_mode": "none", - "inverse_duration_template": "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS.", - "inverse_resolution_template": "This video is not of {height}x{width} resolution.", - "negative_prompt_keep_metadata": true, - "aspect_ratio": "16,9", - "fps": 24, - "num_frames": 189, - "video_save_quality": 10, - "image_save_quality": 95, - "enable_sound": true -} -``` - -- [ ] **Step 2: Extract a conditioning audio clip from a published example** - -Run inside the i4 container (via the `slurm-node` skill), from the repo root: - -```bash -mkdir -p inputs/omni/assets -curl -sS -H "Authorization: Bearer $HF_TOKEN" \ - "https://huggingface.co/nvidia/Cosmos3-Nano/resolve/main/assets/example_t2vs_output.mp4" \ - -o /tmp/example_t2vs_output.mp4 -python - <<'PY' -import av, numpy as np, soundfile as sf -container = av.open("/tmp/example_t2vs_output.mp4") -astream = container.streams.audio[0] -sr = astream.codec_context.sample_rate -frames = [f.to_ndarray() for f in container.decode(astream)] # each [C, n] fltp or [n*C] packed -# Normalize to [N, C] float32. -import numpy as np -chunks = [] -for arr in frames: - if arr.ndim == 2: # planar [C, n] - chunks.append(arr.T) - else: # packed - chunks.append(arr.reshape(-1, astream.channels)) -audio = np.concatenate(chunks, axis=0).astype("float32") -audio = audio[: sr * 4] # keep first 4 seconds to bound file size -sf.write("inputs/omni/assets/a2v_audio.wav", audio, sr) -print("wrote", audio.shape, "@", sr) -PY -``` - -Expected: prints the shape and writes `inputs/omni/assets/a2v_audio.wav` (~1–1.5 MB). - -- [ ] **Step 3: Create the example input file** - -Create `inputs/omni/a2v.json` (image from the existing i2v example; audio is the extracted clip, referenced relative to the input file): - -```json -{ - "model_mode": "audio_image2video", - "name": "a2v", - "prompt": "{\"temporal_caption\": \"A silver robotic arm in a clean lab pours water from a glass jar into a white ceramic cup; soft mechanical whirring and gentle water trickling are audible.\", \"audio_description\": \"Gentle splashing and trickling of water with a faint mechanical whir from servo motors; no speech or music.\", \"resolution\": {\"H\": 720, \"W\": 1280}, \"aspect_ratio\": \"16,9\", \"fps\": 24}", - "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/robot_153.jpg", - "sound_path": "assets/a2v_audio.wav" -} -``` - -- [ ] **Step 4: Validate the example parses into sample args** - -Run inside the dev/slurm env: - -```bash -python -c " -import json, pathlib -d = json.loads(pathlib.Path('inputs/omni/a2v.json').read_text()) -assert d['model_mode'] == 'audio_image2video' -assert d['sound_path'].endswith('.wav') -print('ok') -" -``` -Expected: `ok`. (Full arg-building is exercised by the slurm run in Task 9.) - -- [ ] **Step 5: Commit** - -```bash -git add cosmos_framework/inference/defaults/audio_image2video/sample_args.json inputs/omni/a2v.json inputs/omni/assets/a2v_audio.wav -git commit -m "Add audio_image2video defaults and a2v example input" -``` - ---- - -## Task 8: Documentation - -**Files:** -- Modify: `docs/inference.md` - -- [ ] **Step 1: Add the Modes table row** - -In the Modes table (after the `video2video` row), add: - -```markdown -| `audio_image2video` | text prompt + image + audio | `vision.mp4` (with conditioning audio muxed in) | `prompt`, `vision_path`, `sound_path` | [`inputs/omni/a2v.json`](../inputs/omni/a2v.json) | -``` - -- [ ] **Step 2: Note the checkpoint under the sound sentence** - -Replace the existing sentence (line ~148): - -```markdown -Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To run every example in one batch, use `-i "inputs/omni/*.json"`. -``` - -with: - -```markdown -Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To instead **condition** generation on a real audio clip (audio+image → video), use `model_mode: audio_image2video` with a `sound_path` and the sound-encoder checkpoint `--checkpoint-path Cosmos3-Nano-SoundEncoder` (see [`inputs/omni/a2v.json`](../inputs/omni/a2v.json)). To run every example in one batch, use `-i "inputs/omni/*.json"`. -``` - -- [ ] **Step 3: Add the checkpoint to the Models table** - -In the Models table, add a row: - -```markdown -| Cosmos3-Nano-SoundEncoder | `--checkpoint-path=Cosmos3-Nano-SoundEncoder` | `audio_image2video` (audio+image → video), plus all Nano modes | -``` - -- [ ] **Step 4: Commit** - -```bash -git add docs/inference.md -git commit -m "Document audio_image2video mode and sound-encoder checkpoint" -``` - ---- - -## Task 9: End-to-end A2V run on slurm (verification gate) - -**Files:** none (verification only) - -- [ ] **Step 1: Run A2V inference** via the `slurm-node` skill, from the repo root inside the i4 container: - -```bash -python -m cosmos_framework.scripts.inference \ - --parallelism-preset=latency \ - -i "inputs/omni/a2v.json" \ - -o outputs/a2v \ - --checkpoint-path Cosmos3-Nano-SoundEncoder \ - --seed=0 -``` - -Expected: completes without shape/dtype errors during sound encode/condition; writes `outputs/a2v/a2v/vision.mp4` and `sample_args.json`. - -- [ ] **Step 2: Verify the output** - -```bash -python - <<'PY' -import av -c = av.open("outputs/a2v/a2v/vision.mp4") -assert len(c.streams.video) == 1, "no video stream" -assert len(c.streams.audio) >= 1, "no audio track muxed" -print("video frames:", c.streams.video[0].frames, "| audio streams:", len(c.streams.audio)) -PY -``` -Expected: prints a positive video frame count and ≥1 audio stream. - -- [ ] **Step 3: Inspect `sample_args.json`** to confirm conditioning was applied: - -```bash -python -c " -import json -a = json.load(open('outputs/a2v/a2v/sample_args.json')) -assert a['model_mode'] == 'audio_image2video' -assert a['enable_sound'] is True -assert a['condition_frame_indexes_vision'] == [0] -print('conditioning ok') -" -``` -Expected: `conditioning ok`. - -- [ ] **Step 4: If sound tokens are not frozen during sampling** (risk #1 in the spec — e.g. output audio does not match the input clip), STOP and use `superpowers:systematic-debugging` to trace whether the diffusion loop respects `condition_frame_indexes_sound`. This is the one place a model-side fix might be needed; do not paper over it. - ---- - -## Self-Review - -**Spec coverage:** -- Checkpoint registration → Task 1. ✅ -- `audio_image2video` mode → Task 2. ✅ -- `sound_path` field + validation/download → Task 3. ✅ -- `load_conditioning_audio` → Task 4. ✅ -- `condition_sound` wiring (`ts2v` + preserved image cond) → Tasks 5–6. ✅ -- Defaults + example + audio asset → Task 7. ✅ -- Output decode+mux (no change) → confirmed in Task 9 verification. ✅ -- Tests (colocated) → Tasks 3, 4, 5. ✅ (`sound_data_utils` `ts2v` is exercised indirectly by Task 5's plan assertions; the spec's optional `sound_data_utils_test.py` is dropped as redundant — Task 5 already asserts the resulting `condition_frame_indexes_sound`.) -- Docs → Task 8. ✅ -- E2E slurm verification → Task 9. ✅ - -**Placeholder scan:** No TBD/TODO; every code step shows complete code. ✅ - -**Type consistency:** `load_conditioning_audio(path, *, sample_rate, audio_channels, num_samples) -> Tensor[1,C,N]` is defined in Task 4 and called identically in Task 6. `inject_sound_into_batch(..., *, condition_sound=False)` defined in Task 5, called with `condition_sound=condition_sound` in Task 6. `ModelMode.AUDIO_IMAGE2VIDEO` / `is_sound_condition` / `SOUND_CONDITION_MODEL_MODES` consistent across Tasks 2–3. ✅ - -**Note on `sound_data_utils_test.py`:** The spec listed it as "if absent." Dropped to avoid duplicate coverage; Task 5 asserts the `ts2v` plan's effect end-to-end through `inject_sound_into_batch`. If a direct unit test is preferred, add one asserting `build_sequence_plan_for_sound("ts2v", v, s).condition_frame_indexes_sound == list(range(s))`. From f3bfb2000317267d7df432dbf643f15c043e6998 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 02:31:51 -0700 Subject: [PATCH 16/25] Remove A2V example input (a2v.json + conditioning audio asset) Co-Authored-By: Claude Opus 4.8 (1M context) --- inputs/omni/a2v.json | 7 ------- inputs/omni/assets/a2v_audio.wav | 3 --- 2 files changed, 10 deletions(-) delete mode 100644 inputs/omni/a2v.json delete mode 100644 inputs/omni/assets/a2v_audio.wav diff --git a/inputs/omni/a2v.json b/inputs/omni/a2v.json deleted file mode 100644 index d0da628c..00000000 --- a/inputs/omni/a2v.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "model_mode": "audio_image2video", - "name": "a2v", - "prompt": "{\"subjects\": [{\"description\": \"A claw hammer with a wooden handle and a polished steel head, positioned above a nail embedded in a wooden plank\", \"appearance_details\": \"The hammer has a warm-toned, slightly worn wooden handle with visible wood grain and a smooth, silvery steel head showing minor surface wear and patina\", \"relationship\": \"The hammer is directly above and aligned with the nail, poised to strike it\", \"location\": \"Upper center of frame\", \"relative_size\": \"Large within frame\", \"orientation\": \"Horizontal, head facing downward toward the nail\", \"pose\": \"Suspended above the nail in striking position\", \"action\": \"Repeatedly striking the nail into the wood\", \"state_changes\": \"The hammer rises and falls in a rhythmic motion, making contact with the nail head on each downstroke\", \"clothing\": \"\", \"expression\": \"\", \"gender\": \"\", \"age\": \"\", \"skin_tone_and_texture\": \"\", \"facial_features\": \"\", \"number_of_subjects\": 1, \"number_of_arms\": 0, \"number_of_legs\": 0}, {\"description\": \"A small steel nail partially driven into a wooden plank, with a flat round head and a thin shank showing ring-shank grooves\", \"appearance_details\": \"Silver-colored steel nail with a flat circular head, approximately 2 inches of the shank still visible above the wood surface, with fine ring grooves along the shaft\", \"relationship\": \"The nail is the target of the hammer's strikes, embedded vertically in the wooden surface below\", \"location\": \"Lower center of frame\", \"relative_size\": \"Small within frame\", \"orientation\": \"Vertical, pointing straight up\", \"pose\": \"Standing upright, partially embedded in wood\", \"action\": \"Being driven deeper into the wood with each hammer strike\", \"state_changes\": \"The nail progressively sinks deeper into the wood plank with each successive hammer blow\", \"clothing\": \"\", \"expression\": \"\", \"gender\": \"\", \"age\": \"\", \"skin_tone_and_texture\": \"\", \"facial_features\": \"\", \"number_of_subjects\": 1, \"number_of_arms\": 0, \"number_of_legs\": 0}], \"background_setting\": \"A woodworking workshop or garage workspace. The background is softly blurred, revealing what appears to be other tools, wooden materials, and a window letting in natural light from the right side. The work surface is a thick plank of light-colored natural wood with visible grain patterns.\", \"lighting\": {\"conditions\": \"Warm natural daylight mixed with ambient indoor lighting, creating a soft, diffused illumination\", \"direction\": \"Side-lit from the right, with secondary ambient light from the left\", \"shadows\": \"Soft shadows cast by the hammer onto the wood surface, with gentle shadow beneath the nail; the hammer head casts a subtle shadow on the plank below\", \"illumination_effect\": \"Warm, golden tone that enhances the natural wood textures and gives the metal surfaces a gentle gleam\"}, \"aesthetics\": {\"composition\": \"Centered composition with the hammer and nail aligned vertically in the middle of the frame, creating a strong focal point with the action happening at the intersection\", \"color_scheme\": \"Warm earth tones dominated by honey-brown wood, silver-gray steel, with soft muted background tones of beige and blue-gray\", \"mood_atmosphere\": \"Industrious, focused, craftsmanship, satisfying manual labor\", \"patterns\": \"Horizontal wood grain lines across the plank surface\"}, \"cinematography\": {\"camera_motion\": \"Static with very subtle vibration from hammer impacts\", \"framing\": \"Close-up shot focused tightly on the hammer head and nail\", \"camera_angle\": \"Eye-level, slightly below the hammer looking straight at the point of contact\", \"depth_of_field\": \"Shallow\", \"focus\": \"Sharp focus on the nail and hammer head, with the background and edges of the wood plank softly blurred\", \"lens_focal_length\": \"Macro or short telephoto, approximately 85-100mm equivalent\"}, \"style_medium\": \"Live-action video\", \"artistic_style\": \"Cinematic close-up, documentary-style craftsmanship footage\", \"context\": \"A woodworking or carpentry scene capturing the satisfying act of hammering a nail into wood, suitable for craftsmanship content, ASMR, or instructional material\", \"actions\": [{\"time\": \"0:00-0:01\", \"description\": \"The hammer is poised above the nail, then lifts slightly in preparation for the first strike\"}, {\"time\": \"0:01-0:02\", \"description\": \"The hammer swings down and strikes the nail with a sharp metallic tap, the nail sinks slightly into the wood\"}, {\"time\": \"0:02-0:03\", \"description\": \"The hammer rebounds upward and swings down again for a second strike, driving the nail further\"}, {\"time\": \"0:03-0:04\", \"description\": \"A third strike lands on the nail head, producing another sharp tap; the nail visibly shorter above the wood surface\"}, {\"time\": \"0:04-0:05\", \"description\": \"The hammer rises and delivers a fourth firm blow, the nail now about halfway driven in\"}, {\"time\": \"0:05-0:06\", \"description\": \"A fifth strike drives the nail deeper, with small wood fibers slightly displaced around the entry point\"}, {\"time\": \"0:06-0:07\", \"description\": \"A final strong blow drives the nail nearly flush with the wood surface, the hammer rests momentarily above\"}], \"text_and_signage_elements\": [], \"segments\": [{\"segment_index\": 0, \"time_range\": \"0:00-0:03\", \"description\": \"The hammer begins its rhythmic striking pattern, delivering the first three blows to the nail with measured force\", \"key_changes\": \"The nail progressively sinks from its initial position into the wood; each strike produces a visible impact moment\", \"camera\": \"Static close-up, maintaining focus on the point of contact between hammer and nail\"}, {\"segment_index\": 1, \"time_range\": \"0:03-0:07\", \"description\": \"The hammer continues with stronger, more confident strikes as the nail is driven deeper into the plank\", \"key_changes\": \"The nail becomes significantly shorter above the wood surface; the rhythm of strikes may slightly intensify; small wood displacement visible around the nail\", \"camera\": \"Static close-up, same framing maintained throughout with slight vibration on each impact\"}], \"transitions\": [], \"temporal_caption\": \"The video opens with a close-up of a steel hammer poised just above a nail standing upright in a wooden plank. Within the first second, the hammer lifts slightly and then swings down to deliver the first sharp strike against the nail head. The metallic tap resonates as the nail sinks a small amount into the wood. The hammer rebounds and continues in a steady rhythm, striking the nail approximately once per second. With each successive blow, the nail is driven progressively deeper into the honey-colored wood. By the midpoint of the video, the nail is roughly halfway embedded. The strikes continue with consistent force, and by the final second, the nail is nearly flush with the wood surface. The camera remains static throughout, capturing the satisfying repetitive motion in intimate close-up detail.\", \"audio_description\": \"Sharp, crisp metallic tapping sounds as the steel hammer head contacts the nail with each strike. Each tap is followed by a brief thud as the vibration transfers into the plank. Between strikes, there is a soft whoosh of the hammer moving through air. Ambient workshop sounds are faintly present in the background \\u2014 perhaps a distant hum or subtle room tone. No music or speech is present.\", \"resolution\": {\"H\": 720, \"W\": 1280}, \"aspect_ratio\": \"16,9\", \"duration\": \"7s\", \"fps\": 24}", - "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/01_hammer_nail.png", - "sound_path": "assets/a2v_audio.wav" -} diff --git a/inputs/omni/assets/a2v_audio.wav b/inputs/omni/assets/a2v_audio.wav deleted file mode 100644 index e1efd066..00000000 --- a/inputs/omni/assets/a2v_audio.wav +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:61317f4ce3b3e1e2e23109702d8e175f8742f4b9e818bd8a55d3cdb2789861be -size 1515564 From 3165a9909e9aac9043989317ac52e1b8ceea0966 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 02:32:59 -0700 Subject: [PATCH 17/25] Drop dead a2v.json links from inference docs, keep audio_image2video documented Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/inference.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/inference.md b/docs/inference.md index 23c41b4d..2e11f7be 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -141,12 +141,12 @@ The four `--{dp,cp,cfgp}-*-size` flags override the auto-selected values from `- | `text2video` | text prompt | `vision.mp4` | `prompt` | [`inputs/omni/t2v.json`](../inputs/omni/t2v.json) | | `image2video` | text prompt + image | `vision.mp4` | `prompt`, `vision_path` | [`inputs/omni/i2v.json`](../inputs/omni/i2v.json) | | `video2video` | text prompt + video | `vision.mp4` | `prompt`, `vision_path` | [`inputs/omni/v2v.json`](../inputs/omni/v2v.json) | -| `audio_image2video` | text prompt + image + audio | `vision.mp4` (with the conditioning audio muxed in) | `prompt`, `vision_path`, `sound_path` | [`inputs/omni/a2v.json`](../inputs/omni/a2v.json) | +| `audio_image2video` | text prompt + image + audio | `vision.mp4` (with the conditioning audio muxed in) | `prompt`, `vision_path`, `sound_path` | — | | `forward_dynamics` | observation image/video + prompt + actions | future visual rollout in `vision.mp4` | `domain_name`, `vision_path`, `action_path` | [`inputs/omni/action_forward_dynamics_av.json`](../inputs/omni/action_forward_dynamics_av.json), [`inputs/omni/action_forward_dynamics_camera.json`](../inputs/omni/action_forward_dynamics_camera.json), [`inputs/omni/action_forward_dynamics_robot.json`](../inputs/omni/action_forward_dynamics_robot.json), [`inputs/omni/action_forward_dynamics_batch.jsonl`](../inputs/omni/action_forward_dynamics_batch.jsonl) | | `inverse_dynamics` | observation video + prompt | predicted action sequence in `sample_outputs.json` | `domain_name`, `vision_path` | [`inputs/omni/action_inverse_dynamics_av.json`](../inputs/omni/action_inverse_dynamics_av.json), [`inputs/omni/action_inverse_dynamics_robot.json`](../inputs/omni/action_inverse_dynamics_robot.json), [`inputs/omni/action_inverse_dynamics_batch.jsonl`](../inputs/omni/action_inverse_dynamics_batch.jsonl) | | `policy` | observation image/video + prompt | predicted action sequence in `sample_outputs.json` + future visual rollout in `vision.mp4` | `domain_name`, `vision_path` | [`inputs/omni/action_policy_av.json`](../inputs/omni/action_policy_av.json), [`inputs/omni/action_policy_robot.json`](../inputs/omni/action_policy_robot.json), [`inputs/omni/action_policy_batch.jsonl`](../inputs/omni/action_policy_batch.jsonl) | -Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To instead **condition** generation on a real audio clip (audio+image → video), use `model_mode: audio_image2video` with a `sound_path` (see [`inputs/omni/a2v.json`](../inputs/omni/a2v.json)). To run every example in one batch, use `-i "inputs/omni/*.json"`. +Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To instead **condition** generation on a real audio clip (audio+image → video), use `model_mode: audio_image2video` with an image `vision_path` and a `sound_path` (a `.wav`/`.mp3`/`.flac` clip). To run every example in one batch, use `-i "inputs/omni/*.json"`. ## Parallelism Arguments From a0f1f6665e17a06ebe286f75b1cca8ca2c8681ed Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 02:39:52 -0700 Subject: [PATCH 18/25] Clarify AVAE comments: sound_tokenizer is decoder-only until updated The native encoder.layers.* keys load when present (enabling audio_image2video); the current Cosmos3-Nano sound_tokenizer ships only the decoder, so A2V produces faithful audio only once the checkpoint is updated to include the encoder. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../inference/common/checkpoints.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/cosmos_framework/inference/common/checkpoints.py b/cosmos_framework/inference/common/checkpoints.py index 86877ce2..caf6ca21 100644 --- a/cosmos_framework/inference/common/checkpoints.py +++ b/cosmos_framework/inference/common/checkpoints.py @@ -76,10 +76,11 @@ def _materialize_avae_ckpt(local_dir: str) -> None: ``[C]`` and loads via ``load_state_dict(strict=False)`` — so without remapping the keys, none match and every decoder weight is silently left at init (noise). We invert the forward conversion (key remap + snake reshape) and wrap the result - under ``state_dict``. The encoder ships in the already-native ``encoder.layers.*`` - layout, so it passes through ``_avae_block_key_to_legacy`` unchanged and loads - directly (required for encoding input audio in sound-conditioned modes). - Idempotent. + under ``state_dict``. The decoder alone is sufficient for sound generation. If the + checkpoint also ships the SpecConvNeXt encoder (already in the native + ``encoder.layers.*`` layout), those keys pass through ``_avae_block_key_to_legacy`` + unchanged and load directly — required for encoding input audio in sound-conditioned + modes (audio_image2video). Idempotent. """ import torch from safetensors.torch import load_file @@ -251,13 +252,14 @@ def register_checkpoints(): revision="main", subdirectory="sound_tokenizer", ), - # The sound_tokenizer/ safetensors ship the full AVAE: the diffusers - # OobleckDecoder decoder (decoder.block.*) plus the native SpecConvNeXt - # encoder (encoder.layers.*). _materialize_avae_ckpt remaps the decoder keys - # back to the legacy decoder.layers.* layout and passes the already-native - # encoder keys through unchanged. The encoder is required to ENCODE input - # audio for sound-conditioned modes (audio_image2video); the decoder alone - # only supports sound generation. + # The sound_tokenizer/ safetensors use the diffusers OobleckDecoder layout + # (decoder.block.*); _materialize_avae_ckpt remaps them back to the legacy + # decoder.layers.* layout the native AVAE loader expects. Sound *generation* + # (t2vs/i2vs) needs only the decoder. Sound *conditioning* (audio_image2video) + # additionally needs the SpecConvNeXt encoder (native encoder.layers.* keys), + # which pass through unchanged and load when present. NOTE: the current + # Cosmos3-Nano sound_tokenizer is decoder-only; audio_image2video produces + # faithful audio only once this checkpoint is updated to also ship the encoder. post_download=_materialize_avae_ckpt, ), ) From b62e992dd59cb13751e7d1e17846d8170dcbf288 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 02:45:30 -0700 Subject: [PATCH 19/25] Simplify AVAE materialize comments Co-Authored-By: Claude Opus 4.8 (1M context) --- .../inference/common/checkpoints.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/cosmos_framework/inference/common/checkpoints.py b/cosmos_framework/inference/common/checkpoints.py index caf6ca21..0a10aae3 100644 --- a/cosmos_framework/inference/common/checkpoints.py +++ b/cosmos_framework/inference/common/checkpoints.py @@ -76,11 +76,8 @@ def _materialize_avae_ckpt(local_dir: str) -> None: ``[C]`` and loads via ``load_state_dict(strict=False)`` — so without remapping the keys, none match and every decoder weight is silently left at init (noise). We invert the forward conversion (key remap + snake reshape) and wrap the result - under ``state_dict``. The decoder alone is sufficient for sound generation. If the - checkpoint also ships the SpecConvNeXt encoder (already in the native - ``encoder.layers.*`` layout), those keys pass through ``_avae_block_key_to_legacy`` - unchanged and load directly — required for encoding input audio in sound-conditioned - modes (audio_image2video). Idempotent. + under ``state_dict``. Native ``encoder.layers.*`` keys pass through + ``_avae_block_key_to_legacy`` unchanged. Idempotent. """ import torch from safetensors.torch import load_file @@ -252,14 +249,9 @@ def register_checkpoints(): revision="main", subdirectory="sound_tokenizer", ), - # The sound_tokenizer/ safetensors use the diffusers OobleckDecoder layout - # (decoder.block.*); _materialize_avae_ckpt remaps them back to the legacy - # decoder.layers.* layout the native AVAE loader expects. Sound *generation* - # (t2vs/i2vs) needs only the decoder. Sound *conditioning* (audio_image2video) - # additionally needs the SpecConvNeXt encoder (native encoder.layers.* keys), - # which pass through unchanged and load when present. NOTE: the current - # Cosmos3-Nano sound_tokenizer is decoder-only; audio_image2video produces - # faithful audio only once this checkpoint is updated to also ship the encoder. + # _materialize_avae_ckpt remaps the diffusers OobleckDecoder keys + # (decoder.block.*) back to the legacy decoder.layers.* layout the native AVAE + # loader expects; native encoder.layers.* keys pass through unchanged. post_download=_materialize_avae_ckpt, ), ) From 731d988626f2308edb159d306b328b11717368ea Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 02:53:15 -0700 Subject: [PATCH 20/25] Source AVAE sound_tokenizer from the loaded checkpoint's bundled sound_tokenizer/ By default, use the sound_tokenizer/ co-located in the main model checkpoint (matched to the transformer; uses whatever encoder/decoder that checkpoint ships) instead of the global "AVAE" registry repo. Mirrors vlm_processor_from_checkpoint: after download_checkpoint(), materialize the legacy AVAE .ckpt from the bundled sound_tokenizer/ and point avae_path at it (download_checkpoint_v2 short-circuits local paths). Falls back to the registry when the checkpoint bundles no sound_tokenizer/. So Cosmos3-Super uses Super's sound_tokenizer, Nano uses Nano's. Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/inference.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index abc99311..55af12ea 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -1074,6 +1074,23 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: tokenizer_cfg.pop("revision", None) tokenizer_cfg.pop("subdir", None) tokenizer_cfg["tokenizer_type"] = str(checkpoint_path) + # Source the AVAE sound tokenizer from the loaded checkpoint's own + # bundled sound_tokenizer/ (matched to the transformer; uses whatever + # encoder/decoder that checkpoint ships) instead of the global "AVAE" + # registry repo. Mirrors vlm_processor_from_checkpoint above; falls back + # to the registry when the checkpoint bundles no sound_tokenizer/. + sound_cfg = model_dict["config"].get("sound_tokenizer") + if sound_cfg: + from cosmos_framework.inference.common.checkpoints import ( + _AVAE_LEGACY_CKPT_NAME, + _materialize_avae_ckpt, + ) + + sound_tokenizer_dir = Path(checkpoint_path) / "sound_tokenizer" + if sound_tokenizer_dir.is_dir(): + _materialize_avae_ckpt(str(sound_tokenizer_dir)) + sound_cfg["bucket_name"] = "" + sound_cfg["avae_path"] = str(sound_tokenizer_dir / _AVAE_LEGACY_CKPT_NAME) config = Cosmos3OmniConfig(model=model_dict) model = Cosmos3OmniModel.from_pretrained_dcp( checkpoint_path, From 18ebdcd7b86584081c28071c578d2864ffed80ce Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 03:07:39 -0700 Subject: [PATCH 21/25] Add sound_tokenizer 'from_checkpoint' control (inference-only) The sound_tokenizer node accepts an inference-only 'from_checkpoint' key (default True): True sources the AVAE from the loaded checkpoint's bundled sound_tokenizer/; False keeps the configured avae_path (the registered "AVAE" repo) even when the checkpoint bundles one. The key is popped before AVAEInterface instantiation. Set it in the inference model YAML to override. Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/inference.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index 55af12ea..c7bcaa9d 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -1077,17 +1077,22 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: # Source the AVAE sound tokenizer from the loaded checkpoint's own # bundled sound_tokenizer/ (matched to the transformer; uses whatever # encoder/decoder that checkpoint ships) instead of the global "AVAE" - # registry repo. Mirrors vlm_processor_from_checkpoint above; falls back - # to the registry when the checkpoint bundles no sound_tokenizer/. + # registry repo. sound_cfg = model_dict["config"].get("sound_tokenizer") - if sound_cfg: - from cosmos_framework.inference.common.checkpoints import ( - _AVAE_LEGACY_CKPT_NAME, - _materialize_avae_ckpt, - ) - + if sound_cfg is not None: + # ``from_checkpoint`` is an inference-only routing key on the + # sound_tokenizer node; pop it so it never reaches AVAEInterface. + # True (default): use the checkpoint's bundled sound_tokenizer/. + # False: keep the configured avae_path (the registered "AVAE" repo), + # even when the checkpoint bundles a sound_tokenizer/. + from_checkpoint = sound_cfg.pop("from_checkpoint", True) sound_tokenizer_dir = Path(checkpoint_path) / "sound_tokenizer" - if sound_tokenizer_dir.is_dir(): + if from_checkpoint and sound_tokenizer_dir.is_dir(): + from cosmos_framework.inference.common.checkpoints import ( + _AVAE_LEGACY_CKPT_NAME, + _materialize_avae_ckpt, + ) + _materialize_avae_ckpt(str(sound_tokenizer_dir)) sound_cfg["bucket_name"] = "" sound_cfg["avae_path"] = str(sound_tokenizer_dir / _AVAE_LEGACY_CKPT_NAME) From 15898e6549fc6a1f7e211337ddcbafa556f55d87 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 03:20:26 -0700 Subject: [PATCH 22/25] Make AVAE sound_tokenizer registered-first, bundled fallback Default precedence: use the registered "AVAE" repo (configured avae_path) when one is registered; fall back to the loaded checkpoint's bundled sound_tokenizer/ only when no AVAE is registered. The inference-only sound_tokenizer.from_checkpoint key (default False) forces the bundled one even when an AVAE is registered. Backward-compatible: with the "AVAE" entry registered (default), all checkpoints use the registry AVAE exactly as before; bundled is opt-in (from_checkpoint:true) or used when the registration is removed. Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/inference.py | 27 +++++++++++++++---------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index c7bcaa9d..5822f177 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -1074,20 +1074,25 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: tokenizer_cfg.pop("revision", None) tokenizer_cfg.pop("subdir", None) tokenizer_cfg["tokenizer_type"] = str(checkpoint_path) - # Source the AVAE sound tokenizer from the loaded checkpoint's own - # bundled sound_tokenizer/ (matched to the transformer; uses whatever - # encoder/decoder that checkpoint ships) instead of the global "AVAE" - # registry repo. + # AVAE sound-tokenizer source precedence: prefer the registered "AVAE" + # repo (the configured ``avae_path``); fall back to the loaded + # checkpoint's own bundled ``sound_tokenizer/`` only when no AVAE is + # registered. The inference-only ``from_checkpoint`` key (default False) + # forces the bundled one even when an AVAE is registered; it is popped so + # it never reaches AVAEInterface. sound_cfg = model_dict["config"].get("sound_tokenizer") if sound_cfg is not None: - # ``from_checkpoint`` is an inference-only routing key on the - # sound_tokenizer node; pop it so it never reaches AVAEInterface. - # True (default): use the checkpoint's bundled sound_tokenizer/. - # False: keep the configured avae_path (the registered "AVAE" repo), - # even when the checkpoint bundles a sound_tokenizer/. - from_checkpoint = sound_cfg.pop("from_checkpoint", True) + from cosmos_framework.inference.common.checkpoints import register_checkpoints + from cosmos_framework.utils.checkpoint_db import CheckpointConfig, sanitize_uri + + register_checkpoints() + from_checkpoint = sound_cfg.pop("from_checkpoint", False) + bucket = sound_cfg.get("bucket_name") or "" + avae_path = sound_cfg.get("avae_path") or "" + avae_dir = f"s3://{bucket}/{Path(avae_path).parent}" if (bucket and avae_path) else avae_path + avae_registered = bool(avae_path) and CheckpointConfig.maybe_from_uri(sanitize_uri(avae_dir)) is not None sound_tokenizer_dir = Path(checkpoint_path) / "sound_tokenizer" - if from_checkpoint and sound_tokenizer_dir.is_dir(): + if sound_tokenizer_dir.is_dir() and (from_checkpoint or not avae_registered): from cosmos_framework.inference.common.checkpoints import ( _AVAE_LEGACY_CKPT_NAME, _materialize_avae_ckpt, From 34a7f09239abb826fbc1958def4b6a5dfcd90345 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 03:26:28 -0700 Subject: [PATCH 23/25] Simplify AVAE source check to avae_path presence Use the configured avae_path when set (registered AVAE); fall back to the checkpoint's bundled sound_tokenizer/ only when avae_path is empty, or when from_checkpoint:true forces it. Drops the s3-uri/registry lookup. Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/inference.py | 22 +++++++--------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index 5822f177..3f6b4fe0 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -1074,25 +1074,17 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: tokenizer_cfg.pop("revision", None) tokenizer_cfg.pop("subdir", None) tokenizer_cfg["tokenizer_type"] = str(checkpoint_path) - # AVAE sound-tokenizer source precedence: prefer the registered "AVAE" - # repo (the configured ``avae_path``); fall back to the loaded - # checkpoint's own bundled ``sound_tokenizer/`` only when no AVAE is - # registered. The inference-only ``from_checkpoint`` key (default False) - # forces the bundled one even when an AVAE is registered; it is popped so - # it never reaches AVAEInterface. + # AVAE sound-tokenizer source: use the configured ``avae_path`` (the + # registered "AVAE" repo) when one is set; fall back to the loaded + # checkpoint's own bundled ``sound_tokenizer/`` only when no avae_path is + # configured. The inference-only ``from_checkpoint`` key (default False) + # forces the bundled one even when avae_path is set; it is popped so it + # never reaches AVAEInterface. sound_cfg = model_dict["config"].get("sound_tokenizer") if sound_cfg is not None: - from cosmos_framework.inference.common.checkpoints import register_checkpoints - from cosmos_framework.utils.checkpoint_db import CheckpointConfig, sanitize_uri - - register_checkpoints() from_checkpoint = sound_cfg.pop("from_checkpoint", False) - bucket = sound_cfg.get("bucket_name") or "" - avae_path = sound_cfg.get("avae_path") or "" - avae_dir = f"s3://{bucket}/{Path(avae_path).parent}" if (bucket and avae_path) else avae_path - avae_registered = bool(avae_path) and CheckpointConfig.maybe_from_uri(sanitize_uri(avae_dir)) is not None sound_tokenizer_dir = Path(checkpoint_path) / "sound_tokenizer" - if sound_tokenizer_dir.is_dir() and (from_checkpoint or not avae_registered): + if sound_tokenizer_dir.is_dir() and (from_checkpoint or not sound_cfg.get("avae_path")): from cosmos_framework.inference.common.checkpoints import ( _AVAE_LEGACY_CKPT_NAME, _materialize_avae_ckpt, From 634c83e07671af6dd71b1a0a2912bf7077e2a9a1 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 03:29:52 -0700 Subject: [PATCH 24/25] Simplify comments in A2V sound inference Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/inference.py | 8 +++----- cosmos_framework/inference/sound.py | 4 +--- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index 3f6b4fe0..564da97f 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -1074,11 +1074,9 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: tokenizer_cfg.pop("revision", None) tokenizer_cfg.pop("subdir", None) tokenizer_cfg["tokenizer_type"] = str(checkpoint_path) - # AVAE sound-tokenizer source: use the configured ``avae_path`` (the - # registered "AVAE" repo) when one is set; fall back to the loaded - # checkpoint's own bundled ``sound_tokenizer/`` only when no avae_path is - # configured. The inference-only ``from_checkpoint`` key (default False) - # forces the bundled one even when avae_path is set; it is popped so it + # AVAE source: the configured ``avae_path`` when set, else the loaded + # checkpoint's bundled ``sound_tokenizer/``. The inference-only + # ``from_checkpoint`` key (default False) forces bundled; pop it so it # never reaches AVAEInterface. sound_cfg = model_dict["config"].get("sound_tokenizer") if sound_cfg is not None: diff --git a/cosmos_framework/inference/sound.py b/cosmos_framework/inference/sound.py index 261bea41..59d008e8 100644 --- a/cosmos_framework/inference/sound.py +++ b/cosmos_framework/inference/sound.py @@ -82,9 +82,7 @@ def load_conditioning_audio( data, src_sr = sf.read(str(path), dtype="float32", always_2d=True) # [N, C] waveform = torch.from_numpy(data).transpose(0, 1).contiguous() # [C, N] - # Resample to the tokenizer's rate. Uses scipy (a declared dependency); - # torchaudio is intentionally avoided as it is not a project dependency - # and is absent from the inference container. + # Resample with scipy (torchaudio is not a project dependency). if src_sr != sample_rate: from math import gcd From 3fb1fbee53827dfe7fe6881652a027d923068261 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Fri, 12 Jun 2026 03:44:22 -0700 Subject: [PATCH 25/25] Revert docs/inference.md changes on this branch Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/inference.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/inference.md b/docs/inference.md index 2e11f7be..15801476 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -141,12 +141,11 @@ The four `--{dp,cp,cfgp}-*-size` flags override the auto-selected values from `- | `text2video` | text prompt | `vision.mp4` | `prompt` | [`inputs/omni/t2v.json`](../inputs/omni/t2v.json) | | `image2video` | text prompt + image | `vision.mp4` | `prompt`, `vision_path` | [`inputs/omni/i2v.json`](../inputs/omni/i2v.json) | | `video2video` | text prompt + video | `vision.mp4` | `prompt`, `vision_path` | [`inputs/omni/v2v.json`](../inputs/omni/v2v.json) | -| `audio_image2video` | text prompt + image + audio | `vision.mp4` (with the conditioning audio muxed in) | `prompt`, `vision_path`, `sound_path` | — | | `forward_dynamics` | observation image/video + prompt + actions | future visual rollout in `vision.mp4` | `domain_name`, `vision_path`, `action_path` | [`inputs/omni/action_forward_dynamics_av.json`](../inputs/omni/action_forward_dynamics_av.json), [`inputs/omni/action_forward_dynamics_camera.json`](../inputs/omni/action_forward_dynamics_camera.json), [`inputs/omni/action_forward_dynamics_robot.json`](../inputs/omni/action_forward_dynamics_robot.json), [`inputs/omni/action_forward_dynamics_batch.jsonl`](../inputs/omni/action_forward_dynamics_batch.jsonl) | | `inverse_dynamics` | observation video + prompt | predicted action sequence in `sample_outputs.json` | `domain_name`, `vision_path` | [`inputs/omni/action_inverse_dynamics_av.json`](../inputs/omni/action_inverse_dynamics_av.json), [`inputs/omni/action_inverse_dynamics_robot.json`](../inputs/omni/action_inverse_dynamics_robot.json), [`inputs/omni/action_inverse_dynamics_batch.jsonl`](../inputs/omni/action_inverse_dynamics_batch.jsonl) | | `policy` | observation image/video + prompt | predicted action sequence in `sample_outputs.json` + future visual rollout in `vision.mp4` | `domain_name`, `vision_path` | [`inputs/omni/action_policy_av.json`](../inputs/omni/action_policy_av.json), [`inputs/omni/action_policy_robot.json`](../inputs/omni/action_policy_robot.json), [`inputs/omni/action_policy_batch.jsonl`](../inputs/omni/action_policy_batch.jsonl) | -Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To instead **condition** generation on a real audio clip (audio+image → video), use `model_mode: audio_image2video` with an image `vision_path` and a `sound_path` (a `.wav`/`.mp3`/`.flac` clip). To run every example in one batch, use `-i "inputs/omni/*.json"`. +Set `enable_sound: true` on a `text2video` sample (see [`inputs/omni/t2vs.json`](../inputs/omni/t2vs.json)) to also generate audio. To run every example in one batch, use `-i "inputs/omni/*.json"`. ## Parallelism Arguments