feat: app.py

XY_Tokenizer/.gitignore

@@ -0,0 +1,193 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+__pycache__/
+
+# Pyre type checker
+.pyre/
+
+# windows folder
+*.ini
+
+# models
+*.pkl
+
+*.wav
+*.flac
+*.mp3
+
+# others
+temp/
+
+exp/
+
+slurmlogs/
+slurmlogs_*/
+
+dev/
+.vscode/
+
+config/debug
+.vscode/
+
+
+submit_debug*
+
+random_rep_for_v2.13
+
+exp_eval
+
+data/**/*.txt
+
+tokenize_data/tokenize_result/
+
+
+*.png
+
+reconstruct_evaluation_backup/
+
+semantic_evaluation/scripts/en
+
+reconstruct_evaluation/scripts
+
+*.jsonl
+*.json
+scripts/debug
+
+*.hostfile
+.deepspeed_env
+
+*.idx
+
+backup*
+
+*.ckpt
+
+# Project specific
+output_wavs/
+*.pt
+*.pth
+output.log

XY_Tokenizer/README.md

@@ -0,0 +1,72 @@
+# XY Tokenizer
+
+XY Tokenizer is a speech codec that simultaneously models both semantic and acoustic aspects of speech, converting audio into discrete tokens and decoding them back to high-quality audio. It achieves efficient speech representation at only 1kbps with RVQ8 quantization at 12.5Hz frame rate.
+
+## Features
+
+- **Dual-channel modeling**: Simultaneously captures semantic meaning and acoustic details
+- **Efficient representation**: 1kbps bitrate with RVQ8 quantization at 12.5Hz
+- **High-quality audio tokenization**: Convert speech to discrete tokens and back with minimal quality loss
+- **Long audio support**: Process audio files longer than 30 seconds using chunking with overlap
+- **Batch processing**: Efficiently process multiple audio files in batches
+- **24kHz output**: Generate high-quality 24kHz audio output
+
+## Installation
+
+```bash
+# Create and activate conda environment
+conda create -n xy_tokenizer python=3.10 -y && conda activate xy_tokenizer
+
+# Install dependencies
+pip install -r requirements.txt
+```
+
+## Usage
+
+### Basic Inference
+
+To tokenize audio files and reconstruct them:
+
+```bash
+python inference.py \
+  --config_path ./config/xy_tokenizer_config.yaml \
+  --checkpoint_path ./weights/xy_tokenizer.ckpt \
+  --input_dir ./input_wavs/ \
+  --output_dir ./output_wavs/
+```
+
+### Parameters
+
+- `--config_path`: Path to the model configuration file
+- `--checkpoint_path`: Path to the pre-trained model checkpoint
+- `--input_dir`: Directory containing input WAV files
+- `--output_dir`: Directory to save reconstructed audio files
+- `--device`: Device to run inference on (default: "cuda")
+- `--debug`, `--debug_ip`, `--debug_port`: Debugging options (disabled by default)
+
+## Project Structure
+
+- `xy_tokenizer/`: Core model implementation
+  - `model.py`: Main XY_Tokenizer model class
+  - `nn/`: Neural network components
+- `config/`: Configuration files
+- `utils/`: Utility functions
+- `weights/`: Pre-trained model weights
+- `input_wavs/`: Directory for input audio files
+- `output_wavs/`: Directory for output audio files
+
+## Model Architecture
+
+XY Tokenizer uses a dual-channel architecture that simultaneously models:
+1. **Semantic Channel**: Captures high-level semantic information and linguistic content
+2. **Acoustic Channel**: Preserves detailed acoustic features including speaker characteristics and prosody
+
+The model processes audio through several stages:
+1. Feature extraction (mel-spectrogram)
+2. Parallel semantic and acoustic encoding
+3. Residual Vector Quantization (RVQ8) at 12.5Hz frame rate (1kbps)
+4. Decoding and waveform generation
+
+## License
+
+[Specify your license here]

XY_Tokenizer/config/xy_tokenizer_config.yaml

@@ -0,0 +1,122 @@
+input_sample_rate: &input_sample_rate 16000
+output_sample_rate: &output_sample_rate 24000
+
+generator_params:
+    input_sample_rate: *input_sample_rate
+    output_sample_rate: *output_sample_rate
+
+    feature_extractor_kwargs:
+        chunk_length: 30
+        feature_size: 80
+        hop_length: 160
+        n_fft: 400
+        n_samples: 480000
+        nb_max_frames: 3000
+        padding_side: right
+        padding_value: 0.0
+        return_attention_mask: false
+        sampling_rate: *input_sample_rate
+
+    ## Codec Args
+
+    ## semantic channel
+    semantic_encoder_kwargs:  # 100hz -> 50hz
+        num_mel_bins: 80
+        sampling_rate: *input_sample_rate
+        hop_length: 160
+        stride_size: 2
+        kernel_size: 3
+        d_model: 768
+        scale_embedding: false
+        max_audio_seconds: 30
+        encoder_layers: 12
+        encoder_attention_heads: 12
+        encoder_ffn_dim: 3072
+        activation_function: "gelu"
+
+    semantic_encoder_adapter_kwargs: # 50hz
+        input_dim: 768
+        output_dim: 768
+        d_model: 768
+        max_source_positions: 1500
+        encoder_layers: 4
+        encoder_attention_heads: 12
+        encoder_ffn_dim: 3072
+
+
+    ## acoustic channel
+    acoustic_encoder_kwargs:  # 100hz -> 50hz
+        num_mel_bins: 80
+        sampling_rate: *input_sample_rate
+        hop_length: 160
+        stride_size: 2
+        kernel_size: 3
+        d_model: 768
+        scale_embedding: false
+        max_audio_seconds: 30
+        encoder_layers: 12
+        encoder_attention_heads: 12
+        encoder_ffn_dim: 3072
+        activation_function: "gelu"
+
+
+    ## semantic & acoustic shared parameters 
+    pre_rvq_adapter_kwargs: # 50hz
+        input_dim: 1536
+        output_dim: 768
+        d_model: 768
+        max_source_positions: 1500
+        encoder_layers: 4
+        encoder_attention_heads: 12
+        encoder_ffn_dim: 3072
+        
+    downsample_kwargs:  # 50hz -> 12.5hz
+        d_model: 768
+        avg_pooler: 4
+
+    quantizer_kwargs:  # 12.5hz
+        input_dim: 3072
+        rvq_dim: 512
+        output_dim: 3072
+        num_quantizers: 8
+        codebook_size: 1024
+        codebook_dim: 512
+        quantizer_dropout: 0.0
+        commitment: 1
+    
+    post_rvq_adapter_kwargs: # 12.5hz
+        input_dim: 3072
+        output_dim: 3072
+        d_model: 768
+        max_source_positions: 375
+        encoder_layers: 4
+        encoder_attention_heads: 12
+        encoder_ffn_dim: 3072
+
+    upsample_kwargs:  # 12.5hz -> 50hz
+        d_model: 768
+        stride: 4
+    
+    ## acoustic channel
+    acoustic_decoder_kwargs:  # 50hz -> 100hz
+        num_mel_bins: 80
+        sampling_rate: *input_sample_rate
+        hop_length: 160
+        stride_size: 2
+        kernel_size: 3
+        d_model: 768
+        scale_embedding: false
+        max_audio_seconds: 30
+        decoder_layers: 12
+        decoder_attention_heads: 12
+        decoder_ffn_dim: 3072
+        activation_function: "gelu"
+    
+    vocos_kwargs:  # 100hz -> 24khz
+        input_channels: 80
+        dim: 512
+        intermediate_dim: 4096
+        num_layers: 30
+        n_fft: 960
+        hop_size: 240
+        padding: "same"

XY_Tokenizer/inference.py

@@ -0,0 +1,71 @@
+import os
+import argparse
+import logging
+import torch
+
+from utils.helpers import set_logging, waiting_for_debug, load_audio, save_audio, find_audio_files
+from xy_tokenizer.model import XY_Tokenizer
+
+if __name__ == "__main__":
+    set_logging()
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, default="./config/xy_tokenizer_config.yaml")
+    parser.add_argument("--checkpoint_path", type=str, default="./weights/xy_tokenizer.ckpt")
+    parser.add_argument("--device", type=str, default="cuda")
+    
+    parser.add_argument("--input_dir", type=str, required=True)
+    parser.add_argument("--output_dir", type=str, required=True)
+    
+    
+    parser.add_argument("--debug_ip", type=str)
+    parser.add_argument("--debug_port", type=int)
+    parser.add_argument("--debug", default=0, type=int, nargs="?",
+        help='whether debug or not',
+    )
+    args = parser.parse_args()
+    if args.debug == 1:
+        waiting_for_debug(args.debug_ip, args.debug_port)
+
+    device = torch.device(args.device)
+
+    ## Load codec model
+    generator = XY_Tokenizer.load_from_checkpoint(config_path=args.config_path, ckpt_path=args.checkpoint_path).to(device).eval()
+    
+    ## Find audios
+    audio_paths = find_audio_files(input_dir=args.input_dir)
+    
+    ## Create output directory if not exists
+    os.makedirs(args.output_dir, exist_ok=True)
+    logging.info(f"Processing {len(audio_paths)} audio files, output will be saved to {args.output_dir}")
+
+    with torch.no_grad():
+        ## Process audios in batches
+        batch_size = 8
+        for i in range(0, len(audio_paths), batch_size):
+            batch_paths = audio_paths[i:i + batch_size]
+            logging.info(f"Processing batch {i // batch_size + 1}/{len(audio_paths) // batch_size + 1}, files: {batch_paths}")
+
+            # Load audio files
+            wav_list = [load_audio(path, target_sample_rate=generator.input_sample_rate).squeeze().to(device) for path in batch_paths]
+            logging.info(f"Successfully loaded {len(wav_list)} audio files with lengths {[len(wav) for wav in wav_list]} samples")
+
+            # Encode
+            encode_result = generator.encode(wav_list, overlap_seconds=10)
+            codes_list = encode_result["codes_list"]  # B * (nq, T)
+            logging.info(f"Encoding completed, code lengths: {[codes.shape[-1] for codes in codes_list]}")
+            logging.info(f"{codes_list = }")
+
+            # Decode
+            decode_result = generator.decode(codes_list, overlap_seconds=10)
+            syn_wav_list = decode_result["syn_wav_list"]  # B * (T,)
+            logging.info(f"Decoding completed, generated waveform lengths: {[len(wav) for wav in syn_wav_list]} samples")
+
+            # Save generated audios
+            for path, syn_wav in zip(batch_paths, syn_wav_list):
+                output_path = os.path.join(args.output_dir, os.path.basename(path))
+                save_audio(output_path, syn_wav.cpu().reshape(1, -1), sample_rate=generator.output_sample_rate)
+                logging.info(f"Saved generated audio to {output_path}")
+
+        
+    logging.info("All audio processing completed")

XY_Tokenizer/requirements.txt

@@ -0,0 +1,23 @@
+beartype
+tensorboard
+numpy
+torch
+torchaudio
+einops
+scipy
+huggingface-hub
+soundfile
+matplotlib
+lion_pytorch
+accelerate
+debugpy
+tensorboardX
+librosa
+pesq
+tqdm
+mir_eval
+stopes
+s3prl
+onnxscript
+jiwer
+orjson

XY_Tokenizer/utils/helpers.py

@@ -0,0 +1,117 @@
+import logging
+import torchaudio
+import os
+import sys
+import glob
+import debugpy
+import torch
+import numpy as np
+import re
+
+def count_params_by_module(model_name, model):
+    logging.info(f"Counting num_parameters of {model_name}:")
+    
+    param_stats = {}
+    total_params = 0  # Count total parameters
+    total_requires_grad_params = 0  # Count parameters with requires_grad=True
+    total_no_grad_params = 0  # Count parameters with requires_grad=False
+    
+    for name, param in model.named_parameters():
+        module_name = name.split('.')[0]
+        if module_name not in param_stats:
+            param_stats[module_name] = {'total': 0, 'requires_grad': 0, 'no_grad': 0}
+        
+        param_num = param.numel()
+        param_stats[module_name]['total'] += param_num
+        total_params += param_num
+        
+        if param.requires_grad:
+            param_stats[module_name]['requires_grad'] += param_num
+            total_requires_grad_params += param_num
+        else:
+            param_stats[module_name]['no_grad'] += param_num
+            total_no_grad_params += param_num
+    
+    # Calculate maximum width for each column
+    max_module_name_length = max(len(module) for module in param_stats)
+    max_param_length = max(len(f"{stats['total'] / 1e6:.2f}M") for stats in param_stats.values())
+    
+    # Output parameter statistics for each module
+    for module, stats in param_stats.items():
+        logging.info(f"\t{module:<{max_module_name_length}}: "
+                     f"Total: {stats['total'] / 1e6:<{max_param_length}.2f}M, "
+                     f"Requires Grad: {stats['requires_grad'] / 1e6:<{max_param_length}.2f}M, "
+                     f"No Grad: {stats['no_grad'] / 1e6:<{max_param_length}.2f}M")
+    
+    # Output total parameter statistics
+    logging.info(f"\tTotal parameters: {total_params / 1e6:.2f}M parameters")
+    logging.info(f"\tRequires Grad parameters: {total_requires_grad_params / 1e6:.2f}M parameters")
+    logging.info(f"\tNo Grad parameters: {total_no_grad_params / 1e6:.2f}M parameters")
+    logging.info(f"################################################################")
+
+
+def load_and_resample_audio(audio_path, target_sample_rate):
+    wav, raw_sample_rate = torchaudio.load(audio_path) # (1, T)   tensor 
+    if raw_sample_rate != target_sample_rate:   
+        wav = torchaudio.functional.resample(wav, raw_sample_rate, target_sample_rate) # tensor 
+    return wav.squeeze()
+
+def set_logging():
+    rank = os.environ.get("RANK", 0)
+    logging.basicConfig(
+        level=logging.INFO,
+        stream=sys.stdout,
+        format=f"%(asctime)s [RANK {rank}] (%(module)s:%(lineno)d) %(levelname)s : %(message)s",
+    )
+    
+def waiting_for_debug(ip, port):
+    rank = os.environ.get("RANK", "0")
+    debugpy.listen((ip, port)) # Replace localhost with cluster node IP
+    logging.info(f"[rank = {rank}] Waiting for debugger attach...")
+    debugpy.wait_for_client()
+    logging.info(f"[rank = {rank}] Debugger attached")
+    
+def load_audio(audio_path, target_sample_rate):
+    # Load audio file, wav shape: (channels, time)
+    wav, raw_sample_rate = torchaudio.load(audio_path)
+    
+    # If multi-channel, convert to mono by averaging across channels
+    if wav.shape[0] > 1:
+        wav = torch.mean(wav, dim=0, keepdim=True)  # Average across channels, keep channel dim
+    
+    # Resample if necessary
+    if raw_sample_rate != target_sample_rate:
+        wav = torchaudio.functional.resample(wav, raw_sample_rate, target_sample_rate)
+    
+    # Convert to numpy, add channel dimension, then back to tensor with desired shape
+    wav = np.expand_dims(wav.squeeze(0).numpy(), axis=1)  # Shape: (time, 1)
+    wav = torch.tensor(wav).reshape(1, 1, -1)  # Shape: (1, 1, time)
+    
+    return wav
+
+def save_audio(audio_outpath, audio_out, sample_rate):
+    torchaudio.save(
+        audio_outpath, 
+        audio_out, 
+        sample_rate=sample_rate, 
+        encoding='PCM_S', 
+        bits_per_sample=16
+    )
+    logging.info(f"Successfully saved audio at {audio_outpath}")
+    
+def find_audio_files(input_dir):
+    audio_extensions = ['*.flac', '*.mp3', '*.wav']
+    audios_input = []
+    for ext in audio_extensions:
+        audios_input.extend(glob.glob(os.path.join(input_dir, '**', ext), recursive=True))
+    logging.info(f"Found {len(audios_input)} audio files in {input_dir}")
+    return sorted(audios_input)
+
+def normalize_text(text):
+    # Remove all punctuation (including English and Chinese punctuation)
+    text = re.sub(r'[^\w\s\u4e00-\u9fff]', '', text)
+    # Convert to lowercase (effective for English, no effect on Chinese)
+    text = text.lower()
+    # Remove extra spaces
+    text = ' '.join(text.split())
+    return text

XY_Tokenizer/xy_tokenizer/model.py

