aoiandroid/IndexTTS-Rust

1//! Configuration management for IndexTTS2 3use crate::{Error, Result};4use serde::{Deserialize, Serialize};5use std::path::{Path, PathBuf};6 7/// Main configuration for IndexTTS8#[derive(Debug, Clone, Serialize, Deserialize)]9pub struct Config {10    /// GPT model configuration11    pub gpt: GptConfig,12    /// Vocoder configuration13    pub vocoder: VocoderConfig,14    /// Semantic-to-Mel configuration15    pub s2mel: S2MelConfig,16    /// Dataset/tokenizer configuration17    pub dataset: DatasetConfig,18    /// Emotion configuration19    pub emotions: EmotionConfig,20    /// General inference settings21    pub inference: InferenceConfig,22    /// Model paths23    pub model_dir: PathBuf,24}25 26/// GPT model architecture configuration27#[derive(Debug, Clone, Serialize, Deserialize)]28pub struct GptConfig {29    /// Number of transformer layers30    pub layers: usize,31    /// Model dimension32    pub model_dim: usize,33    /// Number of attention heads34    pub heads: usize,35    /// Maximum text tokens36    pub max_text_tokens: usize,37    /// Maximum mel tokens38    pub max_mel_tokens: usize,39    /// Stop token for mel generation40    pub stop_mel_token: usize,41    /// Start token for text42    pub start_text_token: usize,43    /// Start token for mel44    pub start_mel_token: usize,45    /// Number of mel codes46    pub num_mel_codes: usize,47    /// Number of text tokens in vocabulary48    pub num_text_tokens: usize,49}50 51/// Vocoder configuration52#[derive(Debug, Clone, Serialize, Deserialize)]53pub struct VocoderConfig {54    /// Model name/path55    pub name: String,56    /// Checkpoint path57    pub checkpoint: Option<PathBuf>,58    /// Use FP16 inference59    pub use_fp16: bool,60    /// Use DeepSpeed optimization61    pub use_deepspeed: bool,62}63 64/// Semantic-to-Mel model configuration65#[derive(Debug, Clone, Serialize, Deserialize)]66pub struct S2MelConfig {67    /// Checkpoint path68    pub checkpoint: PathBuf,69    /// Preprocessing parameters70    pub preprocess: PreprocessConfig,71}72 73/// Audio preprocessing configuration74#[derive(Debug, Clone, Serialize, Deserialize)]75pub struct PreprocessConfig {76    /// Sample rate77    pub sr: u32,78    /// FFT size79    pub n_fft: usize,80    /// Hop length81    pub hop_length: usize,82    /// Window length83    pub win_length: usize,84    /// Number of mel bands85    pub n_mels: usize,86    /// Minimum frequency for mel filterbank87    pub fmin: f32,88    /// Maximum frequency for mel filterbank89    pub fmax: f32,90}91 92/// Dataset and tokenizer configuration93#[derive(Debug, Clone, Serialize, Deserialize)]94pub struct DatasetConfig {95    /// BPE model path96    pub bpe_model: PathBuf,97    /// Vocabulary size98    pub vocab_size: usize,99}100 101/// Emotion control configuration102#[derive(Debug, Clone, Serialize, Deserialize)]103pub struct EmotionConfig {104    /// Number of emotion dimensions105    pub num_dims: usize,106    /// Values per dimension107    pub num: Vec<usize>,108    /// Emotion matrix path109    pub matrix_path: Option<PathBuf>,110}111 112/// General inference configuration113#[derive(Debug, Clone, Serialize, Deserialize)]114pub struct InferenceConfig {115    /// Device to use (cpu, cuda:0, etc.)116    pub device: String,117    /// Use FP16 precision118    pub use_fp16: bool,119    /// Batch size120    pub batch_size: usize,121    /// Top-k sampling parameter122    pub top_k: usize,123    /// Top-p (nucleus) sampling parameter124    pub top_p: f32,125    /// Temperature for sampling126    pub temperature: f32,127    /// Repetition penalty128    pub repetition_penalty: f32,129    /// Length penalty130    pub length_penalty: f32,131}132 133impl Default for Config {134    fn default() -> Self {135        Self {136            gpt: GptConfig::default(),137            vocoder: VocoderConfig::default(),138            s2mel: S2MelConfig::default(),139            dataset: DatasetConfig::default(),140            emotions: EmotionConfig::default(),141            inference: InferenceConfig::default(),142            model_dir: PathBuf::from("models"),143        }144    }145}146 147impl Default for GptConfig {148    fn default() -> Self {149        Self {150            layers: 8,151            model_dim: 512,152            heads: 8,153            max_text_tokens: 120,154            max_mel_tokens: 250,155            stop_mel_token: 8193,156            start_text_token: 8192,157            start_mel_token: 8192,158            num_mel_codes: 8194,159            num_text_tokens: 6681,160        }161    }162}163 164impl Default for VocoderConfig {165    fn default() -> Self {166        Self {167            name: "bigvgan_v2_22khz_80band_256x".into(),168            checkpoint: None,169            use_fp16: true,170            use_deepspeed: false,171        }172    }173}174 175impl Default for S2MelConfig {176    fn default() -> Self {177        Self {178            checkpoint: PathBuf::from("models/s2mel.onnx"),179            preprocess: PreprocessConfig::default(),180        }181    }182}183 184impl Default for PreprocessConfig {185    fn default() -> Self {186        Self {187            sr: 22050,188            n_fft: 1024,189            hop_length: 256,190            win_length: 1024,191            n_mels: 80,192            fmin: 0.0,193            fmax: 8000.0,194        }195    }196}197 198impl Default for DatasetConfig {199    fn default() -> Self {200        Self {201            bpe_model: PathBuf::from("models/bpe.model"),202            vocab_size: 6681,203        }204    }205}206 207impl Default for EmotionConfig {208    fn default() -> Self {209        Self {210            num_dims: 8,211            num: vec![5, 6, 8, 6, 5, 4, 7, 6],212            matrix_path: Some(PathBuf::from("models/emotion_matrix.safetensors")),213        }214    }215}216 217impl Default for InferenceConfig {218    fn default() -> Self {219        Self {220            device: "cpu".into(),221            use_fp16: false,222            batch_size: 1,223            top_k: 50,224            top_p: 0.95,225            temperature: 1.0,226            repetition_penalty: 1.0,227            length_penalty: 1.0,228        }229    }230}231 232impl Config {233    /// Load configuration from YAML file234    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self> {235        let path = path.as_ref();236        if !path.exists() {237            return Err(Error::FileNotFound(path.display().to_string()));238        }239 240        let content = std::fs::read_to_string(path)?;241        let config: Config = serde_yaml::from_str(&content)?;242        Ok(config)243    }244 245    /// Save configuration to YAML file246    pub fn save<P: AsRef<Path>>(&self, path: P) -> Result<()> {247        let content = serde_yaml::to_string(self)248            .map_err(|e| Error::Config(format!("Failed to serialize config: {}", e)))?;249        std::fs::write(path, content)?;250        Ok(())251    }252 253    /// Load configuration from JSON file254    pub fn load_json<P: AsRef<Path>>(path: P) -> Result<Self> {255        let path = path.as_ref();256        if !path.exists() {257            return Err(Error::FileNotFound(path.display().to_string()));258        }259 260        let content = std::fs::read_to_string(path)?;261        let config: Config = serde_json::from_str(&content)?;262        Ok(config)263    }264 265    /// Create default configuration and save to file266    pub fn create_default<P: AsRef<Path>>(path: P) -> Result<Self> {267        let config = Config::default();268        config.save(path)?;269        Ok(config)270    }271 272    /// Validate the configuration273    pub fn validate(&self) -> Result<()> {274        // Check model directory exists275        if !self.model_dir.exists() {276            log::warn!(277                "Model directory does not exist: {}",278                self.model_dir.display()279            );280        }281 282        // Validate GPT config283        if self.gpt.layers == 0 {284            return Err(Error::Config("GPT layers must be > 0".into()));285        }286        if self.gpt.model_dim == 0 {287            return Err(Error::Config("GPT model_dim must be > 0".into()));288        }289        if self.gpt.heads == 0 {290            return Err(Error::Config("GPT heads must be > 0".into()));291        }292        if !self.gpt.model_dim.is_multiple_of(self.gpt.heads) {293            return Err(Error::Config(294                "GPT model_dim must be divisible by heads".into(),295            ));296        }297 298        // Validate preprocessing299        if self.s2mel.preprocess.sr == 0 {300            return Err(Error::Config("Sample rate must be > 0".into()));301        }302        if self.s2mel.preprocess.n_fft == 0 {303            return Err(Error::Config("n_fft must be > 0".into()));304        }305        if self.s2mel.preprocess.hop_length == 0 {306            return Err(Error::Config("hop_length must be > 0".into()));307        }308 309        // Validate inference settings310        if self.inference.temperature <= 0.0 {311            return Err(Error::Config("Temperature must be > 0".into()));312        }313        if self.inference.top_p <= 0.0 || self.inference.top_p > 1.0 {314            return Err(Error::Config("top_p must be in (0, 1]".into()));315        }316 317        Ok(())318    }319}320