voxtral_micro/
lib.rs

1//! # Voxtral Micro
2//!
3//! A minimal text-to-speech library using Q4-quantized GGUF models.
4//!
5//! ## Example
6//!
7//! ```rust,no_run
8//! use voxtral_micro::TtsEngine;
9//!
10//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
11//! let mut tts = TtsEngine::new("models/voxtral-tts-q4.gguf").await?;
12//!
13//! let audio = tts.synthesize("Hello world!", None)?;
14//! tts.save_wav("output.wav", &audio)?;
15//! # Ok(())
16//! # }
17//! ```
18
19pub mod audio;
20#[cfg(feature = "wgpu")]
21pub mod gguf;
22pub mod models;
23pub mod tokenizer;
24pub mod tts;
25
26use anyhow::{bail, Context, Result};
27use burn::backend::wgpu::WgpuDevice;
28use burn::backend::Wgpu;
29use burn::tensor::Tensor;
30use std::path::{Path, PathBuf};
31
32// Speed scaling removed: 1.0 = normal speed (no pitch shift compensation)
33use tokenizer::TekkenEncoder;
34
35/// Main TTS engine for speech synthesis from GGUF models.
36pub struct TtsEngine {
37    backbone: gguf::tts_model::Q4TtsBackbone,
38    fm: gguf::tts_model::Q4FmTransformer,
39    codec: tts::codec::CodecDecoder<Wgpu>,
40    tokenizer: TekkenEncoder,
41    voices_dir: PathBuf,
42    device: WgpuDevice,
43    max_frames: usize,
44}
45
46impl TtsEngine {
47    /// Create a new TTS engine from a GGUF model file.
48    ///
49    /// # Arguments
50    /// * `gguf_path` - Path to the Q4 GGUF model file
51    ///
52    /// # Returns
53    /// A TTS engine ready for synthesis, or an error if model loading fails.
54    ///
55    /// # Example
56    /// ```rust,no_run
57    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
58    /// use voxtral_tts::TtsEngine;
59    /// let tts = TtsEngine::new("models/voxtral-tts-q4.gguf").await?;
60    /// # Ok(())
61    /// # }
62    /// ```
63    pub async fn new<P: AsRef<Path>>(gguf_path: P) -> Result<Self> {
64        Self::with_options(gguf_path, None, None).await
65    }
66
67    /// Create a new TTS engine with custom options.
68    ///
69    /// # Arguments
70    /// * `gguf_path` - Path to the Q4 GGUF model file
71    /// * `tokenizer_path` - Optional path to tokenizer JSON (auto-discovered if None)
72    /// * `voices_dir` - Optional path to voices directory (defaults to models/voxtral-tts/voice_embedding)
73    pub async fn with_options<P: AsRef<Path>>(
74        gguf_path: P,
75        tokenizer_path: Option<P>,
76        voices_dir: Option<P>,
77    ) -> Result<Self> {
78        let start = std::time::Instant::now();
79        let gguf_path = gguf_path.as_ref();
80        if !gguf_path.exists() {
81            bail!("GGUF model not found at {}", gguf_path.display());
82        }
83
84        // Resolve tokenizer path
85        let tokenizer_path = match tokenizer_path {
86            Some(p) => p.as_ref().to_path_buf(),
87            None => {
88                let gguf_dir = gguf_path
89                    .parent()
90                    .unwrap_or(&PathBuf::from("."))
91                    .to_path_buf();
92                let candidates = [
93                    gguf_dir.join("tekken.json"),
94                    PathBuf::from("models/tekken.json"),
95                    PathBuf::from("models/voxtral-tts/tekken.json"),
96                ];
97                candidates
98                    .into_iter()
99                    .find(|p| p.exists())
100                    .ok_or_else(|| {
101                        anyhow::anyhow!(
102                            "Tokenizer not found. Provide tokenizer_path or place tekken.json alongside the GGUF file"
103                        )
104                    })?
105            }
106        };
107
108        if !tokenizer_path.exists() {
109            bail!("Tokenizer not found at {}", tokenizer_path.display());
110        }
111
112        let tokenizer =
113            TekkenEncoder::from_file(&tokenizer_path).context("Failed to load tokenizer")?;
114
115        // Resolve voices directory
116        let voices_dir = match voices_dir {
117            Some(d) => d.as_ref().to_path_buf(),
118            None => PathBuf::from("models/voice_embedding"),
119        };
120
121        // Load GGUF model
122        let device = WgpuDevice::default();
123        tracing::info!("Loading Q4 TTS model from {}", gguf_path.display());
124        let load_start = std::time::Instant::now();
125        let mut loader = gguf::Q4TtsModelLoader::from_file(gguf_path)
126            .context("Failed to open GGUF")?;
127        let (backbone, fm, codec) = loader.load(&device).context("Failed to load Q4 model")?;
128        tracing::info!("Model loaded in {:.2}s", load_start.elapsed().as_secs_f32());
129
130        let total_time = start.elapsed().as_secs_f32();
131        tracing::info!("TTS engine initialized in {:.2}s", total_time);
132
133        Ok(Self {
134            backbone,
135            fm,
136            codec,
137            tokenizer,
138            voices_dir,
139            device,
140            max_frames: 2000,
141        })
142    }
143
144    /// Synthesize speech from text using the default voice (casual_female).
145    ///
146    /// # Arguments
147    /// * `text` - The text to synthesize
148    /// * `voice` - Optional voice name (defaults to "casual_female")
149    ///
150    /// # Returns
151    /// An audio buffer containing the synthesized speech at 24kHz.
152    pub fn synthesize(&mut self, text: &str, voice: Option<&str>) -> Result<audio::AudioBuffer> {
153        self.synthesize_with_options(text, voice, 1.0, 1.0, None)
154    }
155
156    /// Synthesize speech with custom options.
157    ///
158    /// # Arguments
159    /// * `text` - The text to synthesize
160    /// * `voice` - Optional voice name (defaults to "casual_female")
161    /// * `speed` - Playback speed multiplier (0.5 to 3.0, where 1.0 is normal)
162    /// * `gain` - Volume gain multiplier (0.1 to 2.0, where 1.0 is normal)
163    /// * `language` - Optional language code (e.g., "en", "fr", "de")
164    ///
165    /// # Returns
166    /// An audio buffer containing the synthesized speech at 24kHz.
167    ///
168    /// # Speed Behavior
169    /// Speed adjustment uses resampling which changes both tempo and pitch.
170    /// Higher speeds result in higher pitch (chipmunk effect), lower speeds
171    /// result in lower pitch.
172    ///
173    /// # Example
174    /// ```no_run
175    /// # use voxtral_micro::TtsEngine;
176    /// # tokio_test::block_on(async {
177    /// let mut tts = TtsEngine::new("models/voxtral-tts-q4.gguf").await?;
178    /// let audio = tts.synthesize_with_options(
179    ///     "Hello world!",
180    ///     None,      // voice: None = default "casual_female"
181    ///     2.0,       // speed: 2.0 = twice as fast
182    ///     0.8,       // gain: 0.8 = 20% quieter
183    ///     Some("en") // language
184    /// )?;
185    /// # Ok::<(), anyhow::Error>(())
186    /// # });
187    /// ```
188    pub fn synthesize_with_options(
189        &mut self,
190        text: &str,
191        voice: Option<&str>,
192        speed: f32,
193        gain: f32,
194        language: Option<&str>,
195    ) -> Result<audio::AudioBuffer> {
196        let synthesis_start = std::time::Instant::now();
197        
198        // Validate parameters
199        if !(0.5..=3.0).contains(&speed) {
200            bail!("Speed must be between 0.5 and 3.0, got {}", speed);
201        }
202        if !(0.1..=2.0).contains(&gain) {
203            bail!("Gain must be between 0.1 and 2.0, got {}", gain);
204        }
205        let voice_name = voice.unwrap_or("casual_female");
206
207        // Tokenize text
208        let tokenize_start = std::time::Instant::now();
209        let token_ids = self.tokenizer.encode(text);
210        tracing::debug!("Tokenization: {:.3}s", tokenize_start.elapsed().as_secs_f32());
211        tracing::info!(
212            text_tokens = token_ids.len(),
213            voice = voice_name,
214            language = ?language,
215            "Synthesizing"
216        );
217
218        // Load voice embedding
219        let voice_path = self
220            .voices_dir
221            .join(format!("{}.safetensors", voice_name));
222        if !voice_path.exists() {
223            bail!(
224                "Voice '{}' not found at {}\n\
225                \n\
226                Voice embeddings are separate files that must be downloaded.\n\
227                Download with:\n\
228                  make download-models\n\
229                Or manually:\n\
230                  uv run --with huggingface_hub hf download \\\n\
231                    TrevorJS/voxtral-tts-q4-gguf \\\n\
232                    --local-dir models\n\
233                \n\
234                Then voices will be available at: models/voxtral-tts/voice_embedding/*.safetensors",
235                voice_name,
236                voice_path.display()
237            );
238        }
239
240        let voice_bytes = std::fs::read(&voice_path)?;
241        let voice_embed: Tensor<Wgpu, 2> = tts::voice::load_voice_from_bytes(
242            &voice_bytes,
243            3072,
244            &self.device,
245        )
246        .context("Failed to load voice")?;
247
248        tracing::info!(
249            voice = voice_name,
250            frames = voice_embed.dims()[0],
251            "Voice loaded"
252        );
253
254        // Build input sequence
255        let special = tts::config::TtsSpecialTokens::default();
256        let bos = self
257            .backbone
258            .embed_tokens_from_ids(&[special.bos_token_id as i32], 1, 1);
259        let begin_audio = self.backbone.embed_tokens_from_ids(
260            &[special.begin_audio_token_id as i32],
261            1,
262            1,
263        );
264        let next_audio_text = self.backbone.embed_tokens_from_ids(
265            &[special.next_audio_text_token_id as i32],
266            1,
267            1,
268        );
269        let repeat_audio_text = self.backbone.embed_tokens_from_ids(
270            &[special.repeat_audio_text_token_id as i32],
271            1,
272            1,
273        );
274        let text_ids_i32: Vec<i32> = token_ids.iter().map(|&id| id as i32).collect();
275        let text_embeds = self
276            .backbone
277            .embed_tokens_from_ids(&text_ids_i32, 1, text_ids_i32.len());
278
279        let input_sequence = Tensor::cat(
280            vec![
281                bos,
282                begin_audio.clone(),
283                voice_embed.unsqueeze_dim::<3>(0),
284                next_audio_text,
285                text_embeds,
286                repeat_audio_text,
287                begin_audio,
288            ],
289            1,
290        );
291
292        let codebook = tts::embeddings::AudioCodebookEmbeddings::new(
293            self.backbone.audio_codebook_embeddings().clone(),
294            tts::config::AudioCodebookLayout::default(),
295        );
296
297        // Generate audio frames
298        let gen_start = std::time::Instant::now();
299        let frames = pollster::block_on(self.backbone.generate_async(
300            input_sequence,
301            &self.fm,
302            &codebook,
303            self.max_frames,
304        ))
305        .map_err(|e| anyhow::anyhow!("Generation failed: {e}"))?;
306        tracing::info!("Frame generation: {:.2}s ({} frames)", gen_start.elapsed().as_secs_f32(), frames.len());
307
308        if frames.is_empty() {
309            bail!("No audio frames generated");
310        }
311
312        // Codec decode
313        let decode_start = std::time::Instant::now();
314        let n_frames = frames.len();
315        let semantic_indices: Vec<usize> = frames.iter().map(|f| f.semantic_idx).collect();
316        let mut acoustic_data = Vec::with_capacity(n_frames * 36);
317        for frame in &frames {
318            for &level in &frame.acoustic_levels {
319                acoustic_data.push(level as f32);
320            }
321        }
322        let acoustic_tensor: Tensor<Wgpu, 2> = Tensor::from_data(
323            burn::tensor::TensorData::new(acoustic_data, [n_frames, 36]),
324            &self.device,
325        );
326        let waveform = self.codec.decode(&semantic_indices, acoustic_tensor);
327        let [_batch, total_samples] = waveform.dims();
328        tracing::info!("Codec decode: {:.2}s", decode_start.elapsed().as_secs_f32());
329
330        let postprocess_start = std::time::Instant::now();
331        let wav_data = waveform.to_data();
332        let mut samples: Vec<f32> = wav_data.as_slice::<f32>().unwrap()[..total_samples].to_vec();
333
334        // Normalize to 0.95 peak
335        let peak = samples.iter().map(|s| s.abs()).fold(0.0f32, f32::max);
336        if peak > 1e-6 {
337            let gain = 0.95 / peak;
338            for s in &mut samples {
339                *s *= gain;
340            }
341        }
342
343        let mut audio = audio::AudioBuffer::new(samples, 24000);
344        
345        // Apply speed adjustment (1.0 = normal speed)
346        if (speed - 1.0).abs() > 0.001 {
347            audio = audio.with_speed(speed);
348            tracing::debug!(speed = speed, "Speed adjusted");
349        }
350        
351        // Apply gain adjustment
352        if (gain - 1.0).abs() > 0.001 {
353            audio = audio.with_gain(gain);
354        }
355        
356        let duration = audio.len() as f64 / audio.sample_rate as f64;
357        tracing::debug!("Post-processing: {:.3}s", postprocess_start.elapsed().as_secs_f32());
358        
359        let total_synthesis = synthesis_start.elapsed().as_secs_f32();
360        tracing::info!(
361            frames = n_frames,
362            duration_sec = format!("{duration:.2}"),
363            speed = speed,
364            gain = gain,
365            total_time_sec = format!("{total_synthesis:.2}"),
366            "Audio generated"
367        );
368
369        Ok(audio)
370    }
371
372    /// List available voice presets in the voices directory.
373    ///
374    /// # Returns
375    /// A vector of voice names, or an error if the directory doesn't exist.
376    pub fn list_voices(&self) -> Result<Vec<String>> {
377        if !self.voices_dir.exists() {
378            bail!(
379                "Voices directory not found at {}\n\
380                \n\
381                Voice embeddings must be downloaded.\n\
382                Download with:\n\
383                  make download-models\n\
384                Or manually:\n\
385                  uv run --with huggingface_hub hf download \\\n\
386                    TrevorJS/voxtral-tts-q4-gguf \\\n\
387                    --local-dir models",
388                self.voices_dir.display()
389            );
390        }
391
392        let mut voices: Vec<String> = std::fs::read_dir(&self.voices_dir)?
393            .filter_map(|e| e.ok())
394            .filter(|e| {
395                e.path()
396                    .extension()
397                    .is_some_and(|ext| ext == "safetensors")
398            })
399            .filter_map(|e| {
400                e.path()
401                    .file_stem()
402                    .map(|s| s.to_string_lossy().into_owned())
403            })
404            .collect();
405        voices.sort();
406        Ok(voices)
407    }
408
409    /// Set the maximum number of audio frames to generate.
410    ///
411    /// Default is 2000 frames. Higher values allow longer audio generation.
412    pub fn set_max_frames(&mut self, max_frames: usize) {
413        self.max_frames = max_frames;
414    }
415
416    /// Set the number of Euler ODE steps for flow matching (quality vs speed tradeoff).
417    ///
418    /// - 3 steps: Real-time performance
419    /// - 4 steps: Balanced (default)
420    /// - 8 steps: Higher quality
421    pub fn set_euler_steps(&mut self, steps: usize) {
422        self.fm.set_euler_steps(steps);
423    }
424
425    /// Save audio buffer to a WAV file.
426    ///
427    /// # Arguments
428    /// * `path` - Output file path
429    /// * `audio` - Audio buffer to save
430    pub fn save_wav<P: AsRef<Path>>(&self, path: P, audio: &audio::AudioBuffer) -> Result<()> {
431        audio.save(path)
432    }
433}
434
435// Re-export commonly used types
436pub use audio::AudioBuffer;
voxtral_micro/lib.rs

voxtral_micro/
lib.rs