@@ -0,0 +1,279 @@
+# -*- coding: utf-8 -*-
+import yaml
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+from .nn.feature_extractor import MelFeatureExtractor
+from .nn.modules import OmniAudioEncoder, OmniAudioDecoder, ResidualDownConv, UpConv, Transformer, Vocos
+from .nn.quantizer import ResidualVQ
+
+class XY_Tokenizer(nn.Module):
+    def __init__(self, generator_params):
+        super().__init__()
+        # Basic parameters
+        self.input_sample_rate = generator_params['input_sample_rate']
+        self.output_sample_rate = generator_params['output_sample_rate']
+        
+        self.encoder_downsample_rate = 1280
+        self.decoder_upsample_rate = 1920
+        self.code_dim = generator_params['quantizer_kwargs']['input_dim']
+        
+        ## Codec part
+
+        ## Semantic channel
+        self.semantic_encoder = OmniAudioEncoder(**generator_params['semantic_encoder_kwargs'])
+        
+        self.semantic_encoder_adapter = Transformer(**generator_params['semantic_encoder_adapter_kwargs'])
+        
+        ## Acoustic channel
+        self.acoustic_encoder = OmniAudioEncoder(**generator_params['acoustic_encoder_kwargs'])
+        
+        ## Semantic & acoustic shared parameters
+        self.pre_rvq_adapter = Transformer(**generator_params['pre_rvq_adapter_kwargs'])
+        
+        self.downsample = ResidualDownConv(**generator_params['downsample_kwargs'])
+        
+        self.quantizer = ResidualVQ(**generator_params['quantizer_kwargs'])
+        self.nq = generator_params['quantizer_kwargs']['num_quantizers']
+
+        self.post_rvq_adapter = Transformer(**generator_params['post_rvq_adapter_kwargs'])
+                    
+        ## Acoustic channel
+        self.upsample = UpConv(**generator_params['upsample_kwargs'])
+
+        self.acoustic_decoder = OmniAudioDecoder(**generator_params['acoustic_decoder_kwargs'])
+
+        self.enhanced_vocos = Vocos(**generator_params['vocos_kwargs'])
+
+        ## Feature extractor
+        self.feature_extractor = MelFeatureExtractor(**generator_params['feature_extractor_kwargs'])
+
+    @torch.inference_mode()
+    def inference_tokenize(self, x, input_lengths):
+        """
+            Input:
+                x: Waveform tensor # (B, 1, T), T <= 30s * sample_rate
+                input_lengths: Valid length for each sample # (B,)
+            Output:
+                dict: Contains the following key-value pairs
+                    "zq": Quantized embeddings # (B, D, T)
+                    "codes": Quantization codes # (nq, B, T)
+                    "codes_lengths": Quantization code lengths # (B,)
+        """
+        list_x = [xi[:, :x_len].reshape(-1).cpu().numpy() for xi, x_len in zip(x, input_lengths)]
+        features = self.feature_extractor(
+            list_x,
+            sampling_rate=self.input_sample_rate,
+            return_tensors="pt",
+            return_attention_mask=True
+        )
+        input_mel = features['input_features'].to(x.device).to(x.dtype) # (B, D, 3000)
+        audio_attention_mask = features['attention_mask'].to(x.device) # (B, 3000)
+        
+        # Get batch size and sequence length of the input
+        mel_output_length = torch.sum(audio_attention_mask, dim=-1).long() # (B,)
+        
+        # Semantic channel
+        semantic_encoder_output, semantic_encoder_output_length = self.semantic_encoder(input_mel, mel_output_length) # (B, D, T), 100hz -> 50hz
+        
+        semantic_encoder_adapter_output, semantic_encoder_adapter_output_length = self.semantic_encoder_adapter(semantic_encoder_output, semantic_encoder_output_length) # (B, D, T), 50hz
+        
+        # Acoustic channel
+        acoustic_encoder_output, acoustic_encoder_output_length = self.acoustic_encoder(input_mel, mel_output_length) # (B, D, T), 100hz -> 50hz
+        
+        # Semantic & acoustic mixing
+        concated_semantic_acoustic_channel = torch.concat([semantic_encoder_adapter_output, acoustic_encoder_output], dim=1) # (B, D, T)
+        concated_semantic_acoustic_channel_length = acoustic_encoder_output_length
+        
+        pre_rvq_adapter_output, pre_rvq_adapter_output_length = self.pre_rvq_adapter(concated_semantic_acoustic_channel, concated_semantic_acoustic_channel_length) # (B, D, T), 50hz
+        
+        downsample_output, downsample_output_length = self.downsample(pre_rvq_adapter_output, pre_rvq_adapter_output_length) # (B, D, T), 50hz -> 12.5hz
+
+        zq, codes, vq_loss, _, quantizer_output_length = self.quantizer(downsample_output, downsample_output_length) # (B, D, T), (nq, B, T), (nq,), (nq, B, D, T), (B,)
+
+        return {
+            "zq": zq, # (B, D, T)
+            "codes": codes, # (nq, B, T)
+            "codes_lengths": quantizer_output_length # (B,)
+        }
+      
+    @torch.inference_mode()  
+    def inference_detokenize(self, codes, codes_lengths):
+        """
+            Input:
+                codes: Quantization codes # (nq, B, T)
+                codes_lengths: Quantization code lengths for each sample # (B,)
+            Output:
+                dict: Contains the following key-value pairs
+                    "y": Synthesized audio waveform # (B, 1, T)
+                    "output_length": Output lengths # (B,)
+        """
+        zq = self.quantizer.decode_codes(codes) # (B, D, T)
+        
+        post_rvq_adapter_output, post_rvq_adapter_output_length = self.post_rvq_adapter(zq, codes_lengths) # (B, D, T), 12.5hz
+        
+        # Acoustic channel            
+        upsample_output, upsample_output_length = self.upsample(post_rvq_adapter_output, post_rvq_adapter_output_length) # (B, D, T), 12.5hz -> 50hz
+
+        acoustic_decoder_output, acoustic_decoder_output_length = self.acoustic_decoder(upsample_output, upsample_output_length) # (B, D, T), 50hz -> 100hz
+
+        y, vocos_output_length = self.enhanced_vocos(acoustic_decoder_output, acoustic_decoder_output_length) # (B, 1, T), 100hz -> 16khz
+        
+        return {
+            "y": y, # (B, 1, T)
+            "output_length": vocos_output_length, # (B,)
+        }
+        
+    @torch.inference_mode()
+    def encode(self, wav_list, overlap_seconds=10, device=torch.device("cuda")):
+        """
+            Input:
+                wav_list: List of audio waveforms, each with potentially different length, may exceed 30 seconds # B * (T,)
+                overlap_seconds: Overlap in seconds, process 30 seconds at a time, keeping (30 - overlap_seconds) seconds of valid output
+            Output:
+                dict: Contains the following key-value pairs
+                    "codes_list": List of quantization codes # B * (nq, T)
+        """
+        duration_seconds = 30 - overlap_seconds
+        chunk_size = int(30 * self.input_sample_rate) # Maximum samples per chunk
+        duration_size = int(duration_seconds * self.input_sample_rate) # Valid output samples per chunk
+        code_duration_length = duration_size // self.encoder_downsample_rate # Valid code length per chunk
+
+        # Get maximum waveform length
+        max_length = max(len(wav) for wav in wav_list)
+        batch_size = len(wav_list)
+        wav_tensor = torch.zeros(batch_size, 1, max_length, device=device)
+        input_lengths = torch.zeros(batch_size, dtype=torch.long, device=device)
+        for i, wav in enumerate(wav_list):
+            wav_tensor[i, 0, :len(wav)] = wav
+            input_lengths[i] = len(wav) # (B,)
+
+        # Calculate number of chunks needed
+        max_chunks = (max_length + duration_size - 1) // duration_size
+        codes_list = []
+
+        # Process the entire batch in chunks
+        for chunk_idx in range(max_chunks):
+            start = chunk_idx * duration_size
+            end = min(start + chunk_size, max_length)
+            chunk = wav_tensor[:, :, start:end] # (B, 1, T')
+            chunk_lengths = torch.clamp(input_lengths - start, 0, end - start) # (B,)
+
+            # Skip empty chunks
+            if chunk_lengths.max() == 0:
+                continue
+
+            # Encode
+            result = self.inference_tokenize(chunk, chunk_lengths) # {"zq": (B, D, T'), "codes": (nq, B, T'), "codes_lengths": (B,)}
+            chunk_codes = result["codes"] # (nq, B, T')
+            chunk_code_lengths = result["codes_lengths"] # (B,)
+
+            # Extract valid portion
+            valid_code_lengths = torch.clamp(chunk_code_lengths, 0, code_duration_length) # (B,)
+            valid_chunk_codes = torch.zeros(self.nq, batch_size, code_duration_length, device=device, dtype=chunk_codes.dtype)
+            for b in range(batch_size):
+                if valid_code_lengths[b] > 0:
+                    valid_chunk_codes[:, b, :valid_code_lengths[b]] = chunk_codes[:, b, :valid_code_lengths[b]] # (nq, B, valid_code_length)
+
+            codes_list.append(valid_chunk_codes) # (nq, B, valid_code_length)
+
+        # Concatenate all chunks
+        if codes_list:
+            codes_tensor = torch.cat(codes_list, dim=-1) # (nq, B, T_total)
+            codes_list = [codes_tensor[:, i, :input_lengths[i] // self.encoder_downsample_rate] for i in range(batch_size)] # B * (nq, T)
+        else:
+            codes_list = [torch.zeros(self.nq, 0, device=device, dtype=torch.long) for _ in range(batch_size)] # B * (nq, 0)
+
+        return {
+            "codes_list": codes_list # B * (nq, T)
+        }
+        
+    @torch.inference_mode()
+    def decode(self, codes_list, overlap_seconds=10, device=torch.device("cuda")):
+        """
+            Input:
+                codes_list: List of quantization codes # B * (nq, T)
+                overlap_seconds: Overlap in seconds, process 30 seconds at a time, keeping (30 - overlap_seconds) seconds of valid output
+            Output:
+                dict: Contains the following key-value pairs
+                    "syn_wav_list": List of synthesized audio waveforms # B * (T,)
+        """
+        duration_seconds = 30 - overlap_seconds
+        chunk_code_length = int(30 * self.input_sample_rate // self.encoder_downsample_rate) # Maximum code length per chunk
+        duration_code_length = int(duration_seconds * self.input_sample_rate // self.encoder_downsample_rate) # Valid code length per chunk
+        duration_wav_length = duration_code_length * self.decoder_upsample_rate # Valid waveform length per chunk
+
+        # Get maximum code length
+        max_code_length = max(codes.shape[-1] for codes in codes_list)
+        batch_size = len(codes_list)
+        codes_tensor = torch.zeros(self.nq, batch_size, max_code_length, device=device, dtype=torch.long)
+        code_lengths = torch.zeros(batch_size, dtype=torch.long, device=device)
+        for i, codes in enumerate(codes_list):
+            codes_tensor[:, i, :codes.shape[-1]] = codes.to(device)
+            code_lengths[i] = codes.shape[-1] # (B,)
+
+        # Calculate number of chunks needed
+        max_chunks = (max_code_length + duration_code_length - 1) // duration_code_length
+        wav_list = []
+
+        # Process the entire batch in chunks
+        for chunk_idx in range(max_chunks):
+            start = chunk_idx * duration_code_length
+            end = min(start + chunk_code_length, max_code_length)
+            chunk_codes = codes_tensor[:, :, start:end] # (nq, B, T')
+            chunk_code_lengths = torch.clamp(code_lengths - start, 0, end - start) # (B,)
+
+            # Skip empty chunks
+            if chunk_code_lengths.max() == 0:
+                continue
+
+            # Decode
+            result = self.inference_detokenize(chunk_codes, chunk_code_lengths) # {"y": (B, 1, T'), "output_length": (B,)}
+            chunk_wav = result["y"] # (B, 1, T')
+            chunk_wav_lengths = result["output_length"] # (B,)
+
+            # Extract valid portion
+            valid_wav_lengths = torch.clamp(chunk_wav_lengths, 0, duration_wav_length) # (B,)
+            valid_chunk_wav = torch.zeros(batch_size, 1, duration_wav_length, device=device)
+            for b in range(batch_size):
+                if valid_wav_lengths[b] > 0:
+                    valid_chunk_wav[b, :, :valid_wav_lengths[b]] = chunk_wav[b, :, :valid_wav_lengths[b]] # (B, 1, valid_wav_length)
+
+            wav_list.append(valid_chunk_wav) # (B, 1, valid_wav_length)
+
+        # Concatenate all chunks
+        if wav_list:
+            wav_tensor = torch.cat(wav_list, dim=-1) # (B, 1, T_total)
+            syn_wav_list = [wav_tensor[i, 0, :code_lengths[i] * self.decoder_upsample_rate] for i in range(batch_size)] # B * (T,)
+        else:
+            syn_wav_list = [torch.zeros(0, device=device) for _ in range(batch_size)] # B * (0,)
+            
+        return {
+            "syn_wav_list": syn_wav_list # B * (T,)
+        }
+    
+    @classmethod
+    def load_from_checkpoint(cls, config_path: str, ckpt_path: str):
+        # Load model from configuration file and checkpoint
+        logging.info(f"Loading model from {config_path} and {ckpt_path}")
+        
+        # Load configuration
+        with open(config_path, 'r') as f:
+            config = yaml.safe_load(f)
+        
+        # Create model instance
+        model = cls(config['generator_params'])
+        
+        # Load checkpoint
+        checkpoint = torch.load(ckpt_path, map_location='cpu')
+        
+        # Check if checkpoint contains 'generator' key
+        if 'generator' in checkpoint:
+            model.load_state_dict(checkpoint['generator'])
+        else:
+            model.load_state_dict(checkpoint)
+        
+        return model

XY_Tokenizer/xy_tokenizer/nn/feature_extractor.py

@@ -0,0 +1,237 @@
+import numpy as np
+import torch
+
+from typing import Union, List, Optional
+from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils import TensorType, logging
+from transformers.utils.import_utils import is_torch_available
+from transformers.audio_utils import mel_filter_bank, spectrogram, window_function
+
+class MelFeatureExtractor(SequenceFeatureExtractor):
+    model_input_names = ["input_features"]
+    
+    def __init__(
+        self,
+        feature_size=80,
+        sampling_rate=16000,
+        hop_length=160,
+        chunk_length=30,
+        n_fft=400,
+        padding_value=0.0,
+        dither=0.0,
+        return_attention_mask=False,
+        max_frequency=None,
+        **kwargs,
+    ):
+        super().__init__(
+            feature_size=feature_size,
+            sampling_rate=sampling_rate,
+            padding_value=padding_value,
+            return_attention_mask=return_attention_mask,
+            **kwargs,
+        )
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.chunk_length = chunk_length
+        self.n_samples = chunk_length * sampling_rate
+        self.nb_max_frames = self.n_samples // hop_length
+        self.sampling_rate = sampling_rate
+        self.dither = dither
+        self.max_frequency = max_frequency if max_frequency is not None else sampling_rate / 2
+        self.mel_filters = mel_filter_bank(
+            num_frequency_bins=1 + n_fft // 2,
+            num_mel_filters=feature_size,
+            min_frequency=0.0,
+            max_frequency=self.max_frequency,
+            sampling_rate=sampling_rate,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+
+    def _np_extract_fbank_features(self, waveform_batch: np.array, device: str) -> np.ndarray:
+        if device != "cpu":
+            raise ValueError(
+                f"Got device `{device}` for feature extraction, but feature extraction on CUDA accelerator "
+                "devices requires torch, which is not installed. Either set `device='cpu'`, or "
+                "install torch according to the official instructions: https://pytorch.org/get-started/locally/"
+            )
+        log_spec_batch = []
+        for waveform in waveform_batch:
+            log_spec = spectrogram(
+                waveform,
+                window_function(self.n_fft, "hann"),
+                frame_length=self.n_fft,
+                hop_length=self.hop_length,
+                power=2.0,
+                dither=self.dither,
+                mel_filters=self.mel_filters,
+                log_mel="log10",
+            )
+            log_spec = log_spec[:, :-1]
+            log_spec = np.maximum(log_spec, log_spec.max() - 8.0)
+            log_spec = (log_spec + 4.0) / 4.0
+            log_spec_batch.append(log_spec)
+        log_spec_batch = np.array(log_spec_batch)
+        return log_spec_batch
+
+    def _torch_extract_fbank_features(self, waveform: np.array, device: str = "cpu") -> np.ndarray:
+        """
+        Compute the log-mel spectrogram of the audio using PyTorch's GPU-accelerated STFT implementation with batching,
+        yielding results similar to cpu computing with 1e-5 tolerance.
+        """
+        waveform = torch.from_numpy(waveform).to(device, torch.float32)
+        window = torch.hann_window(self.n_fft, device=device)
+
+        if self.dither != 0.0:
+            waveform += self.dither * torch.randn(waveform.shape, dtype=waveform.dtype, device=waveform.device)
+
+        stft = torch.stft(waveform, self.n_fft, self.hop_length, window=window, return_complex=True)
+        magnitudes = stft[..., :-1].abs() ** 2
+
+        mel_filters = torch.from_numpy(self.mel_filters).to(device, torch.float32)
+        mel_spec = mel_filters.T @ magnitudes
+
+        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+        if waveform.dim() == 2:
+            max_val = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0]
+            log_spec = torch.maximum(log_spec, max_val - 8.0)
+        else:
+            log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+        log_spec = (log_spec + 4.0) / 4.0
+        if device != "cpu":
+            log_spec = log_spec.detach().cpu()
+        return log_spec.numpy()
+
+    @staticmethod
+    def zero_mean_unit_var_norm(
+        input_values: List[np.ndarray], attention_mask: List[np.ndarray], padding_value: float = 0.0
+    ) -> List[np.ndarray]:
+        """
+        Every array in the list is normalized to have zero mean and unit variance
+        """
+        if attention_mask is not None:
+            attention_mask = np.array(attention_mask, np.int32)
+            normed_input_values = []
+
+            for vector, length in zip(input_values, attention_mask.sum(-1)):
+                normed_slice = (vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7)
+                if length < normed_slice.shape[0]:
+                    normed_slice[length:] = padding_value
+
+                normed_input_values.append(normed_slice)
+        else:
+            normed_input_values = [(x - x.mean()) / np.sqrt(x.var() + 1e-7) for x in input_values]
+
+        return normed_input_values
+
+    def __call__(
+        self,
+        raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],
+        truncation: bool = True,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: Optional[bool] = None,
+        padding: Optional[str] = "max_length",
+        max_length: Optional[int] = None,
+        sampling_rate: Optional[int] = None,
+        do_normalize: Optional[bool] = None,
+        device: Optional[str] = "cpu",
+        return_token_timestamps: Optional[bool] = None,
+        **kwargs,
+    ) -> BatchFeature:
+        """
+        Main method to featurize and prepare for the model one or several sequence(s). Implementation uses PyTorch for
+        the STFT computation if available, otherwise a slower NumPy based one.
+
+        Args:
+            raw_speech (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`):
+                The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float
+                values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not
+                stereo, i.e. single float per timestep.
+            truncation (`bool`, *optional*, default to `True`):
+                Activates truncation to cut input sequences longer than *max_length* to *max_length*.
+            pad_to_multiple_of (`int`, *optional*, defaults to None):
+                If set will pad the sequence to a multiple of the provided value.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors instead of list of python integers. Acceptable values are:
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return Numpy `np.ndarray` objects.
+            sampling_rate (`int`, *optional*):
+                The sampling rate at which the `raw_speech` input was sampled. If provided, it is checked against
+                the extractor's sampling rate.
+            padding_value (`float`, *optional*, defaults to 0.0):
+                The value that is used to fill the padding values / vectors.
+            do_normalize (`bool`, *optional*, defaults to `False`):
+                Whether or not to zero-mean unit-variance normalize the input.
+            device (`str`, *optional*, defaults to `'cpu'`):
+                Specifies the device for computation of the log-mel spectrogram.
+            return_token_timestamps (`bool`, *optional*, defaults to `None`):
+                Whether or not to return the number of frames of the input raw_speech.
+        """
+        if sampling_rate is not None and sampling_rate != self.sampling_rate:
+            logger.warning(
+                f"The provided `raw_speech` input was sampled at {sampling_rate}Hz, but the feature extractor "
+                f"is configured for {self.sampling_rate}Hz. You should resample the audio to match the "
+                f"extractor's sampling rate to ensure correct feature extraction."
+            )
+
+        is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
+        if is_batched_numpy and len(raw_speech.shape) > 2:
+            raise ValueError(f"Only mono-channel audio is supported for input to {self}")
+        is_batched = is_batched_numpy or (
+            isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
+        )
+
+        if is_batched:
+            raw_speech = [np.asarray([speech], dtype=np.float32).T for speech in raw_speech]
+        elif not is_batched and not isinstance(raw_speech, np.ndarray):
+            raw_speech = np.asarray(raw_speech, dtype=np.float32)
+        elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):
+            raw_speech = raw_speech.astype(np.float32)
+
+        if not is_batched:
+            raw_speech = [np.asarray([raw_speech]).T]
+
+        batched_speech = BatchFeature({"input_features": raw_speech})
+
+        padded_inputs = self.pad(
+            batched_speech,
+            padding=padding,
+            max_length=max_length if max_length else self.n_samples,
+            truncation=truncation,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_attention_mask=return_attention_mask or do_normalize,
+        )
+
+        if do_normalize:
+            padded_inputs["input_features"] = self.zero_mean_unit_var_norm(
+                padded_inputs["input_features"],
+                attention_mask=padded_inputs["attention_mask"],
+                padding_value=self.padding_value,
+            )
+            padded_inputs["input_features"] = np.stack(padded_inputs["input_features"], axis=0)
+
+        input_features = padded_inputs.get("input_features").transpose(2, 0, 1)
+
+        extract_fbank_features = (
+            self._torch_extract_fbank_features if is_torch_available() else self._np_extract_fbank_features
+        )
+        input_features = extract_fbank_features(input_features[0], device)
+
+        if isinstance(input_features[0], List):
+            padded_inputs["input_features"] = [np.asarray(feature, dtype=np.float32) for feature in input_features]
+        else:
+            padded_inputs["input_features"] = input_features
+
+        if return_attention_mask:
+            padded_inputs["attention_mask"] = padded_inputs["attention_mask"][:, :: self.hop_length]
+
+        if return_token_timestamps is not None:
+            padded_inputs["num_frames"] = [len(raw_speech_i) // self.hop_length for raw_speech_i in raw_speech]
+
+        if return_tensors is not None:
+            padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
+
+        return padded_inputs

XY_Tokenizer/xy_tokenizer/nn/modules.py

@@ -0,0 +1,1480 @@
+import torch
+import torch.distributed
+import numpy as np
+import logging
+import math
+import copy
+import numpy as np
+import scipy
+import torch
+import librosa
+
+from typing import Optional, Tuple
+from torch import nn, view_as_real, view_as_complex
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import weight_norm, remove_weight_norm
+from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz
+from transformers.activations import ACT2FN
+from dataclasses import dataclass
+from transformers.modeling_outputs import ModelOutput
+from transformers import WhisperModel
+
+
+# Define function to generate positional embeddings using sine and cosine functions to represent sequence position information
+def sinusoids(length, channels, max_timescale=10000):
+    """Returns sinusoidal waves for positional embedding"""
+    assert channels % 2 == 0
+    log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
+    scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+    return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
+
+# Generate sequence mask to distinguish valid sequence and padding parts
+def get_sequence_mask(inputs, inputs_length):
+    if inputs.dim() == 3:
+        bsz, tgt_len, _ = inputs.size()
+    else:
+        bsz, tgt_len = inputs_length.shape[0], torch.max(inputs_length)
+    sequence_mask = torch.arange(0, tgt_len).to(inputs.device)
+    sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view(bsz, tgt_len, 1)
+    return sequence_mask
+
+# Define RMSNorm layer for normalizing hidden states and stabilizing training process
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            hidden_states = hidden_states.to(self.weight.dtype)
+        return self.weight * hidden_states
+
+# Modified variable-length attention mechanism, supporting FP32 with unified interface
+class VarLenAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, causal=False, dropout=0.0):
+        """
+        Initialize variable-length attention module.
+
+        Parameters:
+            embed_dim (int): Embedding dimension (model's hidden dimension)
+            num_heads (int): Number of attention heads
+            causal (bool): Whether to enable causal attention (only attend to current and previous positions)
+            dropout (float): Attention dropout probability
+        """
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.causal = causal
+        self.dropout = nn.Dropout(dropout)
+        self.scaling = self.head_dim ** -0.5  # Scaling factor
+
+        # Linear projection layers for Q, K, V and output
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+    def _create_attention_mask(self, seq_len, max_len, device, dtype):
+        """
+        Create attention mask supporting variable-length sequences and causality.
+
+        Parameters:
+            seq_len (torch.Tensor): Sequence length for each sample, shape [bsz]
+            max_len (int): Maximum sequence length in the batch
+            device: Device for tensor creation
+            dtype: Data type for mask values
+
+        Returns:
+            mask (torch.Tensor): Attention mask, shape [bsz, 1, max_len, max_len], invalid positions set to minimum value
+        """
+        bsz = seq_len.size(0)
+        # Initialize mask as 1 (valid positions)
+        mask = torch.ones(bsz, 1, max_len, max_len, device=device, dtype=dtype)
+
+        # Generate sequence indices
+        seq_indices = torch.arange(max_len, device=device).unsqueeze(0)  # [1, max_len]
+        seq_len_expanded = seq_len.unsqueeze(1)  # [bsz, 1]
+
+        # Mark valid positions (less than seq_len)
+        valid_mask = seq_indices < seq_len_expanded.unsqueeze(-1)  # [bsz, 1, max_len]
+        mask = mask * (valid_mask.unsqueeze(2) & valid_mask.unsqueeze(3)).to(dtype)  # [bsz, 1, max_len, max_len]
+
+        # If causal attention, add upper triangular mask
+        if self.causal:
+            causal_mask = torch.triu(torch.ones(max_len, max_len, device=device, dtype=torch.bool), diagonal=1)
+            mask = mask * (~causal_mask.unsqueeze(0).unsqueeze(1)).to(dtype)  # Keep only lower triangular part
+
+        # Set invalid positions (0) to dtype's minimum value
+        mask = mask + (1.0 - mask) * torch.finfo(dtype).min  # Valid positions unchanged, invalid positions to minimum value
+        return mask
+
+    def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor:
+        """
+        Forward propagation, input and output are [bsz, max_len, embed_dim].
+
+        Parameters:
+            hidden_states (torch.Tensor): Input hidden states, shape [bsz, max_len, embed_dim]
+            seq_len (torch.Tensor): Sequence length for each sample, shape [bsz]
+
+        Returns:
+            attn_output (torch.Tensor): Attention output, shape [bsz, max_len, embed_dim]
+        """
+        bsz, max_len, _ = hidden_states.size()
+
+        # Project to Q, K, V
+        query = self.q_proj(hidden_states) * self.scaling  # [bsz, max_len, embed_dim]
+        key = self.k_proj(hidden_states)                  # [bsz, max_len, embed_dim]
+        value = self.v_proj(hidden_states)                # [bsz, max_len, embed_dim]
+
+        # Reshape to multi-head form
+        query = query.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2)  # [bsz, num_heads, max_len, head_dim]
+        key = key.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2)      # [bsz, num_heads, max_len, head_dim]
+        value = value.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2)  # [bsz, num_heads, max_len, head_dim]
+
+        # Calculate attention scores
+        attn_scores = torch.matmul(query, key.transpose(-1, -2))  # [bsz, num_heads, max_len, max_len]
+
+        # Generate attention mask
+        attn_mask = self._create_attention_mask(seq_len, max_len, hidden_states.device, attn_scores.dtype)  # [bsz, 1, max_len, max_len]
+        # Apply mask (additive form, consistent with HubertEncoder)
+        attn_scores = attn_scores + attn_mask  # Invalid positions set to very small value
+
+        # Softmax calculate attention weights
+        attn_weights = F.softmax(attn_scores, dim=-1)  # [bsz, num_heads, max_len, max_len]
+        attn_weights = self.dropout(attn_weights)
+
+        # Calculate attention output
+        attn_output = torch.matmul(attn_weights, value)  # [bsz, num_heads, max_len, head_dim]
+        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, max_len, self.embed_dim)  # [bsz, max_len, embed_dim]
+
+        # Output projection
+        attn_output = self.out_proj(attn_output)  # [bsz, max_len, embed_dim]
+
+        return attn_output
+
+# Define Transformer layer containing attention mechanism and feedforward network for feature extraction and transformation
+class OmniWhisperTransformerLayer(nn.Module):
+    def __init__(self, activation_function="gelu", d_model=1280, attention_heads=20, ffn_dim=5120, causal=False, ln_type="LayerNorm", attn_type="varlen"):
+        super().__init__()
+        self.embed_dim = d_model
+        # Only keep varlen attention mechanism
+        if attn_type != "varlen":
+            raise ValueError(f"Unknown attn_type: {attn_type}. Only 'varlen' is supported.")
+        self.self_attn = VarLenAttention(self.embed_dim, attention_heads, causal)
+        if ln_type == "LayerNorm":
+            self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        elif ln_type == "RMSNorm":
+            self.self_attn_layer_norm = RMSNorm(self.embed_dim)
+        else:
+            raise ValueError(f"Unknown ln_type: {ln_type}")
+        self.activation_fn = ACT2FN[activation_function]
+        self.fc1 = nn.Linear(self.embed_dim, ffn_dim)
+        self.fc2 = nn.Linear(ffn_dim, self.embed_dim)
+        if ln_type == "LayerNorm":
+            self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+        elif ln_type == "RMSNorm":
+            self.final_layer_norm = RMSNorm(self.embed_dim)
+        else:
+            raise ValueError(f"Unknown ln_type: {ln_type}")
+
+    def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states  # [bsz, max_len, embed_dim]
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        # from torch.cuda.amp import autocast
+        # print(f"{residual.dtype = }")
+        # print(f"Autocast enabled: {torch.is_autocast_enabled():}")
+        # print(f"after layernorm {hidden_states.dtype = }")
+        hidden_states = self.self_attn(hidden_states, seq_len)  # [bsz, max_len, embed_dim]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = residual + hidden_states
+        if (hidden_states.dtype == torch.float16 or hidden_states.dtype == torch.bfloat16) and \
+           (torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()):
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+        return hidden_states
+
+# Define audio encoder to convert input audio features to hidden state representation
+class OmniAudioEncoder(nn.Module):
+    def __init__(
+            self, 
+            num_mel_bins=128,  # Input feature Mel band number, usually the dimension of Mel spectrogram
+            sampling_rate=16000,  # Audio sampling rate, unit Hz
+            hop_length=160,  # Frame shift length (sample number) when calculating Mel spectrogram
+            stride_size=2,  # Convolution layer step, used for downsampling
+            kernel_size=3,  # Convolution kernel size, controlling receptive field
+            d_model=1280,  # Model's hidden state dimension (embedding dimension)
+            scale_embedding=True,  # Whether to scale embedding (usually used for stabilizing training)
+            max_audio_seconds=30,  # Maximum audio duration supported (seconds)
+            encoder_layers=32,  # Transformer encoder layer number
+            encoder_attention_heads=20,  # Attention head number for each Transformer layer
+            encoder_ffn_dim=5120,  # Intermediate dimension for feedforward network
+            activation_function="gelu",  # Activation function type, default GELU
+            attn_type="varlen"  # New parameter, select attention mechanism type
+        ):
+        super().__init__()
+        # Calculate maximum sequence length: Convert sampling rate to frame number after considering downsampling step
+        self.max_source_positions = (max_audio_seconds * sampling_rate // hop_length) // stride_size
+        # Embedding scaling factor, if enabled sqrt(d_model), otherwise 1.0
+        self.embed_scale = math.sqrt(d_model) if scale_embedding else 1.0
+        self.num_mel_bins = num_mel_bins  # Save Mel band number
+        self.d_model = d_model  # Save hidden state dimension
+        self.stride_size = stride_size
+        
+        # First convolution layer: Convert Mel spectrogram features (num_mel_bins) to hidden dimension (d_model)
+        self.conv1 = nn.Conv1d(num_mel_bins, d_model, kernel_size=kernel_size, padding=1)
+        # Second convolution layer: Apply downsampling with stride_size
+        self.conv2 = nn.Conv1d(d_model, d_model, kernel_size=kernel_size, stride=stride_size, padding=1)
+        
+        # Register positional embedding buffer, using sine function to generate, shape (max_source_positions, d_model)
+        self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model))
+        
+        # Create Transformer encoder layer list, each layer contains attention mechanism and feedforward network
+        self.layers = nn.ModuleList([
+            OmniWhisperTransformerLayer(
+                activation_function=activation_function, 
+                d_model=d_model, 
+                attention_heads=encoder_attention_heads, 
+                ffn_dim=encoder_ffn_dim, 
+                causal=False,  # Encoder does not need causal attention
+                attn_type=attn_type  # Pass attention type
+            ) for _ in range(encoder_layers)
+        ])
+        
+        # Last layer normalization for stable output
+        self.layer_norm = nn.LayerNorm(d_model)
+
+    def forward(self, input_features, input_length, output_hidden_states=False):
+        """
+        Forward propagation function to convert input audio features to hidden state representation
+        
+        Parameters:
+            input_features (torch.Tensor): Input Mel spectrogram features, shape [bsz, num_mel_bins, seq_len]
+            input_length (torch.Tensor): Input sequence length for each sample, shape [bsz]
+            output_hidden_states (bool, optional): Whether to return hidden states for each layer, default False
+        
+        Returns:
+            if output_hidden_states is False:
+                hidden_states (torch.Tensor): Encoded hidden states, shape [bsz, d_model, tgt_len]
+                output_length (torch.Tensor): Output sequence length for each sample, shape [bsz]
+            else:
+                hidden_states (torch.Tensor): Encoded hidden states, shape [bsz, d_model, tgt_len]
+                output_length (torch.Tensor): Output sequence length for each sample, shape [bsz]
+                hidden_states_all_layers (tuple): Tuple containing hidden states for each layer, including initial input
+        """
+        # Ensure input feature data type consistent with convolution layer weights
+        input_features = input_features.to(self.conv1.weight.dtype)  # (B, D, T)
+        
+        # First layer convolution + GELU activation, Convert Mel spectrogram to hidden states
+        inputs_embeds = nn.functional.gelu(self.conv1(input_features))  # (B, D, T)
+        
+        # Second layer convolution + GELU activation, Apply downsampling with stride_size
+        inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))  # (B, D, T)
+        
+        # Calculate output length: Result after downsampling with stride_size
+        output_length = (input_length // self.stride_size).long()  # (B,)
+        
+        # Adjust dimension order to [bsz, seq_len, d_model] for Transformer input
+        hidden_states = inputs_embeds.permute(0, 2, 1)  # (B, T, D)
+        
+        # Get batch size and target sequence length
+        bsz, tgt_len, _ = hidden_states.size()
+        
+        # According to current sequence length, take or use complete positional embedding
+        if tgt_len < self.positional_embedding.shape[0]:
+            current_positional_embedding = self.positional_embedding[:tgt_len]
+        else:
+            current_positional_embedding = self.positional_embedding
+        
+        # Add input embedding to positional embedding, convert to float to avoid precision issues
+        hidden_states = (hidden_states.to(torch.float32) + current_positional_embedding).to(hidden_states.dtype)
+        
+        # Generate sequence mask for processing variable-length sequence
+        attention_mask = get_sequence_mask(hidden_states, output_length)  # [bsz, tgt_len, 1]
+        
+        # Initialize hidden states list for storing output for each layer (if needed)
+        hidden_states_all_layers = () if output_hidden_states else None
+        
+        # Process hidden states through Transformer encoder layer by layer
+        for encoder_layer in self.layers:
+            if output_hidden_states:
+                hidden_states_all_layers = hidden_states_all_layers + (hidden_states,)
+            hidden_states = encoder_layer(hidden_states, output_length)  # [bsz, tgt_len, d_model]
+        
+        # Normalize hidden states
+        hidden_states = self.layer_norm(hidden_states)  # [bsz, tgt_len, d_model]
+        if output_hidden_states:
+            hidden_states_all_layers = hidden_states_all_layers + (hidden_states,)
+        
+        # Use mask to zero out padding parts and ensure output only retains valid data
+        hidden_states = torch.where(attention_mask, hidden_states, 0)  # [bsz, tgt_len, d_model]
+        hidden_states = hidden_states.transpose(1, 2)  # [bsz, d_model, tgt_len]
+        
+        if not output_hidden_states:
+            return hidden_states, output_length  
+        else:
+            return hidden_states, output_length, hidden_states_all_layers
+    
+# Define audio decoder to convert hidden states to Mel spectrogram
+class OmniAudioDecoder(nn.Module):
+    def __init__(
+            self, 
+            num_mel_bins=128, 
+            sampling_rate=16000, 
+            hop_length=160, 
+            stride_size=2, 
+            kernel_size=3, 
+            d_model=1280, 
+            scale_embedding=True, 
+            max_audio_seconds=30, 
+            decoder_layers=32, 
+            decoder_attention_heads=20, 
+            decoder_ffn_dim=5120, 
+            activation_function="gelu",
+            attn_type="varlen"  # New parameter, select attention mechanism type
+        ):
+        super().__init__()
+        self.max_source_positions = (max_audio_seconds * sampling_rate // hop_length) // stride_size
+        self.embed_scale = math.sqrt(d_model) if scale_embedding else 1.0
+        self.num_mel_bins = num_mel_bins
+        self.d_model = d_model
+        self.stride_size = stride_size
+        
+        # Correct transpose convolution layer to ensure output length close to stride_size times
+        self.deconv1 = nn.ConvTranspose1d(
+            d_model, 
+            d_model, 
+            kernel_size=kernel_size, 
+            stride=stride_size, 
+            padding=0,  # Do not fill input side
+            output_padding=0  # Can be adjusted to precisely control length
+        )
+        self.deconv2 = nn.ConvTranspose1d(
+            d_model, 
+            num_mel_bins, 
+            kernel_size=kernel_size, 
+            stride=1,  # Only convert channels, do not change length
+            padding=0
+        )
+        
+        # Positional embedding remains consistent
+        self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model)) # (T, D)
+        
+        # Transformer decoder layer
+        self.layers = nn.ModuleList([
+            OmniWhisperTransformerLayer(
+                activation_function=activation_function, 
+                d_model=d_model, 
+                attention_heads=decoder_attention_heads, 
+                ffn_dim=decoder_ffn_dim, 
+                causal=False,  # Decoder uses causal attention
+                attn_type=attn_type  # Pass attention type
+            ) for _ in range(decoder_layers)
+        ])
+        self.layer_norm = nn.LayerNorm(d_model)
+
+    def forward(self, hidden_states, input_length): # (B, D, T)
+        # Input is hidden state output from encoder
+        hidden_states = hidden_states.transpose(1, 2) # (B, T, D)
+        bsz, tgt_len, _ = hidden_states.size()
+        
+        # Add positional embedding
+        if tgt_len < self.positional_embedding.shape[0]:
+            current_positional_embedding = self.positional_embedding[:tgt_len] # (T, D)
+        else:
+            current_positional_embedding = self.positional_embedding
+        hidden_states = (hidden_states.to(torch.float32) + current_positional_embedding).to(hidden_states.dtype) # (B, T, D)
+        
+        # Generate sequence mask
+        attention_mask = get_sequence_mask(hidden_states, input_length)  # [bsz, tgt_len, 1]
+        
+        # Process through decoder layer
+        for decoder_layer in self.layers:
+            hidden_states = decoder_layer(hidden_states, input_length)  # [bsz, tgt_len, d_model]
+        
+        # Final layer normalization
+        hidden_states = self.layer_norm(hidden_states)  # [bsz, tgt_len, d_model]
+        
+        # Use mask to zero out padding parts
+        hidden_states = torch.where(attention_mask, hidden_states, 0)  # [bsz, tgt_len, d_model]
+        
+        # Process through transpose convolution layer to reconstruct audio features
+        hidden_states = hidden_states.permute(0, 2, 1)  # (B, D, T)
+        output_features = nn.functional.gelu(self.deconv1(hidden_states)) # (B, D, T)
+        output_features = nn.functional.gelu(self.deconv2(output_features)) # (B, D, T)
+        
+        # If strictly stride_size times length is needed, can trim extra parts
+        expected_length = tgt_len * self.stride_size
+        if output_features.size(2) > expected_length:
+            output_features = output_features[:, :, :expected_length]
+        
+        output_length = input_length * self.stride_size
+        # Output shape: [bsz, num_mel_bins, seq_len]
+        return output_features, output_length
+
+# The following part remains unchanged
+class ResidualDownConv(nn.Module):
+    def __init__(self, d_model=1280, avg_pooler=4):
+        """
+        Downsampling module containing residual connection and convolution operation
+        
+        Parameters:
+            d_model (int): Input and output hidden dimension
+            avg_pooler (int): Downsampling factor (convolution step)
+        """
+        super().__init__()
+        self.d_model = d_model
+        self.avg_pooler = avg_pooler
+        self.intermediate_dim = d_model * avg_pooler
+        
+        # Convolution layer for downsampling
+        self.gate_proj = nn.Conv1d(d_model, self.intermediate_dim, avg_pooler, avg_pooler, bias=False)
+        self.up_proj = nn.Conv1d(d_model, self.intermediate_dim, avg_pooler, avg_pooler, bias=False)
+        
+        # Downsampled linear projection
+        self.down_proj = nn.Linear(self.intermediate_dim, self.intermediate_dim, bias=False)
+        
+        # Activation function and layer normalization
+        self.act_fn = ACT2FN['silu']
+        self.layer_norm = nn.LayerNorm(self.intermediate_dim)
+
+    def forward(self, x, input_length):
+        """
+        Forward propagation, execute downsampling and residual processing
+        
+        Parameters:
+            x (torch.Tensor): Input tensor, shape [B, D, T]
+        
+        Returns:
+            res (torch.Tensor): Downsampled feature, shape [B, intermediate_dim, seq_len // avg_pooler]
+            valid_mask (torch.Tensor): Valid sequence mask
+        """
+        output_length = input_length // self.avg_pooler
+        x = x.transpose(1, 2) # (B, T, D)
+        batch_size, seq_len, _ = x.shape # (B, T, D)
+        if seq_len % self.avg_pooler != 0:
+            pad_size = self.avg_pooler - seq_len % self.avg_pooler
+            x = F.pad(x, (0, pad_size), "constant", 0)
+        
+        xt = x.permute(0, 2, 1) # (B, D, T)
+        g = self.gate_proj(xt).permute(0, 2, 1)  # (B, T, D)
+        u = self.up_proj(xt).permute(0, 2, 1) # (B, T, D)
+        x = x.reshape(batch_size, -1, self.intermediate_dim)  # (B, T, D)
+
+        c = self.down_proj(self.act_fn(g) * u) # (B, T, D)
+        res = self.layer_norm(c + x) # (B, T, D)
+        res = res.transpose(1, 2) # (B, D, T)
+        return res, output_length # (B, D, T)
+    
+    
+class UpConv(nn.Module):
+    def __init__(self, d_model=1280, stride=4):
+        """
+        Simple upsampling module using transpose convolution
+        
+        Parameters:
+            d_model (int): Input and output hidden dimension
+            stride (int): Upsampling factor (transpose convolution step)
+        """
+        super().__init__()
+        self.d_model = d_model
+        self.stride = stride
+        
+        # Simple transpose convolution layer to keep channel number consistent
+        self.up_conv = nn.ConvTranspose1d(
+            self.stride * d_model, 
+            d_model, 
+            kernel_size=stride, 
+            stride=stride, 
+            bias=False
+        )
+
+    def forward(self, x, input_length):
+        """
+        Forward propagation, execute upsampling
+        
+        Parameters:
+            x (torch.Tensor): Input tensor, shape [B, D * stride, T]
+        
+        Returns:
+            res (torch.Tensor): Upsampled feature, shape [B, D, T * stride]
+        """
+        # Directly apply transpose convolution
+        res = self.up_conv(x)
+        output_length = input_length * self.stride
+        return res, output_length
+    
+
+# Define Transformer encoder containing multiple Transformer layers for feature extraction and transformation
+class Transformer(nn.Module):
+    def __init__(
+            self, 
+            input_dim=1280,  # Input feature dimension
+            d_model=1280,  # Model's hidden state dimension (embedding dimension)
+            output_dim=1280,  # Output feature dimension
+            max_source_positions=1500,  # Maximum sequence length for positional embedding
+            encoder_layers=32,  # Transformer encoder layer number
+            encoder_attention_heads=20,  # Attention head number for each Transformer layer
+            encoder_ffn_dim=5120,  # Intermediate dimension for feedforward network
+            activation_function="gelu",  # Activation function type, default GELU
+            attn_type="varlen"  # Attention mechanism type
+    ):
+        super().__init__()
+        self.input_dim = input_dim  # Save input dimension
+        self.d_model = d_model  # Save hidden state dimension
+        self.output_dim = output_dim  # Save output dimension
+        self.max_source_positions = max_source_positions  # Save maximum sequence length
+
+        # If input dimension and model dimension are not consistent, add input projection layer
+        if input_dim != d_model:
+            self.proj = nn.Linear(input_dim, d_model, bias=True)
+        else:
+            self.proj = None  # No need for input projection layer
+
+        # Register positional embedding buffer, using sine function to generate, shape (max_source_positions, d_model)
+        self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model))
+        
+        # Create Transformer encoder layer list, each layer contains attention mechanism and feedforward network
+        self.layers = nn.ModuleList([
+            OmniWhisperTransformerLayer(
+                activation_function=activation_function, 
+                d_model=d_model, 
+                attention_heads=encoder_attention_heads, 
+                ffn_dim=encoder_ffn_dim, 
+                causal=False,  # Encoder does not need causal attention
+                attn_type=attn_type  # Pass attention type
+            ) for _ in range(encoder_layers)
+        ])
+        
+        # Last layer normalization for stable output
+        self.layer_norm = nn.LayerNorm(d_model)
+
+        # If output dimension and model dimension are not consistent, add output projection layer
+        if output_dim != d_model:
+            self.out_proj = nn.Linear(d_model, output_dim, bias=True)
+        else:
+            self.out_proj = None  # No need for output projection layer
+
+    def forward(self, input_features: torch.Tensor, input_length: torch.Tensor, output_hidden_states: bool = False):
+        """
+        Forward propagation function to convert input features through Transformer layer to hidden state representation
+        
+        Parameters:
+            input_features (torch.Tensor): Input features, shape [bsz, input_dim, seq_len] (B, input_dim, T)
+            input_length (torch.Tensor): Input sequence length for each sample, shape [bsz]
+            output_hidden_states (bool, optional): Whether to return hidden states for each layer, default False
+        
+        Returns:
+            if output_hidden_states is False:
+                hidden_states (torch.Tensor): Encoded hidden states, shape [bsz, output_dim, seq_len] (B, output_dim, T)
+                output_length (torch.Tensor): Output sequence length for each sample, shape [bsz]
+            else:
+                hidden_states (torch.Tensor): Encoded hidden states, shape [bsz, output_dim, seq_len] (B, output_dim, T)
+                output_length (torch.Tensor): Output sequence length for each sample, shape [bsz]
+                hidden_states_all_layers (tuple): Tuple containing hidden states for each layer, each shape [bsz, seq_len, d_model]
+        """
+        # Output length is the same as input length, Transformer does not change sequence length
+        output_length = input_length.long()  # [bsz]
+
+        # If there is input projection layer, map input features from input_dim to d_model
+        if self.proj is not None:
+            hidden_states = self.proj(input_features.permute(0, 2, 1)).permute(0, 2, 1)  # [bsz, d_model, seq_len] (B, D, T)
+        else:
+            hidden_states = input_features  # [bsz, d_model, seq_len] (B, D, T)
+
+        # Adjust input dimension order to [bsz, seq_len, d_model] for Transformer input
+        hidden_states = hidden_states.permute(0, 2, 1)  # [bsz, seq_len, d_model] (B, T, D)
+        
+        # Get batch size and target sequence length
+        bsz, tgt_len, _ = hidden_states.size()
+        
+        # According to current sequence length, take or use complete positional embedding
+        if tgt_len < self.positional_embedding.shape[0]:
+            current_positional_embedding = self.positional_embedding[:tgt_len]  # [tgt_len, d_model]
+        else:
+            current_positional_embedding = self.positional_embedding  # [max_source_positions, d_model]
+        
+        # Add input features to positional embedding, convert to float to avoid precision issues
+        hidden_states = (hidden_states.to(torch.float32) + current_positional_embedding).to(hidden_states.dtype)  # [bsz, seq_len, d_model]
+        
+        # Generate sequence mask for processing variable-length sequence
+        attention_mask = get_sequence_mask(hidden_states, output_length)  # [bsz, tgt_len, 1]
+        
+        # Initialize hidden states list for storing output for each layer (if needed)
+        hidden_states_all_layers = () if output_hidden_states else None
+        
+        # Process hidden states through Transformer encoder layer by layer
+        for encoder_layer in self.layers:
+            if output_hidden_states:
+                hidden_states_all_layers = hidden_states_all_layers + (hidden_states,)
+            hidden_states = encoder_layer(hidden_states, output_length)  # [bsz, seq_len, d_model]
+        
+        # Normalize hidden states
+        hidden_states = self.layer_norm(hidden_states)  # [bsz, seq_len, d_model]
+        if output_hidden_states:
+            hidden_states_all_layers = hidden_states_all_layers + (hidden_states,)
+        
+        # Use mask to zero out padding parts and ensure output only retains valid data
+        hidden_states = torch.where(attention_mask, hidden_states, 0)  # [bsz, seq_len, d_model]
+        
+        # Adjust dimension order to [bsz, d_model, seq_len]
+        hidden_states = hidden_states.transpose(1, 2)  # [bsz, d_model, seq_len] (B, D, T)
+
+        # If there is output projection layer, map hidden states from d_model to output_dim
+        if self.out_proj is not None:
+            hidden_states = self.out_proj(hidden_states.permute(0, 2, 1)).permute(0, 2, 1)  # [bsz, output_dim, seq_len] (B, output_dim, T)
+
+        if not output_hidden_states:
+            return hidden_states, output_length
+        else:
+            return hidden_states, output_length, hidden_states_all_layers
+        
+
+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
+        original_dtype = x.dtype
+        S = mag.float() * (x.float() + 1j * y.float())
+        audio = self.istft(S)
+        audio = audio.to(original_dtype)
+        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 = 128,
+        dim: int = 512,
+        intermediate_dim: int = 4096,
+        num_layers: int = 30,
+        n_fft: int = 640,
+        hop_size: int = 160,
+        padding: str = "same",
+        adanorm_num_embeddings=None,
+    ):
+        super().__init__()
+
+        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)
+        self.hop_size = hop_size
+
+    def forward(self, x, input_length):
+        x = self.backbone(x)
+        x = self.head(x)
+        output_length = input_length * self.hop_size
+        return x[:, None, :], output_length
+

