1use 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#[derive(Args, Debug)]
14pub struct TtsCommand {
15 #[arg(required = true)]
17 pub model: String,
18
19 #[arg(required = true)]
21 pub text: String,
22
23 #[arg(short, long, default_value = "/tmp/ferrum_tts_output.wav")]
25 pub output: String,
26
27 #[arg(short, long, default_value = "auto")]
29 pub language: String,
30
31 #[arg(short, long, default_value = "auto")]
33 pub backend: String,
34
35 #[arg(long)]
37 pub ref_audio: Option<String>,
38
39 #[arg(long)]
41 pub ref_text: Option<String>,
42
43 #[arg(long)]
45 pub streaming: bool,
46
47 #[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 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 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 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 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 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
194fn 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; 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 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 buf.extend_from_slice(b"fmt ");
213 buf.extend_from_slice(&16u32.to_le_bytes()); buf.extend_from_slice(&1u16.to_le_bytes()); 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()); 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 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 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}