Skip to main content

ferrum_cli/commands/
tts.rs

1//! TTS command - Text-to-speech using Qwen3-TTS models
2
3use crate::config::CliConfig;
4use candle_core::{DType, Device as CandleDevice};
5use clap::Args;
6use colored::Colorize;
7use ferrum_models::source::{ModelFormat, ResolvedModelSource};
8use ferrum_models::HfDownloader;
9use ferrum_types::Result;
10use std::path::PathBuf;
11
12/// Synthesize speech from text using Qwen3-TTS models
13#[derive(Args, Debug)]
14pub struct TtsCommand {
15    /// TTS model name (e.g., qwen3-tts, Qwen/Qwen3-TTS-12Hz-0.6B-Base)
16    #[arg(required = true)]
17    pub model: String,
18
19    /// Text to synthesize
20    #[arg(required = true)]
21    pub text: String,
22
23    /// Output WAV file path
24    #[arg(short, long, default_value = "/tmp/ferrum_tts_output.wav")]
25    pub output: String,
26
27    /// Language (e.g., chinese, english, auto)
28    #[arg(short, long, default_value = "auto")]
29    pub language: String,
30
31    /// Backend: auto, cpu, metal (default: auto)
32    #[arg(short, long, default_value = "auto")]
33    pub backend: String,
34
35    /// Reference audio for voice cloning (WAV/M4A/MP3)
36    #[arg(long)]
37    pub ref_audio: Option<String>,
38
39    /// Reference audio transcript (required for ICL voice cloning)
40    #[arg(long)]
41    pub ref_text: Option<String>,
42
43    /// Enable streaming mode (generate audio in chunks)
44    #[arg(long)]
45    pub streaming: bool,
46
47    /// Frames per streaming chunk (default: 10, ~800ms per chunk)
48    #[arg(long, default_value = "10")]
49    pub chunk_frames: usize,
50}
51
52pub async fn execute(cmd: TtsCommand, config: CliConfig) -> Result<()> {
53    let model_id = resolve_tts_alias(&cmd.model);
54    let cache_dir = get_hf_cache_dir(&config);
55
56    eprintln!("{} {}", "Model:".dimmed(), model_id.cyan());
57
58    // Find or download model
59    let source = match find_cached_model(&cache_dir, &model_id) {
60        Some(source) => source,
61        None => {
62            eprintln!(
63                "{} Model '{}' not found locally, downloading...",
64                ">>>".cyan(),
65                model_id
66            );
67            let token = std::env::var("HF_TOKEN")
68                .or_else(|_| std::env::var("HUGGING_FACE_HUB_TOKEN"))
69                .ok();
70            let downloader = HfDownloader::new(cache_dir, token)?;
71            let snapshot_path = downloader.download(&model_id, None).await?;
72            let format = detect_format(&snapshot_path);
73            if format == ModelFormat::Unknown {
74                return Err(ferrum_types::FerrumError::model(
75                    "Downloaded model has unknown format",
76                ));
77            }
78            ResolvedModelSource {
79                original: model_id.clone(),
80                local_path: snapshot_path,
81                format,
82                from_cache: false,
83            }
84        }
85    };
86
87    // Verify architecture — Qwen3-TTS config.json has "talker_config" key
88    let config_path = source.local_path.join("config.json");
89    let config_data = std::fs::read_to_string(&config_path)
90        .map_err(|e| ferrum_types::FerrumError::model(format!("read config.json: {e}")))?;
91    let config_json: serde_json::Value = serde_json::from_str(&config_data)
92        .map_err(|e| ferrum_types::FerrumError::model(format!("parse config.json: {e}")))?;
93
94    if config_json.get("talker_config").is_none() {
95        return Err(ferrum_types::FerrumError::model(format!(
96            "'{}' does not appear to be a Qwen3-TTS model (missing talker_config)",
97            model_id
98        )));
99    }
100
101    let candle_device = select_candle_device(&cmd.backend)?;
102    eprintln!("{} {:?}", "Device:".dimmed(), &candle_device);
103    eprintln!("{}", "Loading TTS model...".dimmed());
104
105    let mut executor = ferrum_models::TtsModelExecutor::from_path(
106        &source.local_path.to_string_lossy(),
107        candle_device,
108        DType::F32,
109    )?;
110    eprintln!("{}", "Model loaded.".green());
111
112    // Synthesize
113    eprintln!(
114        "{} \"{}\"",
115        "Text:".dimmed(),
116        if cmd.text.chars().count() > 80 {
117            let truncated: String = cmd.text.chars().take(77).collect();
118            format!("{truncated}...")
119        } else {
120            cmd.text.clone()
121        }
122        .cyan()
123    );
124
125    let start = std::time::Instant::now();
126    let sample_rate = executor.sample_rate();
127
128    if cmd.streaming && cmd.ref_audio.is_none() {
129        // Streaming mode: generate and save chunks incrementally
130        eprintln!("{}", "Streaming mode enabled".yellow());
131        let mut all_samples = Vec::new();
132        let sr = sample_rate;
133        let t0 = start;
134        let chunks = executor.synthesize_streaming(
135            &cmd.text,
136            &cmd.language,
137            cmd.chunk_frames,
138            |idx, chunk| {
139                let chunk_dur = chunk.len() as f64 / sr as f64;
140                eprintln!(
141                    "  {} chunk {} — {:.2}s audio (at {:.1}s)",
142                    "▶".green(),
143                    idx,
144                    chunk_dur,
145                    t0.elapsed().as_secs_f64(),
146                );
147            },
148        )?;
149        let elapsed = start.elapsed();
150        for chunk in &chunks {
151            all_samples.extend_from_slice(chunk);
152        }
153
154        let duration_secs = all_samples.len() as f64 / sample_rate as f64;
155        save_wav(&cmd.output, &all_samples, sample_rate as u32)?;
156
157        eprintln!("\n{} {}", "Output:".dimmed(), cmd.output.green());
158        eprintln!(
159            "{} {:.2}s audio, {:.2}s elapsed (RTF={:.2}x), {} chunks",
160            "Stats:".dimmed(),
161            duration_secs,
162            elapsed.as_secs_f64(),
163            elapsed.as_secs_f64() / duration_secs.max(0.001),
164            chunks.len(),
165        );
166    } else {
167        // Batch mode
168        let samples = if let Some(ref_audio) = &cmd.ref_audio {
169            let ref_text = cmd.ref_text.as_deref().ok_or_else(|| {
170                ferrum_types::FerrumError::model("--ref-text required for voice cloning")
171            })?;
172            eprintln!("{} {}", "Ref audio:".dimmed(), ref_audio.cyan());
173            executor.synthesize_voice_clone(&cmd.text, &cmd.language, ref_audio, ref_text)?
174        } else {
175            executor.synthesize(&cmd.text, &cmd.language)?
176        };
177        let elapsed = start.elapsed();
178        let duration_secs = samples.len() as f64 / sample_rate as f64;
179        save_wav(&cmd.output, &samples, sample_rate as u32)?;
180
181        eprintln!("\n{} {}", "Output:".dimmed(), cmd.output.green());
182        eprintln!(
183            "{} {:.2}s audio, {:.2}s elapsed (RTF={:.2}x)",
184            "Stats:".dimmed(),
185            duration_secs,
186            elapsed.as_secs_f64(),
187            elapsed.as_secs_f64() / duration_secs.max(0.001),
188        );
189    }
190
191    Ok(())
192}
193
194/// Save PCM samples as a 16-bit mono WAV file (no external dependency).
195fn save_wav(path: &str, samples: &[f32], sample_rate: u32) -> Result<()> {
196    use std::io::Write;
197
198    let num_samples = samples.len() as u32;
199    let bytes_per_sample: u16 = 2; // 16-bit
200    let channels: u16 = 1;
201    let data_size = num_samples * bytes_per_sample as u32;
202    let file_size = 36 + data_size;
203
204    let mut buf: Vec<u8> = Vec::with_capacity(file_size as usize + 8);
205
206    // RIFF header
207    buf.extend_from_slice(b"RIFF");
208    buf.extend_from_slice(&file_size.to_le_bytes());
209    buf.extend_from_slice(b"WAVE");
210
211    // fmt chunk
212    buf.extend_from_slice(b"fmt ");
213    buf.extend_from_slice(&16u32.to_le_bytes()); // chunk size
214    buf.extend_from_slice(&1u16.to_le_bytes()); // PCM format
215    buf.extend_from_slice(&channels.to_le_bytes());
216    buf.extend_from_slice(&sample_rate.to_le_bytes());
217    let byte_rate = sample_rate * channels as u32 * bytes_per_sample as u32;
218    buf.extend_from_slice(&byte_rate.to_le_bytes());
219    let block_align = channels * bytes_per_sample;
220    buf.extend_from_slice(&block_align.to_le_bytes());
221    buf.extend_from_slice(&(bytes_per_sample * 8).to_le_bytes()); // bits per sample
222
223    // data chunk
224    buf.extend_from_slice(b"data");
225    buf.extend_from_slice(&data_size.to_le_bytes());
226
227    for &sample in samples {
228        let clamped = sample.clamp(-1.0, 1.0);
229        let value = (clamped * 32767.0) as i16;
230        buf.extend_from_slice(&value.to_le_bytes());
231    }
232
233    let mut file = std::fs::File::create(path)
234        .map_err(|e| ferrum_types::FerrumError::model(format!("create WAV file: {e}")))?;
235    file.write_all(&buf)
236        .map_err(|e| ferrum_types::FerrumError::model(format!("write WAV file: {e}")))?;
237
238    Ok(())
239}
240
241fn resolve_tts_alias(name: &str) -> String {
242    match name.to_lowercase().as_str() {
243        "qwen3-tts" | "qwen3:tts" | "qwen3-tts:0.6b" => "Qwen/Qwen3-TTS-12Hz-0.6B-Base".to_string(),
244        "qwen3-tts:instruct" | "qwen3-tts-instruct" => {
245            "Qwen/Qwen3-TTS-12Hz-0.6B-Instruct".to_string()
246        }
247        _ => name.to_string(),
248    }
249}
250
251fn get_hf_cache_dir(config: &CliConfig) -> PathBuf {
252    if let Ok(hf_home) = std::env::var("HF_HOME") {
253        return PathBuf::from(hf_home);
254    }
255    let configured = shellexpand::tilde(&config.models.download.hf_cache_dir).to_string();
256    PathBuf::from(configured)
257}
258
259fn find_cached_model(cache_dir: &PathBuf, model_id: &str) -> Option<ResolvedModelSource> {
260    let hub_dir = cache_dir.join("hub");
261    let model_dir_name = format!("models--{}", model_id.replace("/", "--"));
262    let model_dir = hub_dir.join(&model_dir_name);
263
264    if model_dir.exists() {
265        let snapshots_dir = model_dir.join("snapshots");
266        if snapshots_dir.exists() {
267            // Try refs/main first
268            let ref_main = model_dir.join("refs").join("main");
269            if let Ok(rev) = std::fs::read_to_string(&ref_main) {
270                let rev = rev.trim();
271                if !rev.is_empty() {
272                    let snapshot = snapshots_dir.join(rev);
273                    if snapshot.exists() {
274                        let format = detect_format(&snapshot);
275                        if format != ModelFormat::Unknown {
276                            return Some(ResolvedModelSource {
277                                original: model_id.to_string(),
278                                local_path: snapshot,
279                                format,
280                                from_cache: true,
281                            });
282                        }
283                    }
284                }
285            }
286            // Fallback: first snapshot
287            if let Ok(entries) = std::fs::read_dir(&snapshots_dir) {
288                for entry in entries.flatten() {
289                    let path = entry.path();
290                    if path.is_dir() && path.join("config.json").exists() {
291                        let format = detect_format(&path);
292                        if format != ModelFormat::Unknown {
293                            return Some(ResolvedModelSource {
294                                original: model_id.to_string(),
295                                local_path: path,
296                                format,
297                                from_cache: true,
298                            });
299                        }
300                    }
301                }
302            }
303        }
304    }
305
306    None
307}
308
309fn select_candle_device(backend: &str) -> ferrum_types::Result<CandleDevice> {
310    match backend.to_lowercase().as_str() {
311        "cpu" => Ok(CandleDevice::Cpu),
312        "metal" => {
313            #[cfg(all(target_os = "macos", feature = "metal"))]
314            {
315                return CandleDevice::new_metal(0)
316                    .map_err(|error| ferrum_types::FerrumError::device(error.to_string()));
317            }
318            #[allow(unreachable_code)]
319            {
320                Err(ferrum_types::FerrumError::unsupported(
321                    "Metal TTS requires a Metal-enabled Ferrum build",
322                ))
323            }
324        }
325        "cuda" => {
326            #[cfg(feature = "candle-cuda-compat")]
327            {
328                return CandleDevice::new_cuda(0)
329                    .map_err(|error| ferrum_types::FerrumError::device(error.to_string()));
330            }
331            #[allow(unreachable_code)]
332            {
333                Err(ferrum_types::FerrumError::unsupported(
334                    "CUDA TTS requires the candle-cuda-compat feature",
335                ))
336            }
337        }
338        "auto" | _ => {
339            #[cfg(feature = "candle-cuda-compat")]
340            {
341                if let Ok(d) = CandleDevice::new_cuda(0) {
342                    return Ok(d);
343                }
344            }
345            #[cfg(all(target_os = "macos", feature = "metal"))]
346            {
347                if let Ok(device) = CandleDevice::new_metal(0) {
348                    return Ok(device);
349                }
350            }
351            #[allow(unreachable_code)]
352            Ok(CandleDevice::Cpu)
353        }
354    }
355}
356
357fn detect_format(path: &PathBuf) -> ModelFormat {
358    if path.join("model.safetensors").exists() {
359        ModelFormat::SafeTensors
360    } else if std::fs::read_dir(path)
361        .map(|d| {
362            d.filter_map(|e| e.ok())
363                .any(|e| e.path().extension().is_some_and(|ext| ext == "safetensors"))
364        })
365        .unwrap_or(false)
366    {
367        ModelFormat::SafeTensors
368    } else if path.join("pytorch_model.bin").exists() {
369        ModelFormat::PyTorchBin
370    } else {
371        ModelFormat::Unknown
372    }
373}