From 2802038b2eb87140873dc23a2e856ef29774a238 Mon Sep 17 00:00:00 2001 From: jhqxxx <18280426169@163.com> Date: Tue, 31 Mar 2026 21:09:29 +0800 Subject: [PATCH] add server and cli mod --- README.md | 1 + README.zh-CN.md | 1 + docs/changelog.md | 1 + docs/changelog.zh-CN.md | 1 + src/cli/args.rs | 194 ++++++++ src/cli/mod.rs | 360 +++++++++++++++ src/exec/mod.rs | 1 - src/exec/voxcpm.rs | 6 +- src/exec/voxcpm1_5.rs | 63 --- src/lib.rs | 1 - src/main.rs | 710 +----------------------------- src/{api/mod.rs => server/api.rs} | 12 +- src/{api => server}/asr.rs | 6 +- src/{api => server}/asr_types.rs | 0 src/server/mod.rs | 77 ++++ src/{ => server}/process.rs | 4 +- src/utils/mod.rs | 60 +++ tests/test_qwen3.rs | 2 +- tests/test_voxcpm1_5.rs | 2 +- 19 files changed, 717 insertions(+), 785 deletions(-) create mode 100644 src/cli/args.rs create mode 100644 src/cli/mod.rs delete mode 100644 src/exec/voxcpm1_5.rs rename src/{api/mod.rs => server/api.rs} (98%) rename src/{api => server}/asr.rs (98%) rename src/{api => server}/asr_types.rs (100%) create mode 100644 src/server/mod.rs rename src/{ => server}/process.rs (99%) diff --git a/README.md b/README.md index 9ea5982..574b742 100644 --- a/README.md +++ b/README.md @@ -114,6 +114,7 @@ curl http://localhost:10100/v1/chat/completions \ ## Changelog ### 2026-03-31 +- add server adn cli mod - aha model name use modelscope id replace - update WhichModel - Usage add time info diff --git a/README.zh-CN.md b/README.zh-CN.md index 0f7a2d0..88c93b0 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -112,6 +112,7 @@ curl http://localhost:10100/v1/chat/completions \ ## 更新日志 ### 2026-03-31 +- 新增 server 和 cli 模块 - aha模型名称使用 modelscope id 替换 - 更新 WhichModel 枚举 - Usage 增加时间信息 diff --git a/docs/changelog.md b/docs/changelog.md index 37172bd..5cfc0e5 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### 2026-03-31 +- add server adn cli mod - aha model name use modelscope id replace - update WhichModel enum - Usage add time info diff --git a/docs/changelog.zh-CN.md b/docs/changelog.zh-CN.md index 2691938..02d4a77 100644 --- a/docs/changelog.zh-CN.md +++ b/docs/changelog.zh-CN.md @@ -6,6 +6,7 @@ 本项目遵循 [语义化版本](https://semver.org/lang/zh-CN/spec/v2.0.0.html)。 ### 2026-03-31 +- 新增 server 和 cli 模块 - aha模型名称使用 modelscope id 替换 - 更新 WhichModel 枚举 - Usage 增加时间信息 diff --git a/src/cli/args.rs b/src/cli/args.rs new file mode 100644 index 0000000..6f313dc --- /dev/null +++ b/src/cli/args.rs @@ -0,0 +1,194 @@ +use aha::models::common::model_mapping::WhichModel; +use clap::{Args, Parser, Subcommand}; + +#[derive(Parser, Debug)] +#[command(name = "aha")] +#[command(version, about, long_about = None)] +pub(crate) struct Cli { + /// Service listen address + #[arg(short, long, default_value = "127.0.0.1")] + pub address: Option, + + /// Service listen port + #[arg(short, long)] + pub port: Option, + + /// Model type (required for backward compatibility) + #[arg(short, long)] + pub model: Option, + + /// Local model weight path + #[arg(long)] + pub weight_path: Option, + + /// Model download save directory + #[arg(long)] + pub save_dir: Option, + + /// Download retry count + #[arg(long)] + pub download_retries: Option, + + /// Local GGUF model weight path (required for loading models with GGUF). + #[arg(long)] + pub gguf_path: Option, + + /// Local path for mmproj GGUF model weights (required for loading with multimodel GGUF) + #[arg(long)] + pub mmproj_path: Option, + + #[command(subcommand)] + pub command: Option, +} + +#[derive(Subcommand, Debug)] +pub(crate) enum Commands { + /// Download model and start service (default) + Cli(CliArgs), + /// Start service only (--weight-path is optional, defaults to ~/.aha/{model_id}) + Serv(ServArgs), + /// List all running aha services + Ps(ServListArgs), + /// Delete a downloaded model from the default location (~/.aha/{model_id}) + Delete(DeleteArgs), + /// Download model only + Download(DownloadArgs), + /// Run model inference directly + Run(RunArgs), + /// List all supported models + List(ListArgs), +} + +/// Common/shared arguments for server operations +#[derive(Args, Debug)] +pub(crate) struct CommonArgs { + /// Service listen address + #[arg(short, long, default_value = "127.0.0.1")] + pub address: String, + + /// Service listen port + #[arg(short, long, default_value_t = 10100)] + pub port: u16, + + /// Model type (required) + #[arg(short, long)] + pub model: WhichModel, + + /// Allow remote shutdown requests (default: local only, use with caution) + #[arg(long)] + pub allow_remote_shutdown: bool, +} + +/// Arguments for the 'cli' subcommand (download + serve) +#[derive(Args, Debug)] +pub(crate) struct CliArgs { + #[command(flatten)] + pub common: CommonArgs, + + /// Local model weight path (skip download if provided) + #[arg(long)] + pub weight_path: Option, + + /// Model download save directory + #[arg(long)] + pub save_dir: Option, + + /// Download retry count + #[arg(long)] + pub download_retries: Option, + + /// Local GGUF model weight path (required for loading models with GGUF). + #[arg(long)] + pub gguf_path: Option, + + /// Local path for mmproj GGUF model weights (required for loading with multimodel GGUF) + #[arg(long)] + pub mmproj_path: Option, +} + +/// Arguments for the 'serv start' subcommand +#[derive(Args, Debug)] +pub(crate) struct ServArgs { + #[command(flatten)] + pub common: CommonArgs, + + /// Local model weight path (defaults to ~/.aha/{model_id} if not specified) + #[arg(long)] + pub weight_path: Option, + + /// Local GGUF model weight path (required for loading models with GGUF). + #[arg(long)] + pub gguf_path: Option, + + /// Local path for mmproj GGUF model weights (required for loading with multimodel GGUF) + #[arg(long)] + pub mmproj_path: Option, +} + +/// Arguments for the 'serv list' subcommand +#[derive(Args, Debug)] +pub(crate) struct ServListArgs { + /// Compact output format + #[arg(short, long)] + pub compact: bool, +} + +/// Arguments for the 'download' subcommand (download only) +#[derive(Args, Debug)] +pub(crate) struct DownloadArgs { + /// Model type (required) + #[arg(short, long)] + pub model: WhichModel, + + /// Model download save directory + #[arg(short, long)] + pub save_dir: Option, + + /// Download retry count + #[arg(long)] + pub download_retries: Option, +} + +/// Arguments for the 'run' subcommand (direct inference) +#[derive(Args, Debug)] +pub(crate) struct RunArgs { + /// Model type (required) + #[arg(short, long)] + pub model: WhichModel, + + /// Input text or file path + #[arg(short, long, num_args = 1..=2, value_delimiter = ' ')] + pub input: Vec, + + /// Output file path (optional) + #[arg(short, long)] + pub output: Option, + + /// Local model weight path (defaults to ~/.aha/{model_id} if not specified) + #[arg(long)] + pub weight_path: Option, + + /// Local GGUF model weight path (required for loading models with GGUF). + #[arg(long)] + pub gguf_path: Option, + + /// Local path for mmproj GGUF model weights (required for loading with multimodel GGUF) + #[arg(long)] + pub mmproj_path: Option, +} + +/// Arguments for the 'delete' subcommand (delete model from default location) +#[derive(Args, Debug)] +pub(crate) struct DeleteArgs { + /// Model type (required) + #[arg(short, long)] + pub model: WhichModel, +} + +/// Arguments for the 'list' subcommand (list all supported models) +#[derive(Args, Debug)] +pub(crate) struct ListArgs { + /// Output models in JSON format (includes name, model_id, and type fields) + #[arg(short, long)] + pub json: bool, +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs new file mode 100644 index 0000000..c9f0323 --- /dev/null +++ b/src/cli/mod.rs @@ -0,0 +1,360 @@ +use crate::{ + cli::args::{CliArgs, DeleteArgs, DownloadArgs, ListArgs, RunArgs, ServArgs, ServListArgs}, + server::{ + api::init, + process::{ServiceStatus, find_aha_services}, + start_http_server, + }, +}; +use aha::exec::*; +use aha::{ + models::common::model_mapping::WhichModel, + utils::{ + bytes_to_human, dir_size, download_model, get_default_save_dir, get_default_weight_path, + is_model_downloaded, + }, +}; +use anyhow::anyhow; +use serde::Serialize; + +pub mod args; + +/// Model information for JSON output +#[derive(Serialize)] +struct ModelInfo { + model_id: String, + owner: String, + #[serde(rename = "type")] + model_type: String, + downloaded: bool, +} + +/// List all supported models +pub(crate) fn run_list(args: ListArgs) -> anyhow::Result<()> { + let models = WhichModel::model_list(); + + if args.json { + // JSON output + let model_infos: Vec = models + .iter() + .map(|model| ModelInfo { + model_id: model.as_string(), + owner: model.model_owner(), + model_type: model.model_type().to_string(), + downloaded: is_model_downloaded(*model), + }) + .collect(); + println!("{}", serde_json::to_string_pretty(&model_infos)?); + } else { + // Table output (default) + println!("Available models:"); + println!(); + println!( + "{:<40} {:<20} {:<10} {:<10}", + "Model ID", "Owner", "type", "Download" + ); + println!("{}", "-".repeat(80)); + for model in models { + let model_id = model.as_string(); + let owner = model.model_owner(); + let model_type = model.model_type(); + let download_status = if is_model_downloaded(model) { + " ✔" + } else { + "" + }; + println!( + "{:<40} {:<20} {:<10} {:<10}", + model_id, owner, model_type, download_status + ); + } + } + + Ok(()) +} + +/// Run the 'cli' subcommand: download model (if needed) and start service +pub(crate) async fn run_cli(args: CliArgs) -> anyhow::Result<()> { + let CliArgs { + common, + weight_path, + save_dir, + download_retries, + gguf_path, + mmproj_path, + } = args; + let model_id = common.model.as_string(); + + let (model_path, gguf, mmproj) = if model_id.contains("gguf") { + if gguf_path.is_none() { + return Err(anyhow!("gguf model path is required")); + } + ("GGUF".to_string(), gguf_path, mmproj_path) + } else { + let model_path = match weight_path { + Some(path) => path, + None => { + let save_dir = match save_dir { + Some(dir) => dir, + None => get_default_save_dir().expect("Failed to get home directory"), + }; + let max_retries = download_retries.unwrap_or(3); + download_model(&model_id, &save_dir, max_retries).await?; + save_dir + "/" + &model_id + } + }; + (model_path, None, None) + }; + + init(common.model, model_path, gguf, mmproj)?; + start_http_server(common.address, common.port, common.allow_remote_shutdown).await?; + + Ok(()) +} + +/// Run the 'serv' subcommand: start service only (no download) +pub(crate) async fn run_serv(args: ServArgs) -> anyhow::Result<()> { + let ServArgs { + common, + weight_path, + gguf_path, + mmproj_path, + } = args; + let model_id = common.model.as_string(); + let (model_path, gguf, mmproj) = if model_id.contains("gguf") { + if gguf_path.is_none() { + return Err(anyhow!("gguf model path is required")); + } + ("GGUF".to_string(), gguf_path, mmproj_path) + } else { + let model_path = match weight_path { + Some(path) => path, + None => get_default_weight_path(common.model), + }; + (model_path, None, None) + }; + + init(common.model, model_path, gguf, mmproj)?; + start_http_server(common.address, common.port, common.allow_remote_shutdown).await?; + + Ok(()) +} + +/// Run the 'ps' subcommand: list running AHA services +pub(crate) fn run_ps(args: ServListArgs) -> anyhow::Result<()> { + let services = find_aha_services()?; + + if services.is_empty() { + println!("No aha services found running."); + return Ok(()); + } + + if args.compact { + // Compact format: one service per line + for svc in services { + println!("{}", svc.service_id); + } + } else { + // Table format + println!( + "{:<20} {:<10} {:<20} {:<10} {:<15} {:<10}", + "Service ID", "PID", "Model", "Port", "Address", "Status" + ); + println!("{}", "-".repeat(85)); + + for svc in services { + let model = svc.model.as_deref().unwrap_or("N/A"); + let status = match svc.status { + ServiceStatus::Running => "Running", + ServiceStatus::Stopping => "Stopping", + ServiceStatus::Unknown => "Unknown", + }; + println!( + "{:<20} {:<10} {:<20} {:<10} {:<15} {:<10}", + svc.service_id, svc.pid, model, svc.port, svc.address, status, + ); + } + } + + Ok(()) +} + +/// Run the 'download' subcommand: download model only (no server) +pub(crate) async fn run_download(args: DownloadArgs) -> anyhow::Result<()> { + let DownloadArgs { + model, + save_dir, + download_retries, + } = args; + let model_id = model.as_string(); + + let save_dir = match save_dir { + Some(dir) => dir, + None => get_default_save_dir().expect("Failed to get home directory"), + }; + let max_retries = download_retries.unwrap_or(3); + + download_model(&model_id, &save_dir, max_retries).await?; + + Ok(()) +} + +/// Run the 'run' subcommand: direct model inference +pub(crate) fn run_run(args: RunArgs) -> anyhow::Result<()> { + let RunArgs { + model, + input, + output, + weight_path, + gguf_path, + mmproj_path, + } = args; + + // Use default weight path if not specified + let weight_path = match weight_path { + Some(path) => path, + None => get_default_weight_path(model), + }; + match model { + WhichModel::MiniCPM4_0_5B => { + minicpm4::MiniCPM4Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::LFM2_1_2B => { + lfm2::Lfm2Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::LFM2_5_1_2BInstruct => { + lfm2::Lfm2Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::LFM2_5VL1_6B => { + lfm2vl::Lfm2VLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::LFM2VL1_6B => { + lfm2vl::Lfm2VLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen2_5VL3B => { + qwen2_5vl::Qwen2_5VLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen2_5VL7B => { + qwen2_5vl::Qwen2_5VLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3_0_6B => { + qwen3::Qwen3Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3_5_0_8B => { + qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3_5_2B => { + qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3_5_4B => { + qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3_5_9B => { + qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3_5Gguf => { + qwen3_5::Qwen3_5Exec::run_gguf(&input, output.as_deref(), gguf_path, mmproj_path)?; + } + WhichModel::Qwen3ASR0_6B => { + qwen3_asr::Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3ASR1_7B => { + qwen3_asr::Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3VL2B => { + qwen3vl::Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3VL4B => { + qwen3vl::Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3VL8B => { + qwen3vl::Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::Qwen3VL32B => { + qwen3vl::Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::DeepSeekOCR => { + deepseek_ocr::DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::DeepSeekOCR2 => { + deepseek_ocr::DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::HunyuanOCR => { + hunyuan_ocr::HunyuanORExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::PaddleOCRVL => { + paddleocr_vl::PaddleOVLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::PaddleOCRVL1_5 => { + paddleocr_vl::PaddleOVLExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::RMBG2_0 => { + rmbg2_0::RMBG2_0Exec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::VoxCPM => { + voxcpm::VoxCPMExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::VoxCPM1_5 => { + voxcpm::VoxCPMExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::GlmASRNano2512 => { + glm_asr_nano::GlmASRNanoExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::FunASRNano2512 => { + fun_asr_nano::FunASRNanoExec::run(&input, output.as_deref(), &weight_path)?; + } + WhichModel::GlmOCR => { + glm_ocr::GlmOcrExec::run(&input, output.as_deref(), &weight_path)?; + } + } + + Ok(()) +} + +/// Run the 'delete' subcommand: delete model from default location +pub(crate) fn run_delete(args: DeleteArgs) -> anyhow::Result<()> { + let DeleteArgs { model } = args; + let model_id = model.as_string(); + let save_dir = get_default_save_dir().expect("Failed to get home directory"); + let model_path = format!("{}/{}", save_dir, model_id); + + let path = std::path::Path::new(&model_path); + + if !path.exists() { + println!("Model not found: {} does not exist", model_path); + return Ok(()); + } + + // Show model info + println!("Model ID: {}", model_id); + println!("Location: {}", model_path); + + // Calculate size if possible + if let Ok(metadata) = std::fs::metadata(path) + && metadata.is_dir() + && let Ok(total_size) = dir_size(path) + { + println!("Size: {}", bytes_to_human(total_size)); + } + + // Confirm deletion + print!("Are you sure you want to delete this model? (y/N): "); + use std::io::Write; + std::io::stdout().flush()?; + + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + + let input = input.trim().to_lowercase(); + if input != "y" && input != "yes" { + println!("Deletion cancelled."); + return Ok(()); + } + + // Delete the directory + std::fs::remove_dir_all(path)?; + + println!("Model deleted successfully: {}", model_path); + + Ok(()) +} diff --git a/src/exec/mod.rs b/src/exec/mod.rs index a2922c8..4939774 100644 --- a/src/exec/mod.rs +++ b/src/exec/mod.rs @@ -19,7 +19,6 @@ pub mod qwen3_asr; pub mod qwen3vl; pub mod rmbg2_0; pub mod voxcpm; -pub mod voxcpm1_5; use anyhow::Result; diff --git a/src/exec/voxcpm.rs b/src/exec/voxcpm.rs index a9c0895..3673fb0 100644 --- a/src/exec/voxcpm.rs +++ b/src/exec/voxcpm.rs @@ -28,8 +28,10 @@ impl ExecModel for VoxCPMExec { let i_start = Instant::now(); let audio = voxcpm_generate.inference( target_text, - Some("啥子小师叔，打狗还要看主人，你再要继续，我就是你的对手".to_string()), // todo args - Some("file://./assets/audio/voice_01.wav".to_string()), // todo args + // Some("啥子小师叔，打狗还要看主人，你再要继续，我就是你的对手".to_string()), // todo args + // Some("file://./assets/audio/voice_01.wav".to_string()), // todo args + None, + None, 2, 100, // max_len (voxcpm uses 100 vs OpenBMB/VoxCPM1.5's 4096) 10, diff --git a/src/exec/voxcpm1_5.rs b/src/exec/voxcpm1_5.rs deleted file mode 100644 index b6c7d6b..0000000 --- a/src/exec/voxcpm1_5.rs +++ /dev/null @@ -1,63 +0,0 @@ -//! VoxCPM1.5 exec implementation for CLI `run` subcommand -//! -//! This module handles VoxCPM1.5 model inference for direct CLI execution. -//! Input/output parameter interpretation is handled here as per the design: -//! - Input can be text content or a file path (with `file://` prefix) -//! - Output can be a file path or will be auto-generated if not specified - -use std::time::Instant; - -use anyhow::{Ok, Result}; - -use crate::models::voxcpm::generate::VoxCPMGenerate; -use crate::{exec::ExecModel, utils::get_file_path}; - -pub struct VoxCPM1_5Exec; - -impl ExecModel for VoxCPM1_5Exec { - fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> { - let input_text = &input[0]; - let target_text = if input_text.starts_with("file://") { - // let path = &input[7..]; - let path = get_file_path(input_text)?; - std::fs::read_to_string(path)? - } else { - input_text.clone() - }; - - let i_start = Instant::now(); - let mut voxcpm_generate = VoxCPMGenerate::init(weight_path, None, None)?; - let i_duration = i_start.elapsed(); - println!("Time elapsed in load model is: {:?}", i_duration); - - let i_start = Instant::now(); - let audio = voxcpm_generate.inference( - target_text, - Some("啥子小师叔，打狗还要看主人，你再要继续，我就是你的对手".to_string()), // todo args - Some("file://./assets/audio/voice_01.wav".to_string()), // todo args - 2, - 4096, - 10, - 2.0, - 6.0, - )?; - let i_duration = i_start.elapsed(); - println!("Time elapsed in generate is: {:?}", i_duration); - - let output_path = if let Some(out) = output { - out.to_string() - } else { - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH)? - .as_secs(); - format!("voxcpm1_5_{}.wav", timestamp) - }; - - let sample_rate = voxcpm_generate.sample_rate(); - crate::utils::audio_utils::save_wav(&audio, &output_path, sample_rate as u32)?; - - println!("Output saved to: {}", output_path); - - Ok(()) - } -} diff --git a/src/lib.rs b/src/lib.rs index 2fbdb33..4626a8e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,5 @@ pub mod exec; pub mod models; pub mod params; pub mod position_embed; -pub mod process; pub mod tokenizer; pub mod utils; diff --git a/src/main.rs b/src/main.rs index 9b70378..f6ee4c3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,647 +1,12 @@ -use std::sync::atomic::{AtomicBool, Ordering}; -use std::{net::IpAddr, str::FromStr, sync::Arc}; +use clap::Parser; -use aha::{ - models::common::model_mapping::WhichModel, - process::{cleanup_pid_file, create_pid_file}, - utils::{download_model, get_default_save_dir}, +use crate::cli::{ + args::{Cli, CliArgs, Commands, CommonArgs}, + run_cli, run_delete, run_download, run_list, run_ps, run_run, run_serv, }; -use anyhow::anyhow; -use clap::{Args, Parser, Subcommand}; -use rocket::{ - Config, - data::{ByteUnit, Limits}, - routes, -}; -use serde::Serialize; -use crate::api::{init, set_server_port}; -mod api; - -#[derive(Parser, Debug)] -#[command(name = "aha")] -#[command(version, about, long_about = None)] -struct Cli { - /// Service listen address - #[arg(short, long, default_value = "127.0.0.1")] - address: Option, - - /// Service listen port - #[arg(short, long)] - port: Option, - - /// Model type (required for backward compatibility) - #[arg(short, long)] - model: Option, - - /// Local model weight path - #[arg(long)] - weight_path: Option, - - /// Model download save directory - #[arg(long)] - save_dir: Option, - - /// Download retry count - #[arg(long)] - download_retries: Option, - - /// Local GGUF model weight path (required for loading models with GGUF). - #[arg(long)] - gguf_path: Option, - - /// Local path for mmproj GGUF model weights (required for loading with multimodel GGUF) - #[arg(long)] - mmproj_path: Option, - - #[command(subcommand)] - command: Option, -} - -#[derive(Subcommand, Debug)] -enum Commands { - /// Download model and start service (default) - Cli(CliArgs), - /// Start service only (--weight-path is optional, defaults to ~/.aha/{model_id}) - Serv(ServArgs), - /// List all running aha services - Ps(ServListArgs), - /// Delete a downloaded model from the default location (~/.aha/{model_id}) - Delete(DeleteArgs), - /// Download model only - Download(DownloadArgs), - /// Run model inference directly - Run(RunArgs), - /// List all supported models - List(ListArgs), -} - -/// Common/shared arguments for server operations -#[derive(Args, Debug)] -struct CommonArgs { - /// Service listen address - #[arg(short, long, default_value = "127.0.0.1")] - address: String, - - /// Service listen port - #[arg(short, long, default_value_t = 10100)] - port: u16, - - /// Model type (required) - #[arg(short, long)] - model: WhichModel, - - /// Allow remote shutdown requests (default: local only, use with caution) - #[arg(long)] - allow_remote_shutdown: bool, -} - -/// Arguments for the 'cli' subcommand (download + serve) -#[derive(Args, Debug)] -struct CliArgs { - #[command(flatten)] - common: CommonArgs, - - /// Local model weight path (skip download if provided) - #[arg(long)] - weight_path: Option, - - /// Model download save directory - #[arg(long)] - save_dir: Option, - - /// Download retry count - #[arg(long)] - download_retries: Option, - - /// Local GGUF model weight path (required for loading models with GGUF). - #[arg(long)] - gguf_path: Option, - - /// Local path for mmproj GGUF model weights (required for loading with multimodel GGUF) - #[arg(long)] - mmproj_path: Option, -} - -/// Arguments for the 'serv start' subcommand -#[derive(Args, Debug)] -struct ServArgs { - #[command(flatten)] - common: CommonArgs, - - /// Local model weight path (defaults to ~/.aha/{model_id} if not specified) - #[arg(long)] - weight_path: Option, - - /// Local GGUF model weight path (required for loading models with GGUF). - #[arg(long)] - gguf_path: Option, - - /// Local path for mmproj GGUF model weights (required for loading with multimodel GGUF) - #[arg(long)] - mmproj_path: Option, -} - -/// Arguments for the 'serv list' subcommand -#[derive(Args, Debug)] -struct ServListArgs { - /// Compact output format - #[arg(short, long)] - compact: bool, -} - -/// Arguments for the 'download' subcommand (download only) -#[derive(Args, Debug)] -struct DownloadArgs { - /// Model type (required) - #[arg(short, long)] - model: WhichModel, - - /// Model download save directory - #[arg(short, long)] - save_dir: Option, - - /// Download retry count - #[arg(long)] - download_retries: Option, -} - -/// Arguments for the 'run' subcommand (direct inference) -#[derive(Args, Debug)] -struct RunArgs { - /// Model type (required) - #[arg(short, long)] - model: WhichModel, - - /// Input text or file path - #[arg(short, long, num_args = 1..=2, value_delimiter = ' ')] - input: Vec, - - /// Output file path (optional) - #[arg(short, long)] - output: Option, - - /// Local model weight path (defaults to ~/.aha/{model_id} if not specified) - #[arg(long)] - weight_path: Option, - - /// Local GGUF model weight path (required for loading models with GGUF). - #[arg(long)] - gguf_path: Option, - - /// Local path for mmproj GGUF model weights (required for loading with multimodel GGUF) - #[arg(long)] - mmproj_path: Option, -} - -/// Arguments for the 'delete' subcommand (delete model from default location) -#[derive(Args, Debug)] -struct DeleteArgs { - /// Model type (required) - #[arg(short, long)] - model: WhichModel, -} - -/// Arguments for the 'list' subcommand (list all supported models) -#[derive(Args, Debug)] -struct ListArgs { - /// Output models in JSON format (includes name, model_id, and type fields) - #[arg(short, long)] - json: bool, -} - -/// Get the default weight path for a given model -/// Returns ~/.aha/{model_id} e.g., ~/.aha/OpenBMB/VoxCPM1.5 -fn get_default_weight_path(model: WhichModel) -> String { - let model_id = model.as_string(); - let save_dir = get_default_save_dir().expect("Failed to get home directory"); - format!("{}/{}", save_dir, model_id) -} - -/// Check if a model is downloaded by verifying the model directory exists -/// Returns true if ~/.aha/{model_id} directory exists, false otherwise -fn is_model_downloaded(model: WhichModel) -> bool { - let model_id = model.as_string(); - let save_dir = match get_default_save_dir() { - Some(dir) => dir, - None => return false, - }; - let model_path = format!("{}/{}", save_dir, model_id); - std::path::Path::new(&model_path).exists() -} - -/// Model information for JSON output -#[derive(Serialize)] -struct ModelInfo { - model_id: String, - owner: String, - #[serde(rename = "type")] - model_type: String, - downloaded: bool, -} - -/// List all supported models -fn run_list(args: ListArgs) -> anyhow::Result<()> { - let models = WhichModel::model_list(); - - if args.json { - // JSON output - let model_infos: Vec = models - .iter() - .map(|model| ModelInfo { - model_id: model.as_string(), - owner: model.model_owner(), - model_type: model.model_type().to_string(), - downloaded: is_model_downloaded(*model), - }) - .collect(); - println!("{}", serde_json::to_string_pretty(&model_infos)?); - } else { - // Table output (default) - println!("Available models:"); - println!(); - println!( - "{:<40} {:<20} {:<10} {:<10}", - "Model ID", "Owner", "type", "Download" - ); - println!("{}", "-".repeat(80)); - for model in models { - let model_id = model.as_string(); - let owner = model.model_owner(); - let model_type = model.model_type(); - let download_status = if is_model_downloaded(model) { - " ✔" - } else { - "" - }; - println!( - "{:<40} {:<20} {:<10} {:<10}", - model_id, owner, model_type, download_status - ); - } - } - - Ok(()) -} - -/// Run the 'cli' subcommand: download model (if needed) and start service -async fn run_cli(args: CliArgs) -> anyhow::Result<()> { - let CliArgs { - common, - weight_path, - save_dir, - download_retries, - gguf_path, - mmproj_path, - } = args; - let model_id = common.model.as_string(); - - let (model_path, gguf, mmproj) = if model_id.contains("gguf") { - if gguf_path.is_none() { - return Err(anyhow!("gguf model path is required")); - } - ("GGUF".to_string(), gguf_path, mmproj_path) - } else { - let model_path = match weight_path { - Some(path) => path, - None => { - let save_dir = match save_dir { - Some(dir) => dir, - None => get_default_save_dir().expect("Failed to get home directory"), - }; - let max_retries = download_retries.unwrap_or(3); - download_model(&model_id, &save_dir, max_retries).await?; - save_dir + "/" + &model_id - } - }; - (model_path, None, None) - }; - - init(common.model, model_path, gguf, mmproj)?; - start_http_server(common.address, common.port, common.allow_remote_shutdown).await?; - - Ok(()) -} - -/// Run the 'serv' subcommand: start service only (no download) -async fn run_serv(args: ServArgs) -> anyhow::Result<()> { - let ServArgs { - common, - weight_path, - gguf_path, - mmproj_path, - } = args; - let model_id = common.model.as_string(); - let (model_path, gguf, mmproj) = if model_id.contains("gguf") { - if gguf_path.is_none() { - return Err(anyhow!("gguf model path is required")); - } - ("GGUF".to_string(), gguf_path, mmproj_path) - } else { - let model_path = match weight_path { - Some(path) => path, - None => get_default_weight_path(common.model), - }; - (model_path, None, None) - }; - - init(common.model, model_path, gguf, mmproj)?; - start_http_server(common.address, common.port, common.allow_remote_shutdown).await?; - - Ok(()) -} - -/// Run the 'ps' subcommand: list running AHA services -fn run_ps(args: ServListArgs) -> anyhow::Result<()> { - use aha::process::find_aha_services; - - let services = find_aha_services()?; - - if services.is_empty() { - println!("No aha services found running."); - return Ok(()); - } - - if args.compact { - // Compact format: one service per line - for svc in services { - println!("{}", svc.service_id); - } - } else { - // Table format - println!( - "{:<20} {:<10} {:<20} {:<10} {:<15} {:<10}", - "Service ID", "PID", "Model", "Port", "Address", "Status" - ); - println!("{}", "-".repeat(85)); - - for svc in services { - let model = svc.model.as_deref().unwrap_or("N/A"); - let status = match svc.status { - aha::process::ServiceStatus::Running => "Running", - aha::process::ServiceStatus::Stopping => "Stopping", - aha::process::ServiceStatus::Unknown => "Unknown", - }; - println!( - "{:<20} {:<10} {:<20} {:<10} {:<15} {:<10}", - svc.service_id, svc.pid, model, svc.port, svc.address, status, - ); - } - } - - Ok(()) -} - -/// Run the 'download' subcommand: download model only (no server) -async fn run_download(args: DownloadArgs) -> anyhow::Result<()> { - let DownloadArgs { - model, - save_dir, - download_retries, - } = args; - let model_id = model.as_string(); - - let save_dir = match save_dir { - Some(dir) => dir, - None => get_default_save_dir().expect("Failed to get home directory"), - }; - let max_retries = download_retries.unwrap_or(3); - - download_model(&model_id, &save_dir, max_retries).await?; - - Ok(()) -} - -/// Run the 'run' subcommand: direct model inference -fn run_run(args: RunArgs) -> anyhow::Result<()> { - use aha::exec::ExecModel; - - let RunArgs { - model, - input, - output, - weight_path, - gguf_path, - mmproj_path, - } = args; - - // Use default weight path if not specified - let weight_path = match weight_path { - Some(path) => path, - None => get_default_weight_path(model), - }; - match model { - WhichModel::MiniCPM4_0_5B => { - use aha::exec::minicpm4::MiniCPM4Exec; - MiniCPM4Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::LFM2_1_2B => { - use aha::exec::lfm2::Lfm2Exec; - Lfm2Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::LFM2_5_1_2BInstruct => { - use aha::exec::lfm2::Lfm2Exec; - Lfm2Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::LFM2_5VL1_6B => { - use aha::exec::lfm2vl::Lfm2VLExec; - Lfm2VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::LFM2VL1_6B => { - use aha::exec::lfm2vl::Lfm2VLExec; - Lfm2VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen2_5VL3B => { - use aha::exec::qwen2_5vl::Qwen2_5VLExec; - Qwen2_5VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen2_5VL7B => { - use aha::exec::qwen2_5vl::Qwen2_5VLExec; - Qwen2_5VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_0_6B => { - use aha::exec::qwen3::Qwen3Exec; - Qwen3Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5_0_8B => { - use aha::exec::qwen3_5::Qwen3_5Exec; - Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5_2B => { - use aha::exec::qwen3_5::Qwen3_5Exec; - Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5_4B => { - use aha::exec::qwen3_5::Qwen3_5Exec; - Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5_9B => { - use aha::exec::qwen3_5::Qwen3_5Exec; - Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3_5Gguf => { - use aha::exec::qwen3_5::Qwen3_5Exec; - Qwen3_5Exec::run_gguf(&input, output.as_deref(), gguf_path, mmproj_path)?; - } - WhichModel::Qwen3ASR0_6B => { - use aha::exec::qwen3_asr::Qwen3ASRExec; - Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3ASR1_7B => { - use aha::exec::qwen3_asr::Qwen3ASRExec; - Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3VL2B => { - use aha::exec::qwen3vl::Qwen3VLExec; - Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3VL4B => { - use aha::exec::qwen3vl::Qwen3VLExec; - Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3VL8B => { - use aha::exec::qwen3vl::Qwen3VLExec; - Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::Qwen3VL32B => { - use aha::exec::qwen3vl::Qwen3VLExec; - Qwen3VLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::DeepSeekOCR => { - use aha::exec::deepseek_ocr::DeepSeekORExec; - DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::DeepSeekOCR2 => { - use aha::exec::deepseek_ocr::DeepSeekORExec; - DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::HunyuanOCR => { - use aha::exec::hunyuan_ocr::HunyuanORExec; - HunyuanORExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::PaddleOCRVL => { - use aha::exec::paddleocr_vl::PaddleOVLExec; - PaddleOVLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::PaddleOCRVL1_5 => { - use aha::exec::paddleocr_vl::PaddleOVLExec; - PaddleOVLExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::RMBG2_0 => { - use aha::exec::rmbg2_0::RMBG2_0Exec; - RMBG2_0Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::VoxCPM => { - use aha::exec::voxcpm::VoxCPMExec; - VoxCPMExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::VoxCPM1_5 => { - use aha::exec::voxcpm1_5::VoxCPM1_5Exec; - VoxCPM1_5Exec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::GlmASRNano2512 => { - use aha::exec::glm_asr_nano::GlmASRNanoExec; - GlmASRNanoExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::FunASRNano2512 => { - use aha::exec::fun_asr_nano::FunASRNanoExec; - FunASRNanoExec::run(&input, output.as_deref(), &weight_path)?; - } - WhichModel::GlmOCR => { - use aha::exec::glm_ocr::GlmOcrExec; - GlmOcrExec::run(&input, output.as_deref(), &weight_path)?; - } - } - - Ok(()) -} - -/// Run the 'delete' subcommand: delete model from default location -fn run_delete(args: DeleteArgs) -> anyhow::Result<()> { - let DeleteArgs { model } = args; - let model_id = model.as_string(); - let save_dir = get_default_save_dir().expect("Failed to get home directory"); - let model_path = format!("{}/{}", save_dir, model_id); - - let path = std::path::Path::new(&model_path); - - if !path.exists() { - println!("Model not found: {} does not exist", model_path); - return Ok(()); - } - - // Show model info - println!("Model ID: {}", model_id); - println!("Location: {}", model_path); - - // Calculate size if possible - if let Ok(metadata) = std::fs::metadata(path) - && metadata.is_dir() - && let Ok(total_size) = dir_size(path) - { - println!("Size: {}", bytes_to_human(total_size)); - } - - // Confirm deletion - print!("Are you sure you want to delete this model? (y/N): "); - use std::io::Write; - std::io::stdout().flush()?; - - let mut input = String::new(); - std::io::stdin().read_line(&mut input)?; - - let input = input.trim().to_lowercase(); - if input != "y" && input != "yes" { - println!("Deletion cancelled."); - return Ok(()); - } - - // Delete the directory - std::fs::remove_dir_all(path)?; - - println!("Model deleted successfully: {}", model_path); - - Ok(()) -} - -/// Calculate total size of a directory recursively -fn dir_size(path: &std::path::Path) -> anyhow::Result { - let mut total = 0; - if path.is_dir() { - for entry in std::fs::read_dir(path)? { - let entry = entry?; - let entry_path = entry.path(); - if entry_path.is_dir() { - total += dir_size(&entry_path)?; - } else { - total += entry.metadata()?.len(); - } - } - } else { - total = std::fs::metadata(path)?.len(); - } - Ok(total) -} - -/// Convert bytes to human readable format -fn bytes_to_human(bytes: u64) -> String { - const KB: u64 = 1024; - const MB: u64 = KB * 1024; - const GB: u64 = MB * 1024; - const TB: u64 = GB * 1024; - - if bytes >= TB { - format!("{:.2} TB", bytes as f64 / TB as f64) - } else if bytes >= GB { - format!("{:.2} GB", bytes as f64 / GB as f64) - } else if bytes >= MB { - format!("{:.2} MB", bytes as f64 / MB as f64) - } else if bytes >= KB { - format!("{:.2} KB", bytes as f64 / KB as f64) - } else { - format!("{} B", bytes) - } -} +mod cli; +mod server; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -675,66 +40,3 @@ async fn main() -> anyhow::Result<()> { } } } - -pub(crate) async fn start_http_server( - address: String, - port: u16, - allow_remote_shutdown: bool, -) -> anyhow::Result<()> { - // Set server port for shutdown endpoint - set_server_port(port, allow_remote_shutdown); - - // Create PID file for service tracking - let pid = std::process::id(); - create_pid_file(pid, port)?; - - // Set up shutdown flag - let shutdown_flag = Arc::new(AtomicBool::new(false)); - let shutdown_flag_clone = shutdown_flag.clone(); - - // Configure Ctrl+C handler for graceful shutdown - let port_for_cleanup = port; - let shutdown_handler = tokio::spawn(async move { - tokio::signal::ctrl_c().await.ok(); - println!("Received shutdown signal, gracefully shutting down..."); - shutdown_flag_clone.store(true, Ordering::SeqCst); - // Give time for existing requests to complete - tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; - // Cleanup PID file - let _ = cleanup_pid_file(port_for_cleanup); - std::process::exit(0); - }); - - let mut builder = rocket::build().configure(Config { - address: IpAddr::from_str(&address)?, - port, - limits: Limits::default() - .limit("string", ByteUnit::Mebibyte(5)) - .limit("json", ByteUnit::Mebibyte(5)) - .limit("data-form", ByteUnit::Mebibyte(100)) - .limit("file", ByteUnit::Mebibyte(100)), - ..Config::default() - }); - - builder = builder.mount("/v1/chat", routes![api::chat]); - builder = builder.mount("/chat", routes![api::chat]); - // /images/remove_background - builder = builder.mount("/images", routes![api::remove_background]); - // /audio/speech and /audio/transcriptions (ASR transcription endpoint) - builder = builder.mount("/audio", routes![api::speech, api::transcriptions]); - // /v1/audio/transcriptions (OpenAI standard ASR transcription endpoint) - builder = builder.mount("/v1/audio", routes![api::transcriptions]); - // Health check and model info endpoints - builder = builder.mount("/", routes![api::health, api::models]); - // Shutdown endpoint - builder = builder.manage(shutdown_flag); - builder = builder.mount("/", routes![api::shutdown]); - - let _rocket = builder.launch().await?; - - // Cleanup PID file when server exits - cleanup_pid_file(port)?; - shutdown_handler.abort(); - - Ok(()) -} diff --git a/src/api/mod.rs b/src/server/api.rs similarity index 98% rename from src/api/mod.rs rename to src/server/api.rs index 904810a..a91e4a1 100644 --- a/src/api/mod.rs +++ b/src/server/api.rs @@ -4,7 +4,6 @@ use std::sync::{Arc, OnceLock}; use aha::models::{GenerateModel, ModelInstance, common::model_mapping::WhichModel, load_model}; use aha::params::chat::ChatCompletionParameters; -use aha::process::cleanup_pid_file; use aha::utils::string_to_static_str; use rocket::futures::StreamExt; use rocket::serde::{Serialize, json::Json}; @@ -18,17 +17,12 @@ use rocket::{ }; use tokio::sync::RwLock; -// ASR (Automatic Speech Recognition) API module -pub(crate) mod asr; -pub(crate) mod asr_types; - -// Re-export ASR routes -pub(crate) use asr::transcriptions; +use crate::server::process::cleanup_pid_file; /// Wrapper to store model type together with the model instance pub(crate) struct StoredModel { - which_model: WhichModel, - instance: ModelInstance<'static>, + pub which_model: WhichModel, + pub instance: ModelInstance<'static>, } // Export MODEL for use in ASR module diff --git a/src/api/asr.rs b/src/server/asr.rs similarity index 98% rename from src/api/asr.rs rename to src/server/asr.rs index b48f155..26be045 100644 --- a/src/api/asr.rs +++ b/src/server/asr.rs @@ -11,8 +11,10 @@ use rocket::http::Status; use rocket::serde::json::Json; use rocket::{form::Form, post}; -use super::MODEL; -use super::asr_types::{ErrorDetail, ErrorResponse, TranscriptionRequest, TranscriptionResponse}; +use crate::server::api::MODEL; +use crate::server::asr_types::{ + ErrorDetail, ErrorResponse, TranscriptionRequest, TranscriptionResponse, +}; /// Handle audio transcription requests /// diff --git a/src/api/asr_types.rs b/src/server/asr_types.rs similarity index 100% rename from src/api/asr_types.rs rename to src/server/asr_types.rs diff --git a/src/server/mod.rs b/src/server/mod.rs new file mode 100644 index 0000000..d9a81a0 --- /dev/null +++ b/src/server/mod.rs @@ -0,0 +1,77 @@ +use crate::server::api::set_server_port; +use crate::server::process::{cleanup_pid_file, create_pid_file}; +use rocket::data::{ByteUnit, Limits}; +use rocket::{Config, routes}; +use std::net::IpAddr; +use std::str::FromStr; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +// ASR (Automatic Speech Recognition) API module +pub(crate) mod api; +pub(crate) mod asr; +pub(crate) mod asr_types; +pub(crate) mod process; + +pub(crate) async fn start_http_server( + address: String, + port: u16, + allow_remote_shutdown: bool, +) -> anyhow::Result<()> { + // Set server port for shutdown endpoint + set_server_port(port, allow_remote_shutdown); + + // Create PID file for service tracking + let pid = std::process::id(); + create_pid_file(pid, port)?; + + // Set up shutdown flag + let shutdown_flag = Arc::new(AtomicBool::new(false)); + let shutdown_flag_clone = shutdown_flag.clone(); + + // Configure Ctrl+C handler for graceful shutdown + let port_for_cleanup = port; + let shutdown_handler = tokio::spawn(async move { + tokio::signal::ctrl_c().await.ok(); + println!("Received shutdown signal, gracefully shutting down..."); + shutdown_flag_clone.store(true, Ordering::SeqCst); + // Give time for existing requests to complete + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + // Cleanup PID file + let _ = cleanup_pid_file(port_for_cleanup); + std::process::exit(0); + }); + + let mut builder = rocket::build().configure(Config { + address: IpAddr::from_str(&address)?, + port, + limits: Limits::default() + .limit("string", ByteUnit::Mebibyte(5)) + .limit("json", ByteUnit::Mebibyte(5)) + .limit("data-form", ByteUnit::Mebibyte(100)) + .limit("file", ByteUnit::Mebibyte(100)), + ..Config::default() + }); + + builder = builder.mount("/v1/chat", routes![api::chat]); + builder = builder.mount("/chat", routes![api::chat]); + // /images/remove_background + builder = builder.mount("/images", routes![api::remove_background]); + // /audio/speech and /audio/transcriptions (ASR transcription endpoint) + builder = builder.mount("/audio", routes![api::speech, asr::transcriptions]); + // /v1/audio/transcriptions (OpenAI standard ASR transcription endpoint) + builder = builder.mount("/v1/audio", routes![asr::transcriptions]); + // Health check and model info endpoints + builder = builder.mount("/", routes![api::health, api::models]); + // Shutdown endpoint + builder = builder.manage(shutdown_flag); + builder = builder.mount("/", routes![api::shutdown]); + + let _rocket = builder.launch().await?; + + // Cleanup PID file when server exits + cleanup_pid_file(port)?; + shutdown_handler.abort(); + + Ok(()) +} diff --git a/src/process.rs b/src/server/process.rs similarity index 99% rename from src/process.rs rename to src/server/process.rs index 7079185..180dac3 100644 --- a/src/process.rs +++ b/src/server/process.rs @@ -30,6 +30,7 @@ pub struct ServiceInfo { /// Service status #[derive(Debug, Clone, PartialEq)] +#[allow(unused)] pub enum ServiceStatus { Running, Stopping, @@ -102,7 +103,8 @@ pub fn cleanup_pid_file(port: u16) -> Result<()> { /// /// # Arguments /// * `port` - Listen port -pub fn get_pid_from_file(port: u16) -> Option { +#[allow(unused)] +fn get_pid_from_file(port: u16) -> Option { let pid_dir = get_pid_dir().ok()?; let pid_file = pid_dir.join(format!("{}.pid", port)); diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 908f9f9..ddf3d8a 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -9,6 +9,7 @@ use std::io::{Cursor, Read}; use std::time::{SystemTime, UNIX_EPOCH}; use std::{collections::HashMap, fs, path::PathBuf, process::Command, time::Duration}; +use crate::models::common::model_mapping::WhichModel; use crate::params::{ chat::{ AudioUrlType, ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, @@ -975,6 +976,65 @@ pub fn clean_asr_response(raw: &str) -> String { } } +/// Get the default weight path for a given model +/// Returns ~/.aha/{model_id} e.g., ~/.aha/OpenBMB/VoxCPM1.5 +pub fn get_default_weight_path(model: WhichModel) -> String { + let model_id = model.as_string(); + let save_dir = get_default_save_dir().expect("Failed to get home directory"); + format!("{}/{}", save_dir, model_id) +} + +/// Check if a model is downloaded by verifying the model directory exists +/// Returns true if ~/.aha/{model_id} directory exists, false otherwise +pub fn is_model_downloaded(model: WhichModel) -> bool { + let model_id = model.as_string(); + let save_dir = match get_default_save_dir() { + Some(dir) => dir, + None => return false, + }; + let model_path = format!("{}/{}", save_dir, model_id); + std::path::Path::new(&model_path).exists() +} + +/// Calculate total size of a directory recursively +pub fn dir_size(path: &std::path::Path) -> anyhow::Result { + let mut total = 0; + if path.is_dir() { + for entry in std::fs::read_dir(path)? { + let entry = entry?; + let entry_path = entry.path(); + if entry_path.is_dir() { + total += dir_size(&entry_path)?; + } else { + total += entry.metadata()?.len(); + } + } + } else { + total = std::fs::metadata(path)?.len(); + } + Ok(total) +} + +/// Convert bytes to human readable format +pub fn bytes_to_human(bytes: u64) -> String { + const KB: u64 = 1024; + const MB: u64 = KB * 1024; + const GB: u64 = MB * 1024; + const TB: u64 = GB * 1024; + + if bytes >= TB { + format!("{:.2} TB", bytes as f64 / TB as f64) + } else if bytes >= GB { + format!("{:.2} GB", bytes as f64 / GB as f64) + } else if bytes >= MB { + format!("{:.2} MB", bytes as f64 / MB as f64) + } else if bytes >= KB { + format!("{:.2} KB", bytes as f64 / KB as f64) + } else { + format!("{} B", bytes) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/tests/test_qwen3.rs b/tests/test_qwen3.rs index e06cc33..fd94b51 100644 --- a/tests/test_qwen3.rs +++ b/tests/test_qwen3.rs @@ -19,7 +19,7 @@ fn qwen3_0_6b_generate() -> Result<()> { "messages": [ { "role": "user", - "content": "你吃饭了没" + "content": "你好啊，你是谁" } ] } diff --git a/tests/test_voxcpm1_5.rs b/tests/test_voxcpm1_5.rs index 19ef358..569b98a 100644 --- a/tests/test_voxcpm1_5.rs +++ b/tests/test_voxcpm1_5.rs @@ -12,7 +12,7 @@ use anyhow::{Ok, Result}; #[test] fn voxcpm1_5_use_message_generate() -> Result<()> { - // RUST_BACKTRACE=1 cargo test -F cuda voxcpm1_5_use_message_generate -r -- --nocapture + // RUST_BACKTRACE=1 cargo test -F cuda --test test_voxcpm1_5 voxcpm1_5_use_message_generate -r -- --nocapture let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; let model_path = format!("{}/OpenBMB/VoxCPM1.5/", save_dir);