jhqxxx
135 lines
4.3 KiB
Rust
135 lines
4.3 KiB
Rust
use std::{net::IpAddr, str::FromStr, time::Duration};
|
|
|
|
use aha::{models::WhichModel, utils::{download_model, get_default_save_dir}};
|
|
use clap::Parser;
|
|
use modelscope::ModelScope;
|
|
use rocket::{
|
|
Config,
|
|
data::{ByteUnit, Limits},
|
|
routes,
|
|
};
|
|
use tokio::time::sleep;
|
|
|
|
use crate::api::init;
|
|
mod api;
|
|
|
|
#[derive(Parser, Debug)]
|
|
#[command(version, about, long_about = None)]
|
|
struct Args {
|
|
#[arg(short, long, default_value = "127.0.0.1")]
|
|
address: String,
|
|
|
|
#[arg(short, long, default_value_t = 10100)]
|
|
port: u16,
|
|
|
|
#[arg(short, long)]
|
|
model: WhichModel,
|
|
|
|
#[arg(long)]
|
|
weight_path: Option<String>,
|
|
|
|
#[arg(long)]
|
|
save_dir: Option<String>,
|
|
|
|
#[arg(long)]
|
|
download_retries: Option<u32>,
|
|
}
|
|
// async fn download_model(model_id: &str, save_dir: &str, max_retries: u32) -> anyhow::Result<()> {
|
|
// let mut attempts = 0u32;
|
|
// loop {
|
|
// attempts += 1;
|
|
// println!(
|
|
// "Attempting to download model (attempt {}/{})",
|
|
// attempts, max_retries
|
|
// );
|
|
|
|
// match ModelScope::download(model_id, save_dir).await {
|
|
// Ok(()) => {
|
|
// println!("Model downloaded successfully");
|
|
// return Ok(());
|
|
// }
|
|
// Err(e) => {
|
|
// if attempts >= max_retries {
|
|
// return Err(anyhow::anyhow!(
|
|
// "Failed to download model after {} attempts. Last error: {}",
|
|
// max_retries,
|
|
// e
|
|
// ));
|
|
// }
|
|
|
|
// println!(
|
|
// "Download failed (attempt {}): {}. Retrying in 2 seconds...",
|
|
// attempts, e
|
|
// );
|
|
// sleep(Duration::from_secs(2)).await;
|
|
// }
|
|
// }
|
|
// }
|
|
// }
|
|
|
|
#[tokio::main]
|
|
async fn main() -> anyhow::Result<()> {
|
|
let args = Args::parse();
|
|
let model_id = match &args.model {
|
|
WhichModel::MiniCPM4_0_5B => "OpenBMB/MiniCPM4-0.5B",
|
|
WhichModel::Qwen2_5vl3B => "Qwen/Qwen2.5-VL-3B-Instruct",
|
|
WhichModel::Qwen2_5vl7B => "Qwen/Qwen2.5-VL-7B-Instruct",
|
|
WhichModel::Qwen3_0_6B => "Qwen/Qwen3-0.6B",
|
|
WhichModel::Qwen3vl2B => "Qwen/Qwen3-VL-2B-Instruct",
|
|
WhichModel::Qwen3vl4B => "Qwen/Qwen3-VL-4B-Instruct",
|
|
WhichModel::Qwen3vl8B => "Qwen/Qwen3-VL-8B-Instruct",
|
|
WhichModel::Qwen3vl32B => "Qwen/Qwen3-VL-32B-Instruct",
|
|
WhichModel::DeepSeekOCR => "deepseek-ai/DeepSeek-OCR",
|
|
WhichModel::HunyuanOCR => "Tencent-Hunyuan/HunyuanOCR",
|
|
WhichModel::PaddleOCRVL => "PaddlePaddle/PaddleOCR-VL",
|
|
WhichModel::RMBG2_0 => "AI-ModelScope/RMBG-2.0",
|
|
WhichModel::VoxCPM => "OpenBMB/VoxCPM-0.5B",
|
|
WhichModel::VoxCPM1_5 => "OpenBMB/VoxCPM1.5",
|
|
WhichModel::GlmASRNano2512 => "ZhipuAI/GLM-ASR-Nano-2512",
|
|
WhichModel::FunASRNano2512 => "FunAudioLLM/Fun-ASR-Nano-2512",
|
|
};
|
|
let model_path = match &args.weight_path {
|
|
Some(path) => path.clone(),
|
|
None => {
|
|
let save_dir = match &args.save_dir {
|
|
Some(dir) => dir.clone(),
|
|
None => get_default_save_dir().expect("Failed to get home directory"),
|
|
};
|
|
let max_retries = args.download_retries.unwrap_or(3);
|
|
download_model(model_id, &save_dir, max_retries).await?;
|
|
save_dir + "/" + model_id
|
|
}
|
|
};
|
|
// println!("-------------------download path: {}", model_path);
|
|
init(args.model, model_path)?;
|
|
start_http_server(&args).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub(crate) async fn start_http_server(args: &Args) -> anyhow::Result<()> {
|
|
let mut builder = rocket::build().configure(Config {
|
|
address: IpAddr::from_str(&args.address)?,
|
|
port: args.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("/chat", routes![api::chat]);
|
|
// /images/remove_background
|
|
builder = builder.mount("/images", routes![api::remove_background]);
|
|
// /images/speech
|
|
builder = builder.mount("/audio", routes![api::speech]);
|
|
|
|
builder.launch().await?;
|
|
Ok(())
|
|
}
|
|
|
|
// fn main() {
|
|
// println!("Hello, world!");
|
|
// }
|