A2v sound encoder inference

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,15 +522,31 @@ 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):
     """Sound data overrides."""
 
     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:

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"

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,
         ),
     )

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
+}

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,

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,
         )