ThreadAbort/IndexTTS-Rust
77
1#!/usr/bin/env python32"""3Convert IndexTTS-2 PyTorch models to ONNX format for Rust inference!4 5This script converts the three main models:61. GPT model (gpt.pth) - Autoregressive text-to-semantic generation72. S2Mel model (s2mel.pth) - Semantic-to-mel spectrogram conversion83. BigVGAN - Mel-to-waveform vocoder (already available as ONNX from NVIDIA)9 10Usage:11 python tools/convert_to_onnx.py12 13Output:14 models/gpt.onnx15 models/s2mel.onnx16 models/bigvgan.onnx (if needed, otherwise use NVIDIA's)17 18Why ONNX?19 - Cross-platform: Works on Windows, Linux, macOS, M1/M2 Macs20 - Fast: ONNX Runtime is highly optimized21 - Rust-native: ort crate provides excellent ONNX Runtime bindings22 - No Python: Production inference without Python dependency hell!23 24Author: Aye & Hue @ 8b.is25"""26 27import os28import sys29 30# Setup paths31script_dir = os.path.dirname(os.path.abspath(__file__))32project_root = os.path.dirname(script_dir)33os.chdir(project_root)34 35# Set HF cache36os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache'37 38print("=" * 70)39print(" IndexTTS-2 PyTorch to ONNX Converter")40print(" For Rust inference with ort crate!")41print("=" * 70)42print()43 44# Check for models45if not os.path.exists("checkpoints/gpt.pth"):46 print("ERROR: Models not found!")47 print("Run: python tools/download_files.py -s huggingface")48 sys.exit(1)49 50import torch51import torch.onnx52import numpy as np53from pathlib import Path54 55# Add reference code to path56sys.path.insert(0, "indextts - REMOVING - REF ONLY")57 58# Create output directory59output_dir = Path("models")60output_dir.mkdir(exist_ok=True)61 62print(f"PyTorch version: {torch.__version__}")63print(f"Output directory: {output_dir}")64print()65 66 67def export_speaker_encoder():68 """69 Export the CAM++ speaker encoder to ONNX.70 71 This model extracts speaker embeddings from reference audio.72 Input: mel spectrogram [batch, n_mels, time]73 Output: speaker embedding [batch, 192]74 """75 print("\n" + "=" * 50)76 print("Exporting Speaker Encoder (CAM++)")77 print("=" * 50)78 79 try:80 from omegaconf import OmegaConf81 from indextts.s2mel.modules.campplus.DTDNN import CAMPPlus82 83 # Load config84 cfg = OmegaConf.load("checkpoints/config.yaml")85 86 # Create model87 model = CAMPPlus(feat_dim=80, embedding_size=192)88 89 # Load weights90 weights_path = "./checkpoints/hf_cache/models--funasr--campplus/snapshots/fb71fe990cbf6031ae6987a2d76fe64f94377b7e/campplus_cn_common.bin"91 if os.path.exists(weights_path):92 state_dict = torch.load(weights_path, map_location='cpu')93 model.load_state_dict(state_dict)94 print(f"Loaded weights from: {weights_path}")95 96 model.eval()97 98 # CAMPPlus expects [batch, time, n_mels] NOT [batch, n_mels, time]!99 # This is the key insight - the model processes time-series of mel features100 dummy_input = torch.randn(1, 100, 80) # [batch, time, features]101 102 # Verify forward pass works before export103 with torch.no_grad():104 test_output = model(dummy_input)105 print(f"Forward pass works! Output shape: {test_output.shape}")106 107 # Export to ONNX108 output_path = output_dir / "speaker_encoder.onnx"109 torch.onnx.export(110 model,111 dummy_input,112 str(output_path),113 input_names=['mel_spectrogram'],114 output_names=['speaker_embedding'],115 dynamic_axes={116 'mel_spectrogram': {0: 'batch', 1: 'time'}, # time is dim 1!117 'speaker_embedding': {0: 'batch'}118 },119 opset_version=18, # Use 18+ for latest features120 do_constant_folding=True,121 )122 123 # Verify the export124 import onnx125 onnx_model = onnx.load(str(output_path))126 onnx.checker.check_model(onnx_model)127 128 print(f"✓ Exported: {output_path}")129 print(f" Input: mel_spectrogram [batch, time, 80]") # Corrected!130 print(f" Output: speaker_embedding [batch, 192]")131 print(f"✓ ONNX model verified!")132 return True133 134 except Exception as e:135 print(f"✗ Failed to export speaker encoder: {e}")136 import traceback137 traceback.print_exc()138 return False139 140 141def export_gpt_model():142 """143 Export the GPT autoregressive model to ONNX.144 145 This is the most complex model - generates semantic tokens from text.146 We may need to export it in parts due to KV caching.147 148 Input: text_tokens [batch, seq_len], speaker_embedding [batch, 192]149 Output: semantic_codes [batch, code_len]150 """151 print("\n" + "=" * 50)152 print("Exporting GPT Model (Autoregressive)")153 print("=" * 50)154 155 try:156 from omegaconf import OmegaConf157 158 # Load the full model config159 cfg = OmegaConf.load("checkpoints/config.yaml")160 161 # This is tricky - GPT models with KV caching are hard to export162 # We might need to:163 # 1. Export just the forward pass without caching164 # 2. Or export separate encoder/decoder parts165 166 print("GPT model export is complex due to:")167 print(" - Autoregressive generation with KV caching")168 print(" - Dynamic sequence lengths")169 print(" - Multiple internal components")170 print()171 print("Options:")172 print(" A) Export without KV cache (slower but simpler)")173 print(" B) Export encoder + single-step decoder (efficient)")174 print(" C) Use torch.compile + ONNX tracing")175 print()176 177 # For now, let's try the simpler approach178 from infer_v2 import IndexTTS2179 180 # Load model181 tts = IndexTTS2(182 cfg_path="checkpoints/config.yaml",183 model_dir="checkpoints",184 use_fp16=False,185 device="cpu"186 )187 188 # Get the GPT component189 gpt = tts.gpt190 gpt.eval()191 192 print(f"GPT model loaded: {type(gpt)}")193 print(f"Parameters: {sum(p.numel() for p in gpt.parameters()):,}")194 195 # The GPT model architecture:196 # - Text encoder (embeddings + transformer)197 # - Speaker conditioning198 # - Autoregressive decoder199 200 # Let's export the text encoder first201 output_path = output_dir / "gpt_encoder.onnx"202 203 # Create dummy inputs204 text_tokens = torch.randint(0, 30000, (1, 32), dtype=torch.int64)205 206 # This will likely fail due to complex control flow207 # but let's try!208 print(f"Attempting GPT export (may require modifications)...")209 210 # For now, just report what we learned211 print()212 print("Note: Full GPT export requires modifying the model code")213 print("to remove dynamic control flow. Creating a wrapper...")214 215 return False216 217 except Exception as e:218 print(f"✗ Failed to export GPT: {e}")219 import traceback220 traceback.print_exc()221 return False222 223 224def export_s2mel_model():225 """226 Export the Semantic-to-Mel model (flow matching).227 228 This converts semantic codes to mel spectrograms.229 Input: semantic_codes [batch, code_len], speaker_embedding [batch, 192]230 Output: mel_spectrogram [batch, 80, mel_len]231 """232 print("\n" + "=" * 50)233 print("Exporting S2Mel Model (Flow Matching)")234 print("=" * 50)235 236 try:237 from omegaconf import OmegaConf238 239 cfg = OmegaConf.load("checkpoints/config.yaml")240 241 print("S2Mel model (Diffusion/Flow Matching) is also complex:")242 print(" - Multiple denoising steps (iterative)")243 print(" - CFM (Conditional Flow Matching) requires ODE solving")244 print()245 print("Export strategy:")246 print(" 1. Export the single denoising step")247 print(" 2. Run iteration loop in Rust")248 print()249 250 return False251 252 except Exception as e:253 print(f"✗ Failed to export S2Mel: {e}")254 import traceback255 traceback.print_exc()256 return False257 258 259def export_bigvgan():260 """261 Export BigVGAN vocoder to ONNX.262 263 Good news: NVIDIA provides pre-trained BigVGAN models!264 Even better: They're designed for easy ONNX export.265 266 Input: mel_spectrogram [batch, 80, mel_len]267 Output: waveform [batch, 1, wave_len]268 """269 print("\n" + "=" * 50)270 print("Exporting BigVGAN Vocoder")271 print("=" * 50)272 273 try:274 # BigVGAN from NVIDIA is easier to export275 # Let's check if we already have it276 277 print("BigVGAN options:")278 print(" 1. Use NVIDIA's pre-exported ONNX (recommended)")279 print(" https://github.com/NVIDIA/BigVGAN")280 print()281 print(" 2. Export from PyTorch weights (we'll do this)")282 print()283 284 # Try to load BigVGAN285 try:286 from bigvgan import bigvgan287 model = bigvgan.BigVGAN.from_pretrained(288 'nvidia/bigvgan_v2_22khz_80band_256x',289 use_cuda_kernel=False290 )291 model.eval()292 model.remove_weight_norm() # Important for ONNX!293 294 print(f"BigVGAN loaded from HuggingFace")295 print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")296 297 # Create dummy input298 dummy_mel = torch.randn(1, 80, 100)299 300 # Export301 output_path = output_dir / "bigvgan.onnx"302 torch.onnx.export(303 model,304 dummy_mel,305 str(output_path),306 input_names=['mel_spectrogram'],307 output_names=['waveform'],308 dynamic_axes={309 'mel_spectrogram': {0: 'batch', 2: 'mel_length'},310 'waveform': {0: 'batch', 2: 'wave_length'}311 },312 opset_version=18, # Use 18+ for latest features313 do_constant_folding=True,314 )315 316 print(f"✓ Exported: {output_path}")317 print(f" Input: mel_spectrogram [batch, 80, mel_len]")318 print(f" Output: waveform [batch, 1, wave_len]")319 320 # Verify the export321 import onnx322 onnx_model = onnx.load(str(output_path))323 onnx.checker.check_model(onnx_model)324 print(f"✓ ONNX model verified!")325 326 return True327 328 except ImportError:329 print("bigvgan package not installed, installing...")330 os.system("pip install bigvgan")331 print("Please re-run the script.")332 return False333 334 except Exception as e:335 print(f"✗ Failed to export BigVGAN: {e}")336 import traceback337 traceback.print_exc()338 return False339 340 341def main():342 print("\nStarting ONNX conversion...\n")343 344 results = {}345 346 # Export each component347 results['speaker_encoder'] = export_speaker_encoder()348 results['gpt'] = export_gpt_model()349 results['s2mel'] = export_s2mel_model()350 results['bigvgan'] = export_bigvgan()351 352 # Summary353 print("\n" + "=" * 70)354 print(" CONVERSION SUMMARY")355 print("=" * 70)356 357 for name, success in results.items():358 status = "✓ SUCCESS" if success else "✗ NEEDS WORK"359 print(f" {name:20} {status}")360 361 print()362 363 if all(results.values()):364 print("All models converted! Ready for Rust inference.")365 else:366 print("Some models need manual intervention.")367 print()368 print("For complex models (GPT, S2Mel), consider:")369 print(" 1. Modifying the Python code to remove dynamic control flow")370 print(" 2. Using torch.jit.trace with concrete inputs")371 print(" 3. Exporting subcomponents separately")372 print(" 4. Using ONNX Runtime's transformer optimizations")373 374 print()375 print("Output directory:", output_dir.absolute())376 377 378if __name__ == "__main__":379 main()380