zerg-ztc 0.1.11 → 0.1.13

package/packages/ztc-dictation/src/main.rs

@@ -1,430 +0,0 @@
-use clap::Parser;
-use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
-use ringbuf::{HeapRb, traits::{Consumer, Producer, Split}};
-use serde::Serialize;
-use std::io::{self, Write};
-use std::path::PathBuf;
-use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::{Arc, Mutex};
-use std::time::{Duration, Instant};
-use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};
-
-const SAMPLE_RATE: u32 = 16000;
-const CHANNELS: u16 = 1;
-
-#[derive(Parser, Debug)]
-#[command(name = "ztc-audio")]
-#[command(about = "Real-time audio capture and transcription for ZTC")]
-struct Args {
-    /// Whisper model to use: tiny, base, small, medium, large
-    #[arg(short, long, default_value = "base")]
-    model: String,
-
-    /// Just list audio devices and exit
-    #[arg(long)]
-    list_devices: bool,
-
-    /// Download model if not present and exit
-    #[arg(long)]
-    download_model: bool,
-
-    /// Audio input device name (uses default if not specified)
-    #[arg(short, long)]
-    device: Option<String>,
-}
-
-#[derive(Serialize)]
-#[serde(tag = "type")]
-enum OutputMessage {
-    #[serde(rename = "ready")]
-    Ready { device: String, model: String },
-
-    #[serde(rename = "level")]
-    Level { db: f32, rms: f32 },
-
-    #[serde(rename = "text")]
-    Text { content: String, partial: bool },
-
-    #[serde(rename = "error")]
-    Error { message: String },
-
-    #[serde(rename = "device")]
-    Device { name: String, is_default: bool },
-}
-
-fn emit(msg: &OutputMessage) {
-    if let Ok(json) = serde_json::to_string(msg) {
-        println!("{}", json);
-        let _ = io::stdout().flush();
-    }
-}
-
-fn get_model_path(model_name: &str) -> PathBuf {
-    let models_dir = dirs::home_dir()
-        .unwrap_or_else(|| PathBuf::from("."))
-        .join(".ztc")
-        .join("models");
-
-    std::fs::create_dir_all(&models_dir).ok();
-    models_dir.join(format!("ggml-{}.bin", model_name))
-}
-
-fn download_model(model_name: &str) -> Result<PathBuf, String> {
-    let path = get_model_path(model_name);
-
-    if path.exists() {
-        return Ok(path);
-    }
-
-    let url = format!(
-        "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-{}.bin",
-        model_name
-    );
-
-    eprintln!("Downloading whisper model '{}' from {}", model_name, url);
-
-    let response = ureq::get(&url)
-        .call()
-        .map_err(|e| format!("Failed to download model: {}", e))?;
-
-    let total_size = response
-        .header("Content-Length")
-        .and_then(|s| s.parse::<u64>().ok())
-        .unwrap_or(0);
-
-    let pb = indicatif::ProgressBar::new(total_size);
-    pb.set_style(
-        indicatif::ProgressStyle::default_bar()
-            .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})")
-            .unwrap()
-            .progress_chars("#>-"),
-    );
-
-    let mut reader = response.into_reader();
-    let mut file = std::fs::File::create(&path)
-        .map_err(|e| format!("Failed to create model file: {}", e))?;
-
-    let mut buffer = [0u8; 8192];
-    let mut downloaded = 0u64;
-
-    loop {
-        let bytes_read = reader.read(&mut buffer)
-            .map_err(|e| format!("Failed to read from network: {}", e))?;
-
-        if bytes_read == 0 {
-            break;
-        }
-
-        file.write_all(&buffer[..bytes_read])
-            .map_err(|e| format!("Failed to write to file: {}", e))?;
-
-        downloaded += bytes_read as u64;
-        pb.set_position(downloaded);
-    }
-
-    pb.finish_with_message("Download complete");
-    eprintln!("Model saved to {:?}", path);
-
-    Ok(path)
-}
-
-fn list_devices() {
-    let host = cpal::default_host();
-
-    if let Some(device) = host.default_input_device() {
-        if let Ok(name) = device.name() {
-            emit(&OutputMessage::Device {
-                name,
-                is_default: true,
-            });
-        }
-    }
-
-    if let Ok(devices) = host.input_devices() {
-        for device in devices {
-            if let Ok(name) = device.name() {
-                emit(&OutputMessage::Device {
-                    name,
-                    is_default: false,
-                });
-            }
-        }
-    }
-}
-
-fn calculate_db(samples: &[f32]) -> (f32, f32) {
-    if samples.is_empty() {
-        return (-60.0, 0.0);
-    }
-
-    let sum_squares: f32 = samples.iter().map(|s| s * s).sum();
-    let rms = (sum_squares / samples.len() as f32).sqrt();
-
-    // Convert to dB (with floor at -60dB)
-    let db = if rms > 0.0 {
-        20.0 * rms.log10()
-    } else {
-        -60.0
-    };
-
-    (db.max(-60.0), rms)
-}
-
-fn run_audio_capture(
-    device_name: Option<&str>,
-    model_path: PathBuf,
-) -> Result<(), String> {
-    let host = cpal::default_host();
-
-    // Find the audio device
-    let device = if let Some(name) = device_name {
-        host.input_devices()
-            .map_err(|e| format!("Failed to enumerate devices: {}", e))?
-            .find(|d| d.name().map(|n| n.contains(name)).unwrap_or(false))
-            .ok_or_else(|| format!("Device '{}' not found", name))?
-    } else {
-        host.default_input_device()
-            .ok_or_else(|| "No default input device".to_string())?
-    };
-
-    let device_name_str = device.name().unwrap_or_else(|_| "Unknown".to_string());
-
-    // Initialize Whisper
-    let ctx = WhisperContext::new_with_params(
-        model_path.to_str().unwrap(),
-        WhisperContextParameters::default(),
-    )
-    .map_err(|e| format!("Failed to load whisper model: {}", e))?;
-
-    // Get the device's default config and use its sample rate
-    let default_config = device.default_input_config()
-        .map_err(|e| format!("Failed to get default input config: {}", e))?;
-
-    let device_sample_rate = default_config.sample_rate().0;
-    let device_channels = default_config.channels();
-
-    emit(&OutputMessage::Ready {
-        device: device_name_str.clone(),
-        model: model_path.file_stem()
-            .and_then(|s| s.to_str())
-            .unwrap_or("unknown")
-            .to_string(),
-    });
-
-    // Use device's native sample rate and resample to 16kHz for Whisper
-    let config = cpal::StreamConfig {
-        channels: device_channels,
-        sample_rate: cpal::SampleRate(device_sample_rate),
-        buffer_size: cpal::BufferSize::Default,
-    };
-
-    let resample_ratio = SAMPLE_RATE as f32 / device_sample_rate as f32;
-
-    // Shared state
-    let running = Arc::new(AtomicBool::new(true));
-    let running_clone = running.clone();
-
-    // Ring buffer for audio samples (5 seconds of audio at 16kHz)
-    let ring_buffer = HeapRb::<f32>::new(SAMPLE_RATE as usize * 5);
-    let (mut producer, mut consumer) = ring_buffer.split();
-
-    // Buffer for level calculation (100ms chunks)
-    let level_buffer = Arc::new(Mutex::new(Vec::with_capacity(SAMPLE_RATE as usize / 10)));
-    let level_buffer_clone = level_buffer.clone();
-
-    // Resampling state - accumulate fractional samples
-    let resample_accum = Arc::new(Mutex::new(0.0f32));
-    let resample_accum_clone = resample_accum.clone();
-
-    // Handle Ctrl+C
-    ctrlc::set_handler(move || {
-        running_clone.store(false, Ordering::SeqCst);
-    })
-    .map_err(|e| format!("Failed to set signal handler: {}", e))?;
-
-    // Start audio capture
-    let stream = device
-        .build_input_stream(
-            &config,
-            move |data: &[f32], _: &cpal::InputCallbackInfo| {
-                // Convert to mono and resample to 16kHz
-                let num_channels = device_channels as usize;
-
-                // First convert to mono by averaging channels
-                let mono_samples: Vec<f32> = data
-                    .chunks(num_channels)
-                    .map(|frame| frame.iter().sum::<f32>() / num_channels as f32)
-                    .collect();
-
-                // Simple linear resampling to 16kHz
-                // For better quality, could use a proper resampler crate
-                let mut accum = resample_accum_clone.lock().unwrap();
-                let mut resampled = Vec::new();
-
-                for &sample in &mono_samples {
-                    *accum += resample_ratio;
-                    while *accum >= 1.0 {
-                        resampled.push(sample);
-                        *accum -= 1.0;
-                    }
-                }
-
-                // Push resampled samples to ring buffer for transcription
-                for sample in &resampled {
-                    let _ = producer.try_push(*sample);
-                }
-
-                // Also collect for level metering (use original mono samples for responsiveness)
-                if let Ok(mut buf) = level_buffer_clone.lock() {
-                    buf.extend_from_slice(&mono_samples);
-                }
-            },
-            |err| {
-                emit(&OutputMessage::Error {
-                    message: format!("Audio stream error: {}", err),
-                });
-            },
-            None,
-        )
-        .map_err(|e| format!("Failed to build input stream: {}", e))?;
-
-    stream.play().map_err(|e| format!("Failed to start stream: {}", e))?;
-
-    // Main processing loop
-    let mut accumulated_audio: Vec<f32> = Vec::new();
-    let mut last_level_time = Instant::now();
-    let mut last_transcribe_time = Instant::now();
-    let level_interval = Duration::from_millis(50);  // 20 Hz level updates
-    let transcribe_interval = Duration::from_millis(500);  // Transcribe every 500ms
-
-    while running.load(Ordering::SeqCst) {
-        // Read stdin for commands (non-blocking would be better but this works)
-        // For now, just rely on Ctrl+C / signal handling
-
-        // Emit audio levels
-        if last_level_time.elapsed() >= level_interval {
-            if let Ok(mut buf) = level_buffer.lock() {
-                if !buf.is_empty() {
-                    let (db, rms) = calculate_db(&buf);
-                    emit(&OutputMessage::Level { db, rms });
-                    buf.clear();
-                }
-            }
-            last_level_time = Instant::now();
-        }
-
-        // Collect audio for transcription
-        while let Some(sample) = consumer.try_pop() {
-            accumulated_audio.push(sample);
-        }
-
-        // Transcribe periodically
-        if last_transcribe_time.elapsed() >= transcribe_interval && !accumulated_audio.is_empty() {
-            // Run whisper on accumulated audio
-            let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
-            params.set_language(Some("en"));
-            params.set_print_special(false);
-            params.set_print_progress(false);
-            params.set_print_realtime(false);
-            params.set_print_timestamps(false);
-            params.set_single_segment(true);
-            params.set_no_context(true);
-
-            let mut state = ctx.create_state()
-                .map_err(|e| format!("Failed to create whisper state: {}", e))?;
-
-            if state.full(params, &accumulated_audio).is_ok() {
-                let num_segments = state.full_n_segments()
-                    .unwrap_or(0);
-
-                let mut text = String::new();
-                for i in 0..num_segments {
-                    if let Ok(segment) = state.full_get_segment_text(i) {
-                        text.push_str(&segment);
-                    }
-                }
-
-                let text = text.trim().to_string();
-                if !text.is_empty() {
-                    emit(&OutputMessage::Text {
-                        content: text,
-                        partial: true,
-                    });
-                }
-            }
-
-            last_transcribe_time = Instant::now();
-        }
-
-        std::thread::sleep(Duration::from_millis(10));
-    }
-
-    // Final transcription
-    if !accumulated_audio.is_empty() {
-        let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
-        params.set_language(Some("en"));
-        params.set_print_special(false);
-        params.set_print_progress(false);
-        params.set_print_realtime(false);
-        params.set_print_timestamps(false);
-
-        let mut state = ctx.create_state()
-            .map_err(|e| format!("Failed to create whisper state: {}", e))?;
-
-        if state.full(params, &accumulated_audio).is_ok() {
-            let num_segments = state.full_n_segments().unwrap_or(0);
-
-            let mut text = String::new();
-            for i in 0..num_segments {
-                if let Ok(segment) = state.full_get_segment_text(i) {
-                    text.push_str(&segment);
-                }
-            }
-
-            let text = text.trim().to_string();
-            emit(&OutputMessage::Text {
-                content: text,
-                partial: false,
-            });
-        }
-    }
-
-    Ok(())
-}
-
-fn main() {
-    let args = Args::parse();
-
-    if args.list_devices {
-        list_devices();
-        return;
-    }
-
-    if args.download_model {
-        match download_model(&args.model) {
-            Ok(path) => {
-                eprintln!("Model ready at {:?}", path);
-            }
-            Err(e) => {
-                emit(&OutputMessage::Error { message: e });
-                std::process::exit(1);
-            }
-        }
-        return;
-    }
-
-    // Ensure model is downloaded
-    let model_path = match download_model(&args.model) {
-        Ok(path) => path,
-        Err(e) => {
-            emit(&OutputMessage::Error { message: e });
-            std::process::exit(1);
-        }
-    };
-
-    // Run audio capture and transcription
-    if let Err(e) = run_audio_capture(args.device.as_deref(), model_path) {
-        emit(&OutputMessage::Error { message: e });
-        std::process::exit(1);
-    }
-}