voxtral_micro/tts/
config.rs

1//! TTS configuration structs.
2//!
3//! Defines configs for the three TTS pipeline stages (backbone, flow-matching
4//! transformer, codec decoder) plus special token IDs and voice embedding metadata.
5
6use serde::{Deserialize, Serialize};
7
8/// Decoder backbone configuration (Ministral 3B architecture).
9///
10/// Identical to the ASR decoder except: no ADA RMSNorm, no sliding window
11/// specified in weights, and RoPE theta = 1M.
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13pub struct TtsBackboneConfig {
14    /// Number of transformer layers.
15    #[serde(default = "default_backbone_n_layers")]
16    pub n_layers: usize,
17    /// Model hidden dimension.
18    #[serde(default = "default_backbone_dim")]
19    pub dim: usize,
20    /// Number of query attention heads.
21    #[serde(default = "default_backbone_n_heads")]
22    pub n_heads: usize,
23    /// Number of KV attention heads (GQA).
24    #[serde(default = "default_backbone_n_kv_heads")]
25    pub n_kv_heads: usize,
26    /// Per-head dimension.
27    #[serde(default = "default_backbone_head_dim")]
28    pub head_dim: usize,
29    /// SwiGLU FFN hidden dimension.
30    #[serde(default = "default_backbone_ffn_dim")]
31    pub ffn_dim: usize,
32    /// RoPE theta for positional encoding.
33    #[serde(default = "default_backbone_rope_theta")]
34    pub rope_theta: f64,
35    /// Vocabulary size (Tekken tokenizer).
36    #[serde(default = "default_vocab_size")]
37    pub vocab_size: usize,
38    /// Whether token embeddings are tied with the LM head.
39    #[serde(default = "default_true")]
40    pub tied_embeddings: bool,
41    /// RMSNorm epsilon.
42    #[serde(default = "default_norm_eps")]
43    pub norm_eps: f64,
44}
45
46impl Default for TtsBackboneConfig {
47    fn default() -> Self {
48        Self {
49            n_layers: default_backbone_n_layers(),
50            dim: default_backbone_dim(),
51            n_heads: default_backbone_n_heads(),
52            n_kv_heads: default_backbone_n_kv_heads(),
53            head_dim: default_backbone_head_dim(),
54            ffn_dim: default_backbone_ffn_dim(),
55            rope_theta: default_backbone_rope_theta(),
56            vocab_size: default_vocab_size(),
57            tied_embeddings: true,
58            norm_eps: default_norm_eps(),
59        }
60    }
61}
62
63impl TtsBackboneConfig {
64    /// GQA group size (queries per KV head).
65    pub fn gqa_groups(&self) -> usize {
66        self.n_heads / self.n_kv_heads
67    }
68
69    /// Validate config invariants.
70    pub fn validate(&self) -> Result<(), ConfigError> {
71        if self.n_layers == 0 {
72            return Err(ConfigError::InvalidValue("n_layers must be > 0".into()));
73        }
74        if self.dim == 0 {
75            return Err(ConfigError::InvalidValue("dim must be > 0".into()));
76        }
77        if self.n_heads == 0 || self.n_kv_heads == 0 {
78            return Err(ConfigError::InvalidValue(
79                "n_heads and n_kv_heads must be > 0".into(),
80            ));
81        }
82        if !self.n_heads.is_multiple_of(self.n_kv_heads) {
83            return Err(ConfigError::InvalidValue(
84                "n_heads must be divisible by n_kv_heads".into(),
85            ));
86        }
87        if !self.dim.is_multiple_of(self.n_heads) {
88            return Err(ConfigError::InvalidValue(
89                "dim must be divisible by n_heads".into(),
90            ));
91        }
92        Ok(())
93    }
94}
95
96/// Flow-matching transformer configuration.
97///
98/// A small bidirectional (non-causal) transformer that predicts acoustic tokens
99/// from backbone hidden states via an Euler ODE solver.
100#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
101pub struct FmTransformerConfig {
102    /// Number of transformer layers.
103    #[serde(default = "default_fm_n_layers")]
104    pub n_layers: usize,
105    /// Model hidden dimension.
106    #[serde(default = "default_backbone_dim")]
107    pub dim: usize,
108    /// Number of query attention heads.
109    #[serde(default = "default_backbone_n_heads")]
110    pub n_heads: usize,
111    /// Number of KV attention heads (GQA).
112    #[serde(default = "default_backbone_n_kv_heads")]
113    pub n_kv_heads: usize,
114    /// Per-head dimension.
115    #[serde(default = "default_backbone_head_dim")]
116    pub head_dim: usize,
117    /// SwiGLU FFN hidden dimension.
118    #[serde(default = "default_backbone_ffn_dim")]
119    pub ffn_dim: usize,
120    /// RoPE theta (10K, different from backbone's 1M).
121    #[serde(default = "default_fm_rope_theta")]
122    pub rope_theta: f64,
123    /// RMSNorm epsilon.
124    #[serde(default = "default_norm_eps")]
125    pub norm_eps: f64,
126    /// Acoustic state dimensionality (FSQ 36 dims).
127    #[serde(default = "default_acoustic_dim")]
128    pub acoustic_dim: usize,
129    /// Semantic codebook output size (8192 VQ + 128 specials).
130    #[serde(default = "default_semantic_output_size")]
131    pub semantic_output_size: usize,
132    /// Number of Euler ODE steps.
133    #[serde(default = "default_euler_steps")]
134    pub euler_steps: usize,
135    /// Classifier-free guidance scale.
136    #[serde(default = "default_cfg_alpha")]
137    pub cfg_alpha: f32,
138}
139
140impl Default for FmTransformerConfig {
141    fn default() -> Self {
142        Self {
143            n_layers: default_fm_n_layers(),
144            dim: default_backbone_dim(),
145            n_heads: default_backbone_n_heads(),
146            n_kv_heads: default_backbone_n_kv_heads(),
147            head_dim: default_backbone_head_dim(),
148            ffn_dim: default_backbone_ffn_dim(),
149            rope_theta: default_fm_rope_theta(),
150            norm_eps: default_norm_eps(),
151            acoustic_dim: default_acoustic_dim(),
152            semantic_output_size: default_semantic_output_size(),
153            euler_steps: default_euler_steps(),
154            cfg_alpha: default_cfg_alpha(),
155        }
156    }
157}
158
159impl FmTransformerConfig {
160    /// Validate config invariants.
161    pub fn validate(&self) -> Result<(), ConfigError> {
162        if self.n_layers == 0 {
163            return Err(ConfigError::InvalidValue("n_layers must be > 0".into()));
164        }
165        if self.dim == 0 {
166            return Err(ConfigError::InvalidValue("dim must be > 0".into()));
167        }
168        if !self.n_heads.is_multiple_of(self.n_kv_heads) {
169            return Err(ConfigError::InvalidValue(
170                "n_heads must be divisible by n_kv_heads".into(),
171            ));
172        }
173        if self.euler_steps == 0 {
174            return Err(ConfigError::InvalidValue("euler_steps must be > 0".into()));
175        }
176        Ok(())
177    }
178}
179
180/// Codec decoder configuration.
181///
182/// Conv-transformer autoencoder that converts tokens back to a 24 kHz waveform.
183/// Uses ALiBi positional bias, QK-norm, LayerScale, and weight-normed convolutions.
184#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
185pub struct CodecDecoderConfig {
186    /// Hidden dimension for transformer layers.
187    #[serde(default = "default_codec_dim")]
188    pub dim: usize,
189    /// Number of MHA heads (not GQA).
190    #[serde(default = "default_codec_n_heads")]
191    pub n_heads: usize,
192    /// Per-head dimension.
193    #[serde(default = "default_codec_head_dim")]
194    pub head_dim: usize,
195    /// SwiGLU FFN hidden dimension.
196    #[serde(default = "default_codec_ffn_dim")]
197    pub ffn_dim: usize,
198    /// Number of transformer layers per block.
199    #[serde(default = "default_codec_layers_per_block")]
200    pub layers_per_block: usize,
201    /// Sliding window sizes for the 4 transformer block groups.
202    #[serde(default = "default_codec_sliding_windows")]
203    pub sliding_windows: Vec<usize>,
204    /// Input conv: channels in (semantic 256 + acoustic 36 = 292).
205    #[serde(default = "default_codec_input_channels")]
206    pub input_channels: usize,
207    /// Output conv: samples per patch.
208    #[serde(default = "default_codec_output_patch_size")]
209    pub output_patch_size: usize,
210    /// Output sample rate in Hz.
211    #[serde(default = "default_codec_sample_rate")]
212    pub sample_rate: u32,
213    /// Number of semantic VQ codebook entries.
214    #[serde(default = "default_semantic_vq_size")]
215    pub semantic_vq_size: usize,
216    /// Semantic embedding dimension (per entry).
217    #[serde(default = "default_semantic_embed_dim")]
218    pub semantic_embed_dim: usize,
219    /// Number of acoustic FSQ dimensions.
220    #[serde(default = "default_acoustic_dim")]
221    pub acoustic_fsq_dims: usize,
222    /// Number of FSQ levels per dimension.
223    #[serde(default = "default_fsq_levels")]
224    pub fsq_levels: usize,
225    /// RMSNorm epsilon (codec uses 0.01, much larger than backbone's 1e-5).
226    #[serde(default = "default_codec_norm_eps")]
227    pub norm_eps: f64,
228    /// QK-norm epsilon for codec attention.
229    #[serde(default = "default_codec_qk_norm_eps")]
230    pub qk_norm_eps: f64,
231}
232
233impl Default for CodecDecoderConfig {
234    fn default() -> Self {
235        Self {
236            dim: default_codec_dim(),
237            n_heads: default_codec_n_heads(),
238            head_dim: default_codec_head_dim(),
239            ffn_dim: default_codec_ffn_dim(),
240            layers_per_block: default_codec_layers_per_block(),
241            sliding_windows: default_codec_sliding_windows(),
242            input_channels: default_codec_input_channels(),
243            output_patch_size: default_codec_output_patch_size(),
244            sample_rate: default_codec_sample_rate(),
245            semantic_vq_size: default_semantic_vq_size(),
246            semantic_embed_dim: default_semantic_embed_dim(),
247            acoustic_fsq_dims: default_acoustic_dim(),
248            fsq_levels: default_fsq_levels(),
249            norm_eps: default_codec_norm_eps(),
250            qk_norm_eps: default_codec_qk_norm_eps(),
251        }
252    }
253}
254
255impl CodecDecoderConfig {
256    /// Total number of transformer blocks (4 groups of `layers_per_block`).
257    pub fn total_transformer_layers(&self) -> usize {
258        self.sliding_windows.len() * self.layers_per_block
259    }
260
261    /// Validate config invariants.
262    pub fn validate(&self) -> Result<(), ConfigError> {
263        if self.dim == 0 {
264            return Err(ConfigError::InvalidValue("dim must be > 0".into()));
265        }
266        if self.sliding_windows.is_empty() {
267            return Err(ConfigError::InvalidValue(
268                "sliding_windows must not be empty".into(),
269            ));
270        }
271        if self.n_heads == 0 {
272            return Err(ConfigError::InvalidValue("n_heads must be > 0".into()));
273        }
274        if !self.dim.is_multiple_of(self.n_heads) {
275            return Err(ConfigError::InvalidValue(
276                "dim must be divisible by n_heads".into(),
277            ));
278        }
279        Ok(())
280    }
281}
282
283/// Special token IDs for TTS sequence construction.
284///
285/// Input sequence format (per vLLM/mistral-common reference):
286/// ```text
287/// [BOS(1)] [BEGIN_AUDIO(25)] [voice_0..N] [NEXT_AUDIO_TEXT(35)] [text_0..M] [REPEAT_AUDIO_TEXT(36)] [BEGIN_AUDIO(25)]
288/// ```
289#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
290pub struct TtsSpecialTokens {
291    /// Beginning of sequence.
292    pub bos_token_id: u32,
293    /// Marks start of audio section (token 25).
294    pub begin_audio_token_id: u32,
295    /// Transition from audio to text (token 35, `[NEXT_AUDIO_TEXT]`).
296    pub next_audio_text_token_id: u32,
297    /// Transition from text back to audio (token 36, `[REPEAT_AUDIO_TEXT]`).
298    pub repeat_audio_text_token_id: u32,
299    /// Audio placeholder token (token 24, replaced by voice embeddings at embedding level).
300    pub audio_token_id: u32,
301    /// Empty audio sentinel in the audio codebook (index 0 per codebook).
302    pub empty_audio_idx: u32,
303    /// End-of-audio sentinel in the audio codebook (index 1 per codebook).
304    pub end_audio_idx: u32,
305}
306
307impl Default for TtsSpecialTokens {
308    fn default() -> Self {
309        Self {
310            bos_token_id: 1,
311            begin_audio_token_id: 25,       // [BEGIN_AUDIO] = rank 25
312            next_audio_text_token_id: 36,   // [NEXT_AUDIO_TEXT] = rank 36 (text→audio transition)
313            repeat_audio_text_token_id: 35, // [REPEAT_AUDIO_TEXT] = rank 35 (audio→text transition)
314            audio_token_id: 24,             // [AUDIO] = rank 24 (placeholder)
315            empty_audio_idx: 0,
316            end_audio_idx: 1,
317        }
318    }
319}
320
321/// Voice embedding metadata for preset voices.
322#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
323pub struct VoiceEmbeddingConfig {
324    /// Expected embedding dimension (must match backbone dim).
325    #[serde(default = "default_backbone_dim")]
326    pub embed_dim: usize,
327    /// Known preset voice names. An empty list means accept any.
328    #[serde(default = "default_voice_presets")]
329    pub preset_names: Vec<String>,
330}
331
332impl Default for VoiceEmbeddingConfig {
333    fn default() -> Self {
334        Self {
335            embed_dim: default_backbone_dim(),
336            preset_names: default_voice_presets(),
337        }
338    }
339}
340
341/// Audio codebook embedding layout constants.
342#[derive(Debug, Clone, Copy, PartialEq, Eq)]
343pub struct AudioCodebookLayout {
344    /// Number of semantic VQ entries.
345    pub semantic_vq_size: usize,
346    /// Number of acoustic FSQ dimensions (codebooks).
347    pub acoustic_codebooks: usize,
348    /// Number of FSQ levels per acoustic codebook.
349    pub fsq_levels: usize,
350    /// Number of special tokens per codebook (EMPTY_AUDIO, END_AUDIO).
351    pub specials_per_codebook: usize,
352}
353
354impl Default for AudioCodebookLayout {
355    fn default() -> Self {
356        Self {
357            semantic_vq_size: 8192,
358            acoustic_codebooks: 36,
359            fsq_levels: 21,
360            specials_per_codebook: 2,
361        }
362    }
363}
364
365impl AudioCodebookLayout {
366    /// Stride per acoustic codebook (specials + levels).
367    pub fn acoustic_stride(&self) -> usize {
368        self.specials_per_codebook + self.fsq_levels
369    }
370
371    /// Start index of acoustic codebook region in the embedding table.
372    pub fn acoustic_region_start(&self) -> usize {
373        self.specials_per_codebook + self.semantic_vq_size
374    }
375
376    /// Total meaningful entries in the embedding table.
377    pub fn total_entries(&self) -> usize {
378        // 2 semantic specials + 8192 semantic VQ + 36 * 23 acoustic
379        self.specials_per_codebook
380            + self.semantic_vq_size
381            + self.acoustic_codebooks * self.acoustic_stride()
382    }
383
384    /// Global index for a semantic token value.
385    pub fn semantic_global_index(&self, raw_semantic_idx: usize) -> usize {
386        raw_semantic_idx + self.specials_per_codebook
387    }
388
389    /// Global index for an acoustic codebook level.
390    pub fn acoustic_global_index(&self, codebook: usize, level: usize) -> usize {
391        self.acoustic_region_start()
392            + codebook * self.acoustic_stride()
393            + level
394            + self.specials_per_codebook
395    }
396}
397
398/// Top-level TTS configuration combining all stage configs.
399#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
400pub struct TtsConfig {
401    pub backbone: TtsBackboneConfig,
402    pub fm_transformer: FmTransformerConfig,
403    pub codec_decoder: CodecDecoderConfig,
404    pub special_tokens: TtsSpecialTokens,
405    pub voice: VoiceEmbeddingConfig,
406}
407
408impl TtsConfig {
409    /// Validate all sub-configs.
410    pub fn validate(&self) -> Result<(), ConfigError> {
411        self.backbone.validate()?;
412        self.fm_transformer.validate()?;
413        self.codec_decoder.validate()?;
414
415        // Cross-config consistency: FM dim must match backbone dim
416        if self.fm_transformer.dim != self.backbone.dim {
417            return Err(ConfigError::InvalidValue(
418                "fm_transformer.dim must match backbone.dim".into(),
419            ));
420        }
421
422        // Voice embed dim must match backbone dim
423        if self.voice.embed_dim != self.backbone.dim {
424            return Err(ConfigError::InvalidValue(
425                "voice.embed_dim must match backbone.dim".into(),
426            ));
427        }
428
429        Ok(())
430    }
431}
432
433/// Configuration validation error.
434#[derive(Debug, Clone, thiserror::Error)]
435pub enum ConfigError {
436    #[error("invalid config value: {0}")]
437    InvalidValue(String),
438}
439
440// --- Default value functions ---
441
442fn default_backbone_n_layers() -> usize {
443    26
444}
445fn default_backbone_dim() -> usize {
446    3072
447}
448fn default_backbone_n_heads() -> usize {
449    32
450}
451fn default_backbone_n_kv_heads() -> usize {
452    8
453}
454fn default_backbone_head_dim() -> usize {
455    128
456}
457fn default_backbone_ffn_dim() -> usize {
458    9216
459}
460fn default_backbone_rope_theta() -> f64 {
461    1_000_000.0
462}
463fn default_vocab_size() -> usize {
464    131_072
465}
466fn default_norm_eps() -> f64 {
467    1e-5
468}
469fn default_codec_norm_eps() -> f64 {
470    0.01 // from params.json: audio_tokenizer_args.norm_eps
471}
472fn default_codec_qk_norm_eps() -> f64 {
473    1e-6 // from params.json: audio_tokenizer_args.qk_norm_eps
474}
475fn default_true() -> bool {
476    true
477}
478
479fn default_fm_n_layers() -> usize {
480    3
481}
482fn default_fm_rope_theta() -> f64 {
483    10_000.0
484}
485fn default_acoustic_dim() -> usize {
486    36
487}
488fn default_semantic_output_size() -> usize {
489    8320
490}
491fn default_euler_steps() -> usize {
492    5  // Balanced: good quality with reasonable speed (3=fast/artifacts, 8=best quality/slow)
493}
494fn default_cfg_alpha() -> f32 {
495    1.2
496}
497
498fn default_codec_dim() -> usize {
499    1024
500}
501fn default_codec_n_heads() -> usize {
502    8
503}
504fn default_codec_head_dim() -> usize {
505    128
506}
507fn default_codec_ffn_dim() -> usize {
508    4096
509}
510fn default_codec_layers_per_block() -> usize {
511    2
512}
513fn default_codec_sliding_windows() -> Vec<usize> {
514    vec![2, 4, 8, 16]
515}
516fn default_codec_input_channels() -> usize {
517    292
518}
519fn default_codec_output_patch_size() -> usize {
520    240
521}
522fn default_codec_sample_rate() -> u32 {
523    24_000
524}
525fn default_semantic_vq_size() -> usize {
526    8192
527}
528fn default_semantic_embed_dim() -> usize {
529    256
530}
531fn default_fsq_levels() -> usize {
532    21
533}
534
535fn default_voice_presets() -> Vec<String> {
536    [
537        "alloy",
538        "ash",
539        "ballad",
540        "breeze",
541        "casual_female",
542        "casual_male",
543        "coral",
544        "echo",
545        "fable",
546        "nova",
547        "onyx",
548        "professional_female",
549        "professional_male",
550        "sage",
551        "shimmer",
552        "spirit",
553        "verse",
554        "warm_female",
555        "warm_male",
556        "whisper",
557    ]
558    .iter()
559    .map(|s| s.to_string())
560    .collect()
561}
562
563#[cfg(test)]
564mod tests {
565    use super::*;
566
567    #[test]
568    fn test_backbone_defaults() {
569        let config = TtsBackboneConfig::default();
570        assert_eq!(config.n_layers, 26);
571        assert_eq!(config.dim, 3072);
572        assert_eq!(config.n_heads, 32);
573        assert_eq!(config.n_kv_heads, 8);
574        assert_eq!(config.head_dim, 128);
575        assert_eq!(config.ffn_dim, 9216);
576        assert_eq!(config.rope_theta, 1_000_000.0);
577        assert_eq!(config.vocab_size, 131_072);
578        assert!(config.tied_embeddings);
579        assert_eq!(config.gqa_groups(), 4);
580    }
581
582    #[test]
583    fn test_fm_transformer_defaults() {
584        let config = FmTransformerConfig::default();
585        assert_eq!(config.n_layers, 3);
586        assert_eq!(config.dim, 3072);
587        assert_eq!(config.n_heads, 32);
588        assert_eq!(config.n_kv_heads, 8);
589        assert_eq!(config.head_dim, 128);
590        assert_eq!(config.ffn_dim, 9216);
591        assert_eq!(config.rope_theta, 10_000.0);
592        assert_eq!(config.acoustic_dim, 36);
593        assert_eq!(config.semantic_output_size, 8320);
594        assert_eq!(config.euler_steps, 5); // Default balanced quality/speed
595        assert!((config.cfg_alpha - 1.2).abs() < 1e-6);
596    }
597
598    #[test]
599    fn test_codec_decoder_defaults() {
600        let config = CodecDecoderConfig::default();
601        assert_eq!(config.dim, 1024);
602        assert_eq!(config.n_heads, 8);
603        assert_eq!(config.head_dim, 128);
604        assert_eq!(config.layers_per_block, 2);
605        assert_eq!(config.sliding_windows, vec![2, 4, 8, 16]);
606        assert_eq!(config.input_channels, 292);
607        assert_eq!(config.output_patch_size, 240);
608        assert_eq!(config.sample_rate, 24_000);
609        assert_eq!(config.total_transformer_layers(), 8);
610    }
611
612    #[test]
613    fn test_special_tokens_defaults() {
614        let tokens = TtsSpecialTokens::default();
615        assert_eq!(tokens.audio_token_id, 24);
616        assert_eq!(tokens.begin_audio_token_id, 25);
617        assert_eq!(tokens.bos_token_id, 1);
618        assert_eq!(tokens.empty_audio_idx, 0);
619        assert_eq!(tokens.end_audio_idx, 1);
620    }
621
622    #[test]
623    fn test_voice_embedding_defaults() {
624        let config = VoiceEmbeddingConfig::default();
625        assert_eq!(config.embed_dim, 3072);
626        assert_eq!(config.preset_names.len(), 20);
627        assert!(config.preset_names.contains(&"casual_female".to_string()));
628        assert!(config.preset_names.contains(&"whisper".to_string()));
629    }
630
631    #[test]
632    fn test_audio_codebook_layout() {
633        let layout = AudioCodebookLayout::default();
634        assert_eq!(layout.acoustic_stride(), 23);
635        assert_eq!(layout.acoustic_region_start(), 8194);
636        // 2 + 8192 + 36*23 = 2 + 8192 + 828 = 9022
637        assert_eq!(layout.total_entries(), 9022);
638
639        // Semantic index: raw 0 -> global 2
640        assert_eq!(layout.semantic_global_index(0), 2);
641        assert_eq!(layout.semantic_global_index(8191), 8193);
642
643        // Acoustic: codebook 0, level 0 -> 8194 + 0*23 + 0 + 2 = 8196
644        assert_eq!(layout.acoustic_global_index(0, 0), 8196);
645        // Acoustic: codebook 35, level 20 -> 8194 + 35*23 + 20 + 2 = 8194 + 805 + 22 = 9021
646        assert_eq!(layout.acoustic_global_index(35, 20), 9021);
647    }
648
649    #[test]
650    fn test_tts_config_validates() {
651        let config = TtsConfig::default();
652        assert!(config.validate().is_ok());
653    }
654
655    #[test]
656    fn test_validation_rejects_zero_layers() {
657        let mut config = TtsBackboneConfig::default();
658        config.n_layers = 0;
659        assert!(config.validate().is_err());
660    }
661
662    #[test]
663    fn test_validation_rejects_mismatched_heads() {
664        let mut config = TtsBackboneConfig::default();
665        config.n_heads = 7;
666        config.n_kv_heads = 3;
667        assert!(config.validate().is_err());
668    }
669
670    #[test]
671    fn test_validation_rejects_dim_head_mismatch() {
672        let mut config = TtsBackboneConfig::default();
673        config.dim = 100;
674        config.n_heads = 32;
675        assert!(config.validate().is_err());
676    }
677
678    #[test]
679    fn test_validation_rejects_fm_zero_euler_steps() {
680        let mut config = FmTransformerConfig::default();
681        config.euler_steps = 0;
682        assert!(config.validate().is_err());
683    }
684
685    #[test]
686    fn test_validation_rejects_empty_sliding_windows() {
687        let mut config = CodecDecoderConfig::default();
688        config.sliding_windows = vec![];
689        assert!(config.validate().is_err());
690    }
691
692    #[test]
693    fn test_cross_config_validation_dim_mismatch() {
694        let mut config = TtsConfig::default();
695        config.fm_transformer.dim = 1024;
696        assert!(config.validate().is_err());
697    }
698
699    #[test]
700    fn test_cross_config_validation_voice_dim_mismatch() {
701        let mut config = TtsConfig::default();
702        config.voice.embed_dim = 512;
703        assert!(config.validate().is_err());
704    }
705
706    #[test]
707    fn test_serde_roundtrip_backbone() {
708        let config = TtsBackboneConfig::default();
709        let json = serde_json::to_string(&config).unwrap();
710        let parsed: TtsBackboneConfig = serde_json::from_str(&json).unwrap();
711        assert_eq!(config, parsed);
712    }
713
714    #[test]
715    fn test_serde_roundtrip_tts_config() {
716        let config = TtsConfig::default();
717        let json = serde_json::to_string(&config).unwrap();
718        let parsed: TtsConfig = serde_json::from_str(&json).unwrap();
719        assert_eq!(config, parsed);
720    }
721
722    #[test]
723    fn test_serde_partial_deserialize() {
724        // Should fill in defaults for missing fields
725        let json = r#"{"n_layers": 12, "dim": 768}"#;
726        let config: TtsBackboneConfig = serde_json::from_str(json).unwrap();
727        assert_eq!(config.n_layers, 12);
728        assert_eq!(config.dim, 768);
729        // Defaults for the rest
730        assert_eq!(config.n_heads, 32);
731        assert_eq!(config.vocab_size, 131_072);
732    }
733}
voxtral_micro/tts/config.rs

voxtral_micro/tts/
config.rs