XY_Tokenizer/xy_tokenizer/nn/quantizer.py

@@ -0,0 +1,370 @@
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+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))
+
+def sample_vectors(samples, num):
+    # samples: (N, D), num_samples: N, feature dim: D
+    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].float()  # (num, D), ensure fp32
+
+def kmeans(samples, num_clusters, num_iters=10):
+    # samples: (N, D), N samples with D dimensions
+    dim, dtype = samples.shape[-1], torch.float32  # Force fp32
+    means = sample_vectors(samples, num_clusters).float()  # (num_clusters, D), ensure fp32
+    
+    for _ in range(num_iters):
+        dists = -(samples.float().pow(2).sum(1, keepdim=True) -  # (N, 1), ensure fp32
+                 2 * samples.float() @ means.t() +               # (N, num_clusters), ensure fp32
+                 means.t().float().pow(2).sum(0, keepdim=True))  # (1, num_clusters), ensure fp32
+        # dists: (N, num_clusters)
+        buckets = dists.max(dim=-1).indices  # (N)
+        bins = torch.bincount(buckets, minlength=num_clusters)  # (num_clusters)
+        zero_mask = bins == 0  # (num_clusters)
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)  # (num_clusters)
+        
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=torch.float32)  # (num_clusters, D), ensure fp32
+        new_means.scatter_add_(0, buckets.unsqueeze(1).expand(-1, dim), samples.float())  # (num_clusters, D), ensure fp32
+        new_means = new_means / bins_min_clamped[..., None]  # (num_clusters, D)
+        means = torch.where(zero_mask[..., None], means, new_means)  # (num_clusters, D)
+    
+    # Final cluster assignments for returning cluster sizes
+    dists = -(samples.float().pow(2).sum(1, keepdim=True) - 
+             2 * samples.float() @ means.t() + 
+             means.t().float().pow(2).sum(0, keepdim=True))  # (N, num_clusters), ensure fp32
+    buckets = dists.max(dim=-1).indices  # (N)
+    bins = torch.bincount(buckets, minlength=num_clusters).float()  # (num_clusters), ensure fp32
+    
+    return means, bins  # (num_clusters, D), (num_clusters)
+
+class VectorQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+        commitment=1.0,
+        decay=0.99,           # EMA decay
+        epsilon=1e-5,         # Laplace smoothing epsilon
+        threshold_ema_dead=2, # Dead code threshold
+        kmeans_init=True,     # Use kmeans initialization
+        kmeans_iters=10,      # Kmeans iterations
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.commitment = commitment
+        self.decay = decay
+        self.epsilon = epsilon
+        self.threshold_ema_dead = threshold_ema_dead
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        
+        if self.input_dim != self.codebook_dim:
+            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)  # (B, D, T) -> (B, D', T)
+            self.out_project = WNConv1d(self.codebook_dim, self.input_dim, kernel_size=1)  # (B, D', T) -> (B, D, T)
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+
+        # Initialize codebook and EMA buffers
+        init_fn = torch.zeros if kmeans_init else lambda x, y: torch.randn(x, y)
+        self.register_buffer("codebook", init_fn(codebook_size, codebook_dim).float())  # (codebook_size, D'), ensure fp32
+        self.register_buffer("inited", torch.tensor([not kmeans_init], dtype=torch.bool))  # (1)
+        self.register_buffer("cluster_size", torch.zeros(codebook_size).float())  # (codebook_size), ensure fp32
+        self.register_buffer("embed_avg", self.codebook.clone().float())  # (codebook_size, D'), ensure fp32
+
+    def ema_update(self, encodings, embed_onehot):
+        # encodings: (B*T, D'), embed_onehot: (B*T, codebook_size)
+        """Update codebook using EMA"""
+        encodings = encodings.float()  # Ensure fp32
+        embed_onehot = embed_onehot.float()  # Ensure fp32
+        cluster_size_new = embed_onehot.sum(0)  # (codebook_size)
+        embed_sum = encodings.t() @ embed_onehot  # (D', codebook_size)
+        
+        # Distributed reduction
+        if dist.is_initialized():
+            dist.all_reduce(cluster_size_new, op=dist.ReduceOp.SUM)
+            dist.all_reduce(embed_sum, op=dist.ReduceOp.SUM)
+            
+        ema_inplace(self.cluster_size, cluster_size_new, self.decay)  # (codebook_size)
+        ema_inplace(self.embed_avg, embed_sum.t(), self.decay)  # (codebook_size, D')
+        
+        # Laplace smoothing
+        cluster_size = (self.cluster_size + self.epsilon) / (self.cluster_size.sum() + self.codebook_size * self.epsilon)  # (codebook_size)
+        cluster_size = cluster_size * self.cluster_size.sum()  # (codebook_size)
+        self.codebook.copy_(self.embed_avg / cluster_size.unsqueeze(1))  # (codebook_size, D')
+
+    def replace_dead_codes(self, encodings):
+        # encodings: (B*T, D')
+        """Replace dead codes with random samples from current batch"""
+        if self.threshold_ema_dead == 0:
+            return
+        
+        dead_mask = self.cluster_size < self.threshold_ema_dead  # (codebook_size)
+        if dead_mask.any():
+            if dist.is_initialized() and dist.get_rank() == 0:
+                samples = sample_vectors(encodings.float(), self.codebook_size)  # (codebook_size, D'), ensure fp32
+            else:
+                samples = torch.zeros_like(self.codebook).float()  # Placeholder, ensure fp32
+            
+            # Broadcast samples
+            if dist.is_initialized():
+                dist.broadcast(samples, src=0)
+            
+            self.codebook[dead_mask] = samples[:dead_mask.sum()].to(self.codebook.dtype)  # Update dead codes
+
+    def init_codebook(self, encodings):
+        # encodings: (B*T, D')
+        """Initialize codebook with k-means and update cluster_size"""
+        if self.inited.item():
+            return
+        
+        if dist.is_initialized() and dist.get_rank() == 0:
+            embed, cluster_sizes = kmeans(encodings.float(), self.codebook_size, self.kmeans_iters)  # (codebook_size, D'), (codebook_size), ensure fp32
+        else:
+            embed = torch.zeros(self.codebook_size, self.codebook_dim, device=encodings.device).float()  # ensure fp32
+            cluster_sizes = torch.zeros(self.codebook_size, device=encodings.device, dtype=torch.float32)  # ensure fp32
+        
+        # Broadcast results
+        if dist.is_initialized():
+            dist.broadcast(embed, src=0)
+            dist.broadcast(cluster_sizes, src=0)
+            
+        self.codebook.copy_(embed)  # (codebook_size, D')
+        self.embed_avg.copy_(embed.clone())  # (codebook_size, D')
+        self.cluster_size.copy_(cluster_sizes.float())  # (codebook_size)
+        self.inited.fill_(True)
+
+    def forward(self, z):  # z: (B, D, T)
+        # logging.info(f"{self.cluster_size = }, {self.codebook = }, {self.embed_avg = }, {self.inited = }")
+        z = z.float()  # Ensure fp32
+        z_e = self.in_project(z).float()  # (B, D', T), ensure fp32
+        
+        # Rearrange for quantization
+        encodings = rearrange(z_e, "b d t -> (b t) d").float()  # (B*T, D'), ensure fp32
+        
+        # Initialize codebook if needed
+        if self.kmeans_init and not self.inited.item():
+            self.init_codebook(encodings)
+
+        # Quantization
+        dist = (encodings.pow(2).sum(1, keepdim=True) -  # (B*T, 1)
+            2 * encodings @ self.codebook.float().t() +       # (B*T, codebook_size)
+            self.codebook.float().pow(2).sum(1, keepdim=True).t())  # (1, codebook_size)
+        # dist: (B*T, codebook_size)
+        
+        indices = (-dist).max(1)[1]  # (B*T)
+        indices = rearrange(indices, "(b t) -> b t", b=z.size(0))  # (B, T)
+        
+        # Get quantized vectors
+        z_q = self.decode_code(indices).float()  # (B, D', T), ensure fp32
+        
+        # Commitment loss
+        commit_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) * self.commitment  # (B)
+        
+        # EMA updates and dead code replacement during training
+        if self.training and torch.is_grad_enabled():
+            embed_onehot = F.one_hot(indices.view(-1), self.codebook_size).float()  # (B*T, codebook_size), ensure fp32
+            self.ema_update(encodings, embed_onehot)
+            self.replace_dead_codes(encodings)
+
+        # Straight-through estimator
+        z_q = z_e + (z_q - z_e).detach()  # (B, D', T)
+        z_q = self.out_project(z_q).float()  # (B, D, T), ensure fp32
+
+        return z_q, commit_loss, torch.tensor(0.0, device=z.device, dtype=torch.float32), indices, z  # (B, D, T), (B), scalar, (B, T), (B, D', T)
+
+    def decode_code(self, embed_id):  # embed_id: (B, T)
+        return F.embedding(embed_id, self.codebook).transpose(1, 2).float()  # (B, D', T), ensure fp32
+
+class ResidualVQ(nn.Module):
+    def __init__(
+        self,
+        input_dim: int = 1280, # Input dimension, unrelated to RVQ
+        rvq_dim = None, # RVQ dimension. If different from input_dim/output_dim, will add input_dim->rvq_dim/rvq_dim->output_dim projection
+        output_dim: int = None, # Output dimension, unrelated to RVQ
+        num_quantizers: int = 32,
+        codebook_size: int = 1024,
+        codebook_dim: int = 8, # Dimension of each codebook. If different from rvq_dim, will add rvq_dim->codebook_dim and codebook_dim->rvq_dim projections
+        quantizer_dropout: float = 0.5,
+        decay=0.99,
+        epsilon=1e-5,
+        threshold_ema_dead=2,
+        kmeans_init=True,
+        kmeans_iters=10,
+        skip_rvq_ratio: float = 0.0,  # New parameter: probability of skipping RVQ
+        **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_dropout = quantizer_dropout
+        self.skip_rvq_ratio = skip_rvq_ratio  # Store skip probability
+        self.rvq_dim = rvq_dim
+        
+        self.input_proj = WNConv1d(input_dim, rvq_dim, kernel_size=1) if input_dim != rvq_dim else nn.Identity()
+        self.output_proj = WNConv1d(rvq_dim, output_dim, kernel_size=1) if rvq_dim != output_dim else nn.Identity()
+
+        self.quantizers = nn.ModuleList(
+            [
+                VectorQuantize(
+                    input_dim=rvq_dim,
+                    codebook_size=codebook_size,
+                    codebook_dim=codebook_dim,
+                    decay=decay,
+                    epsilon=epsilon,
+                    threshold_ema_dead=threshold_ema_dead,
+                    kmeans_init=kmeans_init,
+                    kmeans_iters=kmeans_iters,
+                    **kwargs,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+
+    def forward(self, z, input_length, n_quantizers: int = None):  # z: (B, D, T), input_length: (B)
+        z = self.input_proj(z)
+
+        with torch.autocast('cuda', enabled = False):    
+            batch_size, _, max_time = z.shape
+            mask = torch.arange(max_time, device=z.device).expand(batch_size, max_time) < input_length.unsqueeze(1)  # (B, T)
+
+            quantized_out = torch.zeros_like(z, dtype=torch.float32)  # (B, D, T), ensure fp32
+            residual = z.clone().float()  # (B, D, T), ensure fp32
+
+            all_commit_losses = []
+            all_indices = []
+            all_quantized = []
+
+            n_quantizers = n_quantizers or self.num_quantizers
+
+            # Randomly decide whether to skip RVQ during training
+            skip_mask = None
+            if self.training and torch.is_grad_enabled() and self.skip_rvq_ratio > 0:
+                # Generate random mask with skip_rvq_ratio probability
+                skip_mask = torch.rand(batch_size, device=z.device) < self.skip_rvq_ratio  # (B,)
+                # If all samples are skipped, force the first sample to be unskipped
+                if skip_mask.all():
+                    skip_mask[0] = False  # Ensure at least one sample (index 0) is not skipped
+
+            if self.training and torch.is_grad_enabled():
+                n_quantizers_tensor = torch.ones((z.shape[0],), dtype=torch.float32, device=z.device) * self.num_quantizers + 1  # (B)
+                dropout = torch.randint(1, self.num_quantizers + 1, (z.shape[0],), dtype=torch.float32, device=z.device)  # (B)
+                n_dropout = int(z.shape[0] * self.quantizer_dropout)
+                n_quantizers_tensor[:n_dropout] = dropout[:n_dropout]  # (B)
+            else:
+                n_quantizers_tensor = torch.full((z.shape[0],), n_quantizers, dtype=torch.float32, device=z.device)  # (B)
+
+            for i, quantizer in enumerate(self.quantizers):
+                if not self.training and i >= n_quantizers:
+                    break
+
+                masked_residual = residual * mask.unsqueeze(1)  # (B, D, T)
+                
+                # If skipping RVQ, directly use input value
+                if self.training and skip_mask is not None and skip_mask.any():
+                    z_q_i = torch.zeros_like(masked_residual, dtype=torch.float32)  # (B, D, T), ensure fp32
+                    commit_loss_i = torch.zeros(batch_size, device=z.device, dtype=torch.float32)  # (B), ensure fp32
+                    indices_i = torch.zeros(batch_size, max_time, device=z.device, dtype=torch.long)  # (B, T)
+                    z_e_i = torch.zeros_like(masked_residual, dtype=torch.float32)  # (B, D, T), ensure fp32
+                    
+                    # Quantize non-skipped samples
+                    non_skipped_mask = ~skip_mask  # (B)
+                    if non_skipped_mask.any():
+                        z_q_i_non_skipped, commit_loss_i_non_skipped, _, indices_i_non_skipped, z_e_i_non_skipped = quantizer(
+                            masked_residual[non_skipped_mask].float()  # Ensure fp32
+                        )
+                        z_q_i[non_skipped_mask] = z_q_i_non_skipped
+                        commit_loss_i[non_skipped_mask] = commit_loss_i_non_skipped
+                        indices_i[non_skipped_mask] = indices_i_non_skipped
+                        z_e_i[non_skipped_mask] = z_e_i_non_skipped
+                else:
+                    z_q_i, commit_loss_i, _, indices_i, z_e_i = quantizer(masked_residual.float())  # (B, D, T), (B), scalar, (B, T), (B, D', T), ensure fp32
+
+                quantizer_mask = (torch.full((z.shape[0],), i, device=z.device, dtype=torch.float32) < n_quantizers_tensor)  # (B)
+                update_mask = (mask & quantizer_mask.unsqueeze(-1)).unsqueeze(1)  # (B, 1, T)
+                
+                # If skipping, output is directly the input
+                if skip_mask is not None:
+                    skip_mask_expanded = skip_mask.unsqueeze(1).unsqueeze(2)  # (B, 1, 1)
+                    z_q_i = torch.where(skip_mask_expanded, masked_residual, z_q_i)  # (B, D, T)
+                    commit_loss_i = torch.where(skip_mask, torch.zeros_like(commit_loss_i), commit_loss_i)  # (B)
+
+                quantized_out = quantized_out + z_q_i * update_mask  # (B, D, T)
+
+                residual_fp32 = residual.to(dtype=torch.float32)  # (B, D, T)
+                z_q_i_fp32 = z_q_i.to(dtype=torch.float32)  # (B, D, T)
+                residual_fp32 = residual_fp32 - z_q_i_fp32 * update_mask  # (B, D, T)
+                residual = residual_fp32.to(dtype=torch.float32)  # (B, D, T), ensure fp32
+
+                valid_mask = mask & quantizer_mask.unsqueeze(-1)  # (B, T)
+                if valid_mask.any():
+                    commit_loss_i = (commit_loss_i * quantizer_mask).sum() / quantizer_mask.sum()  # scalar
+                else:
+                    commit_loss_i = torch.tensor(0.0, device=z.device, dtype=torch.float32)  # scalar, ensure fp32
+
+                all_commit_losses.append(commit_loss_i)  # scalar
+                all_indices.append(indices_i)  # (B, T)
+                all_quantized.append(z_q_i)  # (B, D, T)
+
+            all_commit_losses = torch.stack(all_commit_losses)  # (N)
+            all_indices = torch.stack(all_indices)  # (N, B, T)
+            all_quantized = torch.stack(all_quantized)  # (N, B, D, T)
+
+            output_length = input_length  # (B)
+
+        quantized_out = self.output_proj(quantized_out)
+
+        return (
+            quantized_out,    # (B, D, T)
+            all_indices,      # (N, B, T)
+            all_commit_losses,# (N)
+            all_quantized,    # (N, B, D, T)
+            output_length,    # (B)
+        )
+
+    def decode_codes(self, codes):  # codes: (nq, B, T)
+        """Decode codes from multiple quantizers to embeddings.
+        
+        Args:
+            codes: Tensor of shape (nq, B, T) containing code indices for each quantizer.
+
+        Returns:
+            emb: Tensor of shape (B, D, T) representing the decoded embeddings.
+        """
+        nq, B, T = codes.shape
+        device = codes.device
+        emb = torch.zeros(B, self.rvq_dim, T, device=device, dtype=torch.float32)  # (B, D, T)
+
+        for i, quantizer in enumerate(self.quantizers[:nq]):
+            code_i = codes[i]  # (B, T)
+            quantized_i = quantizer.decode_code(code_i)  # (B, D', T)
+            emb += quantized_i  # Accumulate quantized embeddings
+
+        emb = self.output_proj(emb)  # (B, D, T), apply output projection
+        return emb # (B, D, T)
+
+
+def ema_inplace(moving_avg, new, decay):
+    # moving_avg: (codebook_size) or (codebook_size, D'), new: same as moving_avg
+    """Update exponential moving average in-place"""
+    moving_avg.data.mul_(decay).add_(new.float(), alpha=(1 - decay))  # ensure fp32

app.py

@@ -1,7 +1,495 @@
 import gradio as gr
+import torch
+import torchaudio
+import tempfile
+import json
+import os
+from typing import Optional, Tuple
 
-def greet(name):
-    return "Hello " + name + "!!"
+from generation_utils import load_model, process_batch
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+def load_examples_from_jsonl():
+    """
+    Load examples from examples/examples.jsonl and convert to ROLE_EXAMPLES format
+    """
+    examples = []
+    jsonl_path = "examples/examples.jsonl"
+    
+    if not os.path.exists(jsonl_path):
+        print(f"Warning: {jsonl_path} not found")
+        return []
+
+    with open(jsonl_path, 'r', encoding='utf-8') as f:
+        for line in f:
+            line = line.strip()
+
+            data = json.loads(line)
+            
+            # Extract required fields
+            text = data.get('text', '')
+            base_path = data.get('base_path', 'examples')
+            
+            # Check if this is a role-based example (has speaker1 and speaker2 audio)
+            if 'prompt_audio_speaker1' in data and 'prompt_audio_speaker2' in data:
+                # Role mode example
+                audio_mode = "Role"
+                prompt_audio_1 = os.path.join(base_path, data['prompt_audio_speaker1'])
+                prompt_text_1 = data.get('prompt_text_speaker1', '')
+                prompt_audio_2 = os.path.join(base_path, data['prompt_audio_speaker2'])
+                prompt_text_2 = data.get('prompt_text_speaker2', '')
+                use_normalize = True
+                
+                example = [text, audio_mode, prompt_audio_1, prompt_text_1, prompt_audio_2, prompt_text_2, use_normalize]
+                examples.append(example)
+    
+    print(f"Loaded {len(examples)} examples from {jsonl_path}")
+    return examples
+
+# Load examples from JSONL file
+ROLE_EXAMPLES = load_examples_from_jsonl()
+
+# Language configuration
+LANGUAGES = {
+    "English": {
+        "title": "MOSS-TTSD🪐 Dialogue Generation",
+        "script_input": "### Script Input",
+        "text_to_synthesize": "Text to Synthesize",
+        "text_placeholder": "Text to be synthesized, format: [S1]Role1 text[S2]Role2 text",
+        "use_normalize": "Use text normalization",
+        "normalize_info": "Recommended to enable, improves handling of numbers, punctuation, etc.",
+        "audio_input_mode": "### Audio Input Mode",
+        "select_input_mode": "Select input mode",
+        "mode_info": "Single Audio: Upload one audio with [S1][S2] text; Role Audio: Upload separate audio for Role1 and Role2",
+        "drag_drop_audio": "Drag and drop audio here - or - click to upload",
+        "prompt_text": "Prompt Text",
+        "prompt_placeholder": "Format: [S1]Role1 text[S2]Role2 text",
+        "role1_audio": "**Role1 Audio**",
+        "role1_audio_file": "Role1 Audio File",
+        "role1_text": "Role1 Text",
+        "role1_placeholder": "Role1 text content",
+        "role2_audio": "**Role2 Audio**",
+        "role2_audio_file": "Role2 Audio File",
+        "role2_text": "Role2 Text",
+        "role2_placeholder": "Role2 text content",
+        "generate_audio": "Generate Audio",
+        "generated_audio": "Generated Audio",
+        "status_info": "Status Information",
+        "examples": "### Examples",
+        "examples_desc": "Click on examples below to auto-fill the form",
+        "role_headers": ["Text to Synthesize", "Input Mode", "Role1 Audio File", "Role1 Text", "Role2 Audio File", "Role2 Text", "Use Normalize"]
+    },
+    "中文": {
+        "title": "MOSS-TTSD🪐 对话语音生成",
+        "script_input": "### 文本输入",
+        "text_to_synthesize": "要合成的文本",
+        "text_placeholder": "要合成的文本，格式：[S1]角色1文本[S2]角色2文本",
+        "use_normalize": "使用文本规范化",
+        "normalize_info": "建议启用，改善数字、标点符号等的处理",
+        "audio_input_mode": "### 音频输入模式",
+        "select_input_mode": "选择输入模式",
+        "mode_info": "单音频：上传一个包含[S1][S2]文本的音频；角色音频：分别为角色1和角色2上传音频",
+        "drag_drop_audio": "拖拽音频文件到此处 - 或 - 点击上传",
+        "prompt_text": "提示文本",
+        "prompt_placeholder": "格式：[S1]角色1文本[S2]角色2文本",
+        "role1_audio": "**角色1音频**",
+        "role1_audio_file": "角色1音频文件",
+        "role1_text": "角色1文本",
+        "role1_placeholder": "角色1文本内容",
+        "role2_audio": "**角色2音频**",
+        "role2_audio_file": "角色2音频文件",
+        "role2_text": "角色2文本",
+        "role2_placeholder": "角色2文本内容",
+        "generate_audio": "生成音频",
+        "generated_audio": "生成的音频",
+        "status_info": "状态信息",
+        "examples": "### 示例",
+        "examples_desc": "点击下方示例自动填充表单",
+        "role_headers": ["要合成的文本", "输入模式", "角色1音频文件", "角色1文本", "角色2音频文件", "角色2文本", "使用规范化"]
+    }
+}
+
+# Model configuration
+SYSTEM_PROMPT = "You are a speech synthesizer that generates natural, realistic, and human-like conversational audio from dialogue text."
+MODEL_PATH = "fnlp/MOSS-TTSD-v0"
+SPT_CONFIG_PATH = "XY_Tokenizer/config/xy_tokenizer_config.yaml"
+SPT_CHECKPOINT_PATH = "XY_Tokenizer/weights/xy_tokenizer.ckpt"
+MAX_CHANNELS = 8
+
+from huggingface_hub import hf_hub_download
+
+ckpt_path = hf_hub_download(
+    repo_id="fnlp/XY_Tokenizer_TTSD_V0",
+    filename="xy_tokenizer.ckpt",
+    cache_dir="XY_Tokenizer/weights"
+)
+
+print("Checkpoint downloaded to:", ckpt_path)
+
+# Global variables for caching loaded models
+tokenizer = None
+model = None
+spt = None
+device = None
+
+def initialize_model():
+    """Initialize model (load only on first call)"""
+    global tokenizer, model, spt, device
+    
+    if tokenizer is None:
+        print("Initializing model...")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        tokenizer, model, spt = load_model(MODEL_PATH, SPT_CONFIG_PATH, SPT_CHECKPOINT_PATH)
+        spt = spt.to(device)
+        model = model.to(device)
+        print("Model initialization completed!")
+    
+    return tokenizer, model, spt, device
+
+def process_single_audio_generation(
+    text_input: str,
+    audio_mode: str,
+    prompt_text_single: str,
+    prompt_audio_single: Optional[str] = None,
+    prompt_text_1: str = "",
+    prompt_audio_1: Optional[str] = None,
+    prompt_text_2: str = "",
+    prompt_audio_2: Optional[str] = None,
+    use_normalize: bool = True
+) -> Tuple[Optional[str], str]:
+    """
+    Process single audio generation request
+    
+    Args:
+        text_input: Text to synthesize
+        prompt_text_single: Prompt text for single audio
+        prompt_audio_single: Single audio file path
+        prompt_text_1: Role1 text
+        prompt_audio_1: Role1 audio file path
+        prompt_text_2: Role2 text
+        prompt_audio_2: Role2 audio file path
+        use_normalize: Whether to use text normalization
+    
+    Returns:
+        Generated audio file path and status information
+    """
+    try:
+        # Initialize model
+        tokenizer, model, spt, device = initialize_model()
+        
+        # Build input item
+        item = {
+            "text": text_input,
+        }
+        
+        # Handle different audio input modes (mutually exclusive)
+        if audio_mode == "Single":
+            # Use single audio mode
+            item["prompt_audio"] = prompt_audio_single
+            item["prompt_text"] = prompt_text_single
+        elif audio_mode == "Role" and prompt_audio_1 and prompt_audio_2:
+            # Use role audio mode (requires both audio files)
+            item["prompt_audio_speaker1"] = prompt_audio_1
+            item["prompt_text_speaker1"] = prompt_text_1 if prompt_text_1 else ""
+            item["prompt_audio_speaker2"] = prompt_audio_2
+            item["prompt_text_speaker2"] = prompt_text_2 if prompt_text_2 else ""
+        elif audio_mode == "Role" and prompt_audio_1:
+            # Only Role 1 audio provided, treat as single audio
+            print("Only Role 1 audio provided, treating as single audio.")
+            item["prompt_audio"] = prompt_audio_1
+            item["prompt_text"] = prompt_text_1 if prompt_text_1 else ""
+        elif audio_mode == "Role" and prompt_audio_2:
+            # Only Role 2 audio provided, treat as single audio
+            print("Only Role 2 audio provided, treating as single audio.")
+            item["prompt_audio"] = prompt_audio_2
+            item["prompt_text"] = prompt_text_2 if prompt_text_2 else ""
+        else:
+            return None, "Error: Please select a mode and provide corresponding audio files\n- Single Audio Mode: Provide one audio file and corresponding text\n- Role Mode: Provide audio files for Role1 and Role2"
+        
+        # Set random seed to ensure reproducible results
+        import accelerate
+        accelerate.utils.set_seed(42)
+        
+        # Process batch (single item)
+        actual_texts_data, audio_results = process_batch(
+            batch_items=[item],
+            tokenizer=tokenizer,
+            model=model,
+            spt=spt,
+            device=device,
+            system_prompt=SYSTEM_PROMPT,
+            start_idx=0,
+            use_normalize=use_normalize
+        )
+        
+        # Check results
+        if not audio_results or audio_results[0] is None:
+            return None, "Error: Audio generation failed"
+        
+        audio_result = audio_results[0]
+        
+        # Create temporary output file
+        output_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
+        
+        # Save audio
+        torchaudio.save(output_path, audio_result["audio_data"], audio_result["sample_rate"])
+        
+        # Build status information (using English since this is server-side output)
+        status_info = f"""
+✅ Generation successful!
+📊 Audio Information:
+   - Sample Rate: {audio_result["sample_rate"]} Hz
+   - Audio Length: {audio_result["audio_data"].shape[-1] / audio_result["sample_rate"]:.2f} seconds
+   - Channels: {audio_result["audio_data"].shape[0]}
+
+📝 Text Processing Information:
+   - Original Text: {actual_texts_data[0]['original_text'][:100]}...
+   - Final Text: {actual_texts_data[0]['final_text'][:100]}...
+   - Use Normalize: {actual_texts_data[0]['use_normalize']}
+        """
+        
+        return output_path, status_info
+        
+    except Exception as e:
+        import traceback
+        error_msg = f"Error: Audio generation failed: {str(e)}\n\nDetails:\n{traceback.format_exc()}"
+        return None, error_msg
+
+# Create Gradio interface
+def create_gradio_interface() -> gr.Blocks:
+    with gr.Blocks(title="MOSS-TTSD🪐 Dialogue Generation", theme=gr.themes.Soft()) as demo:
+        
+        # Language selection at the top
+        with gr.Row():
+            language_selector = gr.Radio(
+                choices=["English", "中文"],
+                value="English",
+                label="Language / 语言",
+                info="Select interface language / 选择界面语言"
+            )
+        
+        # Title and header (will be updated based on language)
+        title_md = gr.Markdown("# MOSS-TTSD🪐 Dialogue Generation")
+        github_md = gr.Markdown("### [Github](https://github.com/OpenMOSS/MOSS-TTSD)")
+        
+        with gr.Row():
+            # Left input area
+            with gr.Column(scale=1):
+                script_input_md = gr.Markdown("### Script Input")
+                
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    placeholder="Text to be synthesized, format: [S1]Role1 text[S2]Role2 text",
+                    lines=6,
+                )
+                
+                use_normalize_single = gr.Checkbox(
+                    label="Use text normalization",
+                    value=True,
+                    info="Recommended to enable, improves handling of numbers, punctuation, etc."
+                )
+            
+            # Right audio input area
+            with gr.Column(scale=1):
+                audio_input_mode_md = gr.Markdown("### Audio Input Mode")
+                
+                # Audio input mode selection
+                audio_mode = gr.Radio(
+                    choices=["Single", "Role"],
+                    value="Single",
+                    label="Select input mode",
+                    info="Single Audio: Upload one audio with [S1][S2] text; Role Audio: Upload separate audio for Role1 and Role2"
+                )
+                
+                # Single audio mode
+                with gr.Group(visible=True) as single_mode_group:
+                    prompt_audio_single = gr.File(
+                        label="Drag and drop audio here - or - click to upload",
+                        file_types=["audio"],
+                        type="filepath"
+                    )
+                    prompt_text_single = gr.Textbox(
+                        label="Prompt Text",
+                        placeholder="Format: [S1]Role1 text[S2]Role2 text",
+                        lines=3,
+                    )
+                
+                # Role audio mode
+                with gr.Group(visible=False) as role_mode_group:
+                    with gr.Row():
+                        with gr.Column():
+                            role1_audio_md = gr.Markdown("**Role1 Audio**")
+                            prompt_audio_1 = gr.File(
+                                label="Role1 Audio File",
+                                file_types=["audio"],
+                                type="filepath"
+                            )
+                            prompt_text_1 = gr.Textbox(
+                                label="Role1 Text",
+                                placeholder="Role1 text content",
+                                lines=2
+                            )
+                        
+                        with gr.Column():
+                            role2_audio_md = gr.Markdown("**Role2 Audio**")
+                            prompt_audio_2 = gr.File(
+                                label="Role2 Audio File",
+                                file_types=["audio"],
+                                type="filepath"
+                            )
+                            prompt_text_2 = gr.Textbox(
+                                label="Role2 Text",
+                                placeholder="Role2 text content",
+                                lines=2
+                            )
+        
+        # Generate button
+        with gr.Row():
+            generate_btn = gr.Button("Generate Audio", variant="primary", size="lg")
+        
+        # Output area
+        with gr.Row():
+            with gr.Column():
+                output_audio = gr.Audio(label="Generated Audio", type="filepath")
+                status_info = gr.Textbox(
+                    label="Status Information",
+                    lines=10,
+                    interactive=False
+                )
+        
+        # Examples area
+        with gr.Row():
+            with gr.Column():
+                examples_md = gr.Markdown("### Examples")
+                examples_desc_md = gr.Markdown("Click on examples below to auto-fill the form")
+
+                role_examples = gr.Examples(
+                    examples=ROLE_EXAMPLES,
+                    inputs=[text_input, audio_mode, prompt_audio_1, prompt_text_1, prompt_audio_2, prompt_text_2, use_normalize_single],
+                )
+        
+        # Event handlers
+        
+        # Language change event
+        def update_language(lang):
+            """Update interface language"""
+            texts = LANGUAGES[lang]
+            
+            # Update demo title
+            demo.title = texts["title"]
+            
+            return (
+                gr.Markdown(f"# {texts['title']}"),  # title_md
+                texts["script_input"],  # script_input_md
+                gr.Textbox(
+                    label=texts["text_to_synthesize"],
+                    placeholder=texts["text_placeholder"],
+                    lines=6,
+                ),  # text_input
+                gr.Checkbox(
+                    label=texts["use_normalize"],
+                    value=True,
+                    info=texts["normalize_info"]
+                ),  # use_normalize_single
+                texts["audio_input_mode"],  # audio_input_mode_md
+                gr.Radio(
+                    choices=["Single", "Role"],
+                    value="Single",
+                    label=texts["select_input_mode"],
+                    info=texts["mode_info"]
+                ),  # audio_mode
+                gr.File(
+                    label=texts["drag_drop_audio"],
+                    file_types=["audio"],
+                    type="filepath"
+                ),  # prompt_audio_single
+                gr.Textbox(
+                    label=texts["prompt_text"],
+                    placeholder=texts["prompt_placeholder"],
+                    lines=3,
+                ),  # prompt_text_single
+                texts["role1_audio"],  # role1_audio_md
+                gr.File(
+                    label=texts["role1_audio_file"],
+                    file_types=["audio"],
+                    type="filepath"
+                ),  # prompt_audio_1
+                gr.Textbox(
+                    label=texts["role1_text"],
+                    placeholder=texts["role1_placeholder"],
+                    lines=2
+                ),  # prompt_text_1
+                texts["role2_audio"],  # role2_audio_md
+                gr.File(
+                    label=texts["role2_audio_file"],
+                    file_types=["audio"],
+                    type="filepath"
+                ),  # prompt_audio_2
+                gr.Textbox(
+                    label=texts["role2_text"],
+                    placeholder=texts["role2_placeholder"],
+                    lines=2
+                ),  # prompt_text_2
+                gr.Button(texts["generate_audio"], variant="primary", size="lg"),  # generate_btn
+                gr.Audio(label=texts["generated_audio"], type="filepath"),  # output_audio
+                gr.Textbox(
+                    label=texts["status_info"],
+                    lines=10,
+                    interactive=False
+                ),  # status_info
+                texts["examples"],  # examples_md
+                texts["examples_desc"],  # examples_desc_md
+                gr.Dataset(headers=texts["role_headers"])
+            )
+        
+        language_selector.change(
+            fn=update_language,
+            inputs=[language_selector],
+            outputs=[
+                title_md, script_input_md, text_input, use_normalize_single,
+                audio_input_mode_md, audio_mode, prompt_audio_single, prompt_text_single,
+                role1_audio_md, prompt_audio_1, prompt_text_1,
+                role2_audio_md, prompt_audio_2, prompt_text_2,
+                generate_btn, output_audio, status_info,
+                examples_md, examples_desc_md, role_examples.dataset,
+            ]
+        )
+        
+        # Audio mode toggle event
+        def toggle_audio_mode(mode):
+            if mode == "Single":
+                return gr.Group(visible=True), gr.Group(visible=False)
+            else:
+                return gr.Group(visible=False), gr.Group(visible=True)
+        
+        audio_mode.change(
+            fn=toggle_audio_mode,
+            inputs=[audio_mode],
+            outputs=[single_mode_group, role_mode_group]
+        )
+        
+        # Audio generation event
+        generate_btn.click(
+            fn=process_single_audio_generation,
+            inputs=[
+                text_input,
+                audio_mode,
+                prompt_text_single,
+                prompt_audio_single,
+                prompt_text_1,
+                prompt_audio_1,
+                prompt_text_2,
+                prompt_audio_2,
+                use_normalize_single
+            ],
+            outputs=[output_audio, status_info],
+            show_progress=True
+        )
+    
+    return demo
+
+# Main function
+if __name__ == "__main__":
+    demo = create_gradio_interface()
+    
+    # Launch interface
+    demo.launch()

generation_utils.py

@@ -0,0 +1,452 @@
+import os
+import re
+
+import torch
+import torchaudio
+import numpy as np
+
+from transformers import AutoTokenizer
+from modeling_asteroid import AsteroidTTSInstruct
+from XY_Tokenizer.xy_tokenizer.model import XY_Tokenizer
+
+MAX_CHANNELS = 8
+SILENCE_DURATION = 5.0  # Fixed silence duration: 5 seconds
+
+def load_model(model_path, spt_config_path, spt_checkpoint_path):
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    
+    model = AsteroidTTSInstruct.from_pretrained(model_path, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")
+
+    spt = XY_Tokenizer.load_from_checkpoint(config_path=spt_config_path, ckpt_path=spt_checkpoint_path)
+    
+    model.eval()
+    spt.eval()
+    return tokenizer, model, spt
+
+
+def process_jsonl_item(item):
+    """Process JSONL data items and extract audio and text information according to the new format"""
+    base_path = item.get("base_path", "")
+    text = item.get("text", "")
+    
+    # Process prompt audio and text
+    if "prompt_audio" in item and "prompt_text" in item:
+        print("Using prompt_audio and prompt_text directly from item.")
+        # If prompt_audio and prompt_text exist, use them directly
+        prompt_audio = item["prompt_audio"]
+        prompt_text = item["prompt_text"]
+        
+        # Only perform path joining when prompt_audio is a string path
+        if isinstance(prompt_audio, str) and base_path and prompt_audio:
+            prompt_audio = os.path.join(base_path, prompt_audio)
+    else:
+        print("Using speaker1 and speaker2 information for prompt audio and text.")
+        # Otherwise, merge speaker1 and speaker2 information
+        prompt_audio_speaker1 = item.get("prompt_audio_speaker1", "")
+        prompt_text_speaker1 = item.get("prompt_text_speaker1", "")
+        prompt_audio_speaker2 = item.get("prompt_audio_speaker2", "")
+        prompt_text_speaker2 = item.get("prompt_text_speaker2", "")
+        
+        # Process audio: if it's a string path, perform path joining; if it's a tuple, use directly
+        if isinstance(prompt_audio_speaker1, str):
+            speaker1_audio = os.path.join(base_path, prompt_audio_speaker1) if base_path and prompt_audio_speaker1 else prompt_audio_speaker1
+        else:
+            speaker1_audio = prompt_audio_speaker1  # Use tuple directly
+            
+        if isinstance(prompt_audio_speaker2, str):
+            speaker2_audio = os.path.join(base_path, prompt_audio_speaker2) if base_path and prompt_audio_speaker2 else prompt_audio_speaker2
+        else:
+            speaker2_audio = prompt_audio_speaker2  # Use tuple directly
+        
+        prompt_audio = {
+            "speaker1": speaker1_audio,
+            "speaker2": speaker2_audio
+        }
+        
+        # Merge text
+        prompt_text = ""
+        if prompt_text_speaker1:
+            prompt_text += f"[S1]{prompt_text_speaker1}"
+        if prompt_text_speaker2:
+            prompt_text += f"[S2]{prompt_text_speaker2}"
+        prompt_text = prompt_text.strip()
+    
+    return {
+        "text": text,
+        "prompt_text": prompt_text,
+        "prompt_audio": prompt_audio
+    }
+
+
+def load_audio_data(prompt_audio, target_sample_rate=16000):
+    """Load audio data and return processed audio tensor
+    
+    Args:
+        prompt_audio: Can be in the following formats:
+            - String: audio file path
+            - Tuple: (wav, sr) result from torchaudio.load
+            - Dict: {"speaker1": path_or_tuple, "speaker2": path_or_tuple}
+    """
+    if prompt_audio is None:
+        return None
+    
+    try:
+        # Check if prompt_audio is a dictionary (containing speaker1 and speaker2)
+        if isinstance(prompt_audio, dict) and "speaker1" in prompt_audio and "speaker2" in prompt_audio:
+            # Process audio from both speakers separately
+            wav1, sr1 = _load_single_audio(prompt_audio["speaker1"])
+            wav2, sr2 = _load_single_audio(prompt_audio["speaker2"])
+            # Merge audio from both speakers
+            wav = merge_speaker_audios(wav1, sr1, wav2, sr2, target_sample_rate)
+            if wav is None:
+                return None
+        else:
+            # Single audio
+            wav, sr = _load_single_audio(prompt_audio)
+            # Resample to 16k
+            if sr != target_sample_rate: 
+                wav = torchaudio.functional.resample(wav, sr, target_sample_rate)
+            # Ensure mono channel
+            if wav.shape[0] > 1:
+                wav = wav.mean(dim=0, keepdim=True)  # Convert multi-channel to mono
+            if len(wav.shape) == 1: 
+                wav = wav.unsqueeze(0)
+        
+        return wav
+    except Exception as e:
+        print(f"Error loading audio data: {e}")
+        import traceback
+        traceback.print_exc()
+        return None
+
+
+def _load_single_audio(audio_input):
+    """Load single audio, supports file path or (wav, sr) tuple
+    
+    Args:
+        audio_input: String (file path) or tuple (wav, sr)
+        
+    Returns:
+        tuple: (wav, sr)
+    """
+    if isinstance(audio_input, tuple) and len(audio_input) == 2:
+        # Already a (wav, sr) tuple
+        wav, sr = audio_input
+        return wav, sr
+    elif isinstance(audio_input, str):
+        # Is a file path, needs to be loaded
+        wav, sr = torchaudio.load(audio_input)
+        return wav, sr
+    else:
+        raise ValueError(f"Unsupported audio input format: {type(audio_input)}")
+
+
+def merge_speaker_audios(wav1, sr1, wav2, sr2, target_sample_rate=16000):
+    """Merge audio data from two speakers"""
+    try:
+        # Process first audio
+        if sr1 != target_sample_rate:
+            wav1 = torchaudio.functional.resample(wav1, sr1, target_sample_rate)
+        # Ensure mono channel
+        if wav1.shape[0] > 1:
+            wav1 = wav1.mean(dim=0, keepdim=True)  # Convert multi-channel to mono
+        if len(wav1.shape) == 1:
+            wav1 = wav1.unsqueeze(0)
+        
+        # Process second audio  
+        if sr2 != target_sample_rate:
+            wav2 = torchaudio.functional.resample(wav2, sr2, target_sample_rate)
+        # Ensure mono channel
+        if wav2.shape[0] > 1:
+            wav2 = wav2.mean(dim=0, keepdim=True)  # Convert multi-channel to mono
+        if len(wav2.shape) == 1:
+            wav2 = wav2.unsqueeze(0)
+        
+        # Concatenate audio
+        merged_wav = torch.cat([wav1, wav2], dim=1)
+        return merged_wav
+    except Exception as e:
+        print(f"Error merging audio: {e}")
+        return None
+
+
+def process_inputs(tokenizer, spt, prompt, text, device, audio_data=None, max_channels=8, pad_token=1024):
+    seq = f"<|begin_of_style|>{prompt}<|end_of_style|>\n<|begin_of_text|>{text}<|end_of_text|>\n<|begin_of_speech|>"
+    inputs1 = np.array(tokenizer.encode(seq))
+    input_ids = np.full((inputs1.shape[0], max_channels), pad_token)
+    input_ids[:, 0] = inputs1
+    
+    if audio_data is not None:
+        try:
+            # audio_data should now be a processed audio tensor
+            wav = audio_data
+            
+            # Add fixed 5-second silence at the end of audio (using 16k sample rate)
+            silence_samples = int(SILENCE_DURATION * 16000)
+            silence = torch.zeros(wav.shape[0], silence_samples)
+            wav = torch.cat([wav, silence], dim=1)
+            
+            with torch.no_grad():
+                # Use SPT encoding
+                encode_result = spt.encode([wav.squeeze().to(device)])
+                audio_token = encode_result["codes_list"][0].permute(1, 0).cpu().numpy()  # Adjust dimension order
+                
+            # similar to DAC encoding adjustment
+            audio_token[:, 0] = audio_token[:, 0] + 151665  # Keep this line if offset is needed, otherwise delete
+            input_ids = np.concatenate([input_ids, audio_token])[:-60]
+        except Exception as e:
+            print(f"Error processing audio data: {e}")
+            import traceback
+            traceback.print_exc()
+            # If error occurs, still return input without audio
+    
+    return input_ids
+
+
+def shifting_inputs(input_ids, tokenizer, pad_token=1024, max_channels=8):
+    seq_len = input_ids.shape[0]
+    new_seq_len = seq_len + max_channels - 1
+    shifted_input_ids = np.full((new_seq_len, max_channels), pad_token, dtype=np.int64)
+    shifted_input_ids[:, 0] = np.full(new_seq_len, tokenizer.pad_token_id, dtype=np.int64)
+    for i in range(max_channels):
+        shifted_input_ids[i : (seq_len + i), i] = input_ids[:, i]
+    return shifted_input_ids
+
+
+def rpadding(input_ids, channels, tokenizer):
+    attention_masks = [np.ones(inputs.shape[0]) for inputs in input_ids]
+    max_length = max(ids.shape[0] for ids in input_ids)
+    padded_input_ids, padded_attns = [], []
+        
+    for ids, attn in zip(input_ids, attention_masks):
+        pad_len = max_length - ids.shape[0]
+        input_pad = np.full((pad_len, channels), 1024)
+        input_pad[:, 0] = tokenizer.pad_token_id
+        padded_input_ids.append(np.concatenate([input_pad, ids]))
+        attn_pad = np.zeros(pad_len)
+        padded_attns.append(np.concatenate([attn_pad, attn]))
+
+    input_ids = torch.tensor(np.stack(padded_input_ids))
+    attention_mask = torch.tensor(np.stack(padded_attns))
+
+    return input_ids, attention_mask
+
+
+def find_max_valid_positions(C: torch.Tensor, invalid_value=1024) -> torch.Tensor:
+    values = C[:, :, 1]
+    mask = (values != invalid_value)
+    reversed_mask = mask.flip(dims=[1])
+    reversed_indices = torch.argmax(reversed_mask.int(), dim=1)
+    seq_len = C.size(1)
+    original_indices = seq_len - 1 - reversed_indices
+    has_valid = mask.any(dim=1)
+    original_indices = torch.where(has_valid, original_indices, -1)
+    return original_indices
+
+
+def normalize_text(text: str) -> str:
+    """
+    Normalize multi-speaker script.
+
+    1. Don't preserve line breaks.
+    2. Remove brackets for non-speaker tags (if [] doesn't contain S1/S2...Sx format, remove the brackets themselves).
+    3. Remove decorative symbols: 【】《》（）『』「」"-“” .
+    4. Internal punctuation ！；：、 → ，；only allow ？ and ，。
+    5. Multiple 。 keep only the last one, others → ，。
+    6. Replace consecutive "哈" (>=2) with "(笑)".
+    7. Auto-recognize [S1] / [S2] … tags; if missing, treat as whole segment.
+    """
+    # Replace [1], [2] etc. format with [S1], [S2] etc. format
+    text = re.sub(r'\[(\d+)\]', r'[S\1]', text)
+
+    # Remove decorative characters
+    remove_chars = "【】《》（）『』「」""\"-“”"
+
+
+    # Remove brackets for non-speaker tags (keep content, only remove brackets themselves)
+    text = re.sub(r'\[(?!S\d+\])([^\]]*)\]', r'\1', text)
+
+    # Use positive lookahead to split text by speaker tags (tags themselves are still preserved)
+    segments = re.split(r'(?=\[S\d+\])', text.replace("\n", " "))
+    normalized_lines = []
+
+    for seg in segments:
+        seg = seg.strip()
+        if not seg:
+            continue
+
+        # Extract tags
+        m = re.match(r'^(\[S\d+\])\s*(.*)', seg)
+        tag, content = m.groups() if m else ('', seg)
+
+        # Remove irrelevant symbols
+        content = re.sub(f"[{re.escape(remove_chars)}]", "", content)
+
+        # Handle consecutive "哈" characters: replace 2 or more with "(笑)"
+        content = re.sub(r'哈{2,}', '(笑)', content)
+
+        # First handle multi-character punctuation marks
+        content = content.replace('——', '，')
+        content = content.replace('……', '，')
+
+        # Handle single-character internal punctuation marks
+        internal_punct_map = str.maketrans({
+            '！': '，', '!': ',',
+            '；': '，', ';': ',',
+            '：': '，', ':': ',',
+            '、': '，', 
+            '？': '，', '?': ','
+        })
+        content = content.translate(internal_punct_map)
+        content = content.strip()
+
+        # Keep only the final period
+        if len(content) > 1:
+            last_ch = "。" if content[-1] == "，" else ("." if content[-1] == "," else content[-1])
+            body = content[:-1].replace('。', '，')
+            content = body + last_ch
+
+        normalized_lines.append(f"{tag}{content}".strip())
+
+    return "".join(normalized_lines)
+
+
+def process_batch(batch_items, tokenizer, model, spt, device, system_prompt, start_idx, use_normalize=False):
+    """Process a batch of data items and generate audio, return audio data and metadata"""
+    try:
+        # Prepare batch data
+        batch_size = len(batch_items)
+        texts = []
+        prompts = [system_prompt] * batch_size
+        prompt_audios = []
+        actual_texts_data = []  # Store actual text data used
+        
+        print(f"Processing {batch_size} samples starting from index {start_idx}...")
+        
+        # Extract text and audio from each sample
+        for i, item in enumerate(batch_items):
+            # Use new processing function
+            processed_item = process_jsonl_item(item)
+            
+            text = processed_item["text"]
+            prompt_text = processed_item["prompt_text"]
+            
+            # Merge text
+            full_text = prompt_text + text
+            original_full_text = full_text  # Save original text
+            
+            # Apply text normalization based on parameter
+            if use_normalize:
+                full_text = normalize_text(full_text)
+            
+            # Replace speaker tags
+            final_text = full_text.replace("[S1]", "<speaker1>").replace("[S2]", "<speaker2>")
+            texts.append(final_text)
+            
+            # Save actual text information used
+            actual_texts_data.append({
+                "index": start_idx + i,
+                "original_text": original_full_text,
+                "normalized_text": normalize_text(original_full_text) if use_normalize else None,
+                "final_text": final_text,
+                "use_normalize": use_normalize
+            })
+            
+            # Get reference audio
+            prompt_audios.append(processed_item["prompt_audio"])
+        
+        # Process inputs
+        input_ids_list = []
+        for i, (text, prompt, audio_path) in enumerate(zip(texts, prompts, prompt_audios)):
+            # Load audio data here
+            audio_data = load_audio_data(audio_path) if audio_path else None
+            inputs = process_inputs(tokenizer, spt, prompt, text, device, audio_data)
+            inputs = shifting_inputs(inputs, tokenizer)
+            input_ids_list.append(inputs)
+        
+        # Pad batch inputs
+        input_ids, attention_mask = rpadding(input_ids_list, MAX_CHANNELS, tokenizer)
+        
+        # Batch generation
+        print(f"Starting batch audio generation...")
+        start = input_ids.shape[1] - MAX_CHANNELS + 1
+        
+        # Move inputs to GPU
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+        
+        # Generate model outputs
+        outputs = model.generate(
+            input_ids=input_ids, 
+            attention_mask=attention_mask,
+        )
+        print(f"Original outputs shape: {outputs.shape}")
+        print(f"Start value: {start}")
+        print(f"Shape after slicing: {outputs[:, start:].shape}")
+        print(f"MAX_CHANNELS: {MAX_CHANNELS}")
+        print(f"Calculated seq_len: {outputs.shape[1] - MAX_CHANNELS + 1}")
+        # Process outputs
+        outputs = outputs[:, start:]
+        seq_len = outputs.shape[1] - MAX_CHANNELS + 1
+        speech_ids = torch.full((outputs.shape[0], seq_len, MAX_CHANNELS), 0).to(device)
+        
+        
+        # Adjust output format
+        for j in range(MAX_CHANNELS):
+            speech_ids[..., j] = outputs[:, j : seq_len + j, j]
+            if j == 0: 
+                speech_ids[..., j] = speech_ids[..., j] - 151665
+        
+        # Find valid positions for each sample
+        li = find_max_valid_positions(speech_ids)
+        
+        # Store audio result data
+        audio_results = []
+        
+        # Process batch sample results individually
+        for i in range(batch_size):
+            try:
+                # Extract valid speech tokens
+                end_idx = li[i] + 1
+                if end_idx <= 0:
+                    print(f"Sample {start_idx + i} has no valid speech tokens")
+                    audio_results.append(None)
+                    continue
+                    
+                this_speech_id = speech_ids[i, :end_idx]
+                print(f"Speech token shape for sample {start_idx + i}: {this_speech_id.shape}")
+                
+                # Decode generated audio
+                with torch.no_grad():
+                    codes_list = [this_speech_id.permute(1, 0)]  # Convert to SPT expected format
+                    decode_result = spt.decode(codes_list, overlap_seconds=10)
+                    audio_result = decode_result["syn_wav_list"][0].cpu().detach()
+                    
+                    if audio_result.ndim == 1:  # If 1D [samples]
+                        audio_result = audio_result.unsqueeze(0)  # Convert to 2D [1, samples]
+                
+                # Save audio data instead of file path
+                audio_results.append({
+                    "audio_data": audio_result,
+                    "sample_rate": spt.output_sample_rate,
+                    "index": start_idx + i
+                })
+                print(f"Audio generation completed: sample {start_idx + i}")
+                
+            except Exception as e:
+                print(f"Error processing sample {start_idx + i}: {str(e)}")
+                import traceback
+                traceback.print_exc()
+                audio_results.append(None)
+        
+        # Clean up GPU memory
+        torch.cuda.empty_cache()
+        
+        # Return text data and audio data
+        return actual_texts_data, audio_results
+        
+    except Exception as e:
+        print(f"Error during batch processing: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return [], [None] * len(batch_items)

modeling_asteroid.py

@@ -0,0 +1,399 @@
+import torch
+import torch.nn as nn
+from dataclasses import dataclass
+from transformers.utils import ModelOutput
+from transformers.cache_utils import Cache
+from typing import Optional, List, Tuple, Union
+from transformers.loss.loss_utils import ForCausalLMLoss
+from transformers.generation.streamers import BaseStreamer
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.generation.stopping_criteria import StoppingCriteriaList
+from transformers import PreTrainedModel, GenerationMixin, Qwen3Config, Qwen3Model
+from transformers.generation.logits_process import LogitsProcessorList, RepetitionPenaltyLogitsProcessor, TopKLogitsWarper, TopPLogitsWarper, TemperatureLogitsWarper
+
+
+class AsteroidTTSConfig(Qwen3Config):
+    def __init__(self, 
+                channels = 8,
+                speech_pad_token = 1024,
+                speech_vocab_size = 1025,
+                speech_token_range = [],
+                **kwargs):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.speech_pad_token = speech_pad_token
+        self.speech_vocab_size = speech_vocab_size
+        self.speech_token_range = speech_token_range
+        
+
+@dataclass
+class AsteroidTTSOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    loss_all: Optional[Tuple[torch.FloatTensor]] = None
+    logits_all: Optional[Tuple[torch.FloatTensor]] = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    
+
+@dataclass
+class GenerateDecoderOnlyOutput(ModelOutput):
+    sequences: torch.LongTensor = None
+    scores: Optional[Tuple[torch.FloatTensor]] = None
+    logits: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None
+
+
+class CustomMixin(GenerationMixin):
+    def _sample(
+        self,
+        input_ids: torch.LongTensor,
+        logits_processor: LogitsProcessorList,
+        stopping_criteria: StoppingCriteriaList,
+        generation_config: GenerationConfig,
+        synced_gpus: bool,
+        streamer: Optional["BaseStreamer"],
+        **model_kwargs,
+    ) -> Union[GenerateDecoderOnlyOutput, torch.LongTensor]:
+        # 提取配置参数
+        speech_pad_idx = self.config.speech_pad_token
+        
+        eos_token_id = generation_config.eos_token_id
+        output_attentions = generation_config.output_attentions
+        output_hidden_states = generation_config.output_hidden_states
+        output_scores = generation_config.output_scores
+        output_logits = generation_config.output_logits
+        return_dict_in_generate = generation_config.return_dict_in_generate
+        max_length = generation_config.max_length
+        has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria)
+        do_sample = generation_config.do_sample
+
+        # 初始化输出元组
+        scores = () if (return_dict_in_generate and output_scores) else None
+        raw_logits = () if (return_dict_in_generate and output_logits) else None
+        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+        # 初始化跟踪变量
+        batch_size, cur_len, channels = input_ids.shape  # channels = 8
+        this_peer_finished = False
+        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+        needs_additional_steps = -1 * torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+        tf_inputs = input_ids[:]
+        input_ids = input_ids[:, :-(channels - 1)]
+        model_kwargs["attention_mask"] = model_kwargs["attention_mask"][:, :-(channels - 1)]
+        base_length = input_ids.shape[1]
+        model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+
+        # 定义logits processor
+        if generation_config.do_samples is not None:
+            do_samples = generation_config.do_samples
+            realprocessor = [LogitsProcessorList() for _ in range(channels)]
+            for i, layer_config in enumerate(generation_config.layers):
+                if layer_config.get("repetition_penalty") is not None:
+                    realprocessor[i].append(RepetitionPenaltyLogitsProcessor(penalty=layer_config.get("repetition_penalty")))
+                if layer_config.get("temperature") is not None: 
+                    realprocessor[i].append(TemperatureLogitsWarper(temperature=layer_config.get("temperature")))
+                if layer_config.get("top_k") is not None:
+                    realprocessor[i].append(TopKLogitsWarper(top_k=layer_config.get("top_k")))
+                if layer_config.get("top_p") is not None:
+                    realprocessor[i].append(TopPLogitsWarper(top_p=layer_config.get("top_p")))
+        else:
+            do_samples = [do_sample for _ in range(channels)]
+            realprocessor = [logits_processor for _ in range(channels)]
+        while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+            # 准备模型输入
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+            model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+            model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+            # 前向传递
+            outputs = self(**model_inputs, return_dict=True)
+            model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
+
+            if synced_gpus and this_peer_finished:
+                continue
+
+            # 获取下一个 token 的 logits
+            next_token_logits = [logits[:, -1, :].clone().float().to(input_ids.device) for logits in outputs.logits_all]
+            for i, channel_logits in enumerate(next_token_logits):
+                if i != 0 and input_ids.shape[1] + 1 > tf_inputs.shape[1] - 7 + i: 
+                    channel_logits[:, 1024] = - torch.inf
+                if i == 0 and input_ids.shape[1] + 1 <= tf_inputs.shape[1]: 
+                    channel_logits[:, 152694] = - torch.inf
+            next_token_scores = [realprocessor[i](input_ids[..., i], logits) for i, logits in enumerate(next_token_logits)]
+            # 生成下一个 token
+            next_tokens = []
+            for i, channel_score in enumerate(next_token_scores):
+                if do_samples[i]:
+                    channel_ntk = torch.multinomial(nn.functional.softmax(channel_score, dim=-1), num_samples=1).squeeze(1)
+                elif not do_samples[i]:
+                    channel_ntk = torch.argmax(channel_score, dim=-1)
+                next_tokens.append(channel_ntk)
+            next_tokens = torch.stack(next_tokens, dim=-1)  # [batch_size, channels]
+            # 额外步骤逻辑
+            indices = (~self.is_speech_token(next_tokens[:, 0])) & (needs_additional_steps < 0)
+            needs_additional_steps[indices] = channels - 1  # 对于 8 个通道，需要 6 步
+            
+            if input_ids.shape[1] + 1 <= tf_inputs.shape[1]:
+                i = input_ids.shape[1] + 1 - base_length
+                next_tokens[:, i:] = tf_inputs[:, input_ids.shape[1], i:]
+            
+            # 在额外步骤中替换 token
+            mask = (needs_additional_steps > 0) & (needs_additional_steps < 7)
+            if mask.any().item():
+                next_tokens[mask, 0] = self.config.eos_token_id
+                for i in range(1, channels):
+                    mask_i = mask & (needs_additional_steps < channels - i)
+                    next_tokens[mask_i, i] = speech_pad_idx
+            
+            if has_eos_stopping_criteria:
+                for i in range(channels):
+                    pddp = self.config.eos_token_id if i == 0 else speech_pad_idx
+                    next_tokens[:, i] = next_tokens[:, i] * unfinished_sequences + pddp * (1 - unfinished_sequences)
+                    
+            input_ids = torch.cat([input_ids, next_tokens[:, None, :]], dim=1)
+            if streamer is not None:
+                streamer.put(next_tokens[:, 0].cpu())
+            
+            # 更新 unfinished_sequences
+            needs_additional_steps = torch.where(needs_additional_steps > 0, needs_additional_steps - 1, needs_additional_steps)
+            stopping = stopping_criteria(input_ids[..., 0], scores) | (needs_additional_steps == 0)
+            unfinished_sequences = unfinished_sequences & ~stopping
+            unfinished_sequences = unfinished_sequences | (needs_additional_steps > 0)
+            this_peer_finished = unfinished_sequences.max() == 0
+
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_token_scores,)
+                if output_logits:
+                    raw_logits += (next_token_logits,)
+                if output_attentions:
+                    decoder_attentions += (outputs.attentions,)
+                if output_hidden_states:
+                    decoder_hidden_states += (outputs.hidden_states,)
+
+            cur_len += 1
+            del outputs
+            
+        if streamer is not None:
+            streamer.end()
+
+        if return_dict_in_generate:
+            return GenerateDecoderOnlyOutput(
+                sequences=input_ids,
+                scores=scores,
+                logits=raw_logits,
+                attentions=decoder_attentions,
+                hidden_states=decoder_hidden_states,
+                past_key_values=model_kwargs.get("past_key_values"),
+            )
+        else:
+            return input_ids
+    
+    
+class AsteroidTTSPretrainedModel(PreTrainedModel):
+    config_class = AsteroidTTSConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+
+class AsteroidTTSModel(AsteroidTTSPretrainedModel):
+    def __init__(self, config: AsteroidTTSConfig):
+        super().__init__(config)
+        self.text_pad_idx = config.pad_token_id
+        self.speech_pad_idx = config.speech_pad_token
+        self.embedding_list = nn.ModuleList([])
+        self.embedding_list.append(nn.Embedding(config.vocab_size, config.hidden_size, self.text_pad_idx))
+        # Channels 1 to channels-1: Speech tokens only
+        for _ in range(1, config.channels):
+            self.embedding_list.append(nn.Embedding(config.speech_vocab_size, config.hidden_size, self.speech_pad_idx))
+
+        self.language_model = Qwen3Model(config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embedding_list[0]
+
+    def set_input_embeddings(self, value: nn.Embedding):
+        self.embedding_list[0] = value
+
+    def _prepare_multi_modal_inputs(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+        """
+        Prepares multi-modal embeddings from input_ids of shape (batch_size, channels, sequence_length).
+        For channel 0: text + speech tokens, for channels 1 to channels-1: speech tokens padded with speech_pad_token.
+        """
+        batch_size, seq_length, channels = input_ids.shape
+        if channels != self.config.channels:
+            raise ValueError(f"Expected {self.config.channels} channels, got {channels}")
+        
+        inputs_embeds = torch.zeros(batch_size, seq_length, self.config.hidden_size, device=input_ids.device, dtype=self.embedding_list[0].weight.dtype)
+        for i in range(channels):
+            embed_layer = self.embedding_list[i]
+            channel_input = input_ids[...,i]
+            inputs_embeds += embed_layer(channel_input)
+
+        return inputs_embeds
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,  # Shape: (batch_size, channels, sequence_length)
+        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,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if input_ids is not None:
+            inputs_embeds = self._prepare_multi_modal_inputs(input_ids)
+
+        outputs = self.language_model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+        return outputs
+    
+    
+class AsteroidTTSInstruct(AsteroidTTSPretrainedModel, CustomMixin):
+    _tied_weights_keys = []
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config: AsteroidTTSConfig):
+        super().__init__(config)
+        self.model = AsteroidTTSModel(config)
+        self.channels = config.channels
+        self.weights = [1 for _ in range(self.channels)]
+        self._tied_weights_keys = [f"lm_heads.{i}.weight" for i in range(self.channels)]
+        self.vocab_size = config.vocab_size
+        self.lm_heads = nn.ModuleList([])
+        self.lm_heads.append(nn.Linear(config.hidden_size, config.vocab_size, bias=False))
+        for _ in range(1, config.channels):
+            self.lm_heads.append(nn.Linear(config.hidden_size, config.speech_vocab_size, bias=False))
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embedding_list[0]
+    
+    def can_generate(self):
+        return True
+    
+    def is_speech_token(self, tokens):
+        return (tokens >= self.config.speech_token_range[0]) & (tokens < self.config.speech_token_range[1])
+    
+    def tie_weights(self):
+        for i in range(self.config.channels):
+            self._tie_or_clone_weights(self.lm_heads[i], self.model.embedding_list[i])
+
+    def set_input_embeddings(self, value):
+        self.model.embedding_list[0] = value
+
+    def get_output_embeddings(self):
+        return self.lm_heads[0]
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_heads[0] = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+    
+    def set_weights(self, weights):
+        self.weights = weights
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, AsteroidTTSOutputWithPast]:
+        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
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        logits_all = [lm_head(hidden_states) for lm_head in self.lm_heads]
+
+        loss_all = torch.empty(self.channels, device=input_ids.device if not input_ids is None else inputs_embeds.device)
+
+        if labels is not None:
+            for i in range(self.config.channels):
+                vocab_size = self.config.vocab_size if i == 0 else self.config.speech_vocab_size
+                loss_all[i] = ForCausalLMLoss(logits_all[i], labels[..., i], vocab_size)
+        
+        # total_weight = sum(self.weights)
+        # normalized_weights = [w / total_weight for w in self.weights]
+        normalized_weights = self.weights
+        
+        total_loss = 0
+        for w, loss in zip(normalized_weights, loss_all):
+            total_loss += w * loss
+
+        if not return_dict:
+            output = (logits_all,) + outputs[1:]
+            return (total_loss, loss_all, ) + output if loss is not None else output
+
+        return AsteroidTTSOutputWithPast(
+            loss=total_loss,
+            logits=logits_all[0],
+            loss_all=loss_all,
+            logits_all=logits_all,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

requirements.txt

@@ -0,0 +1,16 @@
+torch>=2.0.0
+torchaudio>=2.0.0
+transformers>=4.30.0
+gradio>=4.0.0
+numpy>=1.21.0
+accelerate>=0.20.0
+PyPDF2
+beautifulsoup4
+soundfile
+librosa
+tqdm
+requests
+openai
+PyYAML
+einops
+huggingface_hub