add demo

app.py

@@ -0,0 +1,147 @@
+import gradio as gr
+import torch
+import soundfile as sf
+import tempfile
+import os
+import warnings
+warnings.filterwarnings("ignore")
+
+try:
+    from models.tts.llm_tts.inference_llm_tts import TTSInferencePipeline
+    MODEL_AVAILABLE = True
+except ImportError:
+    print("Warning: TaDiCodec models not found. Running in demo mode.")
+    MODEL_AVAILABLE = False
+
+class TaDiCodecTTSDemo:
+    def __init__(self):
+        self.pipeline = None
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.load_model()
+    
+    def load_model(self):
+        try:
+            if MODEL_AVAILABLE:
+                print("Loading TaDiCodec-TTS-AR-Qwen2.5-0.5B model...")
+                self.pipeline = TTSInferencePipeline.from_pretrained(
+                    tadicodec_path="amphion/TaDiCodec",
+                    llm_path="amphion/TaDiCodec-TTS-AR-Qwen2.5-0.5B",
+                    device=self.device,
+                )
+                print("Model loaded successfully!")
+            else:
+                print("Running in demo mode - model files not available")
+                self.pipeline = None
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            self.pipeline = None
+    
+    def synthesize_speech(self, text, reference_audio=None, reference_text=""):
+        """
+        Synthesize speech from text using TaDiCodec TTS
+        """
+        if not text.strip():
+            return None, "Please enter some text to synthesize."
+        
+        try:
+            if self.pipeline is not None:
+                # Use actual TaDiCodec inference
+                if reference_audio and reference_text.strip():
+                    audio = self.pipeline(
+                        text=text,
+                        prompt_text=reference_text,
+                        prompt_speech_path=reference_audio,
+                    )
+                else:
+                    audio = self.pipeline(text=text)
+                
+                # Save to temporary file
+                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
+                    sf.write(tmp_file.name, audio, 24000)
+                    return tmp_file.name, "Speech synthesized successfully!"
+            else:
+                # Fallback demo mode - return None to indicate no audio generated
+                return None, "Demo mode - TaDiCodec model not available. Please install the required models."
+                
+        except Exception as e:
+            return None, f"Error during synthesis: {str(e)}"
+
+# Initialize the demo
+demo_instance = TaDiCodecTTSDemo()
+
+def tts_interface(text, reference_audio, reference_text):
+    """Interface function for Gradio"""
+    audio_path, message = demo_instance.synthesize_speech(text, reference_audio, reference_text)
+    return audio_path, message
+
+# Create Gradio interface
+with gr.Blocks(title="TaDiCodec-TTS Demo", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # TaDiCodec-TTS-AR-Qwen2.5-0.5B Demo
+    
+    This is a demo of the TaDiCodec Text-to-Speech model with Qwen2.5-0.5B backbone.
+    
+    **Features:**
+    - Voice cloning with reference audio
+    - Code-switching support (e.g., mixing English and Chinese)
+    - Extremely low bitrate (0.0875 kbps)
+    - High-quality speech generation
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            text_input = gr.Textbox(
+                label="Text to Synthesize",
+                placeholder="Enter the text you want to convert to speech...",
+                lines=3,
+                value="但是 to those who 知道 her well, it was a 标志 of her unwavering 决心 and spirit."
+            )
+            
+            reference_text = gr.Textbox(
+                label="Reference Text",
+                placeholder="Text corresponding to the reference audio for voice cloning...",
+                lines=2,
+                value="In short, we embarked on a mission to make America great again, for all Americans."
+            )
+            
+            reference_audio = gr.Audio(
+                label="Reference Audio",
+                type="filepath",
+            )
+            
+            synthesize_btn = gr.Button("Synthesize Speech", variant="primary")
+        
+        with gr.Column():
+            output_audio = gr.Audio(
+                label="Generated Speech",
+                type="filepath"
+            )
+            
+            status_message = gr.Textbox(
+                label="Status",
+                interactive=False
+            )
+    
+    # Example inputs
+    gr.Markdown("### Example Inputs")
+    examples = [
+        ["Yes, usually people choose to face life with more positive emotions, after all, happy times are always yearning. However, sometimes slowing down and experiencing the details of life can bring deeper joy and satisfaction. What do you think?", "Jittery Jack's jam jars jiggled jauntily, jolting Jack's jumbled jelly-filled jars joyously.", "sample/tongueTwisters_en_018.wav"],
+        ["You think you can just waltz in here and cause chaos? Well, I've got news for you. This time, there's no escaping the consequences. So, one by one, step forward, and let's see who's bold enough to face the music. It's time for a little dose of reality—prepare to be dealt with!","Get in line trouble makers, and I will take care of you.", "sample/en_013.wav"],
+    ]
+    
+    gr.Examples(
+        examples=examples,
+        inputs=[text_input, reference_text, reference_audio],
+        outputs=[output_audio, status_message],
+        fn=tts_interface
+    )
+    
+    # Connect the interface
+    synthesize_btn.click(
+        fn=tts_interface,
+        inputs=[text_input, reference_audio, reference_text],
+        outputs=[output_audio, status_message]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True, server_name="0.0.0.0")

models/codec/amphion_codec/codec.py

@@ -0,0 +1,422 @@
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+from models.codec.amphion_codec.quantize import (
+    ResidualVQ,
+    VectorQuantize,
+    FactorizedVectorQuantize,
+    LookupFreeQuantize,
+)
+
+from models.codec.amphion_codec.vocos import Vocos
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+    if isinstance(m, nn.Linear):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int = 16, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(dim // 2, dilation=1),
+            ResidualUnit(dim // 2, dilation=3),
+            ResidualUnit(dim // 2, dilation=9),
+            Snake1d(dim // 2),
+            WNConv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class CodecEncoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 64,
+        up_ratios: list = [4, 5, 5, 6],
+        out_channels: int = 256,
+        use_tanh: bool = False,
+        cfg=None,
+    ):
+        super().__init__()
+
+        d_model = cfg.d_model if cfg is not None else d_model
+        up_ratios = cfg.up_ratios if cfg is not None else up_ratios
+        out_channels = cfg.out_channels if cfg is not None else out_channels
+        use_tanh = cfg.use_tanh if cfg is not None else use_tanh
+
+        # Create first convolution
+        self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for stride in up_ratios:
+            d_model *= 2
+            self.block += [EncoderBlock(d_model, stride=stride)]
+
+        # Create last convolution
+        self.block += [
+            Snake1d(d_model),
+            WNConv1d(d_model, out_channels, kernel_size=3, padding=1),
+        ]
+
+        if use_tanh:
+            self.block += [nn.Tanh()]
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        return self.block(x)
+
+    def reset_parameters(self):
+        self.apply(init_weights)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(input_dim),
+            WNConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=stride // 2 + stride % 2,
+                output_padding=stride % 2,
+            ),
+            ResidualUnit(output_dim, dilation=1),
+            ResidualUnit(output_dim, dilation=3),
+            ResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class CodecDecoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 256,
+        upsample_initial_channel: int = 1536,
+        up_ratios: list = [5, 5, 4, 2],
+        num_quantizers: int = 8,
+        codebook_size: int = 1024,
+        codebook_dim: int = 256,
+        quantizer_type: str = "vq",
+        quantizer_dropout: float = 0.5,
+        commitment: float = 0.25,
+        codebook_loss_weight: float = 1.0,
+        use_l2_normlize: bool = False,
+        codebook_type: str = "euclidean",
+        kmeans_init: bool = False,
+        kmeans_iters: int = 10,
+        decay: float = 0.8,
+        eps: float = 1e-5,
+        threshold_ema_dead_code: int = 2,
+        weight_init: bool = False,
+        use_vocos: bool = False,
+        vocos_dim: int = 384,
+        vocos_intermediate_dim: int = 1152,
+        vocos_num_layers: int = 8,
+        n_fft: int = 800,
+        hop_size: int = 200,
+        padding: str = "same",
+        cfg=None,
+    ):
+        super().__init__()
+
+        in_channels = (
+            cfg.in_channels
+            if cfg is not None and hasattr(cfg, "in_channels")
+            else in_channels
+        )
+        upsample_initial_channel = (
+            cfg.upsample_initial_channel
+            if cfg is not None and hasattr(cfg, "upsample_initial_channel")
+            else upsample_initial_channel
+        )
+        up_ratios = (
+            cfg.up_ratios
+            if cfg is not None and hasattr(cfg, "up_ratios")
+            else up_ratios
+        )
+        num_quantizers = (
+            cfg.num_quantizers
+            if cfg is not None and hasattr(cfg, "num_quantizers")
+            else num_quantizers
+        )
+        codebook_size = (
+            cfg.codebook_size
+            if cfg is not None and hasattr(cfg, "codebook_size")
+            else codebook_size
+        )
+        codebook_dim = (
+            cfg.codebook_dim
+            if cfg is not None and hasattr(cfg, "codebook_dim")
+            else codebook_dim
+        )
+        quantizer_type = (
+            cfg.quantizer_type
+            if cfg is not None and hasattr(cfg, "quantizer_type")
+            else quantizer_type
+        )
+        quantizer_dropout = (
+            cfg.quantizer_dropout
+            if cfg is not None and hasattr(cfg, "quantizer_dropout")
+            else quantizer_dropout
+        )
+        commitment = (
+            cfg.commitment
+            if cfg is not None and hasattr(cfg, "commitment")
+            else commitment
+        )
+        codebook_loss_weight = (
+            cfg.codebook_loss_weight
+            if cfg is not None and hasattr(cfg, "codebook_loss_weight")
+            else codebook_loss_weight
+        )
+        use_l2_normlize = (
+            cfg.use_l2_normlize
+            if cfg is not None and hasattr(cfg, "use_l2_normlize")
+            else use_l2_normlize
+        )
+        codebook_type = (
+            cfg.codebook_type
+            if cfg is not None and hasattr(cfg, "codebook_type")
+            else codebook_type
+        )
+        kmeans_init = (
+            cfg.kmeans_init
+            if cfg is not None and hasattr(cfg, "kmeans_init")
+            else kmeans_init
+        )
+        kmeans_iters = (
+            cfg.kmeans_iters
+            if cfg is not None and hasattr(cfg, "kmeans_iters")
+            else kmeans_iters
+        )
+        decay = cfg.decay if cfg is not None and hasattr(cfg, "decay") else decay
+        eps = cfg.eps if cfg is not None and hasattr(cfg, "eps") else eps
+        threshold_ema_dead_code = (
+            cfg.threshold_ema_dead_code
+            if cfg is not None and hasattr(cfg, "threshold_ema_dead_code")
+            else threshold_ema_dead_code
+        )
+        weight_init = (
+            cfg.weight_init
+            if cfg is not None and hasattr(cfg, "weight_init")
+            else weight_init
+        )
+        use_vocos = (
+            cfg.use_vocos
+            if cfg is not None and hasattr(cfg, "use_vocos")
+            else use_vocos
+        )
+        vocos_dim = (
+            cfg.vocos_dim
+            if cfg is not None and hasattr(cfg, "vocos_dim")
+            else vocos_dim
+        )
+        vocos_intermediate_dim = (
+            cfg.vocos_intermediate_dim
+            if cfg is not None and hasattr(cfg, "vocos_intermediate_dim")
+            else vocos_intermediate_dim
+        )
+        vocos_num_layers = (
+            cfg.vocos_num_layers
+            if cfg is not None and hasattr(cfg, "vocos_num_layers")
+            else vocos_num_layers
+        )
+        n_fft = cfg.n_fft if cfg is not None and hasattr(cfg, "n_fft") else n_fft
+        hop_size = (
+            cfg.hop_size if cfg is not None and hasattr(cfg, "hop_size") else hop_size
+        )
+        padding = (
+            cfg.padding if cfg is not None and hasattr(cfg, "padding") else padding
+        )
+
+        if quantizer_type == "vq":
+            self.quantizer = ResidualVQ(
+                input_dim=in_channels,
+                num_quantizers=num_quantizers,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_type=quantizer_type,
+                quantizer_dropout=quantizer_dropout,
+                commitment=commitment,
+                codebook_loss_weight=codebook_loss_weight,
+                use_l2_normlize=use_l2_normlize,
+                codebook_type=codebook_type,
+                kmeans_init=kmeans_init,
+                kmeans_iters=kmeans_iters,
+                decay=decay,
+                eps=eps,
+                threshold_ema_dead_code=threshold_ema_dead_code,
+                weight_init=weight_init,
+            )
+        elif quantizer_type == "fvq":
+            self.quantizer = ResidualVQ(
+                input_dim=in_channels,
+                num_quantizers=num_quantizers,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_type=quantizer_type,
+                quantizer_dropout=quantizer_dropout,
+                commitment=commitment,
+                codebook_loss_weight=codebook_loss_weight,
+                use_l2_normlize=use_l2_normlize,
+            )
+        elif quantizer_type == "lfq":
+            self.quantizer = ResidualVQ(
+                input_dim=in_channels,
+                num_quantizers=num_quantizers,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_type=quantizer_type,
+            )
+        else:
+            raise ValueError(f"Unknown quantizer type {quantizer_type}")
+
+        if not use_vocos:
+            # Add first conv layer
+            channels = upsample_initial_channel
+            layers = [WNConv1d(in_channels, channels, kernel_size=7, padding=3)]
+
+            # Add upsampling + MRF blocks
+            for i, stride in enumerate(up_ratios):
+                input_dim = channels // 2**i
+                output_dim = channels // 2 ** (i + 1)
+                layers += [DecoderBlock(input_dim, output_dim, stride)]
+
+            # Add final conv layer
+            layers += [
+                Snake1d(output_dim),
+                WNConv1d(output_dim, 1, kernel_size=7, padding=3),
+                nn.Tanh(),
+            ]
+
+            self.model = nn.Sequential(*layers)
+
+        if use_vocos:
+            self.model = Vocos(
+                input_channels=in_channels,
+                dim=vocos_dim,
+                intermediate_dim=vocos_intermediate_dim,
+                num_layers=vocos_num_layers,
+                adanorm_num_embeddings=None,
+                n_fft=n_fft,
+                hop_size=hop_size,
+                padding=padding,
+            )
+
+        self.reset_parameters()
+
+    def forward(self, x=None, vq=False, eval_vq=False, n_quantizers=None):
+        """
+        if vq is True, x = encoder output, then return quantized output;
+        else, x = quantized output, then return decoder output
+        """
+        if vq is True:
+            if eval_vq:
+                self.quantizer.eval()
+            (
+                quantized_out,
+                all_indices,
+                all_commit_losses,
+                all_codebook_losses,
+                all_quantized,
+            ) = self.quantizer(x, n_quantizers=n_quantizers)
+            return (
+                quantized_out,
+                all_indices,
+                all_commit_losses,
+                all_codebook_losses,
+                all_quantized,
+            )
+
+        return self.model(x)
+
+    def quantize(self, x, n_quantizers=None):
+        self.quantizer.eval()
+        quantized_out, vq, _, _, _ = self.quantizer(x, n_quantizers=n_quantizers)
+        return quantized_out, vq
+
+    # TODO: check consistency of vq2emb and quantize
+    def vq2emb(self, vq, n_quantizers=None):
+        return self.quantizer.vq2emb(vq, n_quantizers=n_quantizers)
+
+    def decode(self, x):
+        return self.model(x)
+
+    def latent2dist(self, x, n_quantizers=None):
+        return self.quantizer.latent2dist(x, n_quantizers=n_quantizers)
+
+    def reset_parameters(self):
+        self.apply(init_weights)

models/codec/amphion_codec/loss.py

@@ -0,0 +1,401 @@
+import librosa
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from torchaudio.transforms import MelSpectrogram
+from einops import rearrange
+from typing import List
+
+
+def stft(x, fft_size, hop_size, win_length, window, use_complex=False):
+    """Perform STFT and convert to magnitude spectrogram.
+    Args:
+        x (Tensor): Input signal tensor (B, T).
+        fft_size (int): FFT size.
+        hop_size (int): Hop size.
+        win_length (int): Window length.
+        window (str): Window function type.
+    Returns:
+        Tensor: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
+    """
+
+    x_stft = torch.stft(
+        x, fft_size, hop_size, win_length, window.to(x.device), return_complex=True
+    )
+
+    # clamp is needed to avoid nan or inf
+    if not use_complex:
+        return torch.sqrt(
+            torch.clamp(x_stft.real**2 + x_stft.imag**2, min=1e-7, max=1e3)
+        ).transpose(2, 1)
+    else:
+        res = torch.cat([x_stft.real.unsqueeze(1), x_stft.imag.unsqueeze(1)], dim=1)
+        res = res.transpose(2, 3)  # [B, 2, T, F]
+        return res
+
+
+def compute_mag_scale(n_fft, sampling_rate):
+    frequencies = librosa.fft_frequencies(sr=sampling_rate, n_fft=n_fft)
+    frequencies = np.where(frequencies > 1e-10, frequencies, -10)
+    db_scale = librosa.frequency_weighting(frequencies).reshape(1, 1, -1)
+    mag_scale = np.sqrt(librosa.db_to_power(db_scale)).astype(np.float32)
+    return torch.from_numpy(mag_scale)
+
+
+class SpectralConvergenceLoss(torch.nn.Module):
+    """Spectral convergence loss module."""
+
+    def __init__(self):
+        """Initialize spectral convergence loss module."""
+        super(SpectralConvergenceLoss, self).__init__()
+
+    def forward(self, x_mag, y_mag):
+        """Calculate forward propagation.
+        Args:
+            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
+            y_mag (Tensor): Magnitude spectrogram of ground-truth signal (B, #frames, #freq_bins).
+        Returns:
+            Tensor: Spectral convergence loss value.
+        """
+        return torch.norm(y_mag - x_mag) / torch.norm(y_mag)
+
+
+class LogSTFTMagnitudeLoss(torch.nn.Module):
+    """Log STFT magnitude loss module."""
+
+    def __init__(self):
+        """Initialize los STFT magnitude loss module."""
+        super(LogSTFTMagnitudeLoss, self).__init__()
+
+    def forward(self, x_mag, y_mag):
+        """Calculate forward propagation.
+        Args:
+            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
+            y_mag (Tensor): Magnitude spectrogram of ground-truth signal (B, #frames, #freq_bins).
+        Returns:
+            Tensor: Log STFT magnitude loss value.
+        """
+        return F.l1_loss(torch.log(y_mag), torch.log(x_mag))
+
+
+class STFTLoss(torch.nn.Module):
+    """STFT loss module."""
+
+    def __init__(
+        self,
+        fft_size=1024,
+        hop_length=120,
+        win_length=600,
+        sampling_rate=16000,
+        window="hann_window",
+        cfg=None,
+    ):
+        """Initialize STFT loss module."""
+        super(STFTLoss, self).__init__()
+
+        fft_size = (
+            cfg.fft_size if cfg is not None and hasattr(cfg, "fft_size") else fft_size
+        )
+        hop_length = (
+            cfg.hop_length
+            if cfg is not None and hasattr(cfg, "hop_length")
+            else hop_length
+        )
+        win_length = (
+            cfg.win_length
+            if cfg is not None and hasattr(cfg, "win_length")
+            else win_length
+        )
+        window = cfg.window if cfg is not None and hasattr(cfg, "window") else window
+
+        self.fft_size = fft_size
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.window = getattr(torch, window)(win_length)
+        self.spectral_convergence_loss = SpectralConvergenceLoss()
+        self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss()
+
+        self.register_buffer("mag_scale", compute_mag_scale(fft_size, sampling_rate))
+
+    def forward(self, x, y):
+        """Calculate forward propagation.
+        Args:
+            x (Tensor): Predicted signal (B, T).
+            y (Tensor): Ground truth signal (B, T).
+        Returns:
+            Tensor: Spectral convergence loss value.
+            Tensor: Log STFT magnitude loss value.
+        """
+        x_mag = (
+            stft(x, self.fft_size, self.hop_length, self.win_length, self.window)
+            * self.mag_scale
+        )
+        y_mag = (
+            stft(y, self.fft_size, self.hop_length, self.win_length, self.window)
+            * self.mag_scale
+        )
+        sc_loss = self.spectral_convergence_loss(x_mag, y_mag)
+        log_mag_loss = self.log_stft_magnitude_loss(x_mag, y_mag)
+
+        return sc_loss, log_mag_loss
+
+
+class MultiResolutionSTFTLoss(torch.nn.Module):
+    """Multi resolution STFT loss module."""
+
+    def __init__(
+        self,
+        fft_sizes=(1024, 2048, 512),
+        hop_sizes=(120, 240, 50),
+        win_lengths=(600, 1200, 240),
+        window="hann_window",
+        sampling_rate=16000,
+        cfg=None,
+    ):
+        """Initialize Multi resolution STFT loss module.
+        Args:
+            fft_sizes (list): List of FFT sizes.
+            hop_sizes (list): List of hop sizes.
+            win_lengths (list): List of window lengths.
+            window (str): Window function type.
+        """
+        super(MultiResolutionSTFTLoss, self).__init__()
+
+        fft_sizes = (
+            cfg.fft_sizes
+            if cfg is not None and hasattr(cfg, "fft_sizes")
+            else fft_sizes
+        )
+        hop_sizes = (
+            cfg.hop_sizes
+            if cfg is not None and hasattr(cfg, "hop_sizes")
+            else hop_sizes
+        )
+        win_lengths = (
+            cfg.win_lengths
+            if cfg is not None and hasattr(cfg, "win_lengths")
+            else win_lengths
+        )
+        window = cfg.window if cfg is not None and hasattr(cfg, "window") else window
+        sampling_rate = (
+            cfg.sampling_rate
+            if cfg is not None and hasattr(cfg, "sampling_rate")
+            else sampling_rate
+        )
+
+        assert len(fft_sizes) == len(hop_sizes) == len(win_lengths)
+        self.stft_losses = torch.nn.ModuleList()
+        for fs, ss, wl in zip(fft_sizes, hop_sizes, win_lengths):
+            self.stft_losses += [
+                STFTLoss(fs, ss, wl, window=window, sampling_rate=sampling_rate)
+            ]
+
+    def forward(self, x, y):
+        """Calculate forward propagation.
+        Args:
+            x (Tensor): Predicted signal (B, T).
+            y (Tensor): GroundTruth signal (B, T).
+        Returns:
+            Tensor: Multi resolution spectral convergence loss value.
+            Tensor: Multi resolution log STFT magnitude loss value.
+        """
+        sc_loss = 0.0
+        mag_loss = 0.0
+        for f in self.stft_losses:
+            sc_l, mag_l = f(x, y)
+            sc_loss += sc_l
+            mag_loss += mag_l
+        sc_loss /= len(self.stft_losses)
+        mag_loss /= len(self.stft_losses)
+
+        return sc_loss, mag_loss
+
+
+class MultiResolutionMelSpectrogramLoss(nn.Module):
+    """Compute distance between mel spectrograms. Can be used
+    in a multi-scale way.
+
+    Parameters
+    ----------
+    n_mels : List[int]
+        Number of mels per STFT, by default [150, 80],
+    window_lengths : List[int], optional
+        Length of each window of each STFT, by default [2048, 512]
+    loss_fn : typing.Callable, optional
+        How to compare each loss, by default nn.L1Loss()
+    clamp_eps : float, optional
+        Clamp on the log magnitude, below, by default 1e-5
+    mag_weight : float, optional
+        Weight of raw magnitude portion of loss, by default 1.0
+    log_weight : float, optional
+        Weight of log magnitude portion of loss, by default 1.0
+    pow : float, optional
+        Power to raise magnitude to before taking log, by default 2.0
+    weight : float, optional
+        Weight of this loss, by default 1.0
+    match_stride : bool, optional
+        Whether to match the stride of convolutional layers, by default False
+
+    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
+    """
+
+    def __init__(
+        self,
+        sample_rate=16000,
+        n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
+        window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
+        clamp_eps: float = 1e-5,
+        mag_weight: float = 0.0,
+        log_weight: float = 1.0,
+        pow: float = 1.0,
+        mel_fmin: List[float] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+        mel_fmax: List[float] = [None, None, None, None, None, None, None],
+        cfg=None,
+    ):
+        super().__init__()
+
+        sample_rate = (
+            cfg.sample_rate
+            if cfg is not None and hasattr(cfg, "sample_rate")
+            else sample_rate
+        )
+        n_mels = cfg.n_mels if cfg is not None and hasattr(cfg, "n_mels") else n_mels
+        window_lengths = (
+            cfg.window_lengths
+            if cfg is not None and hasattr(cfg, "window_lengths")
+            else window_lengths
+        )
+        clamp_eps = (
+            cfg.clamp_eps
+            if cfg is not None and hasattr(cfg, "clamp_eps")
+            else clamp_eps
+        )
+        mag_weight = (
+            cfg.mag_weight
+            if cfg is not None and hasattr(cfg, "mag_weight")
+            else mag_weight
+        )
+        log_weight = (
+            cfg.log_weight
+            if cfg is not None and hasattr(cfg, "log_weight")
+            else log_weight
+        )
+        pow = cfg.pow if cfg is not None and hasattr(cfg, "pow") else pow
+        mel_fmin = (
+            cfg.mel_fmin if cfg is not None and hasattr(cfg, "mel_fmin") else mel_fmin
+        )
+        mel_fmax = (
+            cfg.mel_fmax if cfg is not None and hasattr(cfg, "mel_fmax") else mel_fmax
+        )
+
+        self.mel_transforms = nn.ModuleList(
+            [
+                MelSpectrogram(
+                    sample_rate=sample_rate,
+                    n_fft=window_length,
+                    hop_length=window_length // 4,
+                    n_mels=n_mel,
+                    power=1.0,
+                    center=True,
+                    norm="slaney",
+                    mel_scale="slaney",
+                )
+                for n_mel, window_length in zip(n_mels, window_lengths)
+            ]
+        )
+        self.n_mels = n_mels
+        self.loss_fn = nn.L1Loss()
+        self.clamp_eps = clamp_eps
+        self.log_weight = log_weight
+        self.mag_weight = mag_weight
+        self.mel_fmin = mel_fmin
+        self.mel_fmax = mel_fmax
+        self.pow = pow
+
+    def delta(self, x, k):
+        l = x.shape[1]
+        return x[:, 0 : l - k] - x[:, k:l]
+
+    def forward(self, x, y, mask=None):
+        """Computes mel loss between an estimate and a reference
+        signal.
+
+        Parameters
+        ----------
+        x : AudioSignal
+            Estimate signal
+        y : AudioSignal
+            Reference signal
+
+        Returns
+        -------
+        torch.Tensor
+            Mel loss.
+        """
+        loss = 0.0
+        for mel_transform in self.mel_transforms:
+            x_mel = mel_transform(x)
+            y_mel = mel_transform(y)
+            log_x_mel = x_mel.clamp(self.clamp_eps).pow(self.pow).log10()
+            log_y_mel = y_mel.clamp(self.clamp_eps).pow(self.pow).log10()
+            loss += self.log_weight * self.loss_fn(log_x_mel, log_y_mel)
+            loss += self.mag_weight * self.loss_fn(x_mel, y_mel)
+            # loss += self.loss_fn(self.delta(log_x_mel, 1), self.delta(log_y_mel, 1))
+            # log_x_mel = rearrange(log_x_mel, 'b c t -> b t c')
+            # log_y_mel = rearrange(log_y_mel, 'b c t -> b t c')
+            # for i in range(3):
+            #     loss += self.loss_fn(self.delta(log_x_mel, i), self.delta(log_y_mel, i))
+        # loss /= len(self.mel_transforms)
+        return loss
+
+
+class GANLoss(nn.Module):
+    def __init__(self, mode="lsgan"):
+        super(GANLoss, self).__init__()
+        assert mode in ["lsgan", "lsgan_std", "hinge"]
+        self.mode = mode
+
+    def disc_loss(self, real, fake):
+        if self.mode == "lsgan":
+            real_loss = F.mse_loss(real, torch.ones_like(real))
+            fake_loss = F.mse_loss(fake, torch.zeros_like(fake))
+        elif self.mode == "lsgan_std":
+            real = (real - 1.0).pow(2)
+            fake = (fake - 0.0).pow(2)
+            real_loss = real.mean() + real.std()
+            fake_loss = fake.mean() + fake.std()
+        elif self.mode == "hinge":
+            real_loss = torch.relu(1.0 - real).mean()
+            fake_loss = torch.relu(1.0 + fake).mean()
+        else:
+            raise ValueError(f"no such mode {self.mode}")
+
+        return real_loss, fake_loss
+
+    def disc_loss2(self, fake):
+        if self.mode == "lsgan":
+            fake_loss = F.mse_loss(fake, torch.zeros_like(fake))
+        elif self.mode == "lsgan_std":
+            fake = (fake - 0.0).pow(2)
+            fake_loss = fake.mean() + fake.std()
+        elif self.mode == "hinge":
+            fake_loss = torch.relu(1.0 + fake).mean()
+        else:
+            raise ValueError(f"no such mode {self.mode}")
+
+        return fake_loss
+
+    def gen_loss(self, fake):
+        if self.mode == "lsgan":
+            gen_loss = F.mse_loss(fake, torch.ones_like(fake))
+        elif self.mode == "lsgan_std":
+            fake = (fake - 1.0).pow(2)
+            gen_loss = fake.mean() + fake.std()
+        elif self.mode == "hinge":
+            gen_loss = -fake.mean()
+        else:
+            raise ValueError(f"no such mode {self.mode}")
+
+        return gen_loss

models/codec/amphion_codec/quantize/__init__.py

@@ -0,0 +1,6 @@
+from models.codec.amphion_codec.quantize.factorized_vector_quantize import (
+    FactorizedVectorQuantize,
+)
+from models.codec.amphion_codec.quantize.vector_quantize import VectorQuantize
+from models.codec.amphion_codec.quantize.lookup_free_quantize import LookupFreeQuantize
+from models.codec.amphion_codec.quantize.residual_vq import ResidualVQ

models/codec/amphion_codec/quantize/bsq.py

@@ -0,0 +1,373 @@
+# copy from https://github.com/zhaoyue-zephyrus/bsq-vit/blob/main/transcoder/models/quantizer/bsq.py
+
+from einops import rearrange, reduce
+import torch
+import torch.nn as nn
+from torch.autograd import Function
+import torch.nn.functional as F
+
+
+class DifferentiableEntropyFunction(Function):
+    @staticmethod
+    def forward(ctx, zq, basis, K, eps):
+        zb = (zq + 1) / 2
+        zi = ((zb * basis).sum(-1)).to(torch.int64)
+        cnt = torch.scatter_reduce(
+            torch.zeros(2**K, device=zq.device, dtype=zq.dtype),
+            0,
+            zi.flatten(),
+            torch.ones_like(zi.flatten()).to(zq.dtype),
+            "sum",
+        )
+        prob = (cnt + eps) / (cnt + eps).sum()
+        H = -(prob * torch.log(prob)).sum()
+        ctx.save_for_backward(zq, zi, prob)
+        ctx.K = K
+        return H
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        zq, zi, prob = ctx.saved_tensors
+        grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K
+        reord_grad = grad_array[zi.flatten()].reshape(zi.shape)
+        grad_input = reord_grad.unsqueeze(-1) * zq
+        return grad_input, None, None, None, None
+
+
+def codebook_entropy(zq, basis, K, eps=1e-4):
+    return DifferentiableEntropyFunction.apply(zq, basis, K, eps)
+
+
+class SimpleQuantizer(nn.Module):
+    def __init__(
+        self,
+        embed_dim=14,
+        codebook_size=16384,
+        commitment=0.25,
+        codebook_loss_weight=1.0,
+        use_l2_normlize=True,
+    ):
+        super().__init__()
+
+        self.embed_dim = embed_dim
+        self.codebook_dim = embed_dim
+        self.codebook_size = codebook_size
+        self.commitment = commitment
+        self.codebook_loss_weight = codebook_loss_weight
+        self.use_l2_normlize = use_l2_normlize
+
+        self.codebook = nn.Embedding(self.codebook_size, self.codebook_dim)
+
+    def forward(self, z):
+
+        z_e = z
+
+        z_q, indices = self.decode_latents(z_e)
+
+        if self.training:
+            commit_loss = (
+                F.mse_loss(z_e, z_q.detach(), reduction="mean") * self.commitment
+            )
+            codebook_loss = (
+                F.mse_loss(z_q, z_e.detach(), reduction="mean")
+                * self.codebook_loss_weight
+            )
+        else:
+            commit_loss = torch.zeros(z.shape[0], device=z.device)
+            codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        z_q = z_e + (z_q - z_e).detach()
+
+        return (
+            z_q,
+            commit_loss + codebook_loss,
+            {
+                "commit_loss": commit_loss,
+                "codebook_loss": codebook_loss,
+                "indices": indices,
+            },
+        )
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id)
+
+    def get_codebook_entry(self, indices):
+        return self.embed_code(indices)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b t d -> (b t) d")
+
+        codebook = self.codebook.weight
+
+        if self.use_l2_normlize:
+            # encodings = F.normalize(encodings) we have already normalized the latent before vq
+            codebook = F.normalize(codebook)
+
+        # Compute euclidean distance between encodings and codebook,
+        # if use_l2_normlize is True, the distance is equal to cosine distance
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+
+        return z_q, indices
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = self.decode_code(vq)
+        return emb
+
+    def latent2dist(self, latents):
+
+        encodings = latents
+        codebook = self.codebook.weight
+
+        # L2 normalize encodings and codebook
+        if self.use_l2_normlize:
+            # encodings = F.normalize(encodings)
+            codebook = F.normalize(codebook)
+
+        # Compute euclidean distance between encodings and codebook,
+        # if use_l2_normlize is True, the distance is equal to cosine distance
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )  # (b*t, k)
+
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        dist = rearrange(dist, "(b t) k -> b t k", b=latents.size(0))
+        z_q = self.decode_code(indices)
+
+        return -dist, indices, z_q
+
+
+class BinarySphericalQuantizer(nn.Module):
+    def __init__(
+        self,
+        embed_dim=14,
+        beta=0.0,
+        gamma0=1.0,
+        gamma=1.0,
+        zeta=1.0,
+        input_format="blc",
+        soft_entropy=True,
+        group_size=1,
+        persample_entropy_compute="group",
+        cb_entropy_compute="group",
+        l2_norm=True,
+        inv_temperature=1,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.beta = beta  # loss weight for commit loss
+        self.gamma0 = gamma0  # loss weight for entropy penalty
+        self.gamma = gamma  # loss weight for entropy penalty
+        self.zeta = zeta  # loss weight for entire entropy penalty
+        self.input_format = input_format
+        assert (
+            self.embed_dim % group_size == 0
+        ), "embed_dim must be divisible by group_size"
+        self.num_groups = self.embed_dim // group_size
+        self.group_size = group_size
+        assert persample_entropy_compute in [
+            "group",
+            "analytical",
+        ], "persample_entropy_compute must be either 'group' or 'analytical'"
+        assert cb_entropy_compute in [
+            "group",
+            "nce",
+        ], "cb_entropy_compute must be either 'group' or 'nce'"
+        self.persample_entropy_compute = persample_entropy_compute
+        self.cb_entropy_compute = cb_entropy_compute
+        self.l2_norm = l2_norm
+        self.inv_temperature = inv_temperature
+
+        self.register_buffer("basis", 2 ** torch.arange(embed_dim - 1, -1, -1))
+        self.register_buffer("group_basis", 2 ** torch.arange(group_size - 1, -1, -1))
+
+        self.num_dimensions = 2**embed_dim
+        self.bits_per_index = embed_dim
+
+        # we only need to keep the codebook portion up to the group size
+        # because we approximate the H loss with this subcode
+        group_codes = torch.arange(2**self.group_size)
+        group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]
+        self.register_buffer("group_codebook", group_codebook, persistent=False)
+
+        self.soft_entropy = soft_entropy  # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf
+
+    def quantize(self, z):
+        assert (
+            z.shape[-1] == self.embed_dim
+        ), f"Expected {self.embed_dim} dimensions, got {z.shape[-1]}"
+
+        zhat = torch.where(
+            z > 0,
+            torch.tensor(1, dtype=z.dtype, device=z.device),
+            torch.tensor(-1, dtype=z.dtype, device=z.device),
+        )
+        return z + (zhat - z).detach()
+
+    def forward(self, z):
+        if self.input_format == "bchw":
+            z = rearrange(z, "b c h w -> b h w c")
+        zq = self.quantize(z)
+
+        indices = self.codes_to_indexes(zq.detach())
+        group_indices = self.codes_to_group_indexes(zq.detach())
+        if not self.training:
+            used_codes = torch.unique(indices, return_counts=False)
+        else:
+            used_codes = None
+
+        q_scale = 1.0 / (self.embed_dim**0.5) if self.l2_norm else 1.0
+
+        if self.soft_entropy:
+            persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)
+            entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
+        else:
+            zb_by_sample = (
+                ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)
+            )
+            persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)
+            cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)
+            entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
+
+        zq = zq * q_scale
+
+        # commit loss
+        # commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))
+        commit_loss = torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))
+
+        if self.input_format == "bchw":
+            zq = rearrange(zq, "b h w c -> b c h w")
+
+        return (
+            zq,
+            commit_loss * self.beta
+            + self.zeta * entropy_penalty / self.inv_temperature,
+            {
+                "H": cb_entropy,
+                "used_codes": used_codes,
+                "indices": indices,
+                "group_indices": group_indices,
+                "avg_prob": avg_prob,
+                "commit_loss": commit_loss,
+            },
+        )
+
+    def soft_entropy_loss(self, z):
+        # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size
+        # the sub-code is the last group_size bits of the full code
+        group_code_book = self.group_codebook / (
+            self.embed_dim**0.5 if self.l2_norm else 1
+        )
+        divided_z = rearrange(z, "... (g c) -> ... g c", c=self.group_size)
+
+        # we calculate the distance between the divided_z and the codebook for each subgroup
+        distance = -2 * torch.einsum(
+            "... g c, d c ->... g d", divided_z, group_code_book
+        )
+        prob = (-distance * self.inv_temperature).softmax(dim=-1)
+        if self.persample_entropy_compute == "analytical":
+            if self.l2_norm:
+                p = torch.sigmoid(-4 * z / (self.embed_dim**0.5) * self.inv_temperature)
+            else:
+                p = torch.sigmoid(-4 * z * self.inv_temperature)
+            prob = torch.stack([p, 1 - p], dim=-1)
+            per_sample_entropy = (
+                self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
+            )
+        else:
+            per_sample_entropy = (
+                self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
+            )
+
+        # macro average of the probability of each subgroup
+        avg_prob = reduce(prob, "... g d ->g d", "mean")
+        codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)
+
+        # the approximation of the entropy is the sum of the entropy of each subgroup
+        return per_sample_entropy, codebook_entropy.sum(), avg_prob
+
+    def get_hard_per_sample_entropy(self, zb_by_sample):
+        probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]
+        persample_entropy = -probs_per_dim * torch.log(probs_per_dim + 1e-8) - (
+            1 - probs_per_dim
+        ) * torch.log(1 - probs_per_dim + 1e-8)
+        persample_entropy = persample_entropy.sum(-1)
+        return persample_entropy.mean()
+
+    def codes_to_indexes(self, zhat):
+        """Converts a `code` to an index in the codebook.
+        Args:
+            zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
+        """
+        assert (
+            zhat.shape[-1] == self.embed_dim
+        ), f"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}"
+        return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)
+
+    def codes_to_group_indexes(self, zhat):
+        """Converts a `code` to a list of indexes (in groups) in the codebook.
+        Args:
+            zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
+        """
+        zhat_in_group = rearrange(zhat, "b ... (g c) -> b ... g c", c=self.group_size)
+        return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)
+
+    def indexes_to_codes(self, indices):
+        """Inverse of `indexes_to_codes`."""
+        indices = indices.unsqueeze(-1)
+        codes_non_centered = torch.remainder(torch.floor_divide(indices, self.basis), 2)
+        return codes_non_centered * 2 - 1
+
+    def group_indexes_to_codes(self, group_indices):
+        """Inverse of `group_indexes_to_codes`."""
+        group_indices = group_indices.unsqueeze(-1)
+        codes_non_centered = torch.remainder(
+            torch.floor_divide(group_indices, self.group_basis), 2
+        )
+        codes_non_centered = rearrange(codes_non_centered, "b ... g c -> b ... (g c)")
+        return codes_non_centered * 2 - 1
+
+    def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):
+        if normalize:
+            probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)
+        else:
+            probs = count
+        H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)
+        return H
+
+    def get_group_codebook_entry(self, group_indices):
+        z_q = self.group_indexes_to_codes(group_indices)
+        q_scale = 1.0 / (self.embed_dim**0.5) if self.l2_norm else 1.0
+        z_q = z_q * q_scale
+        if self.input_format == "bchw":
+            h, w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], "Invalid sequence length"
+            z_q = rearrange(z_q, "b (h w) c -> b c h w", h=h)
+        return z_q
+
+    def get_codebook_entry(self, indices):
+        z_q = self.indexes_to_codes(indices)
+        q_scale = 1.0 / (self.embed_dim**0.5) if self.l2_norm else 1.0
+        z_q = z_q * q_scale
+        if self.input_format == "bchw":
+            h, w = int(z_q.shape[1] ** 0.5)
+            assert h * w == z_q.shape[1], "Invalid sequence length"
+            z_q = rearrange(z_q, "b (h w) c -> b c h w", h=h)
+        return z_q
+
+
+if __name__ == "__main__":
+    bsq = BinarySphericalQuantizer()
+    z = torch.randn(3, 20, 14)
+    zq, loss, info = bsq(z)
+    print(zq.shape, loss, info)

