diff options
Diffstat (limited to 'src/main.rs')
| -rw-r--r-- | src/main.rs | 428 |
1 files changed, 425 insertions, 3 deletions
diff --git a/src/main.rs b/src/main.rs index cc8d946..5b2bb0a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,8 +2,430 @@ mod args; pub use args::*; mod output_protocol; -pub use output_protocol::*; +use clap::Parser; +use output_protocol::OutputProtocol; +use std::path::PathBuf; +use std::sync::Arc; +use vtx_engine::EngineBuilder; -fn main() { - println!("Hello, world!"); +/// Output channel enum — wraps each OutputProtocol implementation so we can +/// store a heterogeneous collection and dispatch `send` without trait objects. +enum OutputChannel { + Stdout(Arc<output_protocol::StandardOutputProtocol>), + Stderr(Arc<output_protocol::StandardErrorProtocol>), + Tcp(Arc<output_protocol::TCPOutputProtocol>), + Udp(Arc<output_protocol::UDPOutputProtocol>), + UdpBroadcast(Arc<output_protocol::UDPBroadcastOutputProtocol>), + #[cfg(unix)] + Ipc(Arc<output_protocol::IPCOutputProtocol>), +} + +impl OutputChannel { + async fn init(&self) { + match self { + OutputChannel::Stdout(p) => p.init().await, + OutputChannel::Stderr(p) => p.init().await, + OutputChannel::Tcp(p) => p.init().await, + OutputChannel::Udp(p) => p.init().await, + OutputChannel::UdpBroadcast(p) => p.init().await, + #[cfg(unix)] + OutputChannel::Ipc(p) => p.init().await, + } + } + + async fn send(&self, message: &str) { + match self { + OutputChannel::Stdout(p) => p.clone().send(message).await, + OutputChannel::Stderr(p) => p.clone().send(message).await, + OutputChannel::Tcp(p) => p.clone().send(message).await, + OutputChannel::Udp(p) => p.clone().send(message).await, + OutputChannel::UdpBroadcast(p) => p.clone().send(message).await, + #[cfg(unix)] + OutputChannel::Ipc(p) => p.clone().send(message).await, + } + } +} + +#[tokio::main] +async fn main() { + // Set up tracing so we can see vtx-engine logs (including transcription errors) + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::builder() + .with_default_directive(tracing_subscriber::filter::LevelFilter::WARN.into()) + .from_env_lossy(), + ) + .with_target(false) + .init(); + + let args = DMVOPArguments::parse(); + + // Set global verbose flag + VERBOSE.store(args.verbose, std::sync::atomic::Ordering::Relaxed); + + // --------------------------------------------------------------- + // --list-models: show available models and exit + // --------------------------------------------------------------- + if args.list_models { + list_models(); + return; + } + + // --------------------------------------------------------------- + // --download-model: download a specific model and exit + // --------------------------------------------------------------- + if let Some(ref model_name) = args.download_model { + download_model_cli(model_name).await; + return; + } + + // --------------------------------------------------------------- + // 1. Build the vtx-engine (needed for both listing and capture) + // --------------------------------------------------------------- + eprintln!("[dmvop] Initializing voice engine..."); + + let model = vtx_engine::WhisperModel::parse_identifier(&args.model).unwrap_or_else(|| { + eprintln!( + "[dmvop] Unknown model '{}'. Use --list-models to see available models.", + args.model + ); + std::process::exit(1); + }); + + eprintln!( + "[dmvop] Using model: {} ({})", + model.config_key(), + model.display_name() + ); + + let mut builder = EngineBuilder::new().app_name("dmvop").model(model); + + if let Some(ref lang) = args.lang { + eprintln!("[dmvop] Language hint: {}", lang); + builder = builder.language(lang.as_str()); + } + + let (engine, mut rx) = builder.build().await.expect("Failed to build vtx-engine"); + + // 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(_) => eprintln!("[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 { + eprintln!("[dmvop] Model found: {}", model_status.path); + } + + // --------------------------------------------------------------- + // 2. List devices and exit? + // --------------------------------------------------------------- + let devices = engine.list_input_devices(); + + if args.list_devices { + if devices.is_empty() { + eprintln!("[dmvop] No input devices found."); + } else { + println!("Available input devices:"); + for (i, dev) in devices.iter().enumerate() { + println!( + " [{}] {} (id: {}, type: {:?})", + i, dev.name, dev.id, dev.source_type + ); + } + } + return; + } + + // --------------------------------------------------------------- + // 3. Resolve the format pattern + // --------------------------------------------------------------- + let pattern = resolve_format_pattern(args.format_pattern.as_deref(), args.format_file.as_ref()); + + // --------------------------------------------------------------- + // 4. Create and initialize output channels + // --------------------------------------------------------------- + let mut channels: Vec<OutputChannel> = Vec::new(); + + for mode in &args.output { + match create_output_channel(mode, args.port, args.socket_file.clone(), &args.subnet_mask) { + Some(ch) => { + ch.init().await; + channels.push(ch); + } + None => debug_log!( + "[dmvop] Warning: failed to create output channel {:?}", + mode + ), + } + } + + if channels.is_empty() { + eprintln!("[dmvop] No output channels available. Exiting."); + std::process::exit(1); + } + + // --------------------------------------------------------------- + // 5. Find the requested device and start capture + // --------------------------------------------------------------- + let device_name = match &args.device_name { + Some(n) => n.as_str(), + None => { + eprintln!( + "[dmvop] No device specified. Use --device=<name> or --list-devices to see available devices." + ); + std::process::exit(1); + } + }; + + let device = devices + .iter() + .find(|d| d.id == device_name || d.name == device_name) + .or_else(|| devices.first()); + + match &device { + Some(d) => { + eprintln!("[dmvop] Using input device: {} (id: {})", d.name, d.id); + } + None => { + eprintln!( + "[dmvop] Device '{}' not found and no fallback available.", + device_name + ); + std::process::exit(1); + } + } + + engine + .start_capture(device.map(|d| d.id.clone()), None) + .await + .expect("Failed to start audio capture"); + + eprintln!("[dmvop] Capture started. Waiting for speech..."); + + // --------------------------------------------------------------- + // 5. Event loop — listen for transcription & audio level events + // --------------------------------------------------------------- + let mut last_volume_db: f32 = -60.0; + + loop { + match rx.recv().await { + Ok(event) => match event { + vtx_engine::EngineEvent::TranscriptionComplete(result) => { + let formatted = format_output(&pattern, &result.text, 0.0, last_volume_db); + + for ch in &channels { + ch.send(&formatted).await; + } + } + vtx_engine::EngineEvent::TranscriptionSegment(segment) => { + let formatted = format_output(&pattern, &segment.text, 0.0, last_volume_db); + + for ch in &channels { + ch.send(&formatted).await; + } + } + vtx_engine::EngineEvent::VisualizationData(viz) => { + if let Some(ref metrics) = viz.speech_metrics { + last_volume_db = metrics.amplitude_db; + } + } + vtx_engine::EngineEvent::SpeechStarted => { + debug_log!("[dmvop] Speech started"); + } + vtx_engine::EngineEvent::SpeechEnded { duration_ms } => { + debug_log!("[dmvop] Speech ended ({}ms)", duration_ms); + } + vtx_engine::EngineEvent::CaptureStateChanged { capturing, error } => { + if !capturing { + eprintln!( + "[dmvop] Capture stopped: {}", + error.unwrap_or_else(|| "unknown".to_string()) + ); + break; + } + } + vtx_engine::EngineEvent::ModelDownloadProgress { percent } => { + debug_log!("[dmvop] Downloading model: {}%", percent); + } + vtx_engine::EngineEvent::ModelDownloadComplete { success } => { + debug_log!( + "[dmvop] Model download {}", + if success { "complete" } else { "failed" } + ); + } + _ => {} + }, + Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { + debug_log!("[dmvop] Warning: missed {} events", n); + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + eprintln!("[dmvop] Engine event stream closed"); + break; + } + } + } + + eprintln!("[dmvop] Shutting down."); +} + +/// Resolve the format pattern from either the command-line `--format` value +/// or the contents of `--format-file`. +fn resolve_format_pattern(pattern: Option<&str>, file: Option<&PathBuf>) -> String { + if let Some(path) = file { + match std::fs::read_to_string(path) { + Ok(content) => { + let trimmed = content.trim().to_string(); + if !trimmed.is_empty() { + return trimmed; + } + eprintln!( + "[dmvop] Warning: format file {} is empty, using default pattern", + path.display() + ); + } + Err(e) => { + eprintln!( + "[dmvop] Warning: could not read format file {}: {}", + path.display(), + e + ); + } + } + } + + pattern + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .unwrap_or_else(|| "%{vol},%{word}".to_string()) +} + +/// Create an output channel from an [`OutputMode`]. +fn create_output_channel( + mode: &OutputMode, + port: u16, + socket_file: PathBuf, + subnet_mask: &str, +) -> Option<OutputChannel> { + match mode { + OutputMode::STDOUT => Some(OutputChannel::Stdout(Arc::new( + output_protocol::StandardOutputProtocol, + ))), + OutputMode::STDERR => Some(OutputChannel::Stderr(Arc::new( + output_protocol::StandardErrorProtocol, + ))), + OutputMode::TCP => Some(OutputChannel::Tcp(Arc::new( + output_protocol::TCPOutputProtocol::new(port), + ))), + OutputMode::UDP => Some(OutputChannel::Udp(Arc::new( + output_protocol::UDPOutputProtocol::new(port), + ))), + OutputMode::UDP_BROADCAST => Some(OutputChannel::UdpBroadcast(Arc::new( + output_protocol::UDPBroadcastOutputProtocol::new(port, subnet_mask), + ))), + OutputMode::IPC => { + #[cfg(unix)] + { + Some(OutputChannel::Ipc(Arc::new( + output_protocol::IPCOutputProtocol::new(socket_file), + ))) + } + #[cfg(not(unix))] + { + let _ = socket_file; + eprintln!("[dmvop] IPC (Unix domain socket) is not supported on this platform"); + None + } + } + } +} + +/// Print all available Whisper models and their sizes. +fn list_models() { + println!("Available Whisper models:"); + for model in vtx_engine::WhisperModel::all_in_size_order() { + let size = model.size_mb(); + let size_str = if size >= 1024 { + format!("{:.1} GB", size as f64 / 1024.0) + } else { + format!("{} MB", size) + }; + println!( + " {:20} {} ({})", + model.config_key(), + size_str, + model.display_name() + ); + } +} + +/// Download a specific Whisper model by identifier. +async fn download_model_cli(model_name: &str) { + let model = match vtx_engine::WhisperModel::parse_identifier(model_name) { + Some(m) => m, + None => { + eprintln!( + "[dmvop] Unknown model '{}'. Use --list-models to see available models.", + model_name + ); + std::process::exit(1); + } + }; + + eprintln!( + "[dmvop] Building engine with model '{}'...", + model.config_key() + ); + + let (engine, _rx) = match EngineBuilder::new() + .app_name("dmvop") + .model(model) + .build() + .await + { + Ok(e) => e, + Err(e) => { + eprintln!("[dmvop] Failed to build engine: {}", e); + std::process::exit(1); + } + }; + + let status = engine.check_model_status(); + if status.available { + eprintln!("[dmvop] Model already exists at: {}", status.path); + return; + } + + eprintln!( + "[dmvop] Downloading {} ({} MB)...", + model.config_key(), + model.size_mb() + ); + + match engine.download_model().await { + Ok(_) => { + eprintln!("[dmvop] Model downloaded to: {}", status.path); + } + Err(e) => { + eprintln!("[dmvop] Failed to download model: {}", e); + eprintln!("[dmvop] You can manually download from:"); + eprintln!("[dmvop] {}", model.download_url()); + eprintln!("[dmvop] Place it at: {}", status.path); + std::process::exit(1); + } + } } |