|
| 1 | +//! 说话人模型下载器（TD-2026-08-20-D 清偿 / G1：wespeaker 应用内一键下载）。 |
| 2 | +//! |
| 3 | +//! @ai-context: 对照 structure_models 模式（ureq + .part 原子写 + 1MB 进度事件 + |
| 4 | +//! Content-Length 比对防截断）；单文件（model.onnx，20-70MB）双镜像 |
| 5 | +//! 回退（GitHub release + hf-mirror，与 download-speaker-model.ps1 |
| 6 | +//! 同源）；目标 models/speaker-embedding/model.onnx（speaker_engine |
| 7 | +//! 路径约定）。网络路径不单测（与 model_downloader 同口径）。 |
| 8 | +
|
| 9 | +use std::io::{Read, Write}; |
| 10 | +use std::sync::atomic::{AtomicBool, Ordering}; |
| 11 | +use std::sync::{Arc, Mutex}; |
| 12 | +use std::thread::JoinHandle; |
| 13 | + |
| 14 | +use serde::Serialize; |
| 15 | +use tauri::Emitter; |
| 16 | + |
| 17 | +use crate::error::{AppError, Result}; |
| 18 | + |
| 19 | +/// 下载镜像源（与 scripts/download-speaker-model.ps1 同源；多源回退防单点失效）。 |
| 20 | +const MIRRORS: &[&str] = &[ |
| 21 | + "https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-recongition-models/wespeaker-zh.onnx", |
| 22 | + "https://hf-mirror.com/csukuangfj/sherpa-onnx-speaker-embedding/resolve/main/wespeaker-zh.onnx", |
| 23 | +]; |
| 24 | +/// 下载超时（30 分钟；同 structure_models 口径）。 |
| 25 | +const DOWNLOAD_TIMEOUT_SECS: u64 = 30 * 60; |
| 26 | + |
| 27 | +/// 下载状态（前端进度展示；camelCase 契约）。 |
| 28 | +#[derive(Debug, Clone, Serialize, PartialEq)] |
| 29 | +#[serde(rename_all = "camelCase")] |
| 30 | +pub struct SpeakerDownloadStatus { |
| 31 | + /// idle | downloading | done | failed |
| 32 | + pub state: String, |
| 33 | + pub downloaded_bytes: u64, |
| 34 | + pub total_bytes: u64, |
| 35 | + pub error: Option<String>, |
| 36 | +} |
| 37 | + |
| 38 | +impl Default for SpeakerDownloadStatus { |
| 39 | + fn default() -> Self { |
| 40 | + Self { state: "idle".into(), downloaded_bytes: 0, total_bytes: 0, error: None } |
| 41 | + } |
| 42 | +} |
| 43 | + |
| 44 | +/// 模型已就绪判定（纯函数）：目标文件存在且非空。 |
| 45 | +/// @ai-context: pub 供单测（start 内短路与命令层共用；测试不依赖网络/运行时）。 |
| 46 | +pub fn already_downloaded(dir: &std::path::Path) -> bool { |
| 47 | + dir.join("model.onnx").metadata().is_ok_and(|m| m.len() > 0) |
| 48 | +} |
| 49 | + |
| 50 | +/// 说话人模型下载器（AppState 持有；单文件状态机 + 防并发）。 |
| 51 | +#[derive(Clone)] |
| 52 | +pub struct SpeakerModelDownloader { |
| 53 | + status: Arc<Mutex<SpeakerDownloadStatus>>, |
| 54 | + running: Arc<AtomicBool>, |
| 55 | + /// 线程句柄持有（进程生命周期内完成即弃；join 非必需——状态即完成信号） |
| 56 | + _threads: Arc<Mutex<Vec<JoinHandle<()>>>>, |
| 57 | +} |
| 58 | + |
| 59 | +impl Default for SpeakerModelDownloader { |
| 60 | + fn default() -> Self { |
| 61 | + Self::new() |
| 62 | + } |
| 63 | +} |
| 64 | + |
| 65 | +impl SpeakerModelDownloader { |
| 66 | + pub fn new() -> Self { |
| 67 | + Self { |
| 68 | + status: Arc::new(Mutex::new(SpeakerDownloadStatus::default())), |
| 69 | + running: Arc::new(AtomicBool::new(false)), |
| 70 | + _threads: Arc::new(Mutex::new(Vec::new())), |
| 71 | + } |
| 72 | + } |
| 73 | + |
| 74 | + /// 启动下载（已在下载中 → 拒绝；模型已存在 → done 直接返回）。 |
| 75 | + /// |
| 76 | + /// @param dir - 装配目标目录（models/speaker-embedding，命令层提供） |
| 77 | + pub fn start(&self, dir: std::path::PathBuf, app: tauri::AppHandle) -> Result<()> { |
| 78 | + // 防并发（structure_models 审查 M2 教训：启动标记必须可复位） |
| 79 | + if self.running.swap(true, Ordering::SeqCst) { |
| 80 | + return Err(AppError::Io("说话人模型下载已在进行中".to_string())); |
| 81 | + } |
| 82 | + if already_downloaded(&dir) { |
| 83 | + self.running.store(false, Ordering::SeqCst); |
| 84 | + { |
| 85 | + let mut st = self.status.lock().expect("speaker dl status lock"); |
| 86 | + *st = SpeakerDownloadStatus { state: "done".into(), ..Default::default() }; |
| 87 | + } |
| 88 | + return Ok(()); |
| 89 | + } |
| 90 | + { |
| 91 | + let mut st = self.status.lock().expect("speaker dl status lock"); |
| 92 | + *st = SpeakerDownloadStatus { state: "downloading".into(), ..Default::default() }; |
| 93 | + } |
| 94 | + let status = self.status.clone(); |
| 95 | + let running = self.running.clone(); |
| 96 | + let threads = self._threads.clone(); |
| 97 | + let handle = std::thread::Builder::new() |
| 98 | + .name("entropy-speaker-download".into()) |
| 99 | + .spawn(move || download_impl(dir, app, status, running)) |
| 100 | + .map_err(|e| { |
| 101 | + // spawn 失败必须复位 running（防永久"下载中"） |
| 102 | + self.running.store(false, Ordering::SeqCst); |
| 103 | + AppError::Io(format!("启动说话人模型下载线程失败: {}", e)) |
| 104 | + })?; |
| 105 | + threads.lock().expect("speaker dl threads lock").push(handle); |
| 106 | + Ok(()) |
| 107 | + } |
| 108 | + |
| 109 | + /// 当前状态（无记录 → idle）。 |
| 110 | + pub fn status(&self) -> SpeakerDownloadStatus { |
| 111 | + self.status.lock().expect("speaker dl status lock").clone() |
| 112 | + } |
| 113 | +} |
| 114 | + |
| 115 | +/// 下载实现（后台线程）：双镜像回退 → .part 原子写 → 进度事件 → 落盘校验。 |
| 116 | +fn download_impl( |
| 117 | + dir: std::path::PathBuf, |
| 118 | + app: tauri::AppHandle, |
| 119 | + status: Arc<Mutex<SpeakerDownloadStatus>>, |
| 120 | + running: Arc<AtomicBool>, |
| 121 | +) { |
| 122 | + let _ = std::fs::create_dir_all(&dir); |
| 123 | + let target = dir.join("model.onnx"); |
| 124 | + let part = dir.join("model.onnx.part"); |
| 125 | + let result = (|| -> std::result::Result<(), String> { |
| 126 | + let mut last_err: Option<String> = None; |
| 127 | + for url in MIRRORS { |
| 128 | + match download_one(url, &part, &app, &status) { |
| 129 | + Ok(()) => { |
| 130 | + std::fs::rename(&part, &target) |
| 131 | + .map_err(|e| format!("模型落盘失败: {}", e))?; |
| 132 | + return Ok(()); |
| 133 | + } |
| 134 | + Err(e) => { |
| 135 | + let _ = std::fs::remove_file(&part); |
| 136 | + last_err = Some(format!("{}（{}）", e, url)); |
| 137 | + } |
| 138 | + } |
| 139 | + } |
| 140 | + Err(last_err.unwrap_or_else(|| "无可用镜像".to_string())) |
| 141 | + })(); |
| 142 | + { |
| 143 | + let mut st = status.lock().expect("speaker dl status lock"); |
| 144 | + match result { |
| 145 | + Ok(()) => { |
| 146 | + *st = SpeakerDownloadStatus { state: "done".into(), ..Default::default() }; |
| 147 | + let _ = app.emit("speaker-model:download-done", ()); |
| 148 | + } |
| 149 | + Err(e) => { |
| 150 | + *st = SpeakerDownloadStatus { |
| 151 | + state: "failed".into(), |
| 152 | + error: Some(e.clone()), |
| 153 | + ..Default::default() |
| 154 | + }; |
| 155 | + let _ = app.emit("speaker-model:download-failed", e); |
| 156 | + } |
| 157 | + } |
| 158 | + } |
| 159 | + running.store(false, Ordering::SeqCst); |
| 160 | +} |
| 161 | + |
| 162 | +/// 下载单文件（流式，每 1MB 推送进度；Content-Length 比对防截断——structure_models 同款）。 |
| 163 | +fn download_one( |
| 164 | + url: &str, |
| 165 | + part: &std::path::Path, |
| 166 | + app: &tauri::AppHandle, |
| 167 | + status: &Arc<Mutex<SpeakerDownloadStatus>>, |
| 168 | +) -> std::result::Result<(), String> { |
| 169 | + let agent = ureq::AgentBuilder::new() |
| 170 | + .timeout(std::time::Duration::from_secs(DOWNLOAD_TIMEOUT_SECS)) |
| 171 | + .build(); |
| 172 | + let resp = agent.get(url).call().map_err(|e| format!("请求失败: {}", e))?; |
| 173 | + let total = resp |
| 174 | + .header("Content-Length") |
| 175 | + .and_then(|v| v.parse::<u64>().ok()) |
| 176 | + .unwrap_or(0); |
| 177 | + let mut reader = resp.into_reader(); |
| 178 | + let mut out = std::fs::File::create(part).map_err(|e| format!("创建临时文件失败: {}", e))?; |
| 179 | + let mut buf = [0u8; 64 * 1024]; |
| 180 | + let mut downloaded: u64 = 0; |
| 181 | + let mut last_emit: u64 = 0; |
| 182 | + loop { |
| 183 | + let n = reader.read(&mut buf).map_err(|e| format!("读取响应失败: {}", e))?; |
| 184 | + if n == 0 { |
| 185 | + break; |
| 186 | + } |
| 187 | + out.write_all(&buf[..n]).map_err(|e| format!("写入失败: {}", e))?; |
| 188 | + downloaded += n as u64; |
| 189 | + if downloaded - last_emit >= 1024 * 1024 { |
| 190 | + last_emit = downloaded; |
| 191 | + let _ = app.emit( |
| 192 | + "speaker-model:download-progress", |
| 193 | + crate::model_downloader::DownloadProgress { |
| 194 | + file: "wespeaker-zh.onnx".to_string(), |
| 195 | + downloaded_bytes: downloaded, |
| 196 | + total_bytes: total, |
| 197 | + }, |
| 198 | + ); |
| 199 | + if let Ok(mut st) = status.lock() { |
| 200 | + st.downloaded_bytes = downloaded; |
| 201 | + st.total_bytes = total; |
| 202 | + } |
| 203 | + } |
| 204 | + } |
| 205 | + if total > 0 && downloaded != total { |
| 206 | + return Err(format!("大小不符（期望 {}，实得 {}）", total, downloaded)); |
| 207 | + } |
| 208 | + Ok(()) |
| 209 | +} |
| 210 | + |
| 211 | +#[cfg(test)] |
| 212 | +#[path = "speaker_download_tests.rs"] |
| 213 | +mod tests; |
0 commit comments