vllm-omni/vllm_omni/diffusion/models/magi2/audio_decoder.py

# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2026 SandAI. All Rights Reserved.
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Native Stable Audio Open decoder used by MAGI-2.

The MAGI checkpoint embeds Stable Audio's Oobleck VAE under the
``pretransform.model`` prefix and uses the sequential module names from
``stable-audio-tools``. This module converts those names to Diffusers'
``AutoencoderOobleck`` layout and loads only the decoder tensors.
"""

from __future__ import annotations

import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
from diffusers import AutoencoderOobleck
from safetensors import safe_open

_CHECKPOINT_PREFIX = "pretransform.model.decoder."
_BLOCK_COMPONENTS = {
    0: "snake1",
    1: "conv_t1",
}
_RESIDUAL_COMPONENTS = {
    0: "snake1",
    1: "conv1",
    2: "snake2",
    3: "conv2",
}


def convert_stable_audio_decoder_key(key: str) -> str | None:
    """Map one Stable Audio decoder key to Diffusers' Oobleck decoder.

    Unrelated checkpoint tensors return ``None``. A key inside the decoder
    namespace that does not match the released Stable Audio layout raises so a
    changed checkpoint cannot be loaded partially by accident.
    """

    if not key.startswith(_CHECKPOINT_PREFIX):
        return None
    suffix = key[len(_CHECKPOINT_PREFIX) :]
    parts = suffix.split(".")
    if len(parts) < 3 or parts[0] != "layers":
        raise ValueError(f"Unsupported Stable Audio decoder key: {key}")

    try:
        layer = int(parts[1])
    except ValueError as exc:
        raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
    tail = parts[2:]

    if layer == 0 and len(tail) == 1:
        return f"conv1.{tail[0]}"
    if 1 <= layer <= 5:
        block = layer - 1
        if len(tail) == 3 and tail[0] == "layers":
            try:
                component = _BLOCK_COMPONENTS[int(tail[1])]
            except (KeyError, ValueError) as exc:
                raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
            return f"block.{block}.{component}.{tail[2]}"
        if len(tail) == 5 and tail[0] == "layers" and tail[2] == "layers":
            try:
                residual = int(tail[1]) - 1
                component = _RESIDUAL_COMPONENTS[int(tail[3])]
            except (KeyError, ValueError) as exc:
                raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
            if residual not in (1, 2, 3):
                raise ValueError(f"Unsupported Stable Audio decoder key: {key}")
            return f"block.{block}.res_unit{residual}.{component}.{tail[4]}"
    if layer == 6 and len(tail) == 1:
        return f"snake1.{tail[0]}"
    if layer == 7 and len(tail) == 1:
        return f"conv2.{tail[0]}"
    raise ValueError(f"Unsupported Stable Audio decoder key: {key}")


def convert_stable_audio_decoder_state_dict(
    state_dict: Mapping[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Convert and filter a Stable Audio checkpoint state dictionary."""

    converted: dict[str, torch.Tensor] = {}
    for key, value in state_dict.items():
        target_key = convert_stable_audio_decoder_key(key)
        if target_key is None:
            continue
        if target_key in converted:
            raise ValueError(f"Duplicate converted Stable Audio key: {target_key}")
        if target_key.endswith((".alpha", ".beta")):
            value = value.reshape(1, -1, 1)
        converted[target_key] = value
    return converted


def _load_oobleck_config(model_path: Path) -> tuple[dict[str, Any], int]:
    config_path = model_path / "model_config.json"
    if not config_path.is_file():
        raise FileNotFoundError(f"Stable Audio config does not exist: {config_path}")
    with config_path.open(encoding="utf-8") as file:
        full_config = json.load(file)

    try:
        pretransform = full_config["model"]["pretransform"]
        config = pretransform["config"]
        decoder = config["decoder"]
        decoder_config = decoder["config"]
        sample_rate = int(full_config["sample_rate"])
    except (KeyError, TypeError, ValueError) as exc:
        raise ValueError(f"Invalid Stable Audio model config: {config_path}") from exc

    if pretransform.get("type") != "autoencoder" or decoder.get("type") != "oobleck":
        raise ValueError("MAGI-2 requires an Oobleck Stable Audio pretransform")
    if decoder_config.get("use_nearest_upsample", False):
        raise ValueError("Diffusers Oobleck does not support nearest-neighbor decoder upsampling")
    if decoder_config.get("final_tanh", True):
        raise ValueError("MAGI-2 requires the Stable Audio decoder with final_tanh=false")

    strides = list(decoder_config["strides"])
    downsampling_ratio = int(config["downsampling_ratio"])
    if int(torch.tensor(strides).prod().item()) != downsampling_ratio:
        raise ValueError("Stable Audio strides do not match downsampling_ratio")

    kwargs = {
        "encoder_hidden_size": int(config["encoder"]["config"]["channels"]),
        "downsampling_ratios": strides,
        "channel_multiples": list(decoder_config["c_mults"]),
        "decoder_channels": int(decoder_config["channels"]),
        "decoder_input_channels": int(decoder_config["latent_dim"]),
        "audio_channels": int(decoder_config["out_channels"]),
        "sampling_rate": sample_rate,
    }
    return kwargs, downsampling_ratio


class Magi2AudioDecoder(nn.Module):
    """Decode MAGI-2 audio latents with the bundled Stable Audio VAE weights."""

    def __init__(
        self,
        model_path: str | Path,
        *,
        device: str | torch.device = "cpu",
        dtype: torch.dtype = torch.float32,
    ) -> None:
        super().__init__()
        model_path = Path(model_path)
        if not model_path.is_dir():
            raise ValueError(f"Stable Audio model path must be a local directory: {model_path}")

        oobleck_kwargs, self.downsampling_ratio = _load_oobleck_config(model_path)
        self.sample_rate = int(oobleck_kwargs["sampling_rate"])
        # Construct the Diffusers implementation for architectural compatibility,
        # then retain only the decoder because MAGI-2 has no audio input encoder.
        autoencoder = AutoencoderOobleck(**oobleck_kwargs)
        self.decoder = autoencoder.decoder.to(device=device, dtype=dtype)
        del autoencoder

        weights_path = model_path / "model.safetensors"
        if not weights_path.is_file():
            raise FileNotFoundError(f"Stable Audio weights do not exist: {weights_path}")
        selected: dict[str, torch.Tensor] = {}
        with safe_open(weights_path, framework="pt", device=str(device)) as checkpoint:
            for key in checkpoint.keys():
                if key.startswith(_CHECKPOINT_PREFIX):
                    selected[key] = checkpoint.get_tensor(key)
        converted = convert_stable_audio_decoder_state_dict(selected)
        converted = {key: value.to(dtype=dtype) for key, value in converted.items()}
        self.decoder.load_state_dict(converted, strict=True)
        self.eval().requires_grad_(False)

    @property
    def latent_fps(self) -> float:
        return self.sample_rate / self.downsampling_ratio

    @torch.inference_mode()
    def decode(self, latents: torch.Tensor) -> torch.Tensor:
        parameter = next(self.decoder.parameters())
        return self.decoder(latents.to(device=parameter.device, dtype=parameter.dtype))

    forward = decode


__all__ = [
    "Magi2AudioDecoder",
    "convert_stable_audio_decoder_key",
    "convert_stable_audio_decoder_state_dict",
]