From 2f0c1f3fda6025da64636a4cfd01d4f98bc7d50b Mon Sep 17 00:00:00 2001 From: 魏曹先生 <1992414357@qq.com> Date: Sun, 21 Jun 2026 23:04:20 +0800 Subject: feat(cli): add --models-dir option for custom model path --- Cargo.lock | 2 +- Cargo.toml | 2 +- help.txt | 1 + src/args.rs | 6 ++++- src/main.rs | 83 +++++++++++++++++++++++++++++++++++++++---------------------- 5 files changed, 61 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8b4decf..f7bba13 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2573,7 +2573,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "vtx-engine" version = "0.3.1" -source = "git+https://github.com/Weicao-CatilGrass/vtx-engine?rev=864889b97c52d31098670d33e10d62a96ec8c41e#864889b97c52d31098670d33e10d62a96ec8c41e" +source = "git+https://github.com/Weicao-CatilGrass/vtx-engine?rev=35b954ad50ec6211e6a25ca4777d1685311f696e#35b954ad50ec6211e6a25ca4777d1685311f696e" dependencies = [ "aec3", "block2", diff --git a/Cargo.toml b/Cargo.toml index 69724a4..7821ca1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ pinyin = { version = "0.11.0", features = ["with_tone", "with_tone_num", "with_t [dependencies.vtx-engine] git = "https://github.com/Weicao-CatilGrass/vtx-engine" -rev = "864889b97c52d31098670d33e10d62a96ec8c41e" +rev = "35b954ad50ec6211e6a25ca4777d1685311f696e" [profile.dev] opt-level = 0 diff --git a/help.txt b/help.txt index f4b5507..c318505 100644 --- a/help.txt +++ b/help.txt @@ -7,6 +7,7 @@ Options -m, --model= Whisper model. [default: base_en] --lang= Language hint (zh, ja, en, fr...). Skips detection. --post= Post-process transcribed text, e.g. +pinyin + --models-dir= Directory for model files. Fails if invalid. -f, --format= Output format with %{vol}, %{word}, %{confid}. [default: %{vol},%{word}] -S, --format-file= Read format from file. diff --git a/src/args.rs b/src/args.rs index 7b472c4..9fe76d6 100644 --- a/src/args.rs +++ b/src/args.rs @@ -109,7 +109,11 @@ pub struct DMVOPArguments { )] pub socket_file: PathBuf, - // Subnet mask for UDP broadcast (default: only last octet, e.g., "255.255.255.0") + // Custom directory for model files + #[arg(long = "models-dir", alias = "models", require_equals = true)] + pub models_dir: Option, + + // Subnet mask for UDP broadcast (default: only last octet, e.g., "255.255.255.255.0") #[arg( long = "subnet-mask", alias = "mask", diff --git a/src/main.rs b/src/main.rs index 73b4e06..66af45f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -93,7 +93,7 @@ async fn main() { // --download-model: download a specific model and exit // --------------------------------------------------------------- if let Some(ref model_name) = args.download_model { - download_model_cli(model_name).await; + download_model_cli(model_name, args.models_dir.clone()).await; return; } @@ -116,7 +116,32 @@ async fn main() { model.display_name() ); - let mut builder = EngineBuilder::new().app_name("dmvop").model(model); + let mut builder = EngineBuilder::new().app_name("dmvop"); + + // Apply --models-dir: construct explicit model path, fail on invalid dir + if let Some(ref dir) = args.models_dir { + if dir.is_file() { + eprintln!( + "[dmvop] --models-dir '{}' is a file, not a directory.", + dir.display() + ); + std::process::exit(1); + } + if !dir.exists() { + if let Err(e) = std::fs::create_dir_all(dir) { + eprintln!( + "[dmvop] Failed to create models directory '{}': {}", + dir.display(), + e + ); + std::process::exit(1); + } + } + debug_log!("[dmvop] Models directory: {}", dir.display()); + builder = builder.models_dir(dir); + } else { + builder = builder.model(model); + } if args.instant { debug_log!("[dmvop] Instant mode: aggressive VAD for near-real-time output"); @@ -138,27 +163,6 @@ async fn main() { // Disable PTT mode so VAD drives automatic segmentation engine.set_ptt_mode(false); - // --------------------------------------------------------------- - // Check model availability and download if needed - // --------------------------------------------------------------- - let model_status = engine.check_model_status(); - if !model_status.available { - eprintln!("[dmvop] Model not found at: {}", model_status.path); - eprintln!("[dmvop] Downloading model, please wait..."); - match engine.download_model().await { - Ok(_) => debug_log!("[dmvop] Model downloaded successfully"), - Err(e) => { - eprintln!("[dmvop] Failed to download model: {}", e); - eprintln!("[dmvop] You can manually download a model from:"); - eprintln!("[dmvop] https://huggingface.co/ggerganov/whisper.cpp/tree/main"); - eprintln!("[dmvop] Place it at: {}", model_status.path); - std::process::exit(1); - } - } - } else { - debug_log!("[dmvop] Model found: {}", model_status.path); - } - // --------------------------------------------------------------- // List devices and exit? // --------------------------------------------------------------- @@ -404,7 +408,7 @@ fn list_models() { } /// Download a specific Whisper model by identifier. -async fn download_model_cli(model_name: &str) { +async fn download_model_cli(model_name: &str, models_dir: Option) { let model = match vtx_engine::WhisperModel::parse_identifier(model_name) { Some(m) => m, None => { @@ -421,12 +425,31 @@ async fn download_model_cli(model_name: &str) { model.config_key() ); - let (engine, _rx) = match EngineBuilder::new() - .app_name("dmvop") - .model(model) - .build() - .await - { + let mut builder = EngineBuilder::new().app_name("dmvop"); + if let Some(ref dir) = models_dir { + if dir.is_file() { + eprintln!( + "[dmvop] --models-dir '{}' is a file, not a directory.", + dir.display() + ); + std::process::exit(1); + } + if !dir.exists() { + if let Err(e) = std::fs::create_dir_all(dir) { + eprintln!( + "[dmvop] Failed to create models directory '{}': {}", + dir.display(), + e + ); + std::process::exit(1); + } + } + builder = builder.models_dir(dir); + } else { + builder = builder.model(model); + } + + let (engine, _rx) = match builder.build().await { Ok(e) => e, Err(e) => { eprintln!("[dmvop] Failed to build engine: {}", e); -- cgit