diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 3c234f74..62c3c899 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" @@ -513,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): @@ -520,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: 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: diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index 3bf37032..631f87ef 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,53 @@ 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="https://example.com/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="https://example.com/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) + + 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=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 Path(args.sound_path).name == "clip.wav" diff --git a/cosmos_framework/inference/common/checkpoints.py b/cosmos_framework/inference/common/checkpoints.py index 4041cb93..0a10aae3 100644 --- a/cosmos_framework/inference/common/checkpoints.py +++ b/cosmos_framework/inference/common/checkpoints.py @@ -76,8 +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``. Decoder-only is sufficient: generation only decodes sound - latents to a waveform. 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 @@ -249,9 +249,9 @@ def register_checkpoints(): revision="main", subdirectory="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. + # _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, ), ) 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 +} diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index 83ff6655..564da97f 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 @@ -1062,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) + # 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: + from_checkpoint = sound_cfg.pop("from_checkpoint", False) + sound_tokenizer_dir = Path(checkpoint_path) / "sound_tokenizer" + 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, + ) + + _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, diff --git a/cosmos_framework/inference/sound.py b/cosmos_framework/inference/sound.py index 1ecf3fbe..59d008e8 100644 --- a/cosmos_framework/inference/sound.py +++ b/cosmos_framework/inference/sound.py @@ -59,10 +59,68 @@ 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 with scipy (torchaudio is not a project dependency). + 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, model: Any, + *, + condition_sound: bool = False, ) -> dict[str, Any]: """Add sound data and upgrade the SequencePlan in an existing data batch. @@ -73,6 +131,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. @@ -103,7 +164,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 new file mode 100644 index 00000000..06664294 --- /dev/null +++ b/cosmos_framework/inference/sound_test.py @@ -0,0 +1,87 @@ +# 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 + + +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