models/codec/amphion_codec/quantize/factorized_vector_quantize.py

@@ -0,0 +1,145 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class FactorizedVectorQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+        commitment=0.005,
+        codebook_loss_weight=1.0,
+        use_l2_normlize=True,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.commitment = commitment
+        self.codebook_loss_weight = codebook_loss_weight
+        self.use_l2_normlize = use_l2_normlize
+
+        if self.input_dim != self.codebook_dim:
+            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
+            self.out_project = WNConv1d(
+                self.codebook_dim, self.input_dim, kernel_size=1
+            )
+
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+
+        self.codebook = nn.Embedding(self.codebook_size, self.codebook_dim)
+
+    def forward(self, z):
+        """
+        Parameters
+        ----------
+        z: torch.Tensor[B x D x T]
+
+        Returns
+        -------
+        z_q: torch.Tensor[B x D x T]
+            Quantized continuous representation of input
+        commit_loss: Tensor[B]
+            Commitment loss to train encoder to predict vectors closer to codebook entries
+        codebook_loss: Tensor[B]
+            Codebook loss to update the codebook
+        indices: torch.Tensor[B x T]
+            Codebook indices (quantized discrete representation of input)
+        z_e: torch.Tensor[B x D x T]
+            Projected latents (continuous representation of input before quantization)
+        """
+
+        # Factorized codes project input into low-dimensional space if self.input_dim != self.codebook_dim
+        z_e = self.in_project(z)
+        z_q, indices = self.decode_latents(z_e)
+
+        # Compute commitment loss and codebook loss
+        if self.training:
+            commit_loss = (
+                F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+                * self.commitment
+            )
+            codebook_loss = (
+                F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+                * self.codebook_loss_weight
+            )
+        else:
+            commit_loss = torch.zeros(z.shape[0], device=z.device)
+            codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        z_q = z_e + (z_q - z_e).detach()
+
+        z_q = self.out_project(z_q)
+
+        return z_q, commit_loss, codebook_loss, indices, z_e
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight
+
+        # L2 normalize encodings and codebook
+        if self.use_l2_normlize:
+            encodings = F.normalize(encodings)
+            codebook = F.normalize(codebook)
+
+        # Compute euclidean distance between encodings and codebook,
+        # if use_l2_normlize is True, the distance is equal to cosine distance
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+
+        return z_q, indices
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = self.decode_code(vq)
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb
+
+    def latent2dist(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight
+
+        # L2 normalize encodings and codebook
+        if self.use_l2_normlize:
+            encodings = F.normalize(encodings)
+            codebook = F.normalize(codebook)
+
+        # Compute euclidean distance between encodings and codebook,
+        # if use_l2_normlize is True, the distance is equal to cosine distance
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )  # (b*t, k)
+
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        dist = rearrange(dist, "(b t) k -> b t k", b=latents.size(0))
+        z_q = self.decode_code(indices)
+
+        return -dist, indices, z_q

models/codec/amphion_codec/quantize/lookup_free_quantize.py

@@ -0,0 +1,72 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class LookupFreeQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        assert 2**codebook_dim == codebook_size
+
+        if self.input_dim != self.codebook_dim:
+            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
+            self.out_project = WNConv1d(
+                self.codebook_dim, self.input_dim, kernel_size=1
+            )
+
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+
+    def forward(self, z):
+        z_e = self.in_project(z)
+        z_e = F.sigmoid(z_e)
+
+        z_q = z_e + (torch.round(z_e) - z_e).detach()
+
+        z_q = self.out_project(z_q)
+
+        commit_loss = torch.zeros(z.shape[0], device=z.device)
+        codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        bits = (
+            2
+            ** torch.arange(self.codebook_dim, device=z.device)
+            .unsqueeze(0)
+            .unsqueeze(-1)
+            .long()
+        )  # (1, d, 1)
+        indices = (torch.round(z_e.clone().detach()).long() * bits).sum(1).long()
+
+        return z_q, commit_loss, codebook_loss, indices, z_e
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = torch.zeros(
+            vq.shape[0], self.codebook_dim, vq.shape[-1], device=vq.device
+        )  # (B, d, T)
+        for i in range(self.codebook_dim):
+            emb[:, i, :] = (vq % 2).float()
+            vq = vq // 2
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb

models/codec/amphion_codec/quantize/residual_vq.py

@@ -0,0 +1,172 @@
+from typing import Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+from models.codec.amphion_codec.quantize.factorized_vector_quantize import (
+    FactorizedVectorQuantize,
+)
+from models.codec.amphion_codec.quantize.vector_quantize import VectorQuantize
+from models.codec.amphion_codec.quantize.lookup_free_quantize import LookupFreeQuantize
+
+
+class ResidualVQ(nn.Module):
+    """
+    Introduced in SoundStream: An end2end neural audio codec
+    https://arxiv.org/abs/2107.03312
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 256,
+        num_quantizers: int = 8,
+        codebook_size: int = 1024,
+        codebook_dim: int = 256,
+        quantizer_type: str = "vq",  # "vq" or "fvq" or "lfq"
+        quantizer_dropout: float = 0.5,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.num_quantizers = num_quantizers
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.quantizer_type = quantizer_type
+        self.quantizer_dropout = quantizer_dropout
+
+        if quantizer_type == "vq":
+            VQ = VectorQuantize
+        elif quantizer_type == "fvq":
+            VQ = FactorizedVectorQuantize
+        elif quantizer_type == "lfq":
+            VQ = LookupFreeQuantize
+        else:
+            raise ValueError(f"Unknown quantizer type {quantizer_type}")
+
+        self.quantizers = nn.ModuleList(
+            [
+                VQ(
+                    input_dim=input_dim,
+                    codebook_size=codebook_size,
+                    codebook_dim=codebook_dim,
+                    **kwargs,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+
+    def forward(self, z, n_quantizers: int = None):
+        """
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+        n_quantizers : int, optional
+            No. of quantizers to use
+            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
+            Note: if `self.quantizer_dropout` is True, this argument is ignored
+                when in training mode, and a random number of quantizers is used.
+        Returns
+        -------
+        "quantized_out" : Tensor[B x D x T]
+            Quantized continuous representation of input
+        "all_indices" : Tensor[N x B x T]
+            Codebook indices for each codebook
+            (quantized discrete representation of input)
+        "all_commit_losses" : Tensor[N]
+        "all_codebook_losses" : Tensor[N]
+        "all_quantized" : Tensor[N x B x D x T]
+        """
+
+        quantized_out = 0.0
+        residual = z
+
+        all_commit_losses = []
+        all_codebook_losses = []
+        all_indices = []
+        all_quantized = []
+
+        if n_quantizers is None:
+            n_quantizers = self.num_quantizers
+
+        if self.training:
+            n_quantizers = torch.ones((z.shape[0],)) * self.num_quantizers + 1
+            dropout = torch.randint(1, self.num_quantizers + 1, (z.shape[0],))
+            n_dropout = int(z.shape[0] * self.quantizer_dropout)
+            n_quantizers[:n_dropout] = dropout[:n_dropout]
+            n_quantizers = n_quantizers.to(z.device)
+
+        for i, quantizer in enumerate(self.quantizers):
+            if self.training is False and i >= n_quantizers:
+                break
+
+            z_q_i, commit_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+                residual
+            )
+
+            # Create mask to apply quantizer dropout
+            mask = (
+                torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
+            )
+            quantized_out = quantized_out + z_q_i * mask[:, None, None]
+            residual = residual - z_q_i
+
+            commit_loss_i = (commit_loss_i * mask).mean()
+            codebook_loss_i = (codebook_loss_i * mask).mean()
+
+            all_commit_losses.append(commit_loss_i)
+            all_codebook_losses.append(codebook_loss_i)
+            all_indices.append(indices_i)
+            all_quantized.append(z_q_i)
+
+        all_commit_losses, all_codebook_losses, all_indices, all_quantized = map(
+            torch.stack,
+            (all_commit_losses, all_codebook_losses, all_indices, all_quantized),
+        )
+
+        return (
+            quantized_out,
+            all_indices,
+            all_commit_losses,
+            all_codebook_losses,
+            all_quantized,
+        )
+
+    def vq2emb(self, vq, n_quantizers=None):
+        quantized_out = 0.0
+        if n_quantizers is None:
+            n_quantizers = self.num_quantizers
+        for idx, quantizer in enumerate(self.quantizers):
+            if idx >= n_quantizers:
+                break
+            quantized_out += quantizer.vq2emb(vq[idx])
+        return quantized_out
+
+    def latent2dist(self, z, n_quantizers=None):
+        quantized_out = 0.0
+        residual = z
+
+        all_dists = []
+        all_indices = []
+
+        if n_quantizers is None:
+            n_quantizers = self.num_quantizers
+
+        for i, quantizer in enumerate(self.quantizers):
+            if self.training is False and i >= n_quantizers:
+                break
+            dist_i, indices_i, z_q_i = quantizer.latent2dist(residual)
+            all_dists.append(dist_i)
+            all_indices.append(indices_i)
+
+            quantized_out = quantized_out + z_q_i
+            residual = residual - z_q_i
+
+        all_dists = torch.stack(all_dists)
+        all_indices = torch.stack(all_indices)
+
+        return all_dists, all_indices

models/codec/amphion_codec/quantize/vector_quantize.py

@@ -0,0 +1,396 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+def ema_inplace(moving_avg, new, decay):
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def laplace_smoothing(x, n_categories, eps=1e-5):
+    return (x + eps) / (x.sum() + n_categories * eps)
+
+
+def sample_vectors(samples, num):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+def kmeans(samples, num_clusters, num_iters=10, use_cosine_sim=False):
+    dim, dtype, device = samples.shape[-1], samples.dtype, samples.device
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        if use_cosine_sim:
+            dists = samples @ means.t()
+        else:
+            diffs = rearrange(samples, "n d -> n () d") - rearrange(
+                means, "c d -> () c d"
+            )
+            dists = -(diffs**2).sum(dim=-1)
+
+        buckets = dists.max(dim=-1).indices
+        bins = torch.bincount(buckets, minlength=num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
+        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        if use_cosine_sim:
+            new_means = l2norm(new_means)
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+class EuclideanCodebook(nn.Module):
+    def __init__(
+        self,
+        dim,
+        codebook_size,
+        kmeans_init=False,
+        kmeans_iters=10,
+        decay=0.8,
+        eps=1e-5,
+        threshold_ema_dead_code=2,
+        weight_init=False,
+    ):
+        super().__init__()
+
+        self.decay = decay
+        init_fn = torch.randn if not weight_init else torch.zeros
+        embed = init_fn(codebook_size, dim)
+
+        if weight_init:
+            nn.init.uniform_(embed, -1 / codebook_size, 1 / codebook_size)
+
+        self.codebook_size = codebook_size
+        self.kmeans_iters = kmeans_iters
+        self.eps = eps
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+
+        self.register_buffer(
+            "initted", torch.Tensor([not kmeans_init])
+        )  # if kmeans_init is True, then initted is False; otherwise, initted is True
+        self.register_buffer("cluster_size", torch.zeros(codebook_size))
+        self.register_buffer("embed", embed)
+        self.register_buffer("embed_avg", embed.clone())
+
+    def init_embed_(self, data):
+        embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
+        self.embed.data.copy_(embed)
+        self.embed_avg.data.copy_(embed)
+        self.cluster_size.data.copy_(cluster_size)
+        self.initted.data.copy_(torch.Tensor([True]))
+
+    def replace(self, samples, mask):
+        modified_codebook = torch.where(
+            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
+        )
+        self.embed.data.copy_(modified_codebook)
+
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            return
+
+        expired_codes = self.cluster_size < self.threshold_ema_dead_code
+        if not torch.any(expired_codes):
+            return
+        batch_samples = rearrange(batch_samples, "... d -> (...) d")
+        self.replace(batch_samples, mask=expired_codes)
+
+    def forward(self, x):
+        shape, dtype = x.shape, x.dtype
+        flatten = rearrange(x, "... d -> (...) d")
+        embed = self.embed.t()  # (codebook_size, dim) -> (dim, codebook_size)
+
+        if not self.initted:
+            self.init_embed_(flatten)
+
+        dist = -(
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+
+        embed_ind = dist.max(dim=-1).indices
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
+        embed_ind = embed_ind.view(*shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed)
+
+        if self.training:
+            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
+            embed_sum = (
+                flatten.t() @ embed_onehot
+            )  # (dim, ...) @ (..., codebook_size) -> (dim, codebook_size)
+            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
+            cluster_size = (
+                laplace_smoothing(self.cluster_size, self.codebook_size, self.eps)
+                * self.cluster_size.sum()
+            )
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embed.data.copy_(embed_normalized)
+            self.expire_codes_(x)
+
+        return quantize, embed_ind
+
+    def vq2emb(self, vq):
+        quantize = F.embedding(vq, self.embed)
+        return quantize
+
+    def latent2dist(self, x):
+        shape, dtype = x.shape, x.dtype
+        flatten = rearrange(x, "... d -> (...) d")
+        embed = self.embed.t()  # (codebook_size, dim) -> (dim, codebook_size)
+
+        if not self.initted:
+            self.init_embed_(flatten)
+
+        dist = -(
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+
+        embed_ind = dist.max(dim=-1).indices
+        embed_ind = embed_ind.view(*shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed)
+
+        dist = dist.view(*shape[:-1], -1)
+
+        return dist, embed_ind, quantize
+
+
+class SimpleCodebook(nn.Module):
+    def __init__(
+        self,
+        dim,
+        codebook_size,
+        use_l2_normlize=False,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.codebook_size = codebook_size
+        self.use_l2_normlize = use_l2_normlize
+
+        self.embed = nn.Embedding(self.codebook_size, self.dim)
+
+    def forward(self, x):
+        shape, dtype = x.shape, x.dtype
+        flatten = rearrange(x, "... d -> (...) d")
+        embed = self.embed.weight.t()  # (codebook_size, dim) -> (dim, codebook_size)
+
+        if self.use_l2_normlize:
+            flatten = F.normalize(flatten)
+            embed = F.normalize(embed)
+
+        dist = -(
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+
+        embed_ind = dist.max(dim=-1).indices
+        embed_ind = embed_ind.view(*shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed)
+
+        return quantize, embed_ind
+
+    def vq2emb(self, vq):
+        quantize = F.embedding(vq, self.embed.weight)
+        return quantize
+
+    def latent2dist(self, x):
+        shape, dtype = x.shape, x.dtype
+        flatten = rearrange(x, "... d -> (...) d")
+        embed = self.embed.weight.t()  # (codebook_size, dim) -> (dim, codebook_size)
+
+        if self.use_l2_normlize:
+            flatten = F.normalize(flatten)
+            embed = F.normalize(embed)
+
+        dist = -(
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+
+        embed_ind = dist.max(dim=-1).indices
+        embed_ind = embed_ind.view(*shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed)
+
+        dist = dist.view(*shape[:-1], -1)
+
+        return dist, embed_ind, quantize
+
+
+class VectorQuantize(nn.Module):
+    """Vector quantization and factorized vecotor quantization implementation
+    Args:
+        input_dim (int): Dimension of input.
+        codebook_size (int): Codebook size.
+        codebook_dim (int): Codebook dimension. We suggest use codebook_dim = input_dim
+            if use codebook_type == "euclidean", otherwise, if you want to use
+            factorized vector quantization, use codebook_dim as small number (e.g. 8 or 32).
+        commitment (float): Weight for commitment loss.
+        use_l2_normlize (bool): Whether to use l2 normlized codes for factorized vecotor quantization,
+            we suggest use it as True if you want to use factorized vector quantization
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+        commitment=0.005,
+        codebook_loss_weight=1.0,
+        use_l2_normlize=False,
+        codebook_type="euclidean",  # "euclidean" or "simple"
+        kmeans_init=False,
+        kmeans_iters=10,
+        decay=0.8,
+        eps=1e-5,
+        threshold_ema_dead_code=2,
+        weight_init=False,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.commitment = commitment
+        self.codebook_loss_weight = codebook_loss_weight
+        self.use_l2_normlize = use_l2_normlize
+        self.codebook_type = codebook_type
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.decay = decay
+        self.eps = eps
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.weight_init = weight_init
+
+        if self.input_dim != self.codebook_dim:
+            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
+            self.out_project = WNConv1d(
+                self.codebook_dim, self.input_dim, kernel_size=1
+            )
+
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+
+        if self.codebook_type == "euclidean":
+            self.codebook = EuclideanCodebook(
+                self.codebook_dim,
+                codebook_size=self.codebook_size,
+                kmeans_init=self.kmeans_init,
+                kmeans_iters=self.kmeans_iters,
+                decay=self.decay,
+                eps=self.eps,
+                threshold_ema_dead_code=self.threshold_ema_dead_code,
+                weight_init=self.weight_init,
+            )
+        elif self.codebook_type == "simple":
+            self.codebook = SimpleCodebook(
+                self.codebook_dim,
+                codebook_size=self.codebook_size,
+                use_l2_normlize=self.use_l2_normlize,
+            )
+        else:
+            raise NotImplementedError(
+                f"codebook_type {self.codebook_type} is not implemented!"
+            )
+
+    def forward(self, z):
+        """
+        Parameters
+        ----------
+        z: torch.Tensor[B x D x T]
+
+        Returns
+        -------
+        z_q: torch.Tensor[B x D x T]
+            Quantized continuous representation of input
+        commit_loss: Tensor[B]
+            Commitment loss to train encoder to predict vectors closer to codebook entries
+        codebook_loss: Tensor[B]
+            Codebook loss to update the codebook
+        indices: torch.Tensor[B x T]
+            Codebook indices (quantized discrete representation of input)
+        z_e: torch.Tensor[B x D x T]
+            Projected latents (continuous representation of input before quantization)
+        """
+
+        # Factorized codes project input into low-dimensional space if self.input_dim != self.codebook_dim
+        z_e = self.in_project(z)
+        z_q, indices = self.decode_latents(z_e)
+
+        # Compute commitment loss and codebook loss
+        if self.training:
+            commit_loss = (
+                F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+                * self.commitment
+            )
+            codebook_loss = (
+                F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+                * self.codebook_loss_weight
+            )
+        else:
+            commit_loss = torch.zeros(z.shape[0], device=z.device)
+            codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        z_q = z_e + (z_q - z_e).detach()
+
+        z_q = self.out_project(z_q)
+
+        return z_q, commit_loss, codebook_loss, indices, z_e
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> b t d")
+        z_q, indices = self.codebook(encodings)
+        z_q = z_q.transpose(1, 2)
+        return z_q, indices
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = self.codebook.vq2emb(vq)
+        emb = emb.transpose(1, 2)
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb
+
+    def latent2dist(self, latents):
+        latents = rearrange(latents, "b d t -> b t d")
+        dist, embed_ind, quantize = self.codebook.latent2dist(latents)
+        return dist, embed_ind, quantize.transpose(1, 2)

