voxtral-micro 1.0.0

//! # Voxtral Micro
//!
//! A minimal text-to-speech library using Q4-quantized GGUF models.
//!
//! ## Example
//!
//! ```rust,no_run
//! use voxtral_micro::TtsEngine;
//!
//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
//! let mut tts = TtsEngine::new("models/voxtral-tts-q4.gguf").await?;
//!
//! let audio = tts.synthesize("Hello world!", None)?;
//! tts.save_wav("output.wav", &audio)?;
//! # Ok(())
//! # }
//! ```

pub mod audio;
#[cfg(feature = "wgpu")]
pub mod gguf;
pub mod models;
pub mod tokenizer;
pub mod tts;

use anyhow::{bail, Context, Result};
use burn::backend::wgpu::WgpuDevice;
use burn::backend::Wgpu;
use burn::tensor::Tensor;
use std::path::{Path, PathBuf};

// Speed scaling removed: 1.0 = normal speed (no pitch shift compensation)
use tokenizer::TekkenEncoder;

/// Main TTS engine for speech synthesis from GGUF models.
pub struct TtsEngine {
    backbone: gguf::tts_model::Q4TtsBackbone,
    fm: gguf::tts_model::Q4FmTransformer,
    codec: tts::codec::CodecDecoder<Wgpu>,
    tokenizer: TekkenEncoder,
    voices_dir: PathBuf,
    device: WgpuDevice,
    max_frames: usize,
}

