vllm-omni/vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
OmniVoice Decoder (Stage 1) - Audio token to waveform conversion.

Implements the HiggsAudioV2 decode path using transformers' DacModel decoder
and a custom RVQ quantizer, compatible with transformers 4.x.

Decode path:
  audio_codes [B, 8, T]
    → RVQ codebook lookup + project_out → sum → [B, 1024, T]
    → fc2 Linear(1024, 256) → [B, 256, T]
    → DAC acoustic decoder (conv transpose upsampling) → [B, 1, T*960]
    → 24kHz waveform (25fps × 960 samples/frame)
"""

from __future__ import annotations

import json
import os

import torch
import torch.nn as nn
from vllm.logger import init_logger

from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig

logger = init_logger(__name__)


class HiggsAudioVQLayer(nn.Module):
    """Single VQ layer: codebook lookup + project_out."""

    def __init__(self, codebook_size: int = 1024, codebook_dim: int = 64, hidden_size: int = 1024):
        super().__init__()
        self.codebook = nn.Embedding(codebook_size, codebook_dim)
        self.project_out = nn.Linear(codebook_dim, hidden_size)

    def decode(self, indices: torch.Tensor) -> torch.Tensor:
        """indices: [B, T] → [B, hidden_size, T]"""
        quantized = self.codebook(indices)  # [B, T, codebook_dim]
        quantized = self.project_out(quantized)  # [B, T, hidden_size]
        return quantized.permute(0, 2, 1)  # [B, hidden_size, T]


class HiggsAudioRVQ(nn.Module):
    """Residual Vector Quantizer with 8 codebook layers."""

    def __init__(
        self, num_quantizers: int = 8, codebook_size: int = 1024, codebook_dim: int = 64, hidden_size: int = 1024
    ):
        super().__init__()
        self.quantizers = nn.ModuleList(
            [HiggsAudioVQLayer(codebook_size, codebook_dim, hidden_size) for _ in range(num_quantizers)]
        )

    def decode(self, codes: torch.Tensor) -> torch.Tensor:
        """codes: [num_quantizers, B, T] → [B, hidden_size, T]"""
        result = torch.zeros(
            codes.shape[1],
            self.quantizers[0].project_out.out_features,
            codes.shape[2],
            device=codes.device,
            dtype=torch.float32,
        )
        for i, quantizer in enumerate(self.quantizers):
            result = result + quantizer.decode(codes[i])
        return result


class OmniVoiceDecoder(nn.Module):
    """OmniVoice Stage 1: Token-to-audio decoder.

    Uses DAC acoustic decoder from transformers + custom HiggsAudio RVQ
    quantizer to convert 8-codebook tokens into 24kHz waveform.
    """

    def __init__(self, config: OmniVoiceConfig):
        super().__init__()
        self.config = config
        self.sample_rate = config.sample_rate
        self._loaded = False

        # These are populated by load_weights
        self.quantizer = None
        self.fc2 = None
        self.acoustic_decoder = None

    @torch.inference_mode()
    def forward(self, audio_codes: torch.Tensor) -> torch.Tensor:
        """Decode audio tokens to waveform.

        Args:
            audio_codes: [B, 8, T] - 8-codebook audio token IDs

        Returns:
            waveform: [B, 1, audio_samples] at 24kHz
        """
        if not self._loaded:
            raise RuntimeError("Decoder not loaded. Call load_weights() first.")

        device = audio_codes.device

        # Transpose: [B, 8, T] → [8, B, T]
        codes = audio_codes.transpose(0, 1).long()

        # RVQ decode: sum codebook embeddings → [B, 1024, T]
        quantized = self.quantizer.decode(codes)

        # Project: [B, 1024, T] → fc2 → [B, 256, T]
        # Cast to fc2 weight dtype (may be fp16 when checkpoint stores weights as fp16),
        # then upcast back to float32 — acoustic decoder ConvTranspose1d upsampling
        # produces intermediate values that exceed the fp16 range (~65504), causing NaN.
        quantized = self.fc2(quantized.transpose(1, 2).to(self.fc2.weight.dtype)).transpose(1, 2).float()

        # Acoustic decoder: [B, 256, T] → [B, 1, T*960]
        audio = self.acoustic_decoder(quantized)

        # Ensure [B, 1, samples]
        if audio.dim() == 2:
            audio = audio.unsqueeze(1)

        return audio.to(device)

    def _adjust_output_padding(self, decoder: nn.Module):
        """Adjust ConvTranspose1d output_padding (HiggsAudioV2 modification)."""
        for module in decoder.modules():
            if isinstance(module, nn.ConvTranspose1d):
                stride = module.stride[0] if isinstance(module.stride, tuple) else module.stride
                module.output_padding = (stride % 2,)

    def load_weights(self, model_dir: str, device: torch.device) -> None:
        """Load decoder components from audio_tokenizer/model.safetensors."""
        from safetensors.torch import load_file
        from transformers import DacConfig, DacModel

        audio_tokenizer_path = os.path.join(model_dir, "audio_tokenizer")
        config_path = os.path.join(audio_tokenizer_path, "config.json")
        weights_path = os.path.join(audio_tokenizer_path, "model.safetensors")

        if not os.path.exists(weights_path):
            raise FileNotFoundError(f"Audio tokenizer weights not found at {weights_path}")

        with open(config_path) as f:
            tokenizer_config = json.load(f)

        state_dict = load_file(weights_path, device=str(device))

        # 1. Build RVQ quantizer
        codebook_dim = tokenizer_config.get("codebook_dim", 64)
        codebook_size = tokenizer_config.get("codebook_size", 1024)
        # Hidden size = quantizer project_out output dim
        hidden_size = state_dict["quantizer.quantizers.0.project_out.weight"].shape[0]
        num_quantizers = sum(
            1 for k in state_dict if k.startswith("quantizer.quantizers.") and k.endswith(".codebook.embed")
        )

        self.quantizer = HiggsAudioRVQ(
            num_quantizers=num_quantizers,
            codebook_size=codebook_size,
            codebook_dim=codebook_dim,
            hidden_size=hidden_size,
        ).to(device)

        # Load quantizer weights
        for i in range(num_quantizers):
            prefix = f"quantizer.quantizers.{i}"
            embed_key = f"{prefix}.codebook.embed"
            if embed_key in state_dict:
                self.quantizer.quantizers[i].codebook.weight.data.copy_(state_dict[embed_key])
            proj_out_w = f"{prefix}.project_out.weight"
            proj_out_b = f"{prefix}.project_out.bias"
            if proj_out_w in state_dict:
                self.quantizer.quantizers[i].project_out.weight.data.copy_(state_dict[proj_out_w])
            if proj_out_b in state_dict:
                self.quantizer.quantizers[i].project_out.bias.data.copy_(state_dict[proj_out_b])

        # 2. Build fc2 projection
        fc2_w = state_dict["fc2.weight"]
        fc2_b = state_dict["fc2.bias"]
        self.fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device)
        self.fc2.weight.data.copy_(fc2_w)
        self.fc2.bias.data.copy_(fc2_b)

        # 3. Build DAC acoustic decoder
        dac_cfg = DacConfig(**tokenizer_config["acoustic_model_config"])
        dac_model = DacModel(dac_cfg)
        self.acoustic_decoder = dac_model.decoder.to(device)

        # Load acoustic decoder weights
        loaded = 0
        for name, param in self.acoustic_decoder.named_parameters():
            higgs_name = f"acoustic_decoder.{name}"
            if higgs_name in state_dict:
                param.data.copy_(state_dict[higgs_name])
                loaded += 1

        # Checkpoint weights may be fp16; force float32 so ConvTranspose1d
        # upsampling doesn't overflow (fp16 max ~65504 is too small for
        # intermediate activations in the DAC upsampling chain).
        self.acoustic_decoder.float()

        # Apply HiggsAudioV2 output padding adjustment
        self._adjust_output_padding(self.acoustic_decoder)

        # Remove tanh if present (HiggsAudioV2 uses Identity instead)
        if hasattr(self.acoustic_decoder, "tanh"):
            self.acoustic_decoder.tanh = nn.Identity()

        self.acoustic_decoder.eval()
        self._loaded = True

        logger.info(
            "Loaded OmniVoice decoder: %d quantizers, fc2(%d→%d), acoustic decoder (%d weights)",
            num_quantizers,
            fc2_w.shape[1],
            fc2_w.shape[0],
            loaded,
        )