models/codec/amphion_codec/vocos.py

@@ -0,0 +1,909 @@
+from typing import Optional, Tuple
+
+import numpy as np
+import scipy
+import torch
+from torch import nn, view_as_real, view_as_complex
+from torch import nn
+from torch.nn.utils import weight_norm, remove_weight_norm
+from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz
+from models.codec.melvqgan.melspec import MelSpectrogram
+import librosa
+
+
+def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor:
+    """
+    Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values.
+
+    Args:
+        x (Tensor): Input tensor.
+        clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7.
+
+    Returns:
+        Tensor: Element-wise logarithm of the input tensor with clipping applied.
+    """
+    return torch.log(torch.clip(x, min=clip_val))
+
+
+def symlog(x: torch.Tensor) -> torch.Tensor:
+    return torch.sign(x) * torch.log1p(x.abs())
+
+
+def symexp(x: torch.Tensor) -> torch.Tensor:
+    return torch.sign(x) * (torch.exp(x.abs()) - 1)
+
+
+class STFT(nn.Module):
+    def __init__(
+        self,
+        n_fft: int,
+        hop_length: int,
+        win_length: int,
+        center=True,
+    ):
+        super().__init__()
+        self.center = center
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, T * hop_length)
+
+        if not self.center:
+            pad = self.win_length - self.hop_length
+            x = torch.nn.functional.pad(x, (pad // 2, pad // 2), mode="reflect")
+
+        stft_spec = torch.stft(
+            x,
+            self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=self.window,
+            center=self.center,
+            return_complex=False,
+        )  # (B, n_fft // 2 + 1, T, 2)
+
+        rea = stft_spec[:, :, :, 0]  # (B, n_fft // 2 + 1, T, 2)
+        imag = stft_spec[:, :, :, 1]  # (B, n_fft // 2 + 1, T, 2)
+
+        log_mag = torch.log(
+            torch.abs(torch.sqrt(torch.pow(rea, 2) + torch.pow(imag, 2))) + 1e-5
+        )  # (B, n_fft // 2 + 1, T)
+        phase = torch.atan2(imag, rea)  # (B, n_fft // 2 + 1, T)
+
+        return log_mag, phase
+
+
+class ISTFT(nn.Module):
+    """
+    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+    See issue: https://github.com/pytorch/pytorch/issues/62323
+    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+    The NOLA constraint is met as we trim padded samples anyway.
+
+    Args:
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames.
+        win_length (int): The size of window frame and STFT filter.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(
+        self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"
+    ):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+        Args:
+            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+                            N is the number of frequency bins, and T is the number of time frames.
+
+        Returns:
+            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
+        """
+        if self.padding == "center":
+            # Fallback to pytorch native implementation
+            return torch.istft(
+                spec,
+                self.n_fft,
+                self.hop_length,
+                self.win_length,
+                self.window,
+                center=True,
+            )
+        elif self.padding == "same":
+            pad = (self.win_length - self.hop_length) // 2
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        assert spec.dim() == 3, "Expected a 3D tensor as input"
+        B, N, T = spec.shape
+
+        # Inverse FFT
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        # Overlap and Add
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        )[:, 0, 0, pad:-pad]
+
+        # Window envelope
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = torch.nn.functional.fold(
+            window_sq,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        ).squeeze()[pad:-pad]
+
+        # Normalize
+        assert (window_envelope > 1e-11).all()
+        y = y / window_envelope
+
+        return y
+
+
+class MDCT(nn.Module):
+    """
+    Modified Discrete Cosine Transform (MDCT) module.
+
+    Args:
+        frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, frame_len: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.frame_len = frame_len
+        N = frame_len // 2
+        n0 = (N + 1) / 2
+        window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
+        self.register_buffer("window", window)
+
+        pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len)
+        post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N)
+        # view_as_real: NCCL Backend does not support ComplexFloat data type
+        # https://github.com/pytorch/pytorch/issues/71613
+        self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
+        self.register_buffer("post_twiddle", view_as_real(post_twiddle))
+
+    def forward(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the Modified Discrete Cosine Transform (MDCT) to the input audio.
+
+        Args:
+            audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size
+                and T is the length of the audio.
+
+        Returns:
+            Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames
+                and N is the number of frequency bins.
+        """
+        if self.padding == "center":
+            audio = torch.nn.functional.pad(
+                audio, (self.frame_len // 2, self.frame_len // 2)
+            )
+        elif self.padding == "same":
+            # hop_length is 1/2 frame_len
+            audio = torch.nn.functional.pad(
+                audio, (self.frame_len // 4, self.frame_len // 4)
+            )
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        x = audio.unfold(-1, self.frame_len, self.frame_len // 2)
+        N = self.frame_len // 2
+        x = x * self.window.expand(x.shape)
+        X = torch.fft.fft(
+            x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1
+        )[..., :N]
+        res = X * view_as_complex(self.post_twiddle).expand(X.shape) * np.sqrt(1 / N)
+        return torch.real(res) * np.sqrt(2)
+
+
+class IMDCT(nn.Module):
+    """
+    Inverse Modified Discrete Cosine Transform (IMDCT) module.
+
+    Args:
+        frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, frame_len: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.frame_len = frame_len
+        N = frame_len // 2
+        n0 = (N + 1) / 2
+        window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
+        self.register_buffer("window", window)
+
+        pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N)
+        post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2))
+        self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
+        self.register_buffer("post_twiddle", view_as_real(post_twiddle))
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients.
+
+        Args:
+            X (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size,
+                L is the number of frames, and N is the number of frequency bins.
+
+        Returns:
+            Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio.
+        """
+        B, L, N = X.shape
+        Y = torch.zeros((B, L, N * 2), dtype=X.dtype, device=X.device)
+        Y[..., :N] = X
+        Y[..., N:] = -1 * torch.conj(torch.flip(X, dims=(-1,)))
+        y = torch.fft.ifft(
+            Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1
+        )
+        y = (
+            torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape))
+            * np.sqrt(N)
+            * np.sqrt(2)
+        )
+        result = y * self.window.expand(y.shape)
+        output_size = (1, (L + 1) * N)
+        audio = torch.nn.functional.fold(
+            result.transpose(1, 2),
+            output_size=output_size,
+            kernel_size=(1, self.frame_len),
+            stride=(1, self.frame_len // 2),
+        )[:, 0, 0, :]
+
+        if self.padding == "center":
+            pad = self.frame_len // 2
+        elif self.padding == "same":
+            pad = self.frame_len // 4
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        audio = audio[:, pad:-pad]
+        return audio
+
+
+class FourierHead(nn.Module):
+    """Base class for inverse fourier modules."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class ISTFTHead(FourierHead):
+    """
+    ISTFT Head module for predicting STFT complex coefficients.
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames, which should align with
+                          the resolution of the input features.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
+        super().__init__()
+        out_dim = n_fft + 2
+        self.out = torch.nn.Linear(dim, out_dim)
+        self.istft = ISTFT(
+            n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the ISTFTHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x).transpose(1, 2)
+        mag, p = x.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(
+            mag, max=1e2
+        )  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value
+        S = mag * (x + 1j * y)
+        audio = self.istft(S)
+        return audio
+
+
+class IMDCTSymExpHead(FourierHead):
+    """
+    IMDCT Head module for predicting MDCT coefficients with symmetric exponential function
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        mdct_frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+        sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized
+                                     based on perceptual scaling. Defaults to None.
+        clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        mdct_frame_len: int,
+        padding: str = "same",
+        sample_rate: Optional[int] = None,
+        clip_audio: bool = False,
+    ):
+        super().__init__()
+        out_dim = mdct_frame_len // 2
+        self.out = nn.Linear(dim, out_dim)
+        self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
+        self.clip_audio = clip_audio
+
+        if sample_rate is not None:
+            # optionally init the last layer following mel-scale
+            m_max = _hz_to_mel(sample_rate // 2)
+            m_pts = torch.linspace(0, m_max, out_dim)
+            f_pts = _mel_to_hz(m_pts)
+            scale = 1 - (f_pts / f_pts.max())
+
+            with torch.no_grad():
+                self.out.weight.mul_(scale.view(-1, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the IMDCTSymExpHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x)
+        x = symexp(x)
+        x = torch.clip(
+            x, min=-1e2, max=1e2
+        )  # safeguard to prevent excessively large magnitudes
+        audio = self.imdct(x)
+        if self.clip_audio:
+            audio = torch.clip(x, min=-1.0, max=1.0)
+
+        return audio
+
+
+class IMDCTCosHead(FourierHead):
+    """
+    IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p)
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        mdct_frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+        clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        mdct_frame_len: int,
+        padding: str = "same",
+        clip_audio: bool = False,
+    ):
+        super().__init__()
+        self.clip_audio = clip_audio
+        self.out = nn.Linear(dim, mdct_frame_len)
+        self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the IMDCTCosHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x)
+        m, p = x.chunk(2, dim=2)
+        m = torch.exp(m).clip(
+            max=1e2
+        )  # safeguard to prevent excessively large magnitudes
+        audio = self.imdct(m * torch.cos(p))
+        if self.clip_audio:
+            audio = torch.clip(x, min=-1.0, max=1.0)
+        return audio
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+    Args:
+        dim (int): Number of input channels.
+        intermediate_dim (int): Dimensionality of the intermediate layer.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+            None means non-conditional LayerNorm. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        layer_scale_init_value: float,
+        adanorm_num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(
+            dim, dim, kernel_size=7, padding=3, groups=dim
+        )  # depthwise conv
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(
+            dim, intermediate_dim
+        )  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(
+        self, x: torch.Tensor, cond_embedding_id: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+        if self.adanorm:
+            assert cond_embedding_id is not None
+            x = self.norm(x, cond_embedding_id)
+        else:
+            x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+
+        x = residual + x
+        return x
+
+
+class AdaLayerNorm(nn.Module):
+    """
+    Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes
+
+    Args:
+        num_embeddings (int): Number of embeddings.
+        embedding_dim (int): Dimension of the embeddings.
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.dim = embedding_dim
+        self.scale = nn.Embedding(
+            num_embeddings=num_embeddings, embedding_dim=embedding_dim
+        )
+        self.shift = nn.Embedding(
+            num_embeddings=num_embeddings, embedding_dim=embedding_dim
+        )
+        torch.nn.init.ones_(self.scale.weight)
+        torch.nn.init.zeros_(self.shift.weight)
+
+    def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor:
+        scale = self.scale(cond_embedding_id)
+        shift = self.shift(cond_embedding_id)
+        x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps)
+        x = x * scale + shift
+        return x
+
+
+class ResBlock1(nn.Module):
+    """
+    ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions,
+    but without upsampling layers.
+
+    Args:
+        dim (int): Number of input channels.
+        kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3.
+        dilation (tuple[int], optional): Dilation factors for the dilated convolutions.
+            Defaults to (1, 3, 5).
+        lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function.
+            Defaults to 0.1.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        kernel_size: int = 3,
+        dilation: Tuple[int, int, int] = (1, 3, 5),
+        lrelu_slope: float = 0.1,
+        layer_scale_init_value: Optional[float] = None,
+    ):
+        super().__init__()
+        self.lrelu_slope = lrelu_slope
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=self.get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=self.get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=self.get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=self.get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=self.get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=self.get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+
+        self.gamma = nn.ParameterList(
+            [
+                (
+                    nn.Parameter(
+                        layer_scale_init_value * torch.ones(dim, 1), requires_grad=True
+                    )
+                    if layer_scale_init_value is not None
+                    else None
+                ),
+                (
+                    nn.Parameter(
+                        layer_scale_init_value * torch.ones(dim, 1), requires_grad=True
+                    )
+                    if layer_scale_init_value is not None
+                    else None
+                ),
+                (
+                    nn.Parameter(
+                        layer_scale_init_value * torch.ones(dim, 1), requires_grad=True
+                    )
+                    if layer_scale_init_value is not None
+                    else None
+                ),
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma):
+            xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope)
+            xt = c1(xt)
+            xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope)
+            xt = c2(xt)
+            if gamma is not None:
+                xt = gamma * xt
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+    @staticmethod
+    def get_padding(kernel_size: int, dilation: int = 1) -> int:
+        return int((kernel_size * dilation - dilation) / 2)
+
+
+class Backbone(nn.Module):
+    """Base class for the generator's backbone. It preserves the same temporal resolution across all layers."""
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
+                        C denotes output features, and L is the sequence length.
+
+        Returns:
+            Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
+                    and H denotes the model dimension.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class VocosBackbone(Backbone):
+    """
+    Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+        num_layers (int): Number of ConvNeXtBlock layers.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+                                                None means non-conditional model. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        dim: int,
+        intermediate_dim: int,
+        num_layers: int,
+        layer_scale_init_value: Optional[float] = None,
+        adanorm_num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        layer_scale_init_value = layer_scale_init_value or 1 / num_layers
+        self.convnext = nn.ModuleList(
+            [
+                ConvNeXtBlock(
+                    dim=dim,
+                    intermediate_dim=intermediate_dim,
+                    layer_scale_init_value=layer_scale_init_value,
+                    adanorm_num_embeddings=adanorm_num_embeddings,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        bandwidth_id = kwargs.get("bandwidth_id", None)
+        x = self.embed(x)
+        if self.adanorm:
+            assert bandwidth_id is not None
+            x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id)
+        else:
+            x = self.norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        for conv_block in self.convnext:
+            x = conv_block(x, cond_embedding_id=bandwidth_id)
+        x = self.final_layer_norm(x.transpose(1, 2))
+        return x
+
+
+class VocosResNetBackbone(Backbone):
+    """
+    Vocos backbone module built with ResBlocks.
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        num_blocks (int): Number of ResBlock1 blocks.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        input_channels,
+        dim,
+        num_blocks,
+        layer_scale_init_value=None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.embed = weight_norm(
+            nn.Conv1d(input_channels, dim, kernel_size=3, padding=1)
+        )
+        layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3
+        self.resnet = nn.Sequential(
+            *[
+                ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value)
+                for _ in range(num_blocks)
+            ]
+        )
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        x = self.embed(x)
+        x = self.resnet(x)
+        x = x.transpose(1, 2)
+        return x
+
+
+class Vocos(nn.Module):
+    def __init__(
+        self,
+        input_channels: int = 256,
+        dim: int = 384,
+        intermediate_dim: int = 1152,
+        num_layers: int = 8,
+        n_fft: int = 800,
+        hop_size: int = 200,
+        padding: str = "same",
+        adanorm_num_embeddings=None,
+        cfg=None,
+    ):
+        super().__init__()
+
+        input_channels = (
+            cfg.input_channels
+            if cfg is not None and hasattr(cfg, "input_channels")
+            else input_channels
+        )
+        dim = cfg.dim if cfg is not None and hasattr(cfg, "dim") else dim
+        intermediate_dim = (
+            cfg.intermediate_dim
+            if cfg is not None and hasattr(cfg, "intermediate_dim")
+            else intermediate_dim
+        )
+        num_layers = (
+            cfg.num_layers
+            if cfg is not None and hasattr(cfg, "num_layers")
+            else num_layers
+        )
+        adanorm_num_embeddings = (
+            cfg.adanorm_num_embeddings
+            if cfg is not None and hasattr(cfg, "adanorm_num_embeddings")
+            else adanorm_num_embeddings
+        )
+        n_fft = cfg.n_fft if cfg is not None and hasattr(cfg, "n_fft") else n_fft
+        hop_size = (
+            cfg.hop_size if cfg is not None and hasattr(cfg, "hop_size") else hop_size
+        )
+        padding = (
+            cfg.padding if cfg is not None and hasattr(cfg, "padding") else padding
+        )
+
+        self.backbone = VocosBackbone(
+            input_channels=input_channels,
+            dim=dim,
+            intermediate_dim=intermediate_dim,
+            num_layers=num_layers,
+            adanorm_num_embeddings=adanorm_num_embeddings,
+        )
+        self.head = ISTFTHead(dim, n_fft, hop_size, padding)
+
+    def forward(self, x):
+        x = self.backbone(x)
+        x = self.head(x)
+
+        return x[:, None, :]
+
+
+# if __name__ == "__main__":
+#     vocos_model = Vocos(
+#         input_channels=128,
+#         dim=1024,
+#         intermediate_dim=4096,
+#         adanorm_num_embeddings=None,
+#         num_layers=40,
+#         n_fft=1920,
+#         hop_size=480,
+#         padding="same"
+#     )
+#     mel_model = MelSpectrogram(
+#         1920,
+#         128,
+#         24000,
+#         480,
+#         1920,
+#         0,
+#         12000
+#     )
+#     print(sum(p.numel() for p in vocos_model.parameters())/1e6)
+
+#     speech = librosa.load("/mnt/bn/yuacnwang-speech/dataset/yuanshen/ganyu_en.wav", sr=24000)[0][:36000]
+#     speech = torch.tensor(speech).unsqueeze(0)
+#     print(speech.shape)
+#     mel_feat = mel_model(speech)
+#     print(mel_feat.shape)
+
+#     rec_speech = vocos_model(mel_feat)
+#     print(rec_speech.shape)

models/codec/melvqgan/melspec.py

@@ -0,0 +1,153 @@
+import torch
+import pyworld as pw
+import numpy as np
+import soundfile as sf
+import os
+from torchaudio.functional import pitch_shift
+import librosa
+from librosa.filters import mel as librosa_mel_fn
+import torch.nn as nn
+import torch.nn.functional as F
+import tqdm
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    return np.exp(x) / C
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+    output = dynamic_range_decompression_torch(magnitudes)
+    return output
+
+
+class MelSpectrogram(nn.Module):
+    def __init__(
+        self,
+        n_fft,
+        num_mels,
+        sampling_rate,
+        hop_size,
+        win_size,
+        fmin,
+        fmax,
+        center=False,
+    ):
+        super(MelSpectrogram, self).__init__()
+        self.n_fft = n_fft
+        self.hop_size = hop_size
+        self.win_size = win_size
+        self.sampling_rate = sampling_rate
+        self.num_mels = num_mels
+        self.fmin = fmin
+        self.fmax = fmax
+        self.center = center
+
+        mel_basis = {}
+        hann_window = {}
+
+        mel = librosa_mel_fn(
+            sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
+        )
+        mel_basis = torch.from_numpy(mel).float()
+        hann_window = torch.hann_window(win_size)
+
+        self.register_buffer("mel_basis", mel_basis)
+        self.register_buffer("hann_window", hann_window)
+
+    def forward(self, y):
+        y = torch.nn.functional.pad(
+            y.unsqueeze(1),
+            (
+                int((self.n_fft - self.hop_size) / 2),
+                int((self.n_fft - self.hop_size) / 2),
+            ),
+            mode="reflect",
+        )
+        y = y.squeeze(1)
+        spec = torch.stft(
+            y,
+            self.n_fft,
+            hop_length=self.hop_size,
+            win_length=self.win_size,
+            window=self.hann_window,
+            center=self.center,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=True,
+        )
+        spec = torch.view_as_real(spec)
+
+        spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+        spec = torch.matmul(self.mel_basis, spec)
+        spec = spectral_normalize_torch(spec)
+
+        return spec
+
+
+# if __name__ == "__main__":
+#     mel_model = MelSpectrogram(
+#         sampling_rate=24000,
+#         num_mels=128,
+#         hop_size=480,
+#         n_fft=1920,
+#         win_size=1920,
+#         fmin=0,
+#         fmax=12000,
+#     )
+#     mel_model = mel_model.to("cuda:0")
+
+#     x_len = 0
+#     var = 0
+#     mean = 0
+
+
+#     from utils.util import load_config
+#     cfg = load_config("/storage/wyc/SpeechGeneration/egs/tts/SoundStorm/exp_config_16k_emilia_llama_new_semantic_repcodec_8192_1q_24k.json")
+#     print(cfg)
+#     dataset = SoundStormDataset(AK, SK, bucket_name, cfg=cfg)
+#     print(dataset.__getitem__(0))
+
+#     idx_list = list(range(0, len(dataset)))
+#     np.random.shuffle(idx_list)
+#     for i in tqdm.tqdm(range(10000)):
+#         idx = idx_list[i]
+#         # data_path = dataset[idx]
+#         # speech, _ = librosa.load(data_path, sr=24000)
+#         speech = dataset[idx]["speech"]
+#         speech = torch.tensor(speech).unsqueeze(0).to("cuda")
+#         mel = mel_model(speech)
+#         temp_len = mel.shape[-1]
+#         temp_mean = mel.mean()
+#         temp_var = mel.var()
+
+#         new_mean = (mean * x_len + temp_mean * temp_len) / (x_len + temp_len)
+#         new_var = (var * (x_len - 1) + temp_var * (temp_len - 1) + x_len * (new_mean - mean)**2 + temp_len * (new_mean - temp_mean)**2) / (x_len + temp_len -1)
+
+#         x_len += temp_len
+#         mean = new_mean
+#         var = new_var
+
+#         if i % 100 == 0:
+#             print(mean, var)
+
+#     print(mean)
+#     print(var)

models/tts/llm_tts/chat_template.py

@@ -0,0 +1,96 @@
+def format_chat_prompt_phi3(messages, add_assistant_token=True):
+    """
+    Convert the messages list into the phi-3 chat template format.
+
+    Args:
+        messages: A list of messages containing role and content.
+
+    Returns:
+        str: The formatted prompt string.
+    """
+    prompt = ""
+    for msg in messages:
+        role = msg["role"]
+        content = msg["content"]
+        # Add corresponding tags for system and user messages
+        if role in ["system", "user"]:
+            prompt += f"<|{role}|>\n{content}<|end|>\n"
+        # For assistant messages, add only the start tag if it's the last one
+        elif role == "assistant" and msg != messages[-1]:
+            prompt += f"<|{role}|>\n{content}<|end|>\n"
+        elif role == "assistant" and msg == messages[-1]:
+            prompt += f"<|{role}|>\n{content}"
+
+    # If the last message is not from the assistant, add the assistant tag
+    if messages[-1]["role"] != "assistant" and add_assistant_token:
+        prompt += "<|assistant|>"
+    return prompt
+
+
+def format_chat_prompt_qwen2(messages, add_assistant_token=True):
+    """
+    Custom function to format chat prompts without tool-related logic.
+
+    Args:
+        messages: A list of messages containing role and content.
+        add_generation_prompt: Boolean to add a generation prompt at the end.
+
+    Returns:
+        str: The formatted prompt string.
+    """
+    prompt = ""
+
+    if messages[0]["role"] == "system":
+        prompt += f"<|im_start|>system\n{messages[0]['content']}<|im_end|>\n"
+    else:
+        prompt += "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n"
+
+    for message in messages:
+        if (
+            (message["role"] == "user")
+            or (message["role"] == "system" and not message == messages[0])
+            or (message["role"] == "assistant" and not message.get("tool_calls"))
+        ):
+            prompt += f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>\n"
+        elif message["role"] == "assistant":
+            prompt += f"<|im_start|>{message['role']}"
+            if message.get("content"):
+                prompt += f"\n{message['content']}"
+            prompt += "<|im_end|>\n"
+
+    if add_assistant_token:
+        prompt += "<|im_start|>assistant\n"
+
+    return prompt
+
+
+def gen_chat_prompt_for_tts(text, model_name="phi-3", caption=None):
+    if caption is None:
+        template = [
+            {
+                "role": "system",
+                "content": "You are a powerful AI assistant for speech understanding and generation.",
+            },
+            {
+                "role": "user",
+                "content": f"Please speak the following text out loud: {text}",
+            },
+        ]
+    else:
+        template = [
+            {
+                "role": "system",
+                "content": "You are a powerful AI assistant for speech understanding and generation.",
+            },
+            {
+                "role": "user",
+                "content": f"Please follow the caption: <|start_of_caption|>{caption}<|end_of_caption|> and speak the following text out loud: {text}",
+            },
+        ]
+
+    if model_name == "phi-3":
+        return format_chat_prompt_phi3(template)
+    elif model_name == "qwen2":
+        return format_chat_prompt_qwen2(template)
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")