impl TtsEngine {
    /// Create a new TTS engine from a GGUF model file.
    ///
    /// # Arguments
    /// * `gguf_path` - Path to the Q4 GGUF model file
    ///
    /// # Returns
    /// A TTS engine ready for synthesis, or an error if model loading fails.
    ///
    /// # Example
    /// ```rust,no_run
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// use voxtral_tts::TtsEngine;
    /// let tts = TtsEngine::new("models/voxtral-tts-q4.gguf").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn new<P: AsRef<Path>>(gguf_path: P) -> Result<Self> {
        Self::with_options(gguf_path, None, None).await
    }

    /// Create a new TTS engine with custom options.
    ///
    /// # Arguments
    /// * `gguf_path` - Path to the Q4 GGUF model file
    /// * `tokenizer_path` - Optional path to tokenizer JSON (auto-discovered if None)
    /// * `voices_dir` - Optional path to voices directory (defaults to models/voxtral-tts/voice_embedding)
    pub async fn with_options<P: AsRef<Path>>(
        gguf_path: P,
        tokenizer_path: Option<P>,
        voices_dir: Option<P>,
    ) -> Result<Self> {
        let start = std::time::Instant::now();
        let gguf_path = gguf_path.as_ref();
        if !gguf_path.exists() {
            bail!("GGUF model not found at {}", gguf_path.display());
        }

        // Resolve tokenizer path
        let tokenizer_path = match tokenizer_path {
            Some(p) => p.as_ref().to_path_buf(),
            None => {
                let gguf_dir = gguf_path
                    .parent()
                    .unwrap_or(&PathBuf::from("."))
                    .to_path_buf();
                let candidates = [
                    gguf_dir.join("tekken.json"),
                    PathBuf::from("models/tekken.json"),
                    PathBuf::from("models/voxtral-tts/tekken.json"),
                ];
                candidates
                    .into_iter()
                    .find(|p| p.exists())
                    .ok_or_else(|| {
                        anyhow::anyhow!(
                            "Tokenizer not found. Provide tokenizer_path or place tekken.json alongside the GGUF file"
                        )
                    })?
            }
        };

        if !tokenizer_path.exists() {
            bail!("Tokenizer not found at {}", tokenizer_path.display());
        }

        let tokenizer =
            TekkenEncoder::from_file(&tokenizer_path).context("Failed to load tokenizer")?;

        // Resolve voices directory
        let voices_dir = match voices_dir {
            Some(d) => d.as_ref().to_path_buf(),
            None => PathBuf::from("models/voice_embedding"),
        };

        // Load GGUF model
        let device = WgpuDevice::default();
        tracing::info!("Loading Q4 TTS model from {}", gguf_path.display());
        let load_start = std::time::Instant::now();
        let mut loader = gguf::Q4TtsModelLoader::from_file(gguf_path)
            .context("Failed to open GGUF")?;
        let (backbone, fm, codec) = loader.load(&device).context("Failed to load Q4 model")?;
        tracing::info!("Model loaded in {:.2}s", load_start.elapsed().as_secs_f32());

        let total_time = start.elapsed().as_secs_f32();
        tracing::info!("TTS engine initialized in {:.2}s", total_time);

        Ok(Self {
            backbone,
            fm,
            codec,
            tokenizer,
            voices_dir,
            device,
            max_frames: 2000,
        })
    }

    /// Synthesize speech from text using the default voice (casual_female).
    ///
    /// # Arguments
    /// * `text` - The text to synthesize
    /// * `voice` - Optional voice name (defaults to "casual_female")
    ///
    /// # Returns
    /// An audio buffer containing the synthesized speech at 24kHz.
    pub fn synthesize(&mut self, text: &str, voice: Option<&str>) -> Result<audio::AudioBuffer> {
        self.synthesize_with_options(text, voice, 1.0, 1.0, None)
    }

    /// Synthesize speech with custom options.
    ///
    /// # Arguments
    /// * `text` - The text to synthesize
    /// * `voice` - Optional voice name (defaults to "casual_female")
    /// * `speed` - Playback speed multiplier (0.5 to 3.0, where 1.0 is normal)
    /// * `gain` - Volume gain multiplier (0.1 to 2.0, where 1.0 is normal)
    /// * `language` - Optional language code (e.g., "en", "fr", "de")
    ///
    /// # Returns
    /// An audio buffer containing the synthesized speech at 24kHz.
    ///
    /// # Speed Behavior
    /// Speed adjustment uses resampling which changes both tempo and pitch.
    /// Higher speeds result in higher pitch (chipmunk effect), lower speeds
    /// result in lower pitch.
    ///
    /// # Example
    /// ```no_run
    /// # use voxtral_micro::TtsEngine;
    /// # tokio_test::block_on(async {
    /// let mut tts = TtsEngine::new("models/voxtral-tts-q4.gguf").await?;
    /// let audio = tts.synthesize_with_options(
    ///     "Hello world!",
    ///     None,      // voice: None = default "casual_female"
    ///     2.0,       // speed: 2.0 = twice as fast
    ///     0.8,       // gain: 0.8 = 20% quieter
    ///     Some("en") // language
    /// )?;
    /// # Ok::<(), anyhow::Error>(())
    /// # });
    /// ```
    pub fn synthesize_with_options(
        &mut self,
        text: &str,
        voice: Option<&str>,
        speed: f32,
        gain: f32,
        language: Option<&str>,
    ) -> Result<audio::AudioBuffer> {
        let synthesis_start = std::time::Instant::now();
        
        // Validate parameters
        if !(0.5..=3.0).contains(&speed) {
            bail!("Speed must be between 0.5 and 3.0, got {}", speed);
        }
        if !(0.1..=2.0).contains(&gain) {
            bail!("Gain must be between 0.1 and 2.0, got {}", gain);
        }
        let voice_name = voice.unwrap_or("casual_female");

        // Tokenize text
        let tokenize_start = std::time::Instant::now();
        let token_ids = self.tokenizer.encode(text);
        tracing::debug!("Tokenization: {:.3}s", tokenize_start.elapsed().as_secs_f32());
        tracing::info!(
            text_tokens = token_ids.len(),
            voice = voice_name,
            language = ?language,
            "Synthesizing"
        );

        // Load voice embedding
        let voice_path = self
            .voices_dir
            .join(format!("{}.safetensors", voice_name));
        if !voice_path.exists() {
            bail!(
                "Voice '{}' not found at {}\n\
                \n\
                Voice embeddings are separate files that must be downloaded.\n\
                Download with:\n\
                  make download-models\n\
                Or manually:\n\
                  uv run --with huggingface_hub hf download \\\n\
                    TrevorJS/voxtral-tts-q4-gguf \\\n\
                    --local-dir models\n\
                \n\
                Then voices will be available at: models/voxtral-tts/voice_embedding/*.safetensors",
                voice_name,
                voice_path.display()
            );
        }

        let voice_bytes = std::fs::read(&voice_path)?;
        let voice_embed: Tensor<Wgpu, 2> = tts::voice::load_voice_from_bytes(
            &voice_bytes,
            3072,
            &self.device,
        )
        .context("Failed to load voice")?;

        tracing::info!(
            voice = voice_name,
            frames = voice_embed.dims()[0],
            "Voice loaded"
        );

        // Build input sequence
        let special = tts::config::TtsSpecialTokens::default();
        let bos = self
            .backbone
            .embed_tokens_from_ids(&[special.bos_token_id as i32], 1, 1);
        let begin_audio = self.backbone.embed_tokens_from_ids(
            &[special.begin_audio_token_id as i32],
            1,
            1,
        );
        let next_audio_text = self.backbone.embed_tokens_from_ids(
            &[special.next_audio_text_token_id as i32],
            1,
            1,
        );
        let repeat_audio_text = self.backbone.embed_tokens_from_ids(
            &[special.repeat_audio_text_token_id as i32],
            1,
            1,
        );
        let text_ids_i32: Vec<i32> = token_ids.iter().map(|&id| id as i32).collect();
        let text_embeds = self
            .backbone
            .embed_tokens_from_ids(&text_ids_i32, 1, text_ids_i32.len());

        let input_sequence = Tensor::cat(
            vec![
                bos,
                begin_audio.clone(),
                voice_embed.unsqueeze_dim::<3>(0),
                next_audio_text,
                text_embeds,
                repeat_audio_text,
                begin_audio,
            ],
            1,
        );

        let codebook = tts::embeddings::AudioCodebookEmbeddings::new(
            self.backbone.audio_codebook_embeddings().clone(),
            tts::config::AudioCodebookLayout::default(),
        );

        // Generate audio frames
        let gen_start = std::time::Instant::now();
        let frames = pollster::block_on(self.backbone.generate_async(
            input_sequence,
            &self.fm,
            &codebook,
            self.max_frames,
        ))
        .map_err(|e| anyhow::anyhow!("Generation failed: {e}"))?;
        tracing::info!("Frame generation: {:.2}s ({} frames)", gen_start.elapsed().as_secs_f32(), frames.len());

        if frames.is_empty() {
            bail!("No audio frames generated");
        }

        // Codec decode
        let decode_start = std::time::Instant::now();
        let n_frames = frames.len();
        let semantic_indices: Vec<usize> = frames.iter().map(|f| f.semantic_idx).collect();
        let mut acoustic_data = Vec::with_capacity(n_frames * 36);
        for frame in &frames {
            for &level in &frame.acoustic_levels {
                acoustic_data.push(level as f32);
            }
        }
        let acoustic_tensor: Tensor<Wgpu, 2> = Tensor::from_data(
            burn::tensor::TensorData::new(acoustic_data, [n_frames, 36]),
            &self.device,
        );
        let waveform = self.codec.decode(&semantic_indices, acoustic_tensor);
        let [_batch, total_samples] = waveform.dims();
        tracing::info!("Codec decode: {:.2}s", decode_start.elapsed().as_secs_f32());

        let postprocess_start = std::time::Instant::now();
        let wav_data = waveform.to_data();
        let mut samples: Vec<f32> = wav_data.as_slice::<f32>().unwrap()[..total_samples].to_vec();

        // Normalize to 0.95 peak
        let peak = samples.iter().map(|s| s.abs()).fold(0.0f32, f32::max);
        if peak > 1e-6 {
            let gain = 0.95 / peak;
            for s in &mut samples {
                *s *= gain;
            }
        }

        let mut audio = audio::AudioBuffer::new(samples, 24000);
        
        // Apply speed adjustment (1.0 = normal speed)
        if (speed - 1.0).abs() > 0.001 {
            audio = audio.with_speed(speed);
            tracing::debug!(speed = speed, "Speed adjusted");
        }
        
        // Apply gain adjustment
        if (gain - 1.0).abs() > 0.001 {
            audio = audio.with_gain(gain);
        }
        
        let duration = audio.len() as f64 / audio.sample_rate as f64;
        tracing::debug!("Post-processing: {:.3}s", postprocess_start.elapsed().as_secs_f32());
        
        let total_synthesis = synthesis_start.elapsed().as_secs_f32();
        tracing::info!(
            frames = n_frames,
            duration_sec = format!("{duration:.2}"),
            speed = speed,
            gain = gain,
            total_time_sec = format!("{total_synthesis:.2}"),
            "Audio generated"
        );

        Ok(audio)
    }

    /// List available voice presets in the voices directory.
    ///
    /// # Returns
    /// A vector of voice names, or an error if the directory doesn't exist.
    pub fn list_voices(&self) -> Result<Vec<String>> {
        if !self.voices_dir.exists() {
            bail!(
                "Voices directory not found at {}\n\
                \n\
                Voice embeddings must be downloaded.\n\
                Download with:\n\
                  make download-models\n\
                Or manually:\n\
                  uv run --with huggingface_hub hf download \\\n\
                    TrevorJS/voxtral-tts-q4-gguf \\\n\
                    --local-dir models",
                self.voices_dir.display()
            );
        }

        let mut voices: Vec<String> = std::fs::read_dir(&self.voices_dir)?
            .filter_map(|e| e.ok())
            .filter(|e| {
                e.path()
                    .extension()
                    .is_some_and(|ext| ext == "safetensors")
            })
            .filter_map(|e| {
                e.path()
                    .file_stem()
                    .map(|s| s.to_string_lossy().into_owned())
            })
            .collect();
        voices.sort();
        Ok(voices)
    }

    /// Set the maximum number of audio frames to generate.
    ///
    /// Default is 2000 frames. Higher values allow longer audio generation.
    pub fn set_max_frames(&mut self, max_frames: usize) {
        self.max_frames = max_frames;
    }

    /// Set the number of Euler ODE steps for flow matching (quality vs speed tradeoff).
    ///
    /// - 3 steps: Real-time performance
    /// - 4 steps: Balanced (default)
    /// - 8 steps: Higher quality
    pub fn set_euler_steps(&mut self, steps: usize) {
        self.fm.set_euler_steps(steps);
    }

    /// Save audio buffer to a WAV file.
    ///
    /// # Arguments
    /// * `path` - Output file path
    /// * `audio` - Audio buffer to save
    pub fn save_wav<P: AsRef<Path>>(&self, path: P, audio: &audio::AudioBuffer) -> Result<()> {
        audio.save(path)
    }
}

// Re-export commonly used types
pub use audio::AudioBuffer;