models/tts/llm_tts/inference_llm_tts.py

@@ -0,0 +1,265 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
+from huggingface_hub import snapshot_download
+
+from models.tts.tadicodec.inference_tadicodec import TaDiCodecPipline
+from models.tts.llm_tts.chat_template import gen_chat_prompt_for_tts
+
+
+class TTSInferencePipeline(nn.Module):
+    """
+    TTS inference pipeline that integrates TaDiCodec and LLM models
+    Uses standard LLM for autoregressive generation
+    """
+
+    def __init__(
+        self,
+        tadicodec_path: str,
+        llm_path: str,
+        device: torch.device,
+    ):
+        super().__init__()
+        self.device = device
+        self.llm_path = llm_path
+
+        # Load TaDiCodec pipeline
+        self.tadicodec = TaDiCodecPipline.from_pretrained(
+            ckpt_dir=tadicodec_path, device=device
+        )
+
+        # Load LLM directly from pretrained
+        # Try to use flash attention 2, fallback to default if not available
+        try:
+            self.llm = AutoModelForCausalLM.from_pretrained(
+                llm_path,
+                device_map=device,
+                torch_dtype="auto",
+                trust_remote_code=True,
+                attn_implementation="flash_attention_2",
+            )
+        except Exception as e:
+            print(f"Flash attention 2 not available, using default attention: {e}")
+            self.llm = AutoModelForCausalLM.from_pretrained(
+                llm_path,
+                device_map=device,
+                torch_dtype="auto",
+                trust_remote_code=True,
+            )
+
+        # Load tokenizer directly from pretrained
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            llm_path,
+            trust_remote_code=True,
+        )
+
+    def tensor_to_audio_string(self, tensor):
+        """Convert tensor to audio string format"""
+        if isinstance(tensor, list) and isinstance(tensor[0], list):
+            values = tensor[0]
+        else:
+            values = tensor[0].tolist() if hasattr(tensor, "tolist") else tensor[0]
+
+        result = "<|start_of_audio|>"
+        for value in values:
+            result += f"<|audio_{value}|>"
+        return result
+
+    def extract_audio_ids(self, text):
+        """Extract audio IDs from string containing audio tokens"""
+        import re
+
+        pattern = r"<\|audio_(\d+)\|>"
+        audio_ids = re.findall(pattern, text)
+        return [int(id) for id in audio_ids]
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: str = None,
+        tadicodec_path: str = None,
+        llm_path: str = None,
+        device: Optional[torch.device] = None,
+        auto_download: bool = True,
+    ):
+        """
+        Create pipeline from pretrained models
+
+        Args:
+            model_id: Hugging Face model ID for the LLM (e.g., "amphion/TaDiCodec-TTS-AR-Qwen2.5-3B")
+            tadicodec_path: Path to TaDiCodec model or Hugging Face model ID (defaults to "amphion/TaDiCodec")
+            llm_path: Path to LLM model or Hugging Face model ID (overrides model_id if provided)
+            device: Device to run on
+            auto_download: Whether to automatically download models from Hugging Face if not found locally
+
+        Returns:
+            TTSInferencePipeline instance
+        """
+        resolved_device = (
+            device
+            if device is not None
+            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+
+        # Set default paths if not provided
+        if tadicodec_path is None:
+            tadicodec_path = "amphion/TaDiCodec"
+        
+        if llm_path is None:
+            if model_id is not None:
+                llm_path = model_id
+            else:
+                llm_path = "./ckpt/TaDiCodec-TTS-AR-Qwen2.5-0.5B"
+
+        # Handle TaDiCodec path
+        resolved_tadicodec_path = cls._resolve_model_path(
+            tadicodec_path, auto_download=auto_download, model_type="tadicodec"
+        )
+
+        # Handle LLM path
+        resolved_llm_path = cls._resolve_model_path(
+            llm_path, auto_download=auto_download, model_type="llm"
+        )
+
+        return cls(
+            tadicodec_path=resolved_tadicodec_path,
+            llm_path=resolved_llm_path,
+            device=resolved_device,
+        )
+
+    @staticmethod
+    def _resolve_model_path(
+        model_path: str, auto_download: bool = True, model_type: str = "llm"
+    ) -> str:
+        """
+        Resolve model path, downloading from Hugging Face if necessary
+
+        Args:
+            model_path: Local path or Hugging Face model ID
+            auto_download: Whether to auto-download from HF
+            model_type: Type of model ("llm" or "tadicodec")
+
+        Returns:
+            Resolved local path
+        """
+        # If it's already a local path and exists, return as is
+        if os.path.exists(model_path):
+            return model_path
+
+        # If it looks like a Hugging Face model ID (contains '/')
+        if "/" in model_path and auto_download:
+            print(f"Downloading {model_type} model from Hugging Face: {model_path}")
+            try:
+                # Download to cache directory
+                cache_dir = os.path.join(
+                    os.path.expanduser("~"), ".cache", "huggingface", "hub"
+                )
+                downloaded_path = snapshot_download(
+                    repo_id=model_path,
+                    cache_dir=cache_dir,
+                    local_dir_use_symlinks=False,
+                )
+                print(
+                    f"Successfully downloaded {model_type} model to: {downloaded_path}"
+                )
+                return downloaded_path
+            except Exception as e:
+                print(f"Failed to download {model_type} model from Hugging Face: {e}")
+                raise ValueError(
+                    f"Could not download {model_type} model from {model_path}"
+                )
+
+        # If it's a local path that doesn't exist
+        if not os.path.exists(model_path):
+            if auto_download:
+                raise ValueError(
+                    f"Model path does not exist: {model_path}. Set auto_download=True to download from Hugging Face."
+                )
+            else:
+                raise FileNotFoundError(f"Model path does not exist: {model_path}")
+
+        return model_path
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        text: str,
+        prompt_text: Optional[str] = None,
+        prompt_speech_path: Optional[str] = None,
+        top_k: int = 50,
+        top_p: float = 0.98,
+        temperature: float = 1.0,
+        n_timesteps: int = 25,
+        return_code: bool = False,
+    ):
+        """
+        Perform TTS inference
+
+        Args:
+            text: Target text to synthesize
+            prompt_text: Prompt text for conditioning
+            prompt_speech_path: Path to prompt audio file
+            top_k: Top-k sampling parameter
+            top_p: Top-p sampling parameter
+            temperature: Temperature for sampling
+            n_timesteps: Number of diffusion timesteps
+            return_code: Whether to return audio codes instead of audio
+
+        Returns:
+            Generated audio array or audio codes
+        """
+        # Get prompt audio codes
+        if prompt_speech_path:
+            prompt_speech_code = self.tadicodec(
+                speech_path=prompt_speech_path, return_code=True, text=""
+            )
+        else:
+            raise ValueError("prompt_speech_path is required")
+
+        # Use standard LLM for autoregressive generation
+        # TODO: add a json style chat prompt template
+        prompt = gen_chat_prompt_for_tts(
+            (prompt_text or "") + text,
+            "phi-3" if "Phi" in self.llm_path else "qwen2",
+        ) + self.tensor_to_audio_string(prompt_speech_code)
+
+        input_ids = self.tokenizer.encode(prompt)
+        generate_ids = self.llm.generate(
+            input_ids=torch.tensor(input_ids).unsqueeze(0).to(self.device),
+            min_new_tokens=12,
+            max_new_tokens=400,
+            do_sample=True,
+            top_k=top_k,
+            top_p=top_p,
+            temperature=temperature,
+        )
+
+        output = self.tokenizer.decode(generate_ids[0], skip_special_tokens=False)
+
+        combine_speech_code = self.extract_audio_ids(output)
+        indices = torch.tensor(combine_speech_code).unsqueeze(0).long().to(self.device)
+
+        if return_code:
+            return indices
+
+        # Decode audio
+        text_token_ids = self.tadicodec.tokenize_text(text, prompt_text)
+        prompt_mel = self.tadicodec.extract_mel_feature(prompt_speech_path)
+
+        rec_mel = self.tadicodec.decode(
+            indices=indices,
+            text_token_ids=text_token_ids,
+            prompt_mel=prompt_mel,
+            n_timesteps=n_timesteps,
+        )
+
+        rec_audio = (
+            self.tadicodec.vocoder_model(rec_mel.transpose(1, 2))
+            .detach()
+            .cpu()
+            .numpy()[0][0]
+        )
+
+        return rec_audio

models/tts/llm_tts/inference_mgm_tts.py

@@ -0,0 +1,338 @@
+import os
+import sys
+import json
+import librosa
+from huggingface_hub import snapshot_download
+
+import torch
+import torch.nn as nn
+from typing import Optional
+import safetensors
+from transformers import AutoTokenizer
+from utils.util import load_config
+
+from models.tts.tadicodec.inference_tadicodec import TaDiCodecPipline
+from models.tts.llm_tts.mgm import MGMT2S
+
+from models.tts.llm_tts.chat_template import gen_chat_prompt_for_tts
+
+
+class MGMInferencePipeline(nn.Module):
+    """
+    MGM TTS inference pipeline that integrates TaDiCodec and MGM models
+    Uses diffusion-based generation with mask-guided modeling
+    """
+
+    def __init__(
+        self,
+        tadicodec_path: str,
+        mgm_path: str,
+        device: torch.device,
+    ):
+        super().__init__()
+        self.device = device
+        self.mgm_path = mgm_path
+
+        # Load TaDiCodec pipeline
+        self.tadicodec = TaDiCodecPipline.from_pretrained(
+            ckpt_dir=tadicodec_path, device=device
+        )
+
+        # Load tokenizer directly from pretrained
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            mgm_path,
+            trust_remote_code=True,
+        )
+
+        config_path = os.path.join(mgm_path, "config.json")
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"Config file not found at {config_path}")
+
+        self.cfg = load_config(config_path)
+
+        # Extract MGM config from the loaded config
+        mgm_config = self.cfg.model.mgmt2s
+        if not mgm_config:
+            raise ValueError("MGM config not found in config.json")
+
+        # Load MGM model with config - using the same pattern as llm_infer_eval.py
+        self.mgm = MGMT2S(
+            hidden_size=mgm_config.hidden_size,
+            num_layers=mgm_config.num_layers,
+            num_heads=mgm_config.num_heads,
+            cfg_scale=mgm_config.cfg_scale,
+            cond_codebook_size=mgm_config.cond_codebook_size,
+            cond_dim=mgm_config.cond_dim,
+            phone_vocab_size=mgm_config.phone_vocab_size,
+        )
+
+        # Load model weights
+        model_path = os.path.join(mgm_path, "model.safetensors")
+
+        if os.path.exists(model_path):
+            safetensors.torch.load_model(self.mgm, model_path, strict=True)
+        else:
+            # Try loading from the directory directly
+            safetensors.torch.load_model(self.mgm, mgm_path, strict=True)
+
+        self.mgm.to(device)
+        self.mgm.eval()
+
+    def tensor_to_audio_string(self, tensor):
+        """Convert tensor to audio string format"""
+        if isinstance(tensor, list) and isinstance(tensor[0], list):
+            values = tensor[0]
+        else:
+            values = tensor[0].tolist() if hasattr(tensor, "tolist") else tensor[0]
+
+        result = "<|start_of_audio|>"
+        for value in values:
+            result += f"<|audio_{value}|>"
+        return result
+
+    def extract_audio_ids(self, text):
+        """Extract audio IDs from string containing audio tokens"""
+        import re
+
+        pattern = r"<\|audio_(\d+)\|>"
+        audio_ids = re.findall(pattern, text)
+        return [int(id) for id in audio_ids]
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: str = None,
+        tadicodec_path: str = None,
+        mgm_path: str = None,
+        device: Optional[torch.device] = None,
+        auto_download: bool = True,
+    ):
+        """
+        Create pipeline from pretrained models
+
+        Args:
+            model_id: Hugging Face model ID for the MGM model (e.g., "amphion/TaDiCodec-TTS-MGM")
+            tadicodec_path: Path to TaDiCodec model or Hugging Face model ID (defaults to "amphion/TaDiCodec")
+            mgm_path: Path to MGM model directory or Hugging Face model ID (overrides model_id if provided)
+            device: Device to run on
+            auto_download: Whether to automatically download models from Hugging Face if not found locally
+
+        Returns:
+            MGMInferencePipeline instance
+        """
+        resolved_device = (
+            device
+            if device is not None
+            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+
+        # Set default paths if not provided
+        if tadicodec_path is None:
+            tadicodec_path = "amphion/TaDiCodec"
+        
+        if mgm_path is None:
+            if model_id is not None:
+                mgm_path = model_id
+            else:
+                mgm_path = "./ckpt/TaDiCodec-TTS-MGM"
+
+        # Handle TaDiCodec path
+        resolved_tadicodec_path = cls._resolve_model_path(
+            tadicodec_path, auto_download=auto_download, model_type="tadicodec"
+        )
+
+        # Handle MGM path
+        resolved_mgm_path = cls._resolve_model_path(
+            mgm_path, auto_download=auto_download, model_type="mgm"
+        )
+
+        return cls(
+            tadicodec_path=resolved_tadicodec_path,
+            mgm_path=resolved_mgm_path,
+            device=resolved_device,
+        )
+
+    @staticmethod
+    def _resolve_model_path(
+        model_path: str, auto_download: bool = True, model_type: str = "mgm"
+    ) -> str:
+        """
+        Resolve model path, downloading from Hugging Face if necessary
+
+        Args:
+            model_path: Local path or Hugging Face model ID
+            auto_download: Whether to auto-download from HF
+            model_type: Type of model ("mgm" or "tadicodec")
+
+        Returns:
+            Resolved local path
+        """
+        # If it's already a local path and exists, return as is
+        if os.path.exists(model_path):
+            return model_path
+
+        # If it looks like a Hugging Face model ID (contains '/')
+        if "/" in model_path and auto_download:
+            print(f"Downloading {model_type} model from Hugging Face: {model_path}")
+            try:
+                # Download to cache directory
+                cache_dir = os.path.join(
+                    os.path.expanduser("~"), ".cache", "huggingface", "hub"
+                )
+                downloaded_path = snapshot_download(
+                    repo_id=model_path,
+                    cache_dir=cache_dir,
+                    local_dir_use_symlinks=False,
+                )
+                print(
+                    f"Successfully downloaded {model_type} model to: {downloaded_path}"
+                )
+                return downloaded_path
+            except Exception as e:
+                print(f"Failed to download {model_type} model from Hugging Face: {e}")
+                raise ValueError(
+                    f"Could not download {model_type} model from {model_path}"
+                )
+
+        # If it's a local path that doesn't exist
+        if not os.path.exists(model_path):
+            if auto_download:
+                raise ValueError(
+                    f"Model path does not exist: {model_path}. Set auto_download=True to download from Hugging Face."
+                )
+            else:
+                raise FileNotFoundError(f"Model path does not exist: {model_path}")
+
+        return model_path
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        text: str,
+        prompt_text: Optional[str] = None,
+        prompt_speech_path: Optional[str] = None,
+        n_timesteps_mgm: int = 25,
+        n_timesteps: int = 25,
+        target_len: Optional[int] = None,
+        return_code: bool = False,
+    ):
+        """
+        Perform MGM TTS inference
+
+        Args:
+            text: Target text to synthesize
+            prompt_text: Prompt text for conditioning
+            prompt_speech_path: Path to prompt audio file
+            n_timesteps_mgm: Number of diffusion timesteps for MGM
+            n_timesteps: Number of diffusion timesteps
+            target_len: Target length for audio generation
+            return_code: Whether to return audio codes instead of audio
+
+        Returns:
+            Generated audio array or audio codes
+        """
+        # Get prompt audio codes
+        if prompt_speech_path:
+            prompt_speech_code = self.tadicodec(
+                speech_path=prompt_speech_path, return_code=True, text=""
+            )
+        else:
+            raise ValueError("prompt_speech_path is required")
+
+        # Convert prompt codes to tensor
+        prompt_codes = torch.tensor(prompt_speech_code).to(self.device)
+        prompt_len = prompt_codes.shape[1]
+
+        # Tokenize text for phone conditioning
+        input_text = gen_chat_prompt_for_tts(
+            prompt_text + " " + text,
+            "phi-3" if "phi" in self.cfg.preprocess.tokenizer_path else "qwen2",
+        )
+
+        ##### debug #####
+        print("input_text: ", input_text)
+        ##### debug #####
+
+        text_token_ids = self.tokenizer.encode(input_text)
+        text_token_ids = torch.tensor(text_token_ids).unsqueeze(0).to(self.device)
+
+        # Estimate target length based on text length
+        frame_rate = getattr(self.cfg.preprocess, "frame_rate", 6.25)
+
+        if target_len is None:
+            # If no target_len, estimate based on prompt speech length and text ratio
+            prompt_text_len = len(prompt_text.encode("utf-8"))
+            target_text_len = len(text.encode("utf-8"))
+            prompt_speech_len = librosa.get_duration(filename=prompt_speech_path)
+            target_speech_len = prompt_speech_len * target_text_len / prompt_text_len
+            target_len = int(target_speech_len * frame_rate)
+        else:
+            # If target_len is provided, use it directly
+            target_len = int(target_len * frame_rate)
+
+        ##### debug #####
+        print(f"Prompt length: {prompt_len}, Target length: {target_len}")
+        print(f"Text: {text}")
+        print(f"Prompt text: {prompt_text}")
+        ##### debug #####
+
+        # Generate audio codes using MGM reverse diffusion
+        generated_codes = self.mgm.reverse_diffusion(
+            prompt=prompt_codes,
+            target_len=target_len,
+            phone_id=text_token_ids,
+            n_timesteps=n_timesteps_mgm,
+            cfg=1.5,
+            rescale_cfg=0.75,
+        )
+
+        print(f"Generated codes shape: {generated_codes.shape}")
+
+        combine_codes = torch.cat([prompt_codes, generated_codes], dim=1)
+
+        if return_code:
+            return combine_codes
+
+        # Decode audio using TaDiCodec
+        prompt_mel = self.tadicodec.extract_mel_feature(prompt_speech_path)
+
+        text_token_ids = self.tadicodec.tokenize_text(text, prompt_text)
+        rec_mel = self.tadicodec.decode(
+            indices=combine_codes,
+            text_token_ids=text_token_ids,
+            prompt_mel=prompt_mel,
+            n_timesteps=n_timesteps,
+        )
+
+        rec_audio = (
+            self.tadicodec.vocoder_model(rec_mel.transpose(1, 2))
+            .detach()
+            .cpu()
+            .numpy()[0][0]
+        )
+
+        return rec_audio
+
+
+# Usage example
+if __name__ == "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Create pipeline
+    pipeline = MGMInferencePipeline.from_pretrained(
+        tadicodec_path="./ckpt/TaDiCodec",
+        mgm_path="./ckpt/TaDiCodec-TTS-MGM",
+        device=device,
+    )
+
+    # Inference on single sample
+    audio = pipeline(
+        text="但是 to those who 知道 her well, it was a 标志 of her unwavering 决心 and spirit.",
+        prompt_text="In short, we embarked on a mission to make America great again, for all Americans.",
+        prompt_speech_path="./use_examples/test_audio/trump_0.wav",
+    )
+
+    # Save audio
+    import soundfile as sf
+
+    sf.write("./use_examples/test_audio/mgm_tts_output.wav", audio, 24000)

models/tts/llm_tts/llama_nar_prefix.py

@@ -0,0 +1,457 @@
+from transformers import LlamaConfig, LlamaModel
+import torch
+import torch.nn as nn
+from typing import List, Optional, Tuple, Union
+import math
+
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer
+from transformers.models.llama.modeling_llama import BaseModelOutputWithPast
+
+
+class SinusoidalPosEmb(nn.Module):
+    """
+    Sinusoidal Positional Embedding module.
+
+    This module generates sinusoidal positional embeddings for a given 1D input tensor,
+    which is commonly used for representing timesteps in diffusion models.
+    """
+
+    def __init__(self, dim: int):
+        """
+        Initializes the SinusoidalPosEmb module.
+
+        Args:
+            dim (int): The dimension of the embedding.
+        """
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Generates the positional embedding.
+
+        Args:
+            x (torch.Tensor): A 1D tensor of positions (e.g., timesteps), shape `(batch_size,)`.
+
+        Returns:
+            torch.Tensor: The positional embeddings, shape `(batch_size, dim)`.
+        """
+        device = x.device
+        half_dim = self.dim // 2
+        # Calculate the embedding frequencies based on the log-space formula
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        # Create the embedding matrix by multiplying positions with frequencies
+        emb = x[:, None] * emb[None, :] * 1.0
+        # Concatenate sine and cosine components to form the final embedding
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class LlamaAdaptiveRMSNorm(nn.Module):
+    """
+    Adaptive Root Mean Square Layer Normalization.
+
+    This is a variant of RMSNorm where the scaling factor (weight) is adaptively
+    predicted from a conditional embedding, allowing for conditional modulation
+    of the normalized hidden states.
+    """
+
+    def __init__(
+        self, hidden_size: int = 1024, eps: float = 1e-6, dim_cond: int = 1024
+    ):
+        """
+        Initializes the LlamaAdaptiveRMSNorm module.
+
+        Args:
+            hidden_size (int): The dimension of the hidden states to be normalized.
+            eps (float): A small value added to the variance for numerical stability.
+            dim_cond (int): The dimension of the conditional embedding.
+        """
+        super().__init__()
+        # Linear layer to project the conditional embedding to the hidden size
+        self.to_weight = nn.Linear(dim_cond, hidden_size)
+        # Initialize weights to zero and bias to one for an identity transformation at the start
+        nn.init.zeros_(self.to_weight.weight)
+        nn.init.ones_(self.to_weight.bias)
+        self.variance_epsilon = eps
+        # Disable automatic Hugging Face initialization for this custom module
+        self._is_hf_initialized = True
+
+    def forward(
+        self, hidden_states: torch.Tensor, cond_embedding: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Applies the adaptive RMS normalization.
+
+        Args:
+            hidden_states (torch.Tensor): The input tensor, shape `(batch, seq_len, hidden_size)`.
+            cond_embedding (torch.Tensor): The conditional embedding, shape `(batch, dim_cond)`.
+
+        Returns:
+            torch.Tensor: The normalized and modulated hidden states.
+        """
+        input_dtype = hidden_states.dtype
+        # Calculate variance and normalize the hidden states
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        # Predict the scaling factor from the conditional embedding
+        weight = self.to_weight(cond_embedding)
+        # Unsqueeze if the conditional embedding is per-batch instead of per-token
+        if len(weight.shape) == 2:
+            weight = weight.unsqueeze(1)
+
+        # Apply the learned scaling factor
+        return (weight * hidden_states).to(input_dtype)
+
+
+class LlamaNARDecoderLayer(LlamaDecoderLayer):
+    """
+    A Non-Autoregressive (NAR) Llama Decoder Layer using adaptive layer normalization.
+
+    This class overrides the standard LlamaDecoderLayer to replace its RMSNorm
+    modules with LlamaAdaptiveRMSNorm, allowing it to be conditioned on an external embedding.
+    """
+
+    def __init__(self, config: LlamaConfig, layer_idx: int):
+        """Overrides the LlamaDecoderLayer to use adaptive layer normalization."""
+        super().__init__(config, layer_idx)  # init attention, mlp, etc. from parent
+        self.layer_idx = layer_idx
+        # Override the standard layer norms with our adaptive versions
+        self.input_layernorm = LlamaAdaptiveRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, dim_cond=config.hidden_size
+        )
+        self.post_attention_layernorm = LlamaAdaptiveRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, dim_cond=config.hidden_size
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cond_embedding: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[
+        torch.Tensor,
+        Optional[torch.Tensor],
+        Optional[Tuple[torch.Tensor, torch.Tensor]],
+    ]:
+        """
+        Forward pass for the NAR decoder layer, including conditional embedding.
+
+        Args:
+            hidden_states (torch.Tensor): Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            cond_embedding (torch.Tensor): Conditional embedding for adaptive normalization.
+            attention_mask (Optional[torch.Tensor]): Attention mask of size `(batch, 1, tgt_len, src_len)`.
+            position_ids (Optional[torch.LongTensor]): Indices of positions of each input sequence tokens.
+            past_key_value (Optional[Tuple[torch.Tensor, torch.Tensor]]): Cached past key and value projection states.
+            output_attentions (Optional[bool]): Whether to return the attention tensors.
+            use_cache (Optional[bool]): If True, past key values are returned to speed up decoding.
+
+        Returns:
+            Tuple containing the output hidden states, and optionally attention weights and past key/value states.
+        """
+        residual = hidden_states
+
+        # Apply adaptive pre-attention layer norm
+        hidden_states = self.input_layernorm(
+            hidden_states, cond_embedding=cond_embedding
+        )
+
+        # Self Attention block
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected block
+        residual = hidden_states
+        # Apply adaptive post-attention layer norm
+        hidden_states = self.post_attention_layernorm(
+            hidden_states, cond_embedding=cond_embedding
+        )
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+class DiffLlamaPrefix(LlamaModel):
+    """
+    A Llama-based non-autoregressive transformer model for diffusion (masked generative modeling) tasks.
+
+    This model uses a Llama architecture but modifies it for non-autoregressive generation.
+    Key features:
+    1. Non-causal (fully-visible) attention mask.
+    2. Adaptive layer normalization conditioned on diffusion timesteps.
+    3. Ability to be conditioned on phoneme (text) embeddings, which are prepended as a prefix.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        num_heads: int = 16,
+        num_layers: int = 16,
+        use_phone_cond: bool = True,
+        config: LlamaConfig = LlamaConfig(
+            vocab_size=0, hidden_size=256, num_attention_heads=1, num_hidden_layers=1
+        ),
+    ):
+        """
+        Initializes the DiffLlamaPrefix model.
+
+        Args:
+            hidden_size (int): The hidden dimension of the transformer.
+            num_heads (int): The number of attention heads.
+            num_layers (int): The number of transformer layers.
+            use_phone_cond (bool): Whether to use phoneme embeddings as a conditional prefix.
+            config (LlamaConfig): A LlamaConfig object. A default is provided for convenience.
+        """
+        super().__init__(config)
+
+        self.use_phone_cond = use_phone_cond
+
+        # Create a stack of non-autoregressive Llama layers
+        self.layers = nn.ModuleList(
+            [
+                LlamaNARDecoderLayer(
+                    LlamaConfig(
+                        hidden_size=hidden_size,
+                        num_attention_heads=num_heads,
+                        max_position_embeddings=4096,
+                        intermediate_size=hidden_size * 4,
+                    ),
+                    layer_idx=i,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        # Final adaptive layer norm
+        self.norm = LlamaAdaptiveRMSNorm(hidden_size, dim_cond=hidden_size)
+
+        # Modules for diffusion step conditioning
+        self.diff_step_embedding = SinusoidalPosEmb(hidden_size)
+        self.diff_step_mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(hidden_size * 4, hidden_size),
+        )
+
+        # MLP for processing phoneme embedding condition
+        if self.use_phone_cond:
+            self.cond_mlp = nn.Sequential(
+                nn.Linear(hidden_size, hidden_size * 4),
+                nn.SiLU(),
+                nn.Linear(hidden_size * 4, hidden_size),
+            )
+
+        # This loop is redundant if layers are initialized correctly, but ensures config consistency.
+        for layer in self.layers:
+            layer.input_layernorm = LlamaAdaptiveRMSNorm(
+                hidden_size, dim_cond=hidden_size
+            )
+            layer.post_attention_layernorm = LlamaAdaptiveRMSNorm(
+                hidden_size, dim_cond=hidden_size
+            )
+
+        # We handle embeddings manually, so disable the default token embedder
+        self.embed_tokens = None
+
+        self.post_init()
+
+    def _prepare_decoder_attention_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_shape: Tuple[int, int],
+        inputs_embeds: torch.Tensor,
+        past_key_values_length: int,
+    ) -> Optional[torch.Tensor]:
+        """
+        Creates a non-causal (fully-visible) attention mask.
+
+        This method overrides the default causal mask creation. It converts a 2D padding mask
+        `[bsz, seq_len]` into a 4D attention mask `[bsz, 1, tgt_seq_len, src_seq_len]`
+        suitable for self-attention, without applying a causal triangle.
+
+        Args:
+            attention_mask (torch.Tensor): The 2D padding mask.
+            input_shape (Tuple[int, int]): The shape of the input (`batch_size`, `seq_len`).
+            inputs_embeds (torch.Tensor): The input embeddings tensor.
+            past_key_values_length (int): The length of any cached key-values.
+
+        Returns:
+            Optional[torch.Tensor]: The 4D attention mask, or None if the input mask is None.
+        """
+        combined_attention_mask = None
+
+        def _expand_mask(
+            mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None
+        ) -> torch.Tensor:
+            """Expands a 2D attention mask to a 4D attention mask."""
+            bsz, src_len = mask.size()
+            tgt_len = tgt_len if tgt_len is not None else src_len
+            expanded_mask = (
+                mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+            )
+            # Invert the mask and convert to additive format (-inf for masked positions)
+            inverted_mask = 1.0 - expanded_mask
+            return inverted_mask.masked_fill(
+                inverted_mask.to(torch.bool), torch.finfo(dtype).min
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(
+                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+            ).to(inputs_embeds.device)
+            combined_attention_mask = (
+                expanded_attn_mask
+                if combined_attention_mask is None
+                else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        diffusion_step: torch.Tensor,
+        x_mask: torch.Tensor,
+        phone_embedding: Optional[torch.Tensor] = None,
+        phone_mask: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.Tensor:
+        """
+        Forward pass of the DiffLlamaPrefix model.
+
+        Args:
+            x (torch.Tensor): The primary input tensor, e.g., noisy data (`batch, seq_len, hidden_size`).
+            diffusion_step (torch.Tensor): Diffusion timesteps, shape `(batch,)`.
+            x_mask (torch.Tensor): The padding mask for `x`, shape `(batch, seq_len)`.
+            phone_embedding (Optional[torch.Tensor]): Phoneme embeddings prefix, shape `(batch, phone_len, hidden_size)`.
+            phone_mask (Optional[torch.Tensor]): The padding mask for `phone_embedding`, shape `(batch, phone_len)`.
+            input_ids, etc.: Standard Hugging Face arguments, mostly for compatibility.
+
+        Returns:
+            torch.Tensor: The final output tensor of shape `(batch, seq_len, hidden_size)`.
+        """
+        # 1. Prepend conditional prefix (phoneme embeddings)
+        if self.use_phone_cond and phone_embedding is not None:
+            # Process condition through an MLP
+            phone_embedding = self.cond_mlp(phone_embedding)  # (B, T_phone, C)
+            phone_length = phone_embedding.shape[1]
+            # Concatenate prefix and main input
+            inputs_embeds = torch.cat([phone_embedding, x], dim=1)
+            attention_mask = torch.cat([phone_mask, x_mask], dim=1)
+        else:
+            inputs_embeds = x
+            attention_mask = x_mask
+            phone_length = 0
+
+        # 2. Process diffusion step embedding for adaptive normalization
+        diffusion_step_emb = self.diff_step_embedding(diffusion_step).to(x.device)
+        diffusion_step_emb = self.diff_step_mlp(diffusion_step_emb)  # (B, C)
+
+        # 3. Standard Transformer Preamble (adapted from LlamaModel)
+        batch_size, seq_length, _ = inputs_embeds.shape
+
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        past_key_values_length = 0
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+
+        if position_ids is None:
+            device = inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length,
+                seq_length + past_key_values_length,
+                dtype=torch.long,
+                device=device,
+            )
+            position_ids = position_ids.unsqueeze(0)
+
+        # Create the non-causal attention mask
+        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask,
+            (batch_size, seq_length),
+            inputs_embeds,
+            past_key_values_length,
+        )
+
+        hidden_states = inputs_embeds
+
+        # 4. Transformer Decoder Layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = (
+                past_key_values[idx] if past_key_values is not None else None
+            )
+
+            # Pass the processed diffusion step embedding to the adaptive layer
+            layer_outputs = decoder_layer(
+                hidden_states,
+                cond_embedding=diffusion_step_emb,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        # 5. Final Normalization and Output Processing
+        hidden_states = self.norm(hidden_states, cond_embedding=diffusion_step_emb)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        # Remove the conditional prefix from the final output sequence
+        return hidden_states[:, phone_length:]

models/tts/llm_tts/mgm.py

@@ -0,0 +1,385 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import math
+from einops import rearrange
+from models.tts.llm_tts.llama_nar_prefix import DiffLlamaPrefix
+
+
+def top_k(logits, thres=0.9):
+    k = math.ceil((1 - thres) * logits.shape[-1])
+    val, ind = logits.topk(k, dim=-1)
+    probs = torch.full_like(logits, float("-inf"))
+    probs.scatter_(2, ind, val)
+    return probs
+
+
+def log(t, eps=1e-10):
+    return torch.log(t + eps)
+
+
+def gumbel_noise(t):
+    noise = torch.zeros_like(t).uniform_(0, 1)
+    return -log(-log(noise))
+
+
+def gumbel_sample(t, temperature=1.0, dim=-1):
+    return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim=dim)
+
+
+class MGMT2S(nn.Module):
+    def __init__(
+        self,
+        hidden_size=1024,
+        num_layers=16,
+        num_heads=16,
+        cfg_scale=0.2,
+        cond_codebook_size=8192,
+        cond_dim=1024,
+        use_phone_cond=True,
+        phone_vocab_size=32100,
+        cfg=None,
+    ):
+        super().__init__()
+
+        hidden_size = (
+            cfg.hidden_size
+            if cfg is not None and hasattr(cfg, "hidden_size")
+            else hidden_size
+        )
+        num_layers = (
+            cfg.num_layers
+            if cfg is not None and hasattr(cfg, "num_layers")
+            else num_layers
+        )
+        num_heads = (
+            cfg.num_heads
+            if cfg is not None and hasattr(cfg, "num_heads")
+            else num_heads
+        )
+        cfg_scale = (
+            cfg.cfg_scale
+            if cfg is not None and hasattr(cfg, "cfg_scale")
+            else cfg_scale
+        )
+        cond_codebook_size = (
+            cfg.cond_codebook_size
+            if cfg is not None and hasattr(cfg, "cond_codebook_size")
+            else cond_codebook_size
+        )
+        cond_dim = (
+            cfg.cond_dim if cfg is not None and hasattr(cfg, "cond_dim") else cond_dim
+        )
+        use_phone_cond = (
+            cfg.use_phone_cond
+            if cfg is not None and hasattr(cfg, "use_phone_cond")
+            else use_phone_cond
+        )
+        phone_vocab_size = (
+            cfg.phone_vocab_size
+            if cfg is not None and hasattr(cfg, "phone_vocab_size")
+            else phone_vocab_size
+        )
+
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.cfg_scale = cfg_scale
+        self.cond_codebook_size = cond_codebook_size
+        self.cond_dim = cond_dim
+        self.use_phone_cond = use_phone_cond
+        self.phone_vocab_size = phone_vocab_size
+
+        self.mask_emb = nn.Embedding(1, self.hidden_size)
+
+        self.to_logit = nn.Linear(self.hidden_size, self.cond_codebook_size)
+
+        self.cond_emb = nn.Embedding(cond_codebook_size, self.hidden_size)
+
+        if self.use_phone_cond:
+            self.phone_emb = nn.Embedding(self.phone_vocab_size, hidden_size)
+            torch.nn.init.normal_(self.phone_emb.weight, mean=0.0, std=0.02)
+
+        self.reset_parameters()
+
+        self.diff_estimator = DiffLlamaPrefix(
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            use_phone_cond=use_phone_cond,
+        )
+
+    def mask_prob(self, t):
+        return torch.sin(t * np.pi / 2).to(t.device)
+
+    def forward_diffusion(self, x0, t):
+        # x0: semantic tokens (B, T)
+        new_t = t
+        mask_prob = self.mask_prob(new_t)  # (B,)
+        # if mask_prob[i] < 0.2, mask_prob[i] = 0.2
+        mask_prob = torch.where(
+            mask_prob < 0.2, torch.ones_like(mask_prob) * 0.2, mask_prob
+        )
+        mask_token = self.mask_emb(
+            torch.LongTensor([0]).to(x0.device)
+        )  # (1, hidden_size)
+
+        xt = torch.zeros(x0.shape[0], x0.shape[1], self.hidden_size).to(x0.device)
+
+        cfg_scale = self.cfg_scale
+
+        #  a segment of r% sequence length is masked, where r ~ U[60, 100]
+        if torch.rand(1) > cfg_scale:
+            prompt_len = torch.randint(
+                min(x0.shape[1] // 4, 5), int(x0.shape[1] * 0.4), (x0.shape[0],)
+            ).to(
+                x0.device
+            )  # (B,)
+        else:
+            prompt_len = torch.zeros(x0.shape[0]).to(x0)  # (B,)
+
+        # get is prompt
+        is_prompt = torch.zeros_like(x0[:, :])  # (B, T)
+        col_indices = (
+            torch.arange(is_prompt.shape[1])
+            .repeat(is_prompt.shape[0], 1)
+            .to(prompt_len)
+        )  # (B, T)
+        is_prompt[col_indices < prompt_len.unsqueeze(1)] = 1  # (B, T) 1 if prompt
+
+        # Add mask
+        mask = torch.bernoulli(torch.ones_like(x0[:, :]) * mask_prob[..., None])
+        mask[is_prompt.bool()] = 0
+        mask_num = mask[:,].sum(dim=1, keepdim=False)
+        all_zero_mask = (mask_num == 0).bool()
+        row_indices_to_modify = torch.nonzero(all_zero_mask)
+        mask[row_indices_to_modify, prompt_len[row_indices_to_modify]] = 1
+        mask = mask[..., None]  # (B, T, 1)
+        xt = (
+            xt + mask * mask_token[:, None, :] + (1 - mask) * self.cond_emb(x0[:, :])
+        )  # (B, T, hidden_size)
+
+        return xt, new_t, mask, prompt_len, mask_prob
+
+    def loss_t(self, x0, x_mask, t, phone_embedding=None, phone_mask=None):
+        xt, new_t, mask, prompt_len, mask_prob = self.forward_diffusion(x0, t)
+        # xt: (B, T, hidden_size)
+        # new_t: (B,)
+        # mask: (B, T, 1)   mask if 1, not mask if 0
+        # prompt_len: (B,)
+        # mask_prob: (B,)
+
+        # # drop all condition for cfg, so if prompt_len is 0, we also drop phone_embedding
+        # if self.use_phone_cond and phone_embedding != None:
+        #     phone_embedding = phone_embedding * torch.where(prompt_len > 0, torch.ones_like(prompt_len), torch.zeros_like(prompt_len)).to(x0.device).unsqueeze(-1).unsqueeze(-1)
+
+        embeds = self.diff_estimator(
+            xt, new_t, x_mask, phone_embedding=phone_embedding, phone_mask=phone_mask
+        )  # (B, T, hidden_size)
+        logits = self.to_logit(embeds)  # (B, T, codebook_size)
+
+        # final mask used for loss calculation
+        final_mask = mask * x_mask[..., None]  # (B, T, 1)
+
+        return logits, final_mask, x0, prompt_len, mask_prob
+
+    def compute_loss(self, x0, x_mask, phone_embedding=None, phone_mask=None):
+        # x0: (B, T)
+        # x_mask: (B, T) mask is 0 for padding
+        t = torch.rand(x0.shape[0], device=x0.device, requires_grad=False)
+        t = torch.clamp(t, 1e-5, 1.0)
+        return self.loss_t(x0, x_mask, t, phone_embedding, phone_mask)
+
+    def reset_parameters(self):
+        def _reset_parameters(m):
+            if isinstance(m, nn.MultiheadAttention):
+                if m._qkv_same_embed_dim:
+                    nn.init.normal_(m.in_proj_weight, std=0.02)
+                else:
+                    nn.init.normal_(m.q_proj_weight, std=0.02)
+                    nn.init.normal_(m.k_proj_weight, std=0.02)
+                    nn.init.normal_(m.v_proj_weight, std=0.02)
+
+                if m.in_proj_bias is not None:
+                    nn.init.constant_(m.in_proj_bias, 0.0)
+                    nn.init.constant_(m.out_proj.bias, 0.0)
+                if m.bias_k is not None:
+                    nn.init.xavier_normal_(m.bias_k)
+                if m.bias_v is not None:
+                    nn.init.xavier_normal_(m.bias_v)
+
+            elif (
+                isinstance(m, nn.Conv1d)
+                or isinstance(m, nn.ConvTranspose1d)
+                or isinstance(m, nn.Conv2d)
+                or isinstance(m, nn.ConvTranspose2d)
+            ):
+                m.weight.data.normal_(0.0, 0.02)
+
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(mean=0.0, std=0.02)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+
+            elif isinstance(m, nn.Embedding):
+                m.weight.data.normal_(mean=0.0, std=0.02)
+                if m.padding_idx is not None:
+                    m.weight.data[m.padding_idx].zero_()
+
+        self.apply(_reset_parameters)
+
+    @torch.no_grad()
+    def reverse_diffusion(
+        self,
+        prompt,
+        target_len,
+        phone_id,
+        prompt_mask=None,
+        temp=0.9,
+        filter_thres=0.98,
+        n_timesteps=25,
+        cfg=1.0,
+        rescale_cfg=1.0,
+    ):
+        # prompt: (B, T)
+        if self.use_phone_cond and phone_id != None:
+            phone_embedding = self.phone_emb(phone_id)
+        else:
+            phone_embedding = None
+
+        prompt_code = prompt  # (B, prompt_len)
+        prompt_len = prompt_code.shape[1]
+
+        x_mask = torch.ones(prompt_code.shape[0], target_len).to(
+            prompt_code.device
+        )  # (B, target_len)
+        phone_mask = torch.ones_like(phone_id)
+
+        if prompt_mask == None:
+            prompt_mask = torch.ones(prompt_code.shape[0], prompt_len).to(
+                prompt_code.device
+            )  # (B, prompt_len)
+
+        cum = torch.zeros(x_mask.shape[0], x_mask.shape[1], self.hidden_size).to(
+            x_mask.device
+        )  # (B, T, hidden_size)
+
+        bsz, seq_len, _ = cum.shape
+
+        choice_temp = 1.0
+        start_temp = temp  # temperature for sampling
+        start_choice_temp = choice_temp  # temperature for choicing mask tokens
+
+        xt = torch.LongTensor(bsz, seq_len).to(x_mask.device)
+
+        steps = n_timesteps
+        to_logit = self.to_logit
+        cond_emb = self.cond_emb
+
+        mask_token = self.mask_emb(torch.LongTensor([0]).to(xt.device))
+        mask = torch.full((bsz, seq_len, 1), True).to(x_mask.device)  # (B, T, 1)
+        seq = torch.full((bsz, seq_len), 0).to(x_mask.device)
+        h = 1.0 / steps
+
+        cur_prompt = 0
+        cur_prompt = cur_prompt + cond_emb(prompt_code)
+
+        t_list = [1.0 - i * h for i in range(steps)]
+        t_list.append(0.0)
+        for i in range(steps):
+            t = t_list[i] * torch.ones(bsz).to(x_mask.device)
+            token = cond_emb(seq)  # (B, T, hidden_size)
+            cur = cum + mask * mask_token[:, None, :] + (~mask) * token
+
+            xt_input = torch.cat([cur_prompt, cur], dim=1)  # (B, T, hidden_size)
+            xt_mask = torch.cat(
+                [prompt_mask, x_mask], dim=1
+            )  # (B, T), mask is 0 for padding
+
+            embeds = self.diff_estimator(
+                xt_input,
+                t,
+                xt_mask,
+                phone_embedding=phone_embedding,
+                phone_mask=phone_mask,
+            )
+            embeds = embeds[:, prompt_len:, :]
+
+            # classifier free guidance
+            # phone_embedding=phone_embedding[:,phone_embedding.shape[1]:,:] means phone_embedding is None
+            if cfg > 0:
+                mask_embeds = self.diff_estimator(
+                    cur,
+                    t,
+                    x_mask,
+                    phone_embedding=phone_embedding[:, phone_embedding.shape[1] :, :],
+                    phone_mask=phone_mask[:, prompt_len:],
+                )
+                pos_emb_std = embeds.std()  # std(g_cond)
+                embeds = embeds + cfg * (embeds - mask_embeds)  # g_cfg
+                rescale_embeds = embeds * pos_emb_std / embeds.std()  # g_final
+                embeds = rescale_cfg * rescale_embeds + (1 - rescale_cfg) * embeds
+
+            logits = to_logit(embeds)  # (B, T, codebook_size)
+            annealing_scale = t_list[i]
+
+            choice_temp = start_choice_temp * annealing_scale
+            temp = start_temp * annealing_scale
+            logits = top_k(logits, filter_thres)
+
+            if i == steps - 1:
+                # greedy
+                if steps == 1:
+                    temp = 0.2
+                    sampled_ids = gumbel_sample(logits, temperature=max(temp, 1e-3))
+                else:
+                    sampled_ids = logits.argmax(dim=-1)
+
+            else:
+                # sampling
+                sampled_ids = gumbel_sample(logits, temperature=max(temp, 1e-3))
+
+            seq = torch.where(mask.squeeze(-1), sampled_ids, seq)
+
+            scores = logits.softmax(dim=-1)
+            scores = scores.gather(2, rearrange(sampled_ids, "b n -> b n 1"))
+            scores = rearrange(scores, "b n 1 -> b n")
+
+            scores = choice_temp * gumbel_noise(scores) + scores
+            scores = 1 - scores
+
+            next_t = t_list[i + 1] * torch.ones(bsz).to(x_mask.device)
+
+            next_mask_num = (self.mask_prob(next_t) * seq_len).long()[0].item()
+
+            if next_mask_num == 0:
+                break
+            scores = scores.masked_fill(
+                ~mask.squeeze(-1), -torch.finfo(scores.dtype).max
+            )
+
+            mask_indices = scores.topk(next_mask_num, dim=-1).indices
+            mask = torch.zeros_like(scores, dtype=torch.bool).scatter(
+                1, mask_indices, True
+            )
+            seq = seq.masked_fill(mask, 0)
+
+            mask = mask.unsqueeze(-1)
+
+        cum = cum + cond_emb(seq)
+        xt = seq
+
+        return xt
+
+    def forward(self, x0, x_mask, phone_id=None, phone_mask=None):
+        # x0: (B, T)
+        # x_mask: (B, T) mask is 0 for padding
+        if self.use_phone_cond and phone_id != None:
+            phone_embedding = self.phone_emb(phone_id)
+        else:
+            phone_embedding = None
+
+        logits, final_mask, x0, prompt_len, mask_prob = self.compute_loss(
+            x0, x_mask, phone_embedding, phone_mask=phone_mask
+        )
+        return logits, final_mask, x0, prompt_len, mask_prob

models/tts/tadicodec/infer_utils.py

@@ -0,0 +1,24 @@
+from models.codec.melvqgan.melspec import MelSpectrogram
+from models.codec.amphion_codec.vocos import Vocos
+
+
+def build_vocoder_model(cfg, device):
+    vocoder_model = Vocos(cfg=cfg.model.vocos)
+    vocoder_model.eval()
+    vocoder_model.to(device)
+    return vocoder_model
+
+
+def build_mel_model(cfg, device):
+    mel_model = MelSpectrogram(
+        sampling_rate=cfg.preprocess.sample_rate,
+        n_fft=cfg.preprocess.n_fft,
+        num_mels=cfg.preprocess.num_mels,
+        hop_size=cfg.preprocess.hop_size,
+        win_size=cfg.preprocess.win_size,
+        fmin=cfg.preprocess.fmin,
+        fmax=cfg.preprocess.fmax,
+    )
+    mel_model.eval()
+    mel_model.to(device)
+    return mel_model

models/tts/tadicodec/inference_tadicodec.py

@@ -0,0 +1,279 @@
+import os
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import librosa
+import safetensors
+import accelerate
+from huggingface_hub import snapshot_download
+
+from transformers import AutoTokenizer
+
+from models.tts.tadicodec.infer_utils import build_vocoder_model, build_mel_model
+from models.tts.tadicodec.modeling_tadicodec import TaDiCodec
+
+
+class TaDiCodecPipline(nn.Module):
+    def __init__(
+        self,
+        cfg,
+        model_path: str,
+        device: torch.device,
+        tokenizer_path: Optional[str] = None,
+        vocoder_ckpt_path: Optional[str] = None,
+    ):
+        super().__init__()
+        self.cfg = cfg
+        self.model_path = model_path
+        self.device = device
+
+        # tokenizer
+        self.tokenizer = (
+            AutoTokenizer.from_pretrained(tokenizer_path)
+            if tokenizer_path is not None
+            else os.path.join(model_path, "text_tokenizer")
+        )
+
+        # mel feature extractor
+        self.mel_model = build_mel_model(cfg, device)
+
+        # main model
+        tadiconfig = cfg.model.tadicodec
+
+        self.tadicodec = TaDiCodec(cfg=tadiconfig)
+        safetensors.torch.load_model(self.tadicodec, model_path, strict=False)
+        self.tadicodec.to(torch.float32)
+        self.tadicodec.to(device)
+        self.tadicodec.eval()
+
+        # vocoder
+        self.vocoder_model = build_vocoder_model(cfg, device)
+        v_path = (
+            vocoder_ckpt_path
+            if vocoder_ckpt_path
+            else os.path.join(model_path, "vocoder")
+        )
+        accelerate.load_checkpoint_and_dispatch(self.vocoder_model, v_path)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        ckpt_dir: str = "./ckpt/TaDiCodec",
+        device: Optional[torch.device] = None,
+        auto_download: bool = True,
+    ):
+        """Create a pipeline from a checkpoint directory or Hugging Face model ID.
+
+        Expected structure under `ckpt_dir`:
+          - config.json                # model and preprocess config
+          - model.safetensors          # TaDiCodec weights
+          - vocoder/                   # directory containing vocoder weights
+              model.safetensors or other *.safetensors
+          - text_tokenizer/               # directory containing text tokenizer
+              tokenizer.json
+              tokenizer_config.json
+
+        Args:
+            ckpt_dir: Directory containing `config.json`, `model.safetensors`, and `vocoder/`, or Hugging Face model ID.
+            device: Device to place models on. Defaults to CUDA if available else CPU.
+            auto_download: Whether to automatically download models from Hugging Face if not found locally
+
+        Returns:
+            TaDiCodecPipline
+        """
+        import os
+        import glob
+        from utils.util import load_config
+
+        resolved_device = (
+            device
+            if device is not None
+            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+
+        # Resolve checkpoint directory
+        resolved_ckpt_dir = cls._resolve_model_path(
+            ckpt_dir, auto_download=auto_download, model_type="tadicodec"
+        )
+
+        # Load config
+        config_path = os.path.join(resolved_ckpt_dir, "config.json")
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"Config not found: {config_path}")
+        cfg = load_config(config_path, lowercase=False)
+
+        # Resolve main model weights
+        model_weights_path = os.path.join(resolved_ckpt_dir, "model.safetensors")
+
+        # Resolve vocoder weights
+        vocoder_ckpt_path = os.path.join(resolved_ckpt_dir, "vocoder")
+
+        text_tokenizer_dir = os.path.join(resolved_ckpt_dir, "text_tokenizer")
+
+        return cls(
+            cfg=cfg,
+            model_path=model_weights_path,
+            device=resolved_device,
+            vocoder_ckpt_path=vocoder_ckpt_path,
+            tokenizer_path=text_tokenizer_dir,
+        )
+
+    @staticmethod
+    def _resolve_model_path(
+        model_path: str, auto_download: bool = True, model_type: str = "tadicodec"
+    ) -> str:
+        """
+        Resolve model path, downloading from Hugging Face if necessary
+
+        Args:
+            model_path: Local path or Hugging Face model ID
+            auto_download: Whether to auto-download from HF
+            model_type: Type of model ("tadicodec")
+
+        Returns:
+            Resolved local path
+        """
+        # If it's already a local path and exists, return as is
+        if os.path.exists(model_path):
+            return model_path
+
+        # If it looks like a Hugging Face model ID (contains '/')
+        if "/" in model_path and auto_download:
+            print(f"Downloading {model_type} model from Hugging Face: {model_path}")
+            try:
+                # Download to cache directory
+                cache_dir = os.path.join(
+                    os.path.expanduser("~"), ".cache", "huggingface", "hub"
+                )
+                downloaded_path = snapshot_download(
+                    repo_id=model_path,
+                    cache_dir=cache_dir,
+                    local_dir_use_symlinks=False,
+                )
+                print(
+                    f"Successfully downloaded {model_type} model to: {downloaded_path}"
+                )
+                return downloaded_path
+            except Exception as e:
+                print(f"Failed to download {model_type} model from Hugging Face: {e}")
+                raise ValueError(
+                    f"Could not download {model_type} model from {model_path}"
+                )
+
+        # If it's a local path that doesn't exist
+        if not os.path.exists(model_path):
+            if auto_download:
+                raise ValueError(
+                    f"Model path does not exist: {model_path}. Set auto_download=True to download from Hugging Face."
+                )
+            else:
+                raise FileNotFoundError(f"Model path does not exist: {model_path}")
+
+        return model_path
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        text: Optional[str] = None,
+        speech_path: Optional[str] = None,
+        prompt_text: Optional[str] = None,
+        prompt_speech_path: Optional[str] = None,
+        n_timesteps: int = 32,
+        return_code: bool = False,
+        cfg_scale: float = 2.0,
+    ):
+        # tokenize text
+        text_input_ids = self.tokenize_text(text, prompt_text)
+
+        # extract mel features
+        target_mel = self.extract_mel_feature(speech_path)
+        prompt_mel = (
+            self.extract_mel_feature(prompt_speech_path) if prompt_speech_path else None
+        )
+
+        # encode to codes from mel
+        if prompt_mel is not None:
+            vq_emb, indices = self.encode(torch.cat([prompt_mel, target_mel], dim=1))
+        else:
+            vq_emb, indices = self.encode(target_mel)
+
+        if return_code:
+            return indices
+
+        # decode mel from codes + optional text/prompt
+        rec_mel = self.decode(
+            vq_emb=vq_emb,
+            text_token_ids=text_input_ids,
+            prompt_mel=(
+                prompt_mel if prompt_mel is not None else target_mel[:, : 50 * 3]
+            ),
+            n_timesteps=n_timesteps,
+            cfg=cfg_scale,
+            rescale_cfg=0.75,
+        )
+
+        # vocoder
+        rec_audio = (
+            self.vocoder_model(rec_mel.transpose(1, 2)).detach().cpu().numpy()[0][0]
+        )
+        return rec_audio
+
+    def tokenize_text(
+        self, text: Optional[str] = None, prompt_text: Optional[str] = None
+    ):
+        if self.tokenizer is None or text is None:
+            return None
+        if prompt_text is not None:
+            text_token_ids = self.tokenizer(
+                prompt_text + text, return_tensors="pt", add_special_tokens=False
+            ).input_ids.to(self.device)
+        else:
+            text_token_ids = self.tokenizer(
+                text, return_tensors="pt", add_special_tokens=False
+            ).input_ids.to(self.device)
+        return text_token_ids
+
+    @torch.no_grad()
+    def extract_mel_feature(self, speech_path: Optional[str]):
+        assert speech_path is not None and os.path.exists(
+            speech_path
+        ), f"Invalid speech_path: {speech_path}"
+        speech = librosa.load(speech_path, sr=24000)[0]
+        speech = torch.tensor(speech).to(self.device).unsqueeze(0)
+        mel_feature = self.mel_model(speech)  # (B, n_mels, T)
+        mel_feature = mel_feature.transpose(1, 2)  # (B, T, n_mels)
+        mel_feature = (mel_feature - self.cfg.preprocess.mel_mean) / math.sqrt(
+            self.cfg.preprocess.mel_var
+        )
+        return mel_feature
+
+    @torch.no_grad()
+    def encode(self, mel_feat: torch.Tensor):
+        vq_emb, indices = self.tadicodec.encode(
+            mel_feat, torch.ones(mel_feat.shape[0], mel_feat.shape[1]).to(self.device)
+        )
+        return vq_emb, indices
+
+    @torch.no_grad()
+    def decode(
+        self,
+        vq_emb: Optional[torch.Tensor] = None,
+        indices: Optional[torch.Tensor] = None,
+        text_token_ids: Optional[torch.Tensor] = None,
+        prompt_mel: Optional[torch.Tensor] = None,
+        n_timesteps: int = 32,
+        cfg: float = 1.0,
+        rescale_cfg: float = 0.75,
+    ):
+        rec_mel = self.tadicodec.reverse_diffusion(
+            vq_emb=vq_emb,
+            indices=indices,
+            text_ids=text_token_ids,
+            prompt_mel=prompt_mel,
+            n_timesteps=n_timesteps,
+            cfg=cfg,
+            rescale_cfg=rescale_cfg,
+        )
+        return rec_mel

models/tts/tadicodec/llama_nar_prefix.py

@@ -0,0 +1,572 @@
+from transformers import LlamaConfig, LlamaForCausalLM, LlamaModel
+import torch
+import torch.nn.functional as F
+import numpy as np
+import os
+import torch.nn as nn
+from typing import List, Optional, Tuple, Union
+import math
+
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer
+from transformers.models.llama.modeling_llama import BaseModelOutputWithPast
+
+
+class SinusoidalPosEmb(nn.Module):
+    """
+    Sinusoidal Positional Embedding module.
+
+    This module generates sinusoidal positional embeddings for a given 1D input tensor,
+    which is commonly used for representing timesteps in diffusion models.
+    """
+
+    def __init__(self, dim: int):
+        """
+        Initializes the SinusoidalPosEmb module.
+
+        Args:
+            dim (int): The dimension of the embedding.
+        """
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Generates the positional embedding.
+
+        Args:
+            x (torch.Tensor): A 1D tensor of positions (e.g., timesteps), shape `(batch_size,)`.
+
+        Returns:
+            torch.Tensor: The positional embeddings, shape `(batch_size, dim)`.
+        """
+        device = x.device
+        half_dim = self.dim // 2
+        # Calculate the embedding frequencies
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        # Create the embedding matrix by multiplying positions with frequencies
+        emb = x[:, None] * emb[None, :] * 1.0
+        # Concatenate sine and cosine components to form the final embedding
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class LlamaAdaptiveRMSNorm(nn.Module):
+    """
+    Adaptive Root Mean Square Layer Normalization.
+
+    This is a variant of RMSNorm where the scaling factor (weight) is adaptively
+    predicted from a conditional embedding, allowing for conditional modulation
+    of the normalized hidden states.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        eps: float = 1e-6,
+        dim_cond: int = 1024,
+        use_cond: bool = True,
+    ):
+        """
+        Initializes the LlamaAdaptiveRMSNorm module.
+
+        Args:
+            hidden_size (int): The dimension of the hidden states to be normalized.
+            eps (float): A small value added to the variance for numerical stability.
+            dim_cond (int): The dimension of the conditional embedding.
+            use_cond (bool): If True, use conditional embedding to modulate the output.
+        """
+        super().__init__()
+        self.use_cond = use_cond
+        if self.use_cond:
+            # Linear layer to project the conditional embedding to the hidden size
+            self.to_weight = nn.Linear(dim_cond, hidden_size)
+            # Initialize weights to zero and bias to one for an identity transformation at the start
+            nn.init.zeros_(self.to_weight.weight)
+            nn.init.ones_(self.to_weight.bias)
+        self.variance_epsilon = eps
+        self._is_hf_initialized = True  # Disable automatic Hugging Face initialization
+
+    def forward(
+        self, hidden_states: torch.Tensor, cond_embedding: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Applies the adaptive RMS normalization.
+
+        Args:
+            hidden_states (torch.Tensor): The input tensor, shape `(batch, seq_len, hidden_size)`.
+            cond_embedding (Optional[torch.Tensor]): The conditional embedding, shape `(batch, dim_cond)` or `(batch, seq_len, dim_cond)`.
+
+        Returns:
+            torch.Tensor: The normalized and modulated hidden states.
+        """
+        input_dtype = hidden_states.dtype
+        # Calculate variance and normalize the hidden states
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        # Apply conditional modulation if enabled
+        if self.use_cond:
+            # Project conditional embedding to get the scaling weight
+            weight = self.to_weight(cond_embedding)
+            # Unsqueeze if the conditional embedding is per-batch instead of per-token
+            if len(weight.shape) == 2:
+                weight = weight.unsqueeze(1)
+
+            # Apply the learned scaling factor
+            return (weight * hidden_states).to(input_dtype)
+        else:
+            return hidden_states
+
+
+class LlamaNARDecoderLayer(LlamaDecoderLayer):
+    """
+    A Non-Autoregressive (NAR) Llama Decoder Layer using adaptive layer normalization.
+
+    This class overrides the standard LlamaDecoderLayer to replace its RMSNorm
+    modules with LlamaAdaptiveRMSNorm, allowing it to be conditioned on an external embedding
+    (e.g., from diffusion timesteps).
+    """
+
+    def __init__(self, config: LlamaConfig, layer_idx: int, use_cond: bool = True):
+        """
+        Overrides the LlamaDecoderLayer to use adaptive layer normalization.
+
+        Args:
+            config (LlamaConfig): The configuration object for the Llama model.
+            layer_idx (int): The index of the layer.
+            use_cond (bool): Whether to use adaptive conditioning in the layer norms.
+        """
+        super().__init__(config, layer_idx)  # init attention, mlp, etc.
+        # Override the standard layer norms with our adaptive versions
+        self.input_layernorm = LlamaAdaptiveRMSNorm(
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+            dim_cond=config.hidden_size,
+            use_cond=use_cond,
+        )
+        self.post_attention_layernorm = LlamaAdaptiveRMSNorm(
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+            dim_cond=config.hidden_size,
+            use_cond=use_cond,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cond_embedding: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        """
+        Forward pass for the NAR decoder layer, including conditional embedding.
+
+        Args:
+            hidden_states (torch.FloatTensor): Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            cond_embedding (torch.Tensor): Conditional embedding for adaptive normalization.
+            attention_mask (Optional[torch.Tensor]): Attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            position_ids (Optional[torch.LongTensor]): Indices of positions of each input sequence tokens.
+            past_key_value (Optional[Tuple[torch.Tensor]]): Cached past key and value projection states.
+            output_attentions (Optional[bool]): Whether or not to return the attentions tensors of all attention layers.
+            use_cache (Optional[bool]): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding.
+
+        Returns:
+            Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+            A tuple containing the output hidden states and optionally the attention weights and past key/value states.
+        """
+        residual = hidden_states
+
+        # Apply adaptive pre-attention layer norm
+        hidden_states = self.input_layernorm(
+            hidden_states, cond_embedding=cond_embedding
+        )
+
+        # Self Attention block
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected block
+        residual = hidden_states
+        # Apply adaptive post-attention layer norm
+        hidden_states = self.post_attention_layernorm(
+            hidden_states, cond_embedding=cond_embedding
+        )
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+class DiffLlamaPrefix(LlamaModel):
+    """
+    A Llama-based non-autoregressive transformer model for diffusion tasks.
+
+    This model uses a Llama architecture but modifies it for non-autoregressive generation.
+    Key features:
+    1.  Non-causal (fully-visible) attention mask.
+    2.  Adaptive layer normalization conditioned on diffusion timesteps.
+    3.  Ability to be conditioned on text embeddings, which are prepended as a prefix.
+    4.  Input and output linear projection layers for feature mapping.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        num_heads: int = 16,
+        num_layers: int = 16,
+        in_dim: Optional[int] = None,
+        out_dim: Optional[int] = None,
+        use_text_emb: bool = True,
+        use_diff_step: bool = True,
+        use_cond: bool = True,
+        config: LlamaConfig = LlamaConfig(
+            vocab_size=0, hidden_size=256, num_attention_heads=1, num_hidden_layers=1
+        ),
+    ):
+        """
+        Initializes the DiffLlamaPrefix model.
+
+        Args:
+            hidden_size (int): The hidden dimension of the transformer.
+            num_heads (int): The number of attention heads.
+            num_layers (int): The number of transformer layers.
+            in_dim (Optional[int]): Dimension of the input features. If set, an input linear layer is created.
+            out_dim (Optional[int]): Dimension of the output features. If set, an output linear layer is created.
+            use_text_emb (bool): Whether to use text embeddings as a conditional prefix.
+            use_diff_step (bool): Whether to use diffusion timestep conditioning.
+            use_cond (bool): Whether to use an additional per-token conditional input `cond`.
+            config (LlamaConfig): A LlamaConfig object. A default is provided for convenience.
+        """
+        super().__init__(config)
+
+        self.use_text_emb = use_text_emb
+        self.use_diff_step = use_diff_step
+        self.use_cond = use_cond
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # Create a stack of non-autoregressive Llama layers
+        self.layers = nn.ModuleList(
+            [
+                LlamaNARDecoderLayer(
+                    LlamaConfig(
+                        hidden_size=hidden_size,
+                        num_attention_heads=num_heads,
+                        max_position_embeddings=4096,
+                        intermediate_size=hidden_size * 4,
+                    ),
+                    layer_idx=i,
+                    use_cond=use_cond,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        # Final adaptive layer norm
+        self.norm = LlamaAdaptiveRMSNorm(
+            hidden_size,
+            dim_cond=hidden_size,
+            use_cond=(use_cond or use_diff_step),
+        )
+
+        # MLP for processing text embedding condition
+        if self.use_text_emb:
+            self.text_mlp = nn.Sequential(
+                nn.Linear(hidden_size, hidden_size * 4),
+                nn.SiLU(),
+                nn.Linear(hidden_size * 4, hidden_size),
+            )
+
+        # Modules for processing diffusion timestep condition
+        if self.use_diff_step:
+            self.diff_step_embedding = SinusoidalPosEmb(hidden_size)
+            self.diff_step_mlp = nn.Sequential(
+                nn.Linear(hidden_size, hidden_size * 4),
+                nn.SiLU(),
+                nn.Linear(hidden_size * 4, hidden_size),
+            )
+
+        # MLP for processing the general per-token `cond` condition
+        if self.use_cond:
+            self.cond_mlp = nn.Sequential(
+                nn.Linear(hidden_size, hidden_size * 4),
+                nn.SiLU(),
+                nn.Linear(hidden_size * 4, hidden_size),
+            )
+
+        # Ensure all layers use the correct adaptive norm configuration
+        for layer in self.layers:
+            layer.input_layernorm = LlamaAdaptiveRMSNorm(
+                hidden_size,
+                dim_cond=hidden_size,
+                use_cond=(use_cond or use_diff_step),
+            )
+            layer.post_attention_layernorm = LlamaAdaptiveRMSNorm(
+                hidden_size,
+                dim_cond=hidden_size,
+                use_cond=(use_cond or use_diff_step),
+            )
+
+        # We handle embeddings manually, so disable the default token embedder
+        self.embed_tokens = None
+
+        # Optional input and output projection layers
+        if self.in_dim is not None:
+            self.in_linear = nn.Linear(self.in_dim, hidden_size)
+        if self.out_dim is not None:
+            self.out_linear = nn.Linear(hidden_size, self.out_dim)
+
+        self.post_init()
+
+    def _prepare_decoder_attention_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_shape: Tuple[int, int, int],
+        inputs_embeds: torch.Tensor,
+        past_key_values_length: int,
+    ) -> Optional[torch.Tensor]:
+        """
+        Creates a non-causal (fully-visible) attention mask.
+
+        This method overrides the default causal mask creation in LlamaModel.
+        It converts a 2D padding mask `[bsz, seq_len]` into a 4D attention mask
+        `[bsz, 1, tgt_seq_len, src_seq_len]` suitable for self-attention,
+        without applying a causal triangle.
+
+        Args:
+            attention_mask (torch.Tensor): The 2D padding mask.
+            input_shape (Tuple[int, int, int]): The shape of the input embeddings.
+            inputs_embeds (torch.Tensor): The input embeddings tensor.
+            past_key_values_length (int): The length of any cached key-values.
+
+        Returns:
+            Optional[torch.Tensor]: The 4D attention mask, or None if the input mask is None.
+        """
+        combined_attention_mask = None
+
+        def _expand_mask(
+            mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None
+        ) -> torch.Tensor:
+            """
+            Expands a 2D attention mask to a 4D attention mask.
+            """
+            bsz, src_len = mask.size()
+            tgt_len = tgt_len if tgt_len is not None else src_len
+
+            expanded_mask = (
+                mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+            )
+
+            # Invert the mask: 1.0 for valid tokens, 0.0 for padded tokens
+            # and then convert to additive mask format (-inf for masked positions)
+            inverted_mask = 1.0 - expanded_mask
+            return inverted_mask.masked_fill(
+                inverted_mask.to(torch.bool), torch.finfo(dtype).min
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(
+                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+            ).to(inputs_embeds.device)
+            combined_attention_mask = (
+                expanded_attn_mask
+                if combined_attention_mask is None
+                else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        cond: Optional[torch.Tensor] = None,
+        diffusion_step: Optional[torch.Tensor] = None,
+        text_embedding: Optional[torch.Tensor] = None,
+        text_mask: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[torch.Tensor, Tuple]:
+        """
+        Forward pass of the DiffLlamaPrefix model.
+
+        Args:
+            x (torch.Tensor): The primary input tensor, e.g., noisy data in diffusion `(batch, seq_len, in_dim)`.
+            x_mask (torch.Tensor): The padding mask for `x`, shape `(batch, seq_len)`.
+            cond (Optional[torch.Tensor]): Additional per-token conditional input, shape `(batch, seq_len, hidden_size)`.
+            diffusion_step (Optional[torch.Tensor]): Diffusion timesteps, shape `(batch,)`.
+            text_embedding (Optional[torch.Tensor]): Text embeddings to be used as a prefix, shape `(batch, text_len, hidden_size)`.
+            text_mask (Optional[torch.Tensor]): The padding mask for `text_embedding`, shape `(batch, text_len)`.
+            input_ids, attention_mask, etc.: Standard Hugging Face arguments (mostly unused here but kept for compatibility).
+
+        Returns:
+            Union[torch.Tensor, Tuple]: The final output tensor `(batch, seq_len, out_dim)`.
+            If `output_hidden_states` is True, returns a tuple `(output_tensor, all_hidden_states)`.
+        """
+        # 1. Input and Conditional Projections
+        if self.in_dim is not None:
+            x = self.in_linear(x)
+
+        if self.use_cond and cond is not None:
+            cond_embedding = self.cond_mlp(cond)  # (B, T, C)
+            x = x + cond_embedding
+
+        # 2. Prepend Text Embedding Prefix
+        if self.use_text_emb and text_embedding is not None:
+            text_embedding = self.text_mlp(text_embedding)  # (B, T, C)
+            text_length = text_embedding.shape[1]
+            # Concatenate text prefix and main input
+            inputs_embeds = torch.cat([text_embedding, x], dim=1)
+            attention_mask = torch.cat([text_mask, x_mask], dim=1)
+        else:
+            inputs_embeds = x
+            attention_mask = x_mask
+            text_length = 0
+
+        # 3. Process Diffusion Step Embedding for Adaptive Norm
+        if self.use_diff_step and diffusion_step is not None:
+            # Convert scalar timesteps to vector embeddings and project with an MLP
+            diffusion_step_emb = self.diff_step_embedding(diffusion_step).to(x.device)
+            diffusion_step_emb = self.diff_step_mlp(diffusion_step_emb)  # (B, C)
+        else:
+            diffusion_step_emb = None
+
+        # 4. Standard Transformer Preamble (adapted from LlamaModel)
+        batch_size, seq_length, _ = inputs_embeds.shape
+
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length,
+                seq_length + past_key_values_length,
+                dtype=torch.long,
+                device=device,
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past),
+                dtype=torch.bool,
+                device=inputs_embeds.device,
+            )
+        # Create the non-causal attention mask
+        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask,
+            (batch_size, seq_length),
+            inputs_embeds,
+            past_key_values_length,
+        )
+
+        hidden_states = inputs_embeds
+
+        # 5. Transformer Decoder Layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                # Store hidden states before the layer, excluding the text prefix part
+                all_hidden_states += (hidden_states[:, text_length:],)
+
+            past_key_value = (
+                past_key_values[idx] if past_key_values is not None else None
+            )
+
+            # Pass the processed diffusion step embedding to the adaptive layer
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cond_embedding=diffusion_step_emb,
+            )
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        # 6. Final Normalization and Output Processing
+        hidden_states = self.norm(hidden_states, cond_embedding=diffusion_step_emb)
+
+        if output_hidden_states:
+            # Add hidden states from the last decoder layer
+            all_hidden_states += (hidden_states[:, text_length:],)
+
+        next_cache = next_decoder_cache if use_cache else None
+
+        # Remove the text prefix from the final output sequence
+        hidden_states = hidden_states[
+            :,
+            text_length:,
+        ]
+
+        # Apply final output projection if specified
+        if self.out_dim is not None:
+            hidden_states = self.out_linear(hidden_states)
+
+        # 7. Return results (simplified for this snippet)
+        if output_hidden_states:
+            return hidden_states, all_hidden_states
+        else:
+            return hidden_states

models/tts/tadicodec/modeling_tadicodec.py

@@ -0,0 +1,641 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import math
+from einops import rearrange
+from typing import Optional, Dict, Any
+
+from models.tts.tadicodec.llama_nar_prefix import DiffLlamaPrefix
+from models.codec.amphion_codec.quantize.bsq import (
+    BinarySphericalQuantizer,
+    SimpleQuantizer,
+)
+import torch.nn.functional as F
+
+
+class TaDiCodec(nn.Module):
+    """
+    TaDiCodec: A diffusion-based codec model for Text-to-Speech (TTS)
+    that uses a non-autoregressive Llama-style transformer architecture.
+
+    It consists of:
+    1. An Encoder that processes input features (e.g., mel-spectrograms or SSL features)
+       into a latent representation.
+    2. A Vector Quantizer (VQ) that discretizes the latent representation into codes.
+    3. A Decoder that generates the output feature (e.g., mel-spectrogram) from the codes,
+       optional text conditioning, and a prompt, using a flow-matching diffusion process.
+    """
+
+    def __init__(
+        self,
+        mel_dim: int = 128,
+        in_dim: int = 128,
+        hidden_size: int = 1024,
+        encoder_num_layers: int = 8,
+        decoder_num_layers: int = 16,
+        num_heads: int = 16,
+        cond_drop_p: float = 0.2,  # drop code for decoder
+        context_drop_p: float = 0.2,  # drop context (mel) for decoder
+        down_sample_factor: int = 8,  # down sample factor for vq
+        vq_emb_dim: int = 14,  # codebook size 2^vq_emb_dim, 2^14 = 16384
+        use_text_cond: bool = True,  # use text cond for decoder
+        text_vocab_size: int = 32100,  # vocab size
+        cond_dim: int = 1024,
+        cond_scale_factor: int = 1,
+        sigma: float = 1e-5,
+        time_scheduler: str = "linear",
+        use_vq: bool = True,
+        vq_type: str = "bsq",
+        use_repa_loss: bool = False,
+        cfg: Optional[Any] = None,
+    ):
+        """
+        Initializes the TaDiCodec model.
+
+        Args:
+            mel_dim (int): Dimension of the mel-spectrogram.
+            in_dim (int): Dimension of the encoder's input features.
+            hidden_size (int): Hidden size of the transformer models.
+            encoder_num_layers (int): Number of layers in the encoder transformer.
+            decoder_num_layers (int): Number of layers in the decoder transformer.
+            num_heads (int): Number of attention heads in the transformers.
+            cond_drop_p (float): Dropout probability for the VQ code condition in the decoder.
+            context_drop_p (float): Dropout probability for the prompt context in the decoder.
+            down_sample_factor (int): Factor by which to downsample the latent representation before VQ.
+            vq_emb_dim (int): Dimension of the vector quantizer's embedding space.
+            use_text_cond (bool): Whether to use text embeddings as a condition in the decoder.
+            text_vocab_size (int): Size of the text vocabulary.
+            cond_dim (int): Dimension of the conditional input.
+            cond_scale_factor (int): Scaling factor for the condition.
+            sigma (float): Small constant used in the flow matching formula.
+            time_scheduler (str): Type of time scheduler for diffusion (e.g., 'linear').
+            use_vq (bool): Whether to use vector quantization.
+            vq_type (str): Type of vector quantizer ('bsq' or 'simple').
+            use_repa_loss (bool): Whether to use the representational alignment loss.
+            cfg (Optional[Any]): A configuration object that can override the default parameters.
+        """
+        super().__init__()
+
+        # Override parameters with values from the config object if provided
+        mel_dim = (
+            cfg.mel_dim if cfg is not None and hasattr(cfg, "mel_dim") else mel_dim
+        )
+        in_dim = cfg.in_dim if cfg is not None and hasattr(cfg, "in_dim") else in_dim
+        hidden_size = (
+            cfg.hidden_size
+            if cfg is not None and hasattr(cfg, "hidden_size")
+            else hidden_size
+        )
+        encoder_num_layers = (
+            cfg.encoder_num_layers
+            if cfg is not None and hasattr(cfg, "encoder_num_layers")
+            else encoder_num_layers
+        )
+        decoder_num_layers = (
+            cfg.decoder_num_layers
+            if cfg is not None and hasattr(cfg, "decoder_num_layers")
+            else decoder_num_layers
+        )
+        num_heads = (
+            cfg.num_heads
+            if cfg is not None and hasattr(cfg, "num_heads")
+            else num_heads
+        )
+        cond_drop_p = (
+            cfg.cond_drop_p
+            if cfg is not None and hasattr(cfg, "cond_drop_p")
+            else cond_drop_p
+        )
+        context_drop_p = (
+            cfg.context_drop_p
+            if cfg is not None and hasattr(cfg, "context_drop_p")
+            else context_drop_p
+        )
+        down_sample_factor = (
+            cfg.down_sample_factor
+            if cfg is not None and hasattr(cfg, "down_sample_factor")
+            else down_sample_factor
+        )
+        vq_emb_dim = (
+            cfg.vq_emb_dim
+            if cfg is not None and hasattr(cfg, "vq_emb_dim")
+            else vq_emb_dim
+        )
+        use_text_cond = (
+            cfg.use_text_cond
+            if cfg is not None and hasattr(cfg, "use_text_cond")
+            else use_text_cond
+        )
+        text_vocab_size = (
+            cfg.text_vocab_size
+            if cfg is not None and hasattr(cfg, "text_vocab_size")
+            else text_vocab_size
+        )
+        cond_dim = (
+            cfg.cond_dim if cfg is not None and hasattr(cfg, "cond_dim") else cond_dim
+        )
+        cond_scale_factor = (
+            cfg.cond_scale_factor
+            if cfg is not None and hasattr(cfg, "cond_scale_factor")
+            else cond_scale_factor
+        )
+        sigma = cfg.sigma if cfg is not None and hasattr(cfg, "sigma") else sigma
+        time_scheduler = (
+            cfg.time_scheduler
+            if cfg is not None and hasattr(cfg, "time_scheduler")
+            else time_scheduler
+        )
+        use_vq = cfg.use_vq if cfg is not None and hasattr(cfg, "use_vq") else use_vq
+        vq_type = (
+            cfg.vq_type if cfg is not None and hasattr(cfg, "vq_type") else vq_type
+        )
+        use_repa_loss = (
+            cfg.use_repa_loss
+            if cfg is not None and hasattr(cfg, "use_repa_loss")
+            else use_repa_loss
+        )
+
+        self.mel_dim = mel_dim
+        self.in_dim = in_dim
+        self.hidden_size = hidden_size
+        self.encoder_num_layers = encoder_num_layers
+        self.decoder_num_layers = decoder_num_layers
+        self.num_heads = num_heads
+        self.cond_drop_p = cond_drop_p
+        self.context_drop_p = context_drop_p
+        self.vq_emb_dim = vq_emb_dim
+        self.down_sample_factor = down_sample_factor
+        self.use_text_cond = use_text_cond
+        self.text_vocab_size = text_vocab_size
+        self.cond_dim = cond_dim
+        self.cond_scale_factor = cond_scale_factor
+        self.sigma = sigma
+        self.time_scheduler = time_scheduler
+        self.use_vq = use_vq
+        self.vq_type = vq_type
+        self.use_repa_loss = use_repa_loss
+
+        # Text embedding layer
+        if self.use_text_cond:
+            self.text_emb = nn.Embedding(text_vocab_size, hidden_size)
+
+        # VQ related layers
+        self.vq_in_linear = nn.Linear(hidden_size, vq_emb_dim)
+        if self.use_vq:
+            if self.vq_type == "bsq":
+                self.bsq = BinarySphericalQuantizer(embed_dim=vq_emb_dim)
+            else:
+                self.bsq = SimpleQuantizer(embed_dim=vq_emb_dim)
+        self.vq_out_linear = nn.Linear(vq_emb_dim, hidden_size)
+
+        # Repa (Representational Alignment) MLP for auxiliary loss
+        if self.use_repa_loss:
+            self.repa_mlp = nn.Sequential(
+                nn.Linear(hidden_size, hidden_size * 2),
+                nn.GELU(),
+                nn.Linear(hidden_size * 2, 1024),
+            )
+            self.repa_layer_idx = (
+                6  # The decoder layer from which to extract hidden states
+            )
+
+        self.reset_parameters()
+
+        # Encoder: A non-autoregressive Llama-style model without time/text conditioning
+        self.encoder = DiffLlamaPrefix(
+            hidden_size=hidden_size,
+            num_layers=encoder_num_layers,
+            num_heads=num_heads,
+            in_dim=self.in_dim,
+            out_dim=None,  # Outputs hidden states for VQ
+            use_text_emb=False,
+            use_diff_step=False,
+            use_cond=False,
+        )
+
+        # Decoder: A non-autoregressive Llama-style model with time, text, and code conditioning
+        self.decoder = DiffLlamaPrefix(
+            hidden_size=hidden_size,
+            num_layers=decoder_num_layers,
+            num_heads=num_heads,
+            in_dim=self.mel_dim,
+            out_dim=self.mel_dim,
+            use_text_emb=use_text_cond,
+            use_diff_step=True,
+            use_cond=True,
+        )
+
+    @torch.no_grad()
+    def forward_diffusion(
+        self, x: torch.Tensor, t: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Performs the forward diffusion process based on flow matching.
+        It takes the clean data `x` and a timestep `t` to produce a noisy sample `xt`.
+        It also creates a prompt/target mask.
+
+        Args:
+            x (torch.Tensor): The clean input data (e.g., mel-spectrogram), shape `(B, T, mel_dim)`.
+            t (torch.Tensor): The diffusion timestep for each sample in the batch, shape `(B,)`.
+
+        Returns:
+            Tuple[torch.Tensor, ...]:
+                - xt (torch.Tensor): The noisy sample at time `t`, shape `(B, T, mel_dim)`.
+                - z (torch.Tensor): The noise vector used, drawn from N(0, I).
+                - new_t (torch.Tensor): The original `t` tensor.
+                - prompt_len (torch.Tensor): The length of the prompt for each sample.
+                - mask (torch.Tensor): A mask where 1 indicates the target (noisy) region and 0 indicates the prompt (clean) region.
+        """
+        new_t = t
+        t = t.unsqueeze(-1).unsqueeze(-1)  # Reshape for broadcasting
+        z = torch.randn(
+            x.shape, dtype=x.dtype, device=x.device, requires_grad=False
+        )  # (B, T, mel_dim)
+
+        context_drop_p = self.context_drop_p
+
+        # Randomly decide the length of the prompt (un-noised context)
+        if torch.rand(1) > context_drop_p:
+            prompt_len = torch.randint(
+                min(x.shape[1] // 4, 5), int(x.shape[1] * 0.4), (x.shape[0],)
+            ).to(x.device)
+        else:
+            # Drop the context entirely by setting prompt length to 0
+            prompt_len = torch.zeros(x.shape[0], device=x.device)
+
+        # Create a mask to distinguish prompt from target
+        is_prompt = torch.zeros_like(x[:, :, 0])  # (B, T)
+        col_indices = torch.arange(is_prompt.shape[1], device=prompt_len.device).repeat(
+            is_prompt.shape[0], 1
+        )  # (B, T)
+        is_prompt[col_indices < prompt_len.unsqueeze(1)] = 1  # 1 if it's a prompt frame
+
+        mask = torch.ones_like(x[:, :, 0])  # Mask is 1 for target, 0 for prompt
+        mask[is_prompt.bool()] = 0
+        mask = mask.unsqueeze(-1)  # (B, T, 1)
+
+        # Flow matching formula: xt = (1 - (1 - sigma) * t) * x0 + t * x
+        # where x0 ~ N(0, 1) and x is the clean data sample.
+        # The equation is applied only to the target region.
+        xt = ((1 - (1 - self.sigma) * t) * z + t * x) * mask + x * (1 - mask)
+
+        return xt, z, new_t, prompt_len, mask
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        x_in: Optional[torch.Tensor] = None,
+        text_ids: Optional[torch.Tensor] = None,
+        text_mask: Optional[torch.Tensor] = None,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        """
+        The main training-time forward pass of the model.
+
+        Args:
+            x (torch.Tensor): The target mel-spectrogram, shape `(B, T, mel_dim)`.
+            x_mask (torch.Tensor): Padding mask for `x`, shape `(B, T)`.
+            x_in (Optional[torch.Tensor]): Optional input for the encoder (e.g., SSL features). If None, `x` is used.
+            text_ids (Optional[torch.Tensor]): Input text token IDs for conditioning, shape `(B, text_len)`.
+            text_mask (Optional[torch.Tensor]): Padding mask for `text_ids`.
+
+        Returns:
+            Dict[str, Optional[torch.Tensor]]: A dictionary containing various tensors for loss computation,
+            such as the predicted flow, the final predicted mel, the noise target, etc.
+        """
+        # 1. Encoder pass
+        if x_in is None:
+            # Use target mel as encoder input
+            vq_emb_pre = self.encoder(x=x, x_mask=x_mask)
+        else:
+            # Use provided SSL features as encoder input
+            vq_emb_pre = self.encoder(x=x_in, x_mask=x_mask)
+
+        # 2. Downsampling before VQ
+        _, T, _ = vq_emb_pre.shape
+        vq_emb_pre = vq_emb_pre.transpose(1, 2)
+        vq_emb_pre = F.interpolate(
+            vq_emb_pre, size=T // self.down_sample_factor, mode="linear"
+        )
+        vq_emb_pre = vq_emb_pre.transpose(1, 2)
+
+        # 3. Vector Quantization
+        vq_emb_pre = self.vq_in_linear(vq_emb_pre)
+        vq_emb_pre = F.normalize(vq_emb_pre, dim=-1)  # L2 normalize before quantization
+
+        if self.use_vq:
+            vq_emb, vq_loss, info = self.bsq(vq_emb_pre)
+            commit_loss = info["commit_loss"]
+        else:
+            vq_emb = vq_emb_pre
+            vq_loss = torch.tensor(0.0, device=x.device)
+            commit_loss = torch.tensor(0.0, device=x.device)
+            info = None
+
+        vq_emb_post = self.vq_out_linear(vq_emb)
+
+        # 4. Upsampling after VQ
+        vq_emb_post = vq_emb_post.transpose(1, 2)
+        vq_emb_post = F.interpolate(vq_emb_post, size=T, mode="linear")
+        vq_emb_post = vq_emb_post.transpose(1, 2)
+
+        # 5. Decoder with flow matching
+        # Sample a random timestep t for each item in the batch
+        t = torch.rand(x.shape[0], device=x.device, requires_grad=False)
+        t = torch.clamp(t, 1e-5, 1.0)  # Clamp to avoid numerical issues at boundaries
+
+        # Perform forward diffusion to get the noisy input `xt` and the noise `z`
+        xt, z, new_t, prompt_len, mask = self.forward_diffusion(x, t)
+        noise = z
+
+        # 6. Prepare conditions for the decoder
+        if self.use_text_cond:
+            text_emb = self.text_emb(text_ids)
+        else:
+            text_emb = None
+
+        # Use the upsampled VQ embedding as the primary condition
+        cond_emb = vq_emb_post
+        # Apply condition dropout for classifier-free guidance training
+        if torch.rand(1) < self.cond_drop_p:
+            cond_emb = torch.zeros_like(cond_emb)
+
+        # 7. Decoder pass
+        if self.use_repa_loss:
+            # If using Repa loss, we need to output intermediate hidden states
+            flow_pred, hidden_states = self.decoder(
+                x=xt,
+                x_mask=x_mask,
+                text_embedding=text_emb,
+                text_mask=text_mask,
+                cond=cond_emb,
+                diffusion_step=new_t,
+                output_hidden_states=True,
+            )
+            ssl_feat_pred = self.repa_mlp(hidden_states[self.repa_layer_idx])
+        else:
+            flow_pred = self.decoder(
+                x=xt,
+                x_mask=x_mask,
+                text_embedding=text_emb,
+                text_mask=text_mask,
+                cond=cond_emb,
+                diffusion_step=new_t,
+            )
+            ssl_feat_pred = None
+
+        # Predict the clean data `x0_pred` from the noisy input `xt` and the predicted flow
+        # x_pred = xt + (1 - t) * flow_pred
+        x_pred = xt + (1 - t.unsqueeze(-1).unsqueeze(-1)) * flow_pred
+
+        # Final mask should consider both the prompt/target mask and the original padding mask
+        final_mask = mask * x_mask.unsqueeze(-1)
+
+        return {
+            "noise": noise,
+            "x": x,
+            "flow_pred": flow_pred,
+            "x_pred": x_pred,
+            "final_mask": final_mask,
+            "prompt_len": prompt_len,
+            "vq_loss": vq_loss,
+            "commit_loss": commit_loss,
+            "ssl_feat_pred": ssl_feat_pred,
+        }
+
+    @torch.no_grad()
+    def encode(
+        self, x: torch.Tensor, x_mask: torch.Tensor
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Encodes an input feature `x` into discrete VQ codes. (Inference)
+
+        Args:
+            x (torch.Tensor): Input feature, shape `(B, T, in_dim)`.
+            x_mask (torch.Tensor): Padding mask for `x`, shape `(B, T)`.
+
+        Returns:
+            Tuple[torch.Tensor, Optional[torch.Tensor]]:
+                - vq_emb (torch.Tensor): The quantized continuous embeddings.
+                - indices (Optional[torch.Tensor]): The discrete code indices, if VQ is used.
+        """
+        # Encoder pass
+        vq_emb_pre = self.encoder(x=x, x_mask=x_mask)
+
+        # Downsampling
+        _, T, _ = vq_emb_pre.shape
+        vq_emb_pre = vq_emb_pre.transpose(1, 2)
+        vq_emb_pre = F.interpolate(
+            vq_emb_pre, size=T // self.down_sample_factor, mode="linear"
+        )
+        vq_emb_pre = vq_emb_pre.transpose(1, 2)
+
+        # VQ
+        vq_emb_pre = self.vq_in_linear(vq_emb_pre)
+        vq_emb_pre = F.normalize(vq_emb_pre, dim=-1)  # L2 norm
+
+        if self.use_vq:
+            vq_emb, _, info = self.bsq(vq_emb_pre)
+            indices = info["indices"]
+        else:
+            vq_emb = vq_emb_pre
+            indices = None
+
+        return vq_emb, indices
+
+    @torch.no_grad()
+    def index2vq(self, indices: torch.Tensor) -> torch.Tensor:
+        """
+        Converts VQ code indices back to continuous embeddings.
+
+        Args:
+            indices (torch.Tensor): The discrete code indices.
+
+        Returns:
+            torch.Tensor: The corresponding continuous codebook embeddings.
+        """
+        return self.bsq.get_codebook_entry(indices).float()
+
+    @torch.no_grad()
+    def reverse_diffusion(
+        self,
+        vq_emb: Optional[torch.Tensor] = None,
+        indices: Optional[torch.Tensor] = None,
+        text_ids: Optional[torch.Tensor] = None,
+        prompt_mel: Optional[torch.Tensor] = None,
+        x_mask: Optional[torch.Tensor] = None,
+        prompt_mask: Optional[torch.Tensor] = None,
+        text_mask: Optional[torch.Tensor] = None,
+        n_timesteps: int = 32,
+        cfg: float = 1.0,
+        rescale_cfg: float = 0.75,
+    ) -> torch.Tensor:
+        """
+        Performs the reverse diffusion process to generate mel-spectrograms from conditions. (Inference)
+
+        Args:
+            vq_emb (Optional[torch.Tensor]): Pre-quantized embeddings.
+            indices (Optional[torch.Tensor]): Discrete VQ code indices. If provided, `vq_emb` is ignored.
+            text_ids (Optional[torch.Tensor]): Text token IDs for conditioning.
+            prompt_mel (Optional[torch.Tensor]): A mel-spectrogram prompt.
+            x_mask (Optional[torch.Tensor]): Padding mask for the target generation length.
+            prompt_mask (Optional[torch.Tensor]): Padding mask for the prompt.
+            text_mask (Optional[torch.Tensor]): Padding mask for the text.
+            n_timesteps (int): Number of steps in the reverse diffusion process.
+            cfg (float): Classifier-Free Guidance scale.
+            rescale_cfg (float): Rescaling factor for CFG to prevent saturation.
+
+        Returns:
+            torch.Tensor: The generated mel-spectrogram.
+        """
+        if vq_emb is None:
+            assert indices is not None, "Either vq_emb or indices must be provided"
+            vq_emb = self.index2vq(indices.long())
+
+        # Upsample VQ embeddings to match the target mel length
+        vq_emb_post = self.vq_out_linear(vq_emb)
+        vq_emb_post = vq_emb_post.transpose(1, 2)
+        vq_emb_post = F.interpolate(
+            vq_emb_post, scale_factor=self.down_sample_factor, mode="linear"
+        )
+        vq_emb_post = vq_emb_post.transpose(1, 2)
+
+        # Prepare text embeddings
+        if self.use_text_cond:
+            text_emb = self.text_emb(text_ids)
+            if text_mask is None:
+                text_mask = torch.ones_like(text_ids)
+        else:
+            text_emb, text_mask = None, None
+
+        cond_emb = vq_emb_post
+
+        # Handle prompt
+        if prompt_mel is None:
+            prompt_mel = torch.zeros(
+                cond_emb.shape[0], 0, self.mel_dim, device=cond_emb.device
+            )
+
+        prompt_len = prompt_mel.shape[1]
+        target_len = cond_emb.shape[1] - prompt_len
+
+        # Prepare masks
+        if x_mask is None:
+            x_mask = torch.ones(cond_emb.shape[0], target_len, device=cond_emb.device)
+        if prompt_mask is None:
+            prompt_mask = torch.ones(
+                cond_emb.shape[0], prompt_len, device=cond_emb.device
+            )
+
+        xt_mask = torch.cat([prompt_mask, x_mask], dim=1)
+
+        # Initialize with random noise
+        z = torch.randn(
+            (cond_emb.shape[0], target_len, self.mel_dim),
+            dtype=cond_emb.dtype,
+            device=cond_emb.device,
+        )
+        xt = z
+        h = 1.0 / n_timesteps
+
+        # Iterative denoising loop (Euler method)
+        for i in range(n_timesteps):
+            # Concatenate prompt and current noisy sample
+            xt_input = torch.cat([prompt_mel, xt], dim=1)
+            # Calculate current timestep
+            t = (0 + (i + 0.5) * h) * torch.ones(
+                z.shape[0], dtype=z.dtype, device=z.device
+            )
+
+            # Get conditional flow prediction
+            flow_pred = self.decoder(
+                x=xt_input,
+                x_mask=xt_mask,
+                text_embedding=text_emb,
+                text_mask=text_mask,
+                cond=cond_emb,
+                diffusion_step=t,
+            )
+            flow_pred = flow_pred[
+                :, prompt_len:, :
+            ]  # Extract flow for the target region
+
+            # Classifier-Free Guidance (CFG)
+            if cfg > 0 and self.use_text_cond:
+                # Get unconditional flow prediction by dropping conditions
+                uncond_flow_pred = self.decoder(
+                    x=xt_input,
+                    x_mask=xt_mask,
+                    text_embedding=None,  # Drop text
+                    text_mask=None,
+                    cond=torch.zeros_like(cond_emb),  # Drop code
+                    diffusion_step=t,
+                )
+                uncond_flow_pred = uncond_flow_pred[:, prompt_len:, :]
+
+                # Combine conditional and unconditional predictions
+                flow_pred_cfg = uncond_flow_pred + cfg * (flow_pred - uncond_flow_pred)
+
+                # Rescale to prevent saturation, as in Stable Diffusion
+                if rescale_cfg > 0:
+                    flow_pred_std = flow_pred.std()
+                    cfg_std = flow_pred_cfg.std()
+                    # Avoid division by zero
+                    if cfg_std > 1e-6:
+                        rescale_flow_pred = flow_pred_cfg * (flow_pred_std / cfg_std)
+                        flow_pred = (
+                            rescale_cfg * rescale_flow_pred
+                            + (1 - rescale_cfg) * flow_pred_cfg
+                        )
+                    else:
+                        flow_pred = flow_pred_cfg
+                else:
+                    flow_pred = flow_pred_cfg
+
+            # Update the noisy sample
+            dxt = flow_pred * h
+            xt = xt + dxt
+
+        return xt
+
+    def reset_parameters(self):
+        """
+        Applies custom weight initialization to the model's submodules.
+        """
+
+        def _reset_parameters(m: nn.Module):
+            if isinstance(m, nn.MultiheadAttention):
+                if m._qkv_same_embed_dim:
+                    nn.init.normal_(m.in_proj_weight, std=0.02)
+                else:
+                    nn.init.normal_(m.q_proj_weight, std=0.02)
+                    nn.init.normal_(m.k_proj_weight, std=0.02)
+                    nn.init.normal_(m.v_proj_weight, std=0.02)
+
+                if m.in_proj_bias is not None:
+                    nn.init.constant_(m.in_proj_bias, 0.0)
+                    nn.init.constant_(m.out_proj.bias, 0.0)
+                if m.bias_k is not None:
+                    nn.init.xavier_normal_(m.bias_k)
+                if m.bias_v is not None:
+                    nn.init.xavier_normal_(m.bias_v)
+
+            elif (
+                isinstance(m, nn.Conv1d)
+                or isinstance(m, nn.ConvTranspose1d)
+                or isinstance(m, nn.Conv2d)
+                or isinstance(m, nn.ConvTranspose2d)
+            ):
+                m.weight.data.normal_(0.0, 0.02)
+
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(mean=0.0, std=0.02)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+
+            elif isinstance(m, nn.Embedding):
+                m.weight.data.normal_(mean=0.0, std=0.02)
+                if m.padding_idx is not None:
+                    m.weight.data[m.padding_idx].zero_()
+
+        self.apply(_reset_parameters)

requirements.txt

@@ -0,0 +1,9 @@
+gradio
+torch
+torchaudio
+numpy
+huggingface-hub
+transformers
+datasets
+librosa
+soundfile

utils/hparam.py

@@ -0,0 +1,659 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# This code is modified from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/training/python/training/hparam.py  pylint: disable=line-too-long
+"""Hyperparameter values."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+import numbers
+import re
+import six
+
+# Define the regular expression for parsing a single clause of the input
+# (delimited by commas).  A legal clause looks like:
+#   <variable name>[<index>]? = <rhs>
+# where <rhs> is either a single token or [] enclosed list of tokens.
+# For example:  "var[1] = a" or "x = [1,2,3]"
+PARAM_RE = re.compile(
+    r"""
+  (?P<name>[a-zA-Z][\w\.]*)      # variable name: "var" or "x"
+  (\[\s*(?P<index>\d+)\s*\])?  # (optional) index: "1" or None
+  \s*=\s*
+  ((?P<val>[^,\[]*)            # single value: "a" or None
+   |
+   \[(?P<vals>[^\]]*)\])       # list of values: None or "1,2,3"
+  ($|,\s*)""",
+    re.VERBOSE,
+)
+
+
+def _parse_fail(name, var_type, value, values):
+    """Helper function for raising a value error for bad assignment."""
+    raise ValueError(
+        "Could not parse hparam '%s' of type '%s' with value '%s' in %s"
+        % (name, var_type.__name__, value, values)
+    )
+
+
+def _reuse_fail(name, values):
+    """Helper function for raising a value error for reuse of name."""
+    raise ValueError("Multiple assignments to variable '%s' in %s" % (name, values))
+
+
+def _process_scalar_value(name, parse_fn, var_type, m_dict, values, results_dictionary):
+    """Update results_dictionary with a scalar value.
+
+    Used to update the results_dictionary to be returned by parse_values when
+    encountering a clause with a scalar RHS (e.g.  "s=5" or "arr[0]=5".)
+
+    Mutates results_dictionary.
+
+    Args:
+      name: Name of variable in assignment ("s" or "arr").
+      parse_fn: Function for parsing the actual value.
+      var_type: Type of named variable.
+      m_dict: Dictionary constructed from regex parsing.
+        m_dict['val']: RHS value (scalar)
+        m_dict['index']: List index value (or None)
+      values: Full expression being parsed
+      results_dictionary: The dictionary being updated for return by the parsing
+        function.
+
+    Raises:
+      ValueError: If the name has already been used.
+    """
+    try:
+        parsed_value = parse_fn(m_dict["val"])
+    except ValueError:
+        _parse_fail(name, var_type, m_dict["val"], values)
+
+    # If no index is provided
+    if not m_dict["index"]:
+        if name in results_dictionary:
+            _reuse_fail(name, values)
+        results_dictionary[name] = parsed_value
+    else:
+        if name in results_dictionary:
+            # The name has already been used as a scalar, then it
+            # will be in this dictionary and map to a non-dictionary.
+            if not isinstance(results_dictionary.get(name), dict):
+                _reuse_fail(name, values)
+        else:
+            results_dictionary[name] = {}
+
+        index = int(m_dict["index"])
+        # Make sure the index position hasn't already been assigned a value.
+        if index in results_dictionary[name]:
+            _reuse_fail("{}[{}]".format(name, index), values)
+        results_dictionary[name][index] = parsed_value
+
+
+def _process_list_value(name, parse_fn, var_type, m_dict, values, results_dictionary):
+    """Update results_dictionary from a list of values.
+
+    Used to update results_dictionary to be returned by parse_values when
+    encountering a clause with a list RHS (e.g.  "arr=[1,2,3]".)
+
+    Mutates results_dictionary.
+
+    Args:
+      name: Name of variable in assignment ("arr").
+      parse_fn: Function for parsing individual values.
+      var_type: Type of named variable.
+      m_dict: Dictionary constructed from regex parsing.
+        m_dict['val']: RHS value (scalar)
+      values: Full expression being parsed
+      results_dictionary: The dictionary being updated for return by the parsing
+        function.
+
+    Raises:
+      ValueError: If the name has an index or the values cannot be parsed.
+    """
+    if m_dict["index"] is not None:
+        raise ValueError("Assignment of a list to a list index.")
+    elements = filter(None, re.split("[ ,]", m_dict["vals"]))
+    # Make sure the name hasn't already been assigned a value
+    if name in results_dictionary:
+        raise _reuse_fail(name, values)
+    try:
+        results_dictionary[name] = [parse_fn(e) for e in elements]
+    except ValueError:
+        _parse_fail(name, var_type, m_dict["vals"], values)
+
+
+def _cast_to_type_if_compatible(name, param_type, value):
+    """Cast hparam to the provided type, if compatible.
+
+    Args:
+      name: Name of the hparam to be cast.
+      param_type: The type of the hparam.
+      value: The value to be cast, if compatible.
+
+    Returns:
+      The result of casting `value` to `param_type`.
+
+    Raises:
+      ValueError: If the type of `value` is not compatible with param_type.
+        * If `param_type` is a string type, but `value` is not.
+        * If `param_type` is a boolean, but `value` is not, or vice versa.
+        * If `param_type` is an integer type, but `value` is not.
+        * If `param_type` is a float type, but `value` is not a numeric type.
+    """
+    fail_msg = "Could not cast hparam '%s' of type '%s' from value %r" % (
+        name,
+        param_type,
+        value,
+    )
+
+    # Some callers use None, for which we can't do any casting/checking. :(
+    if issubclass(param_type, type(None)):
+        return value
+
+    # Avoid converting a non-string type to a string.
+    if issubclass(param_type, (six.string_types, six.binary_type)) and not isinstance(
+        value, (six.string_types, six.binary_type)
+    ):
+        raise ValueError(fail_msg)
+
+    # Avoid converting a number or string type to a boolean or vice versa.
+    if issubclass(param_type, bool) != isinstance(value, bool):
+        raise ValueError(fail_msg)
+
+    # Avoid converting float to an integer (the reverse is fine).
+    if issubclass(param_type, numbers.Integral) and not isinstance(
+        value, numbers.Integral
+    ):
+        raise ValueError(fail_msg)
+
+    # Avoid converting a non-numeric type to a numeric type.
+    if issubclass(param_type, numbers.Number) and not isinstance(value, numbers.Number):
+        raise ValueError(fail_msg)
+
+    return param_type(value)
+
+
+def parse_values(values, type_map, ignore_unknown=False):
+    """Parses hyperparameter values from a string into a python map.
+
+    `values` is a string containing comma-separated `name=value` pairs.
+    For each pair, the value of the hyperparameter named `name` is set to
+    `value`.
+
+    If a hyperparameter name appears multiple times in `values`, a ValueError
+    is raised (e.g. 'a=1,a=2', 'a[1]=1,a[1]=2').
+
+    If a hyperparameter name in both an index assignment and scalar assignment,
+    a ValueError is raised.  (e.g. 'a=[1,2,3],a[0] = 1').
+
+    The hyperparameter name may contain '.' symbols, which will result in an
+    attribute name that is only accessible through the getattr and setattr
+    functions.  (And must be first explicit added through add_hparam.)
+
+    WARNING: Use of '.' in your variable names is allowed, but is not well
+    supported and not recommended.
+
+    The `value` in `name=value` must follows the syntax according to the
+    type of the parameter:
+
+    *  Scalar integer: A Python-parsable integer point value.  E.g.: 1,
+       100, -12.
+    *  Scalar float: A Python-parsable floating point value.  E.g.: 1.0,
+       -.54e89.
+    *  Boolean: Either true or false.
+    *  Scalar string: A non-empty sequence of characters, excluding comma,
+       spaces, and square brackets.  E.g.: foo, bar_1.
+    *  List: A comma separated list of scalar values of the parameter type
+       enclosed in square brackets.  E.g.: [1,2,3], [1.0,1e-12], [high,low].
+
+    When index assignment is used, the corresponding type_map key should be the
+    list name.  E.g. for "arr[1]=0" the type_map must have the key "arr" (not
+    "arr[1]").
+
+    Args:
+      values: String.  Comma separated list of `name=value` pairs where
+        'value' must follow the syntax described above.
+      type_map: A dictionary mapping hyperparameter names to types.  Note every
+        parameter name in values must be a key in type_map.  The values must
+        conform to the types indicated, where a value V is said to conform to a
+        type T if either V has type T, or V is a list of elements of type T.
+        Hence, for a multidimensional parameter 'x' taking float values,
+        'x=[0.1,0.2]' will parse successfully if type_map['x'] = float.
+      ignore_unknown: Bool. Whether values that are missing a type in type_map
+        should be ignored. If set to True, a ValueError will not be raised for
+        unknown hyperparameter type.
+
+    Returns:
+      A python map mapping each name to either:
+      * A scalar value.
+      * A list of scalar values.
+      * A dictionary mapping index numbers to scalar values.
+      (e.g. "x=5,L=[1,2],arr[1]=3" results in {'x':5,'L':[1,2],'arr':{1:3}}")
+
+    Raises:
+      ValueError: If there is a problem with input.
+      * If `values` cannot be parsed.
+      * If a list is assigned to a list index (e.g. 'a[1] = [1,2,3]').
+      * If the same rvalue is assigned two different values (e.g. 'a=1,a=2',
+        'a[1]=1,a[1]=2', or 'a=1,a=[1]')
+    """
+    results_dictionary = {}
+    pos = 0
+    while pos < len(values):
+        m = PARAM_RE.match(values, pos)
+        if not m:
+            raise ValueError("Malformed hyperparameter value: %s" % values[pos:])
+        # Check that there is a comma between parameters and move past it.
+        pos = m.end()
+        # Parse the values.
+        m_dict = m.groupdict()
+        name = m_dict["name"]
+        if name not in type_map:
+            if ignore_unknown:
+                continue
+            raise ValueError("Unknown hyperparameter type for %s" % name)
+        type_ = type_map[name]
+
+        # Set up correct parsing function (depending on whether type_ is a bool)
+        if type_ == bool:
+
+            def parse_bool(value):
+                if value in ["true", "True"]:
+                    return True
+                elif value in ["false", "False"]:
+                    return False
+                else:
+                    try:
+                        return bool(int(value))
+                    except ValueError:
+                        _parse_fail(name, type_, value, values)
+
+            parse = parse_bool
+        else:
+            parse = type_
+
+        # If a singe value is provided
+        if m_dict["val"] is not None:
+            _process_scalar_value(
+                name, parse, type_, m_dict, values, results_dictionary
+            )
+
+        # If the assigned value is a list:
+        elif m_dict["vals"] is not None:
+            _process_list_value(name, parse, type_, m_dict, values, results_dictionary)
+
+        else:  # Not assigned a list or value
+            _parse_fail(name, type_, "", values)
+
+    return results_dictionary
+
+
+class HParams(object):
+    """Class to hold a set of hyperparameters as name-value pairs.
+
+    A `HParams` object holds hyperparameters used to build and train a model,
+    such as the number of hidden units in a neural net layer or the learning rate
+    to use when training.
+
+    You first create a `HParams` object by specifying the names and values of the
+    hyperparameters.
+
+    To make them easily accessible the parameter names are added as direct
+    attributes of the class.  A typical usage is as follows:
+
+    ```python
+    # Create a HParams object specifying names and values of the model
+    # hyperparameters:
+    hparams = HParams(learning_rate=0.1, num_hidden_units=100)
+
+    # The hyperparameter are available as attributes of the HParams object:
+    hparams.learning_rate ==> 0.1
+    hparams.num_hidden_units ==> 100
+    ```
+
+    Hyperparameters have type, which is inferred from the type of their value
+    passed at construction type.   The currently supported types are: integer,
+    float, boolean, string, and list of integer, float, boolean, or string.
+
+    You can override hyperparameter values by calling the
+    [`parse()`](#HParams.parse) method, passing a string of comma separated
+    `name=value` pairs.  This is intended to make it possible to override
+    any hyperparameter values from a single command-line flag to which
+    the user passes 'hyper-param=value' pairs.  It avoids having to define
+    one flag for each hyperparameter.
+
+    The syntax expected for each value depends on the type of the parameter.
+    See `parse()` for a description of the syntax.
+
+    Example:
+
+    ```python
+    # Define a command line flag to pass name=value pairs.
+    # For example using argparse:
+    import argparse
+    parser = argparse.ArgumentParser(description='Train my model.')
+    parser.add_argument('--hparams', type=str,
+                        help='Comma separated list of "name=value" pairs.')
+    args = parser.parse_args()
+    ...
+    def my_program():
+      # Create a HParams object specifying the names and values of the
+      # model hyperparameters:
+      hparams = tf.HParams(learning_rate=0.1, num_hidden_units=100,
+                           activations=['relu', 'tanh'])
+
+      # Override hyperparameters values by parsing the command line
+      hparams.parse(args.hparams)
+
+      # If the user passed `--hparams=learning_rate=0.3` on the command line
+      # then 'hparams' has the following attributes:
+      hparams.learning_rate ==> 0.3
+      hparams.num_hidden_units ==> 100
+      hparams.activations ==> ['relu', 'tanh']
+
+      # If the hyperparameters are in json format use parse_json:
+      hparams.parse_json('{"learning_rate": 0.3, "activations": "relu"}')
+    ```
+    """
+
+    _HAS_DYNAMIC_ATTRIBUTES = True  # Required for pytype checks.
+
+    def __init__(self, model_structure=None, **kwargs):
+        """Create an instance of `HParams` from keyword arguments.
+
+        The keyword arguments specify name-values pairs for the hyperparameters.
+        The parameter types are inferred from the type of the values passed.
+
+        The parameter names are added as attributes of `HParams` object, so they
+        can be accessed directly with the dot notation `hparams._name_`.
+
+        Example:
+
+        ```python
+        # Define 3 hyperparameters: 'learning_rate' is a float parameter,
+        # 'num_hidden_units' an integer parameter, and 'activation' a string
+        # parameter.
+        hparams = tf.HParams(
+            learning_rate=0.1, num_hidden_units=100, activation='relu')
+
+        hparams.activation ==> 'relu'
+        ```
+
+        Note that a few names are reserved and cannot be used as hyperparameter
+        names.  If you use one of the reserved name the constructor raises a
+        `ValueError`.
+
+        Args:
+          model_structure: An instance of ModelStructure, defining the feature
+            crosses to be used in the Trial.
+          **kwargs: Key-value pairs where the key is the hyperparameter name and
+            the value is the value for the parameter.
+
+        Raises:
+          ValueError: If both `hparam_def` and initialization values are provided,
+            or if one of the arguments is invalid.
+
+        """
+        # Register the hyperparameters and their type in _hparam_types.
+        # This simplifies the implementation of parse().
+        # _hparam_types maps the parameter name to a tuple (type, bool).
+        # The type value is the type of the parameter for scalar hyperparameters,
+        # or the type of the list elements for multidimensional hyperparameters.
+        # The bool value is True if the value is a list, False otherwise.
+        self._hparam_types = {}
+        self._model_structure = model_structure
+        for name, value in six.iteritems(kwargs):
+            self.add_hparam(name, value)
+
+    def add_hparam(self, name, value):
+        """Adds {name, value} pair to hyperparameters.
+
+        Args:
+          name: Name of the hyperparameter.
+          value: Value of the hyperparameter. Can be one of the following types:
+            int, float, string, int list, float list, or string list.
+
+        Raises:
+          ValueError: if one of the arguments is invalid.
+        """
+        # Keys in kwargs are unique, but 'name' could the name of a pre-existing
+        # attribute of this object.  In that case we refuse to use it as a
+        # hyperparameter name.
+        if getattr(self, name, None) is not None:
+            raise ValueError("Hyperparameter name is reserved: %s" % name)
+        if isinstance(value, (list, tuple)):
+            if not value:
+                raise ValueError(
+                    "Multi-valued hyperparameters cannot be empty: %s" % name
+                )
+            self._hparam_types[name] = (type(value[0]), True)
+        else:
+            self._hparam_types[name] = (type(value), False)
+        setattr(self, name, value)
+
+    def set_hparam(self, name, value):
+        """Set the value of an existing hyperparameter.
+
+        This function verifies that the type of the value matches the type of the
+        existing hyperparameter.
+
+        Args:
+          name: Name of the hyperparameter.
+          value: New value of the hyperparameter.
+
+        Raises:
+          KeyError: If the hyperparameter doesn't exist.
+          ValueError: If there is a type mismatch.
+        """
+        param_type, is_list = self._hparam_types[name]
+        if isinstance(value, list):
+            if not is_list:
+                raise ValueError(
+                    "Must not pass a list for single-valued parameter: %s" % name
+                )
+            setattr(
+                self,
+                name,
+                [_cast_to_type_if_compatible(name, param_type, v) for v in value],
+            )
+        else:
+            if is_list:
+                raise ValueError(
+                    "Must pass a list for multi-valued parameter: %s." % name
+                )
+            setattr(self, name, _cast_to_type_if_compatible(name, param_type, value))
+
+    def del_hparam(self, name):
+        """Removes the hyperparameter with key 'name'.
+
+        Does nothing if it isn't present.
+
+        Args:
+          name: Name of the hyperparameter.
+        """
+        if hasattr(self, name):
+            delattr(self, name)
+            del self._hparam_types[name]
+
+    def parse(self, values):
+        """Override existing hyperparameter values, parsing new values from a string.
+
+        See parse_values for more detail on the allowed format for values.
+
+        Args:
+          values: String.  Comma separated list of `name=value` pairs where 'value'
+            must follow the syntax described above.
+
+        Returns:
+          The `HParams` instance.
+
+        Raises:
+          ValueError: If `values` cannot be parsed or a hyperparameter in `values`
+          doesn't exist.
+        """
+        type_map = {}
+        for name, t in self._hparam_types.items():
+            param_type, _ = t
+            type_map[name] = param_type
+
+        values_map = parse_values(values, type_map)
+        return self.override_from_dict(values_map)
+
+    def override_from_dict(self, values_dict):
+        """Override existing hyperparameter values, parsing new values from a dictionary.
+
+        Args:
+          values_dict: Dictionary of name:value pairs.
+
+        Returns:
+          The `HParams` instance.
+
+        Raises:
+          KeyError: If a hyperparameter in `values_dict` doesn't exist.
+          ValueError: If `values_dict` cannot be parsed.
+        """
+        for name, value in values_dict.items():
+            self.set_hparam(name, value)
+        return self
+
+    def set_model_structure(self, model_structure):
+        self._model_structure = model_structure
+
+    def get_model_structure(self):
+        return self._model_structure
+
+    def to_json(self, indent=None, separators=None, sort_keys=False):
+        """Serializes the hyperparameters into JSON.
+
+        Args:
+          indent: If a non-negative integer, JSON array elements and object members
+            will be pretty-printed with that indent level. An indent level of 0, or
+            negative, will only insert newlines. `None` (the default) selects the
+            most compact representation.
+          separators: Optional `(item_separator, key_separator)` tuple. Default is
+            `(', ', ': ')`.
+          sort_keys: If `True`, the output dictionaries will be sorted by key.
+
+        Returns:
+          A JSON string.
+        """
+
+        def remove_callables(x):
+            """Omit callable elements from input with arbitrary nesting."""
+            if isinstance(x, dict):
+                return {
+                    k: remove_callables(v)
+                    for k, v in six.iteritems(x)
+                    if not callable(v)
+                }
+            elif isinstance(x, list):
+                return [remove_callables(i) for i in x if not callable(i)]
+            return x
+
+        return json.dumps(
+            remove_callables(self.values()),
+            indent=indent,
+            separators=separators,
+            sort_keys=sort_keys,
+        )
+
+    def parse_json(self, values_json):
+        """Override existing hyperparameter values, parsing new values from a json object.
+
+        Args:
+          values_json: String containing a json object of name:value pairs.
+
+        Returns:
+          The `HParams` instance.
+
+        Raises:
+          KeyError: If a hyperparameter in `values_json` doesn't exist.
+          ValueError: If `values_json` cannot be parsed.
+        """
+        values_map = json.loads(values_json)
+        return self.override_from_dict(values_map)
+
+    def values(self):
+        """Return the hyperparameter values as a Python dictionary.
+
+        Returns:
+          A dictionary with hyperparameter names as keys.  The values are the
+          hyperparameter values.
+        """
+        return {n: getattr(self, n) for n in self._hparam_types.keys()}
+
+    def get(self, key, default=None):
+        """Returns the value of `key` if it exists, else `default`."""
+        if key in self._hparam_types:
+            # Ensure that default is compatible with the parameter type.
+            if default is not None:
+                param_type, is_param_list = self._hparam_types[key]
+                type_str = "list<%s>" % param_type if is_param_list else str(param_type)
+                fail_msg = (
+                    "Hparam '%s' of type '%s' is incompatible with "
+                    "default=%s" % (key, type_str, default)
+                )
+
+                is_default_list = isinstance(default, list)
+                if is_param_list != is_default_list:
+                    raise ValueError(fail_msg)
+
+                try:
+                    if is_default_list:
+                        for value in default:
+                            _cast_to_type_if_compatible(key, param_type, value)
+                    else:
+                        _cast_to_type_if_compatible(key, param_type, default)
+                except ValueError as e:
+                    raise ValueError("%s. %s" % (fail_msg, e))
+
+            return getattr(self, key)
+
+        return default
+
+    def __contains__(self, key):
+        return key in self._hparam_types
+
+    def __str__(self):
+        return str(sorted(self.values().items()))
+
+    def __repr__(self):
+        return "%s(%s)" % (type(self).__name__, self.__str__())
+
+    @staticmethod
+    def _get_kind_name(param_type, is_list):
+        """Returns the field name given parameter type and is_list.
+
+        Args:
+          param_type: Data type of the hparam.
+          is_list: Whether this is a list.
+
+        Returns:
+          A string representation of the field name.
+
+        Raises:
+          ValueError: If parameter type is not recognized.
+        """
+        if issubclass(param_type, bool):
+            # This check must happen before issubclass(param_type, six.integer_types),
+            # since Python considers bool to be a subclass of int.
+            typename = "bool"
+        elif issubclass(param_type, six.integer_types):
+            # Setting 'int' and 'long' types to be 'int64' to ensure the type is
+            # compatible with both Python2 and Python3.
+            typename = "int64"
+        elif issubclass(param_type, (six.string_types, six.binary_type)):
+            # Setting 'string' and 'bytes' types to be 'bytes' to ensure the type is
+            # compatible with both Python2 and Python3.
+            typename = "bytes"
+        elif issubclass(param_type, float):
+            typename = "float"
+        else:
+            raise ValueError("Unsupported parameter type: %s" % str(param_type))
+
+        suffix = "list" if is_list else "value"
+        return "_".join([typename, suffix])

utils/util.py

@@ -0,0 +1,687 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import collections
+import glob
+import os
+import random
+import time
+import argparse
+from collections import OrderedDict
+
+import json5
+import numpy as np
+import glob
+from torch.nn import functional as F
+
+
+try:
+    from ruamel.yaml import YAML as yaml
+except:
+    from ruamel_yaml import YAML as yaml  # type: ignore
+
+import torch
+
+from utils.hparam import HParams
+import logging
+from logging import handlers
+
+
+def str2bool(v):
+    """Used in argparse.ArgumentParser.add_argument to indicate
+    that a type is a bool type and user can enter
+
+        - yes, true, t, y, 1, to represent True
+        - no, false, f, n, 0, to represent False
+
+    See https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse  # noqa
+    """
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+def find_checkpoint_of_mapper(mapper_ckpt_dir):
+    mapper_ckpts = glob.glob(os.path.join(mapper_ckpt_dir, "ckpts/*.pt"))
+
+    # Select the max steps
+    mapper_ckpts.sort()
+    mapper_weights_file = mapper_ckpts[-1]
+    return mapper_weights_file
+
+
+def pad_f0_to_tensors(f0s, batched=None):
+    # Initialize
+    tensors = []
+
+    if batched == None:
+        # Get the max frame for padding
+        size = -1
+        for f0 in f0s:
+            size = max(size, f0.shape[-1])
+
+        tensor = torch.zeros(len(f0s), size)
+
+        for i, f0 in enumerate(f0s):
+            tensor[i, : f0.shape[-1]] = f0[:]
+
+        tensors.append(tensor)
+    else:
+        start = 0
+        while start + batched - 1 < len(f0s):
+            end = start + batched - 1
+
+            # Get the max frame for padding
+            size = -1
+            for i in range(start, end + 1):
+                size = max(size, f0s[i].shape[-1])
+
+            tensor = torch.zeros(batched, size)
+
+            for i in range(start, end + 1):
+                tensor[i - start, : f0s[i].shape[-1]] = f0s[i][:]
+
+            tensors.append(tensor)
+
+            start = start + batched
+
+        if start != len(f0s):
+            end = len(f0s)
+
+            # Get the max frame for padding
+            size = -1
+            for i in range(start, end):
+                size = max(size, f0s[i].shape[-1])
+
+            tensor = torch.zeros(len(f0s) - start, size)
+
+            for i in range(start, end):
+                tensor[i - start, : f0s[i].shape[-1]] = f0s[i][:]
+
+            tensors.append(tensor)
+
+    return tensors
+
+
+def pad_mels_to_tensors(mels, batched=None):
+    """
+    Args:
+        mels: A list of mel-specs
+    Returns:
+        tensors: A list of tensors containing the batched mel-specs
+        mel_frames: A list of tensors containing the frames of the original mel-specs
+    """
+    # Initialize
+    tensors = []
+    mel_frames = []
+
+    # Split mel-specs into batches to avoid cuda memory exceed
+    if batched == None:
+        # Get the max frame for padding
+        size = -1
+        for mel in mels:
+            size = max(size, mel.shape[-1])
+
+        tensor = torch.zeros(len(mels), mels[0].shape[0], size)
+        mel_frame = torch.zeros(len(mels), dtype=torch.int32)
+
+        for i, mel in enumerate(mels):
+            tensor[i, :, : mel.shape[-1]] = mel[:]
+            mel_frame[i] = mel.shape[-1]
+
+        tensors.append(tensor)
+        mel_frames.append(mel_frame)
+    else:
+        start = 0
+        while start + batched - 1 < len(mels):
+            end = start + batched - 1
+
+            # Get the max frame for padding
+            size = -1
+            for i in range(start, end + 1):
+                size = max(size, mels[i].shape[-1])
+
+            tensor = torch.zeros(batched, mels[0].shape[0], size)
+            mel_frame = torch.zeros(batched, dtype=torch.int32)
+
+            for i in range(start, end + 1):
+                tensor[i - start, :, : mels[i].shape[-1]] = mels[i][:]
+                mel_frame[i - start] = mels[i].shape[-1]
+
+            tensors.append(tensor)
+            mel_frames.append(mel_frame)
+
+            start = start + batched
+
+        if start != len(mels):
+            end = len(mels)
+
+            # Get the max frame for padding
+            size = -1
+            for i in range(start, end):
+                size = max(size, mels[i].shape[-1])
+
+            tensor = torch.zeros(len(mels) - start, mels[0].shape[0], size)
+            mel_frame = torch.zeros(len(mels) - start, dtype=torch.int32)
+
+            for i in range(start, end):
+                tensor[i - start, :, : mels[i].shape[-1]] = mels[i][:]
+                mel_frame[i - start] = mels[i].shape[-1]
+
+            tensors.append(tensor)
+            mel_frames.append(mel_frame)
+
+    return tensors, mel_frames
+
+
+def load_model_config(args):
+    """Load model configurations (in args.json under checkpoint directory)
+
+    Args:
+        args (ArgumentParser): arguments to run bins/preprocess.py
+
+    Returns:
+        dict: dictionary that stores model configurations
+    """
+    if args.checkpoint_dir is None:
+        assert args.checkpoint_file is not None
+        checkpoint_dir = os.path.split(args.checkpoint_file)[0]
+    else:
+        checkpoint_dir = args.checkpoint_dir
+    config_path = os.path.join(checkpoint_dir, "args.json")
+    print("config_path: ", config_path)
+
+    config = load_config(config_path)
+    return config
+
+
+def remove_and_create(dir):
+    if os.path.exists(dir):
+        os.system("rm -r {}".format(dir))
+    os.makedirs(dir, exist_ok=True)
+
+
+def has_existed(path, warning=False):
+    if not warning:
+        return os.path.exists(path)
+
+    if os.path.exists(path):
+        answer = input(
+            "The path {} has existed. \nInput 'y' (or hit Enter) to skip it, and input 'n' to re-write it [y/n]\n".format(
+                path
+            )
+        )
+        if not answer == "n":
+            return True
+
+    return False
+
+
+def remove_older_ckpt(saved_model_name, checkpoint_dir, max_to_keep=5):
+    if os.path.exists(os.path.join(checkpoint_dir, "checkpoint")):
+        with open(os.path.join(checkpoint_dir, "checkpoint"), "r") as f:
+            ckpts = [x.strip() for x in f.readlines()]
+    else:
+        ckpts = []
+    ckpts.append(saved_model_name)
+    for item in ckpts[:-max_to_keep]:
+        if os.path.exists(os.path.join(checkpoint_dir, item)):
+            os.remove(os.path.join(checkpoint_dir, item))
+    with open(os.path.join(checkpoint_dir, "checkpoint"), "w") as f:
+        for item in ckpts[-max_to_keep:]:
+            f.write("{}\n".format(item))
+
+
+def set_all_random_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.random.manual_seed(seed)
+
+
+def save_checkpoint(
+    args,
+    generator,
+    g_optimizer,
+    step,
+    discriminator=None,
+    d_optimizer=None,
+    max_to_keep=5,
+):
+    saved_model_name = "model.ckpt-{}.pt".format(step)
+    checkpoint_path = os.path.join(args.checkpoint_dir, saved_model_name)
+
+    if discriminator and d_optimizer:
+        torch.save(
+            {
+                "generator": generator.state_dict(),
+                "discriminator": discriminator.state_dict(),
+                "g_optimizer": g_optimizer.state_dict(),
+                "d_optimizer": d_optimizer.state_dict(),
+                "global_step": step,
+            },
+            checkpoint_path,
+        )
+    else:
+        torch.save(
+            {
+                "generator": generator.state_dict(),
+                "g_optimizer": g_optimizer.state_dict(),
+                "global_step": step,
+            },
+            checkpoint_path,
+        )
+
+    print("Saved checkpoint: {}".format(checkpoint_path))
+
+    if os.path.exists(os.path.join(args.checkpoint_dir, "checkpoint")):
+        with open(os.path.join(args.checkpoint_dir, "checkpoint"), "r") as f:
+            ckpts = [x.strip() for x in f.readlines()]
+    else:
+        ckpts = []
+    ckpts.append(saved_model_name)
+    for item in ckpts[:-max_to_keep]:
+        if os.path.exists(os.path.join(args.checkpoint_dir, item)):
+            os.remove(os.path.join(args.checkpoint_dir, item))
+    with open(os.path.join(args.checkpoint_dir, "checkpoint"), "w") as f:
+        for item in ckpts[-max_to_keep:]:
+            f.write("{}\n".format(item))
+
+
+def attempt_to_restore(
+    generator, g_optimizer, checkpoint_dir, discriminator=None, d_optimizer=None
+):
+    checkpoint_list = os.path.join(checkpoint_dir, "checkpoint")
+    if os.path.exists(checkpoint_list):
+        checkpoint_filename = open(checkpoint_list).readlines()[-1].strip()
+        checkpoint_path = os.path.join(checkpoint_dir, "{}".format(checkpoint_filename))
+        print("Restore from {}".format(checkpoint_path))
+        checkpoint = torch.load(checkpoint_path, map_location="cpu")
+        if generator:
+            if not list(generator.state_dict().keys())[0].startswith("module."):
+                raw_dict = checkpoint["generator"]
+                clean_dict = OrderedDict()
+                for k, v in raw_dict.items():
+                    if k.startswith("module."):
+                        clean_dict[k[7:]] = v
+                    else:
+                        clean_dict[k] = v
+                generator.load_state_dict(clean_dict)
+            else:
+                generator.load_state_dict(checkpoint["generator"])
+        if g_optimizer:
+            g_optimizer.load_state_dict(checkpoint["g_optimizer"])
+        global_step = 100000
+        if discriminator and "discriminator" in checkpoint.keys():
+            discriminator.load_state_dict(checkpoint["discriminator"])
+            global_step = checkpoint["global_step"]
+            print("restore discriminator")
+        if d_optimizer and "d_optimizer" in checkpoint.keys():
+            d_optimizer.load_state_dict(checkpoint["d_optimizer"])
+            print("restore d_optimizer...")
+    else:
+        global_step = 0
+    return global_step
+
+
+class ExponentialMovingAverage(object):
+    def __init__(self, decay):
+        self.decay = decay
+        self.shadow = {}
+
+    def register(self, name, val):
+        self.shadow[name] = val.clone()
+
+    def update(self, name, x):
+        assert name in self.shadow
+        update_delta = self.shadow[name] - x
+        self.shadow[name] -= (1.0 - self.decay) * update_delta
+
+
+def apply_moving_average(model, ema):
+    for name, param in model.named_parameters():
+        if name in ema.shadow:
+            ema.update(name, param.data)
+
+
+def register_model_to_ema(model, ema):
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            ema.register(name, param.data)
+
+
+class YParams(HParams):
+    def __init__(self, yaml_file):
+        if not os.path.exists(yaml_file):
+            raise IOError("yaml file: {} is not existed".format(yaml_file))
+        super().__init__()
+        self.d = collections.OrderedDict()
+        with open(yaml_file) as fp:
+            for _, v in yaml().load(fp).items():
+                for k1, v1 in v.items():
+                    try:
+                        if self.get(k1):
+                            self.set_hparam(k1, v1)
+                        else:
+                            self.add_hparam(k1, v1)
+                        self.d[k1] = v1
+                    except Exception:
+                        import traceback
+
+                        print(traceback.format_exc())
+
+    # @property
+    def get_elements(self):
+        return self.d.items()
+
+
+def override_config(base_config, new_config):
+    """Update new configurations in the original dict with the new dict
+
+    Args:
+        base_config (dict): original dict to be overridden
+        new_config (dict): dict with new configurations
+
+    Returns:
+        dict: updated configuration dict
+    """
+    for k, v in new_config.items():
+        if type(v) == dict:
+            if k not in base_config.keys():
+                base_config[k] = {}
+            base_config[k] = override_config(base_config[k], v)
+        else:
+            base_config[k] = v
+    return base_config
+
+
+def get_lowercase_keys_config(cfg):
+    """Change all keys in cfg to lower case
+
+    Args:
+        cfg (dict): dictionary that stores configurations
+
+    Returns:
+        dict: dictionary that stores configurations
+    """
+    updated_cfg = dict()
+    for k, v in cfg.items():
+        if type(v) == dict:
+            v = get_lowercase_keys_config(v)
+        updated_cfg[k.lower()] = v
+    return updated_cfg
+
+
+def _load_config(config_fn, lowercase=False):
+    """Load configurations into a dictionary
+
+    Args:
+        config_fn (str): path to configuration file
+        lowercase (bool, optional): whether changing keys to lower case. Defaults to False.
+
+    Returns:
+        dict: dictionary that stores configurations
+    """
+    with open(config_fn, "r") as f:
+        data = f.read()
+    config_ = json5.loads(data)
+    if "base_config" in config_:
+        # load configurations from new path
+        p_config_path = os.path.join(os.getenv("WORK_DIR"), config_["base_config"])
+        p_config_ = _load_config(p_config_path)
+        config_ = override_config(p_config_, config_)
+    if lowercase:
+        # change keys in config_ to lower case
+        config_ = get_lowercase_keys_config(config_)
+    return config_
+
+
+def load_config(config_fn, lowercase=False):
+    """Load configurations into a dictionary
+
+    Args:
+        config_fn (str): path to configuration file
+        lowercase (bool, optional): _description_. Defaults to False.
+
+    Returns:
+        JsonHParams: an object that stores configurations
+    """
+    config_ = _load_config(config_fn, lowercase=lowercase)
+    # create an JsonHParams object with configuration dict
+    cfg = JsonHParams(**config_)
+    return cfg
+
+
+def save_config(save_path, cfg):
+    """Save configurations into a json file
+
+    Args:
+        save_path (str): path to save configurations
+        cfg (dict): dictionary that stores configurations
+    """
+    with open(save_path, "w") as f:
+        json5.dump(
+            cfg, f, ensure_ascii=False, indent=4, quote_keys=True, sort_keys=True
+        )
+
+
+class JsonHParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if type(v) == dict:
+                v = JsonHParams(**v)
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+
+class ValueWindow:
+    def __init__(self, window_size=100):
+        self._window_size = window_size
+        self._values = []
+
+    def append(self, x):
+        self._values = self._values[-(self._window_size - 1) :] + [x]
+
+    @property
+    def sum(self):
+        return sum(self._values)
+
+    @property
+    def count(self):
+        return len(self._values)
+
+    @property
+    def average(self):
+        return self.sum / max(1, self.count)
+
+    def reset(self):
+        self._values = []
+
+
+class Logger(object):
+    def __init__(
+        self,
+        filename,
+        level="info",
+        when="D",
+        backCount=10,
+        fmt="%(asctime)s : %(message)s",
+    ):
+        self.level_relations = {
+            "debug": logging.DEBUG,
+            "info": logging.INFO,
+            "warning": logging.WARNING,
+            "error": logging.ERROR,
+            "crit": logging.CRITICAL,
+        }
+        if level == "debug":
+            fmt = "%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s"
+        self.logger = logging.getLogger(filename)
+        format_str = logging.Formatter(fmt)
+        self.logger.setLevel(self.level_relations.get(level))
+        sh = logging.StreamHandler()
+        sh.setFormatter(format_str)
+        th = handlers.TimedRotatingFileHandler(
+            filename=filename, when=when, backupCount=backCount, encoding="utf-8"
+        )
+        th.setFormatter(format_str)
+        self.logger.addHandler(sh)
+        self.logger.addHandler(th)
+        self.logger.info(
+            "==========================New Starting Here=============================="
+        )
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def slice_segments(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = ids_str[i]
+        idx_end = idx_str + segment_size
+        ret[i] = x[i, :, idx_str:idx_end]
+    return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def subsequent_mask(length):
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def convert_pad_shape(pad_shape):
+    l = pad_shape[::-1]
+    pad_shape = [item for sublist in l for item in sublist]
+    return pad_shape
+
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def generate_path(duration, mask):
+    """
+    duration: [b, 1, t_x]
+    mask: [b, 1, t_y, t_x]
+    """
+    device = duration.device
+
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm
+
+
+def get_current_time():
+    pass
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """
+    Args:
+      lengths:
+        A 1-D tensor containing sentence lengths.
+      max_len:
+        The length of masks.
+    Returns:
+      Return a 2-D bool tensor, where masked positions
+      are filled with `True` and non-masked positions are
+      filled with `False`.
+
+    >>> lengths = torch.tensor([1, 3, 2, 5])
+    >>> make_pad_mask(lengths)
+    tensor([[False,  True,  True,  True,  True],
+            [False, False, False,  True,  True],
+            [False, False,  True,  True,  True],
+            [False, False, False, False, False]])
+    """
+    assert lengths.ndim == 1, lengths.ndim
+    max_len = max(max_len, lengths.max())
+    n = lengths.size(0)
+    seq_range = torch.arange(0, max_len, device=lengths.device)
+    expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
+
+    return expaned_lengths >= lengths.unsqueeze(-1)