tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update WhichModel enum

Browse Source
This commit is contained in:
jhqxxx
2026-03-31 18:45:01 +08:00
parent c5551e7fb8
commit 741869d6dd
44 changed files with 691 additions and 812 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/api/mod.rs
+31 -102
View File
@@ -2,7 +2,7 @@ use std::pin::pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, OnceLock};
use aha::models::{GenerateModel, ModelInstance, WhichModel, load_model};
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;
@@ -243,70 +243,34 @@ struct ErrorResponse {
error: String,
}
/// Convert WhichModel to a display-friendly model ID (kebab-case)
fn which_model_to_id(which_model: WhichModel) -> &'static str {
match which_model {
WhichModel::MiniCPM4_0_5B => "minicpm4-0.5b",
WhichModel::LFM2_1_2B => "lfm2-1.2b",
WhichModel::LFM2_5_1_2BInstruct => "lfm2.5-1.2b-instruct",
WhichModel::LFM2_5VL1_6B => "lfm2.5-vl-1.6b",
WhichModel::LFM2VL1_6B => "lfm2-vl-1.6b",
WhichModel::Qwen2_5VL3B => "qwen2.5vl-3b",
WhichModel::Qwen2_5VL7B => "qwen2.5vl-7b",
WhichModel::Qwen3_0_6B => "qwen3-0.6b",
WhichModel::Qwen3_5_0_8B => "qwen3.5-0.8b",
WhichModel::Qwen3_5_2B => "qwen3.5-2b",
WhichModel::Qwen3_5_4B => "qwen3.5-4b",
WhichModel::Qwen3_5_9B => "qwen3.5-9b",
WhichModel::Qwen3_5Gguf => "qwen3.5-gguf",
WhichModel::Qwen3ASR0_6B => "qwen3asr-0.6b",
WhichModel::Qwen3ASR1_7B => "qwen3asr-1.7b",
WhichModel::Qwen3VL2B => "qwen3vl-2b",
WhichModel::Qwen3VL4B => "qwen3vl-4b",
WhichModel::Qwen3VL8B => "qwen3vl-8b",
WhichModel::Qwen3VL32B => "qwen3vl-32b",
WhichModel::DeepSeekOCR => "deepseek-ocr",
WhichModel::DeepSeekOCR2 => "deepseek-ocr2",
WhichModel::HunyuanOCR => "hunyuan-ocr",
WhichModel::PaddleOCRVL => "paddleocr-vl",
WhichModel::PaddleOCRVL1_5 => "paddleocr-vl1.5",
WhichModel::RMBG2_0 => "rmbg2.0",
WhichModel::VoxCPM => "voxcpm",
WhichModel::VoxCPM1_5 => "voxcpm1.5",
WhichModel::GlmASRNano2512 => "glm-asr-nano-2512",
WhichModel::FunASRNano2512 => "fun-asr-nano-2512",
WhichModel::GlmOCR => "glm-ocr",
}
}
/// Get the owner/organization name for a model
fn which_model_to_owner(which_model: WhichModel) -> &'static str {
match which_model {
WhichModel::MiniCPM4_0_5B => "OpenBMB",
WhichModel::Qwen2_5VL3B | WhichModel::Qwen2_5VL7B => "Qwen",
WhichModel::Qwen3_0_6B | WhichModel::Qwen3ASR0_6B | WhichModel::Qwen3ASR1_7B => "Qwen",
WhichModel::Qwen3VL2B
| WhichModel::Qwen3VL4B
| WhichModel::Qwen3VL8B
| WhichModel::Qwen3VL32B
| WhichModel::Qwen3_5Gguf => "Qwen",
WhichModel::Qwen3_5_0_8B
| WhichModel::Qwen3_5_2B
| WhichModel::Qwen3_5_4B
| WhichModel::Qwen3_5_9B => "Qwen",
WhichModel::DeepSeekOCR | WhichModel::DeepSeekOCR2 => "deepseek-ai",
WhichModel::HunyuanOCR => "Tencent-Hunyuan",
WhichModel::PaddleOCRVL | WhichModel::PaddleOCRVL1_5 => "PaddlePaddle",
WhichModel::RMBG2_0 => "AI-ModelScope",
WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => "OpenBMB",
WhichModel::GlmASRNano2512 | WhichModel::GlmOCR => "ZhipuAI",
WhichModel::FunASRNano2512 => "FunAudioLLM",
WhichModel::LFM2_1_2B
| WhichModel::LFM2_5_1_2BInstruct
| WhichModel::LFM2_5VL1_6B
| WhichModel::LFM2VL1_6B => "LiquidAI",
}
}
// /// Get the owner/organization name for a model
// fn which_model_to_owner(which_model: WhichModel) -> &'static str {
// match which_model {
// WhichModel::MiniCPM4_0_5B => "OpenBMB",
// WhichModel::Qwen2_5VL3B | WhichModel::Qwen2_5VL7B => "Qwen",
// WhichModel::Qwen3_0_6B | WhichModel::Qwen3ASR0_6B | WhichModel::Qwen3ASR1_7B => "Qwen",
// WhichModel::Qwen3VL2B
// | WhichModel::Qwen3VL4B
// | WhichModel::Qwen3VL8B
// | WhichModel::Qwen3VL32B
// | WhichModel::Qwen3_5Gguf => "Qwen",
// WhichModel::Qwen3_5_0_8B
// | WhichModel::Qwen3_5_2B
// | WhichModel::Qwen3_5_4B
// | WhichModel::Qwen3_5_9B => "Qwen",
// WhichModel::DeepSeekOCR | WhichModel::DeepSeekOCR2 => "deepseek-ai",
// WhichModel::HunyuanOCR => "Tencent-Hunyuan",
// WhichModel::PaddleOCRVL | WhichModel::PaddleOCRVL1_5 => "PaddlePaddle",
// WhichModel::RMBG2_0 => "AI-ModelScope",
// WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => "OpenBMB",
// WhichModel::GlmASRNano2512 | WhichModel::GlmOCR => "ZhipuAI",
// WhichModel::FunASRNano2512 => "FunAudioLLM",
// WhichModel::LFM2_1_2B
// | WhichModel::LFM2_5_1_2BInstruct
// | WhichModel::LFM2_5VL1_6B
// | WhichModel::LFM2VL1_6B => "LiquidAI",
// }
// }
#[get("/models")]
pub(crate) async fn models() -> (Status, (ContentType, Json<serde_json::Value>)) {
@@ -315,10 +279,10 @@ pub(crate) async fn models() -> (Status, (ContentType, Json<serde_json::Value>))
let which_model = guard.which_model;
let model_obj = ModelObject {
id: which_model_to_id(which_model).to_string(),
id: which_model.as_string(),
object: "model".to_string(),
created: None, // We don't track creation time
owned_by: which_model_to_owner(which_model).to_string(),
owned_by: which_model.model_owner(),
};
drop(guard);
@@ -419,41 +383,6 @@ mod tests {
assert_eq!(WhichModel::VoxCPM.model_type(), "image");
assert_eq!(WhichModel::VoxCPM1_5.model_type(), "image");
}
// Test model_id retrieval
#[test]
fn test_get_model_id() {
assert_eq!(WhichModel::Qwen3_0_6B.model_id(), "Qwen/Qwen3-0.6B");
assert_eq!(
WhichModel::DeepSeekOCR.model_id(),
"deepseek-ai/DeepSeek-OCR"
);
assert_eq!(WhichModel::VoxCPM1_5.model_id(), "OpenBMB/VoxCPM1.5");
}
// Test OpenAI-compatible model ID conversion
#[test]
fn test_which_model_to_id() {
assert_eq!(which_model_to_id(WhichModel::Qwen3_0_6B), "qwen3-0.6b");
assert_eq!(which_model_to_id(WhichModel::DeepSeekOCR), "deepseek-ocr");
assert_eq!(which_model_to_id(WhichModel::VoxCPM1_5), "voxcpm1.5");
assert_eq!(
which_model_to_id(WhichModel::MiniCPM4_0_5B),
"minicpm4-0.5b"
);
}
// Test owner/organization mapping
#[test]
fn test_which_model_to_owner() {
assert_eq!(which_model_to_owner(WhichModel::Qwen3_0_6B), "Qwen");
assert_eq!(which_model_to_owner(WhichModel::DeepSeekOCR), "deepseek-ai");
assert_eq!(which_model_to_owner(WhichModel::VoxCPM1_5), "OpenBMB");
assert_eq!(
which_model_to_owner(WhichModel::HunyuanOCR),
"Tencent-Hunyuan"
);
}
}
// Shutdown endpoint
src/exec/deepseek_ocr.rs
+1 -1
View File
@@ -28,7 +28,7 @@ impl ExecModel for DeepSeekORExec {
let message = format!(
r#"{{
"model": "deepseek-ocr",
"model": "deepseek-ai/DeepSeek-OCR",
"messages": [
{{
"role": "user",
src/exec/glm_ocr.rs
+1 -1
View File
@@ -27,7 +27,7 @@ impl ExecModel for GlmOcrExec {
let message = format!(
r#"{{
"model": "glm-ocr",
"model": "ZhipuAI/GLM-OCR",
"messages": [
{{
"role": "user",
src/exec/hunyuan_ocr.rs
+1 -1
View File
@@ -28,7 +28,7 @@ impl ExecModel for HunyuanORExec {
let message = format!(
r#"{{
"model": "hunyuan-ocr",
"model": "Tencent-Hunyuan/HunyuanOCR",
"messages": [
{{
"role": "user",
src/exec/paddleocr_vl.rs
+1 -1
View File
@@ -28,7 +28,7 @@ impl ExecModel for PaddleOVLExec {
let message = format!(
r#"{{
"model": "paddleocr-vl",
"model": "PaddlePaddle/PaddleOCR-VL",
"messages": [
{{
"role": "user",
src/exec/rmbg2_0.rs
+1 -1
View File
@@ -29,7 +29,7 @@ impl ExecModel for RMBG2_0Exec {
// Create ChatCompletionParameters for image background removal
let message = format!(
r#"{{
"model": "rmbg2.0",
"model": "AI-ModelScope/RMBG-2.0",
"messages": [
{{
"role": "user",
src/exec/voxcpm.rs
+1 -1
View File
@@ -31,7 +31,7 @@ impl ExecModel for VoxCPMExec {
Some("啥子小师叔，打狗还要看主人，你再要继续，我就是你的对手".to_string()), // todo args
Some("file://./assets/audio/voice_01.wav".to_string()), // todo args
2,
100, // max_len (voxcpm uses 100 vs voxcpm1.5's 4096)
100, // max_len (voxcpm uses 100 vs OpenBMB/VoxCPM1.5's 4096)
10,
2.0,
6.0,
src/main.rs
+29 -57
View File
@@ -2,12 +2,12 @@ use std::sync::atomic::{AtomicBool, Ordering};
use std::{net::IpAddr, str::FromStr, sync::Arc};
use aha::{
models::WhichModel,
models::common::model_mapping::WhichModel,
process::{cleanup_pid_file, create_pid_file},
utils::{download_model, get_default_save_dir},
};
use anyhow::anyhow;
use clap::{Args, Parser, Subcommand, ValueEnum};
use clap::{Args, Parser, Subcommand};
use rocket::{
Config,
data::{ByteUnit, Limits},
@@ -213,7 +213,7 @@ struct ListArgs {
/// 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.model_id();
let model_id = model.as_string();
let save_dir = get_default_save_dir().expect("Failed to get home directory");
format!("{}/{}", save_dir, model_id)
}
@@ -221,7 +221,7 @@ fn get_default_weight_path(model: WhichModel) -> String {
/// 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.model_id();
let model_id = model.as_string();
let save_dir = match get_default_save_dir() {
Some(dir) => dir,
None => return false,
@@ -233,8 +233,8 @@ fn is_model_downloaded(model: WhichModel) -> bool {
/// Model information for JSON output
#[derive(Serialize)]
struct ModelInfo {
name: String,
model_id: String,
owner: String,
#[serde(rename = "type")]
model_type: String,
downloaded: bool,
@@ -242,48 +242,17 @@ struct ModelInfo {
/// List all supported models
fn run_list(args: ListArgs) -> anyhow::Result<()> {
let models = [
WhichModel::MiniCPM4_0_5B,
WhichModel::LFM2_1_2B,
WhichModel::LFM2_5_1_2BInstruct,
WhichModel::Qwen2_5VL3B,
WhichModel::Qwen2_5VL7B,
WhichModel::Qwen3_0_6B,
WhichModel::Qwen3_5_0_8B,
WhichModel::Qwen3_5_2B,
WhichModel::Qwen3_5_4B,
WhichModel::Qwen3_5_9B,
WhichModel::Qwen3ASR0_6B,
WhichModel::Qwen3ASR1_7B,
WhichModel::Qwen3VL2B,
WhichModel::Qwen3VL4B,
WhichModel::Qwen3VL8B,
WhichModel::Qwen3VL32B,
WhichModel::DeepSeekOCR,
WhichModel::DeepSeekOCR2,
WhichModel::HunyuanOCR,
WhichModel::PaddleOCRVL,
WhichModel::PaddleOCRVL1_5,
WhichModel::RMBG2_0,
WhichModel::VoxCPM,
WhichModel::VoxCPM1_5,
WhichModel::GlmASRNano2512,
WhichModel::FunASRNano2512,
WhichModel::GlmOCR,
];
let models = WhichModel::model_list();
if args.json {
// JSON output
let model_infos: Vec<ModelInfo> = models
.iter()
.map(|model| {
let possible_value = model.to_possible_value().unwrap();
ModelInfo {
name: possible_value.get_name().to_string(),
model_id: model.model_id().to_string(),
model_type: model.model_type().to_string(),
downloaded: is_model_downloaded(*model),
}
.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)?);
@@ -292,20 +261,23 @@ fn run_list(args: ListArgs) -> anyhow::Result<()> {
println!("Available models:");
println!();
println!(
"{:<30} {:<40} {:<10}",
"Model Name", "ModelScope ID", "Download"
"{:<40} {:<20} {:<10} {:<10}",
"Model ID", "Owner", "type", "Download"
);
println!("{}", "-".repeat(80));
for model in models {
let possible_value = model.to_possible_value().unwrap();
let name = possible_value.get_name();
let id = model.model_id();
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!("{:<30} {:<40} {:<10}", name, id, download_status);
println!(
"{:<40} {:<20} {:<10} {:<10}",
model_id, owner, model_type, download_status
);
}
}
@@ -322,9 +294,9 @@ async fn run_cli(args: CliArgs) -> anyhow::Result<()> {
gguf_path,
mmproj_path,
} = args;
let model_id = common.model.model_id();
let model_id = common.model.as_string();
let (model_path, gguf, mmproj) = if model_id.eq("GGUF") {
let (model_path, gguf, mmproj) = if model_id.contains("gguf") {
if gguf_path.is_none() {
return Err(anyhow!("gguf model path is required"));
}
@@ -338,8 +310,8 @@ async fn run_cli(args: CliArgs) -> anyhow::Result<()> {
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
download_model(&model_id, &save_dir, max_retries).await?;
save_dir + "/" + &model_id
}
};
(model_path, None, None)
@@ -359,8 +331,8 @@ async fn run_serv(args: ServArgs) -> anyhow::Result<()> {
gguf_path,
mmproj_path,
} = args;
let model_id = common.model.model_id();
let (model_path, gguf, mmproj) = if model_id.eq("GGUF") {
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"));
}
@@ -427,7 +399,7 @@ async fn run_download(args: DownloadArgs) -> anyhow::Result<()> {
save_dir,
download_retries,
} = args;
let model_id = model.model_id();
let model_id = model.as_string();
let save_dir = match save_dir {
Some(dir) => dir,
@@ -435,7 +407,7 @@ async fn run_download(args: DownloadArgs) -> anyhow::Result<()> {
};
let max_retries = download_retries.unwrap_or(3);
download_model(model_id, &save_dir, max_retries).await?;
download_model(&model_id, &save_dir, max_retries).await?;
Ok(())
}
@@ -587,7 +559,7 @@ fn run_run(args: RunArgs) -> anyhow::Result<()> {
/// Run the 'delete' subcommand: delete model from default location
fn run_delete(args: DeleteArgs) -> anyhow::Result<()> {
let DeleteArgs { model } = args;
let model_id = model.model_id();
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);
src/models/common/mod.rs
+1
View File
@@ -8,6 +8,7 @@ use candle_nn::{
};
pub mod gguf;
pub mod model_mapping;
use crate::{
position_embed::rope::{RoPE, apply_rotary_pos_emb, apply_rotary_pos_emb_roformer},
src/models/common/model_mapping.rs
+128
View File
@@ -0,0 +1,128 @@
use clap::ValueEnum;
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum WhichModel {
#[value(name = "LiquidAI/LFM2-1.2B")]
LFM2_1_2B,
#[value(name = "LiquidAI/LFM2.5-1.2B-Instruct")]
LFM2_5_1_2BInstruct,
#[value(name = "LiquidAI/LFM2.5-VL-1.6B")]
LFM2_5VL1_6B,
#[value(name = "LiquidAI/LFM2-VL-1.6B")]
LFM2VL1_6B,
#[value(name = "OpenBMB/MiniCPM4-0.5B")]
MiniCPM4_0_5B,
#[value(name = "Qwen/Qwen2.5-VL-3B-Instruct")]
Qwen2_5VL3B,
#[value(name = "Qwen/Qwen2.5-VL-7B-Instruct")]
Qwen2_5VL7B,
#[value(name = "Qwen/Qwen3-0.6B")]
Qwen3_0_6B,
#[value(name = "Qwen/Qwen3.5-0.8B")]
Qwen3_5_0_8B,
#[value(name = "Qwen/Qwen3.5-2B")]
Qwen3_5_2B,
#[value(name = "Qwen/Qwen3.5-4B")]
Qwen3_5_4B,
#[value(name = "Qwen/Qwen3.5-9B")]
Qwen3_5_9B,
#[value(name = "qwen3.5-gguf")] // todo
Qwen3_5Gguf,
#[value(name = "Qwen/Qwen3-ASR-0.6B")]
Qwen3ASR0_6B,
#[value(name = "Qwen/Qwen3-ASR-1.7B")]
Qwen3ASR1_7B,
#[value(name = "Qwen/Qwen3-VL-2B-Instruct")]
Qwen3VL2B,
#[value(name = "Qwen/Qwen3-VL-4B-Instruct")]
Qwen3VL4B,
#[value(name = "Qwen/Qwen3-VL-8B-Instruct")]
Qwen3VL8B,
#[value(name = "Qwen/Qwen3-VL-32B-Instruct")]
Qwen3VL32B,
#[value(name = "deepseek-ai/DeepSeek-OCR")]
DeepSeekOCR,
#[value(name = "deepseek-ai/DeepSeek-OCR-2")]
DeepSeekOCR2,
#[value(name = "Tencent-Hunyuan/HunyuanOCR")]
HunyuanOCR,
#[value(name = "PaddlePaddle/PaddleOCR-VL")]
PaddleOCRVL,
#[value(name = "PaddlePaddle/PaddleOCR-VL-1.5")]
PaddleOCRVL1_5,
#[value(name = "AI-ModelScope/RMBG-2.0")]
RMBG2_0,
#[value(name = "OpenBMB/VoxCPM-0.5B")]
VoxCPM,
#[value(name = "OpenBMB/VoxCPM1.5")]
VoxCPM1_5,
#[value(name = "ZhipuAI/GLM-ASR-Nano-2512")]
GlmASRNano2512,
#[value(name = "FunAudioLLM/Fun-ASR-Nano-2512")]
FunASRNano2512,
#[value(name = "ZhipuAI/GLM-OCR")]
GlmOCR,
}
impl WhichModel {
/// Get the ModelScope model ID for this model variant
pub fn as_string(&self) -> String {
self.to_possible_value()
.expect("not exists")
.get_name()
.to_string()
}
/// Get the WhichModel enum list
pub fn model_list() -> Vec<Self> {
WhichModel::value_variants().to_vec()
}
pub fn model_owner(&self) -> String {
let name = self.as_string();
let names: Vec<&str> = name.split("/").collect();
if names.len() < 2 {
"none".to_string()
} else {
names.first().map_or("none", |&s| s).to_string()
}
}
/// Get the model type category for this model variant
pub fn model_type(self) -> &'static str {
match self {
// LLM models
WhichModel::MiniCPM4_0_5B
| WhichModel::Qwen3_0_6B
| WhichModel::LFM2_1_2B
| WhichModel::LFM2_5_1_2BInstruct => "llm",
WhichModel::Qwen2_5VL3B
| WhichModel::Qwen2_5VL7B
| WhichModel::Qwen3VL2B
| WhichModel::Qwen3VL4B
| WhichModel::Qwen3VL8B
| WhichModel::Qwen3VL32B
| WhichModel::Qwen3_5_0_8B
| WhichModel::Qwen3_5_2B
| WhichModel::Qwen3_5_4B
| WhichModel::Qwen3_5_9B
| WhichModel::Qwen3_5Gguf
| WhichModel::LFM2_5VL1_6B
| WhichModel::LFM2VL1_6B => "vlm",
// OCR models
WhichModel::DeepSeekOCR
| WhichModel::DeepSeekOCR2
| WhichModel::HunyuanOCR
| WhichModel::GlmOCR
| WhichModel::PaddleOCRVL
| WhichModel::PaddleOCRVL1_5 => "ocr",
// ASR models
WhichModel::Qwen3ASR0_6B
| WhichModel::Qwen3ASR1_7B
| WhichModel::GlmASRNano2512
| WhichModel::FunASRNano2512 => "asr",
// Image models
WhichModel::RMBG2_0 => "image",
WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => "tts",
}
}
}
src/models/deepseek_ocr/generate.rs
+1 -1
View File
@@ -44,7 +44,7 @@ impl DeepseekOCRGenerateModel {
let model_name = std::path::Path::new(path)
.file_name()
.and_then(|s| s.to_str())
.unwrap_or("deepseek-ocr")
.unwrap_or("deepseek-ai/DeepSeek-OCR")
.to_string();
let version = if model_name.contains("2") || cfg.vision_config.width.qwen2_0_5b.is_some() {
2usize
src/models/glm_ocr/generate.rs
+1 -1
View File
@@ -60,7 +60,7 @@ impl GlmOcrGenerateModel {
let model_name = std::path::Path::new(path)
.file_name()
.and_then(|s| s.to_str())
.unwrap_or("glm-ocr")
.unwrap_or("ZhipuAI/GLM-OCR")
.to_string();
Ok(Self {
// chat_template,
src/models/mod.rs
+3 -143
View File
@@ -22,8 +22,9 @@ pub mod rmbg2_0;
pub mod voxcpm;
pub mod w2v_bert_2_0;
use crate::params::chat::{
ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
use crate::{
models::common::model_mapping::WhichModel,
params::chat::{ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse},
};
use anyhow::{Result, anyhow};
use rocket::futures::Stream;
@@ -40,147 +41,6 @@ use crate::models::{
rmbg2_0::generate::RMBG2_0Model, voxcpm::generate::VoxCPMGenerate,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum WhichModel {
#[value(name = "lfm2-1.2b")]
LFM2_1_2B,
#[value(name = "lfm2.5-1.2b-instruct")]
LFM2_5_1_2BInstruct,
#[value(name = "lfm2.5-vl-1.6b")]
LFM2_5VL1_6B,
#[value(name = "lfm2-vl-1.6b")]
LFM2VL1_6B,
#[value(name = "minicpm4-0.5b")]
MiniCPM4_0_5B,
#[value(name = "qwen2.5vl-3b")]
Qwen2_5VL3B,
#[value(name = "qwen2.5vl-7b")]
Qwen2_5VL7B,
#[value(name = "qwen3-0.6b")]
Qwen3_0_6B,
#[value(name = "qwen3.5-0.8b")]
Qwen3_5_0_8B,
#[value(name = "qwen3.5-2b")]
Qwen3_5_2B,
#[value(name = "qwen3.5-4b")]
Qwen3_5_4B,
#[value(name = "qwen3.5-9b")]
Qwen3_5_9B,
#[value(name = "qwen3.5-gguf")]
Qwen3_5Gguf,
#[value(name = "qwen3asr-0.6b")]
Qwen3ASR0_6B,
#[value(name = "qwen3asr-1.7b")]
Qwen3ASR1_7B,
#[value(name = "qwen3vl-2b")]
Qwen3VL2B,
#[value(name = "qwen3vl-4b")]
Qwen3VL4B,
#[value(name = "qwen3vl-8b")]
Qwen3VL8B,
#[value(name = "qwen3vl-32b")]
Qwen3VL32B,
#[value(name = "deepseek-ocr")]
DeepSeekOCR,
#[value(name = "deepseek-ocr2")]
DeepSeekOCR2,
#[value(name = "hunyuan-ocr")]
HunyuanOCR,
#[value(name = "paddleocr-vl")]
PaddleOCRVL,
#[value(name = "paddleocr-vl1.5")]
PaddleOCRVL1_5,
#[value(name = "rmbg2.0")]
RMBG2_0,
#[value(name = "voxcpm")]
VoxCPM,
#[value(name = "voxcpm1.5")]
VoxCPM1_5,
#[value(name = "glm-asr-nano-2512")]
GlmASRNano2512,
#[value(name = "fun-asr-nano-2512")]
FunASRNano2512,
#[value(name = "glm-ocr")]
GlmOCR,
}
impl WhichModel {
/// Get the ModelScope model ID for this model variant
pub fn model_id(self) -> &'static str {
match self {
WhichModel::LFM2_1_2B => "LiquidAI/LFM2-1.2B",
WhichModel::LFM2_5_1_2BInstruct => "LiquidAI/LFM2.5-1.2B-Instruct",
WhichModel::LFM2_5VL1_6B => "LiquidAI/LFM2.5-VL-1.6B",
WhichModel::LFM2VL1_6B => "LiquidAI/LFM2-VL-1.6B",
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::Qwen3_5_0_8B => "Qwen/Qwen3.5-0.8B",
WhichModel::Qwen3_5_2B => "Qwen/Qwen3.5-2B",
WhichModel::Qwen3_5_4B => "Qwen/Qwen3.5-4B",
WhichModel::Qwen3_5_9B => "Qwen/Qwen3.5-9B",
WhichModel::Qwen3_5Gguf => "GGUF",
WhichModel::Qwen3ASR0_6B => "Qwen/Qwen3-ASR-0.6B",
WhichModel::Qwen3ASR1_7B => "Qwen/Qwen3-ASR-1.7B",
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::DeepSeekOCR2 => "deepseek-ai/DeepSeek-OCR-2",
WhichModel::HunyuanOCR => "Tencent-Hunyuan/HunyuanOCR",
WhichModel::PaddleOCRVL => "PaddlePaddle/PaddleOCR-VL",
WhichModel::PaddleOCRVL1_5 => "PaddlePaddle/PaddleOCR-VL-1.5",
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",
WhichModel::GlmOCR => "ZhipuAI/GLM-OCR",
}
}
/// Get the model type category for this model variant
pub fn model_type(self) -> &'static str {
match self {
// LLM models
WhichModel::MiniCPM4_0_5B
| WhichModel::Qwen3_0_6B
| WhichModel::LFM2_1_2B
| WhichModel::LFM2_5_1_2BInstruct => "llm",
WhichModel::Qwen2_5VL3B
| WhichModel::Qwen2_5VL7B
| WhichModel::Qwen3VL2B
| WhichModel::Qwen3VL4B
| WhichModel::Qwen3VL8B
| WhichModel::Qwen3VL32B
| WhichModel::Qwen3_5_0_8B
| WhichModel::Qwen3_5_2B
| WhichModel::Qwen3_5_4B
| WhichModel::Qwen3_5_9B
| WhichModel::Qwen3_5Gguf
| WhichModel::LFM2_5VL1_6B
| WhichModel::LFM2VL1_6B => "vlm",
// OCR models
WhichModel::DeepSeekOCR
| WhichModel::DeepSeekOCR2
| WhichModel::HunyuanOCR
| WhichModel::GlmOCR
| WhichModel::PaddleOCRVL
| WhichModel::PaddleOCRVL1_5 => "ocr",
// ASR models
WhichModel::Qwen3ASR0_6B
| WhichModel::Qwen3ASR1_7B
| WhichModel::GlmASRNano2512
| WhichModel::FunASRNano2512 => "asr",
// Image models
WhichModel::RMBG2_0 => "image",
WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => "tts",
}
}
}
pub trait GenerateModel {
fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse>;
fn generate_stream(
src/models/rmbg2_0/generate.rs
+1 -1
View File
@@ -44,7 +44,7 @@ impl RMBG2_0Model {
let model_name = std::path::Path::new(path)
.file_name()
.and_then(|s| s.to_str())
.unwrap_or("rmbg2.0")
.unwrap_or("AI-ModelScope/RMBG-2.0")
.to_string();
Ok(Self {
model,
src/models/voxcpm/generate.rs
+1 -1
View File
@@ -86,7 +86,7 @@ impl VoxCPMGenerate {
let m_dtype = get_dtype(dtype, cfg_dtype);
let model_list = find_type_files(path, "bin")?;
// voxcpm0.5B模型文件是.bin类型， voxcpm1.5模型文件是.safetensors类型
// voxcpm0.5B模型文件是.bin类型， OpenBMB/VoxCPM1.5模型文件是.safetensors类型
let vb_voxcpm = if model_list.is_empty() {
let model_list = find_type_files(path, "safetensors")?;
unsafe { VarBuilder::from_mmaped_safetensors(&model_list, m_dtype, device)? }
src/params/shared.rs
+9
View File
@@ -5,9 +5,18 @@ pub struct Usage {
/// Number of tokens in the prompt.
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_tokens: Option<u32>,
/// Number of tokens in the prompt.
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_ms: Option<f64>,
/// Number of tokens in the completion.
#[serde(skip_serializing_if = "Option::is_none")]
pub completion_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub completion_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub completion_per_token_ms: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub completion_tps: Option<f64>,
/// Number of tokens in the entire response.
pub total_tokens: u32,
/// Breakdown of tokens used in the prompt.
src/process.rs
+2 -2
View File
@@ -258,7 +258,7 @@ mod tests {
let service = ServiceInfo {
service_id: "12345@10100".to_string(),
pid: 12345,
model: Some("qwen3-0.6b".to_string()),
model: Some("Qwen/Qwen3-0.6B".to_string()),
port: 10100,
address: "127.0.0.1".to_string(),
status: ServiceStatus::Running,
@@ -266,7 +266,7 @@ mod tests {
let service_clone = service.clone();
assert_eq!(service_clone.service_id, "12345@10100");
assert_eq!(service_clone.pid, 12345);
assert_eq!(service_clone.model, Some("qwen3-0.6b".to_string()));
assert_eq!(service_clone.model, Some("Qwen/Qwen3-0.6B".to_string()));
assert_eq!(service_clone.port, 10100);
}
src/utils/mod.rs
+62 -18
View File
@@ -494,25 +494,8 @@ pub fn build_audio_completion_response(
response
}
pub fn build_completion_response(
res: String,
model_name: &str,
completion_tokens: Option<u32>,
prompt_tokens: Option<u32>,
) -> ChatCompletionResponse {
fn build_response(res: String, model_name: &str, usage: Option<Usage>) -> ChatCompletionResponse {
let id = uuid::Uuid::new_v4().to_string();
let usage = if prompt_tokens.is_none() && completion_tokens.is_none() {
None
} else {
Some(Usage {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens.unwrap_or(0) + completion_tokens.unwrap_or(0),
prompt_tokens_details: None,
completion_tokens_details: None,
})
};
let mut response = ChatCompletionResponse {
id: Some(id),
choices: vec![],
@@ -588,6 +571,67 @@ pub fn build_completion_response(
response
}
pub fn build_completion_response(
res: String,
model_name: &str,
completion_tokens: Option<u32>,
prompt_tokens: Option<u32>,
) -> ChatCompletionResponse {
let usage = if prompt_tokens.is_none() && completion_tokens.is_none() {
None
} else {
Some(Usage {
prompt_tokens,
prompt_ms: None,
completion_tokens,
completion_ms: None,
completion_per_token_ms: None,
completion_tps: None,
total_tokens: prompt_tokens.unwrap_or(0) + completion_tokens.unwrap_or(0),
prompt_tokens_details: None,
completion_tokens_details: None,
})
};
build_response(res, model_name, usage)
}
pub fn build_completion_response_with_time(
res: String,
model_name: &str,
completion_tokens: Option<u32>,
completion_ms: Option<f64>,
prompt_tokens: Option<u32>,
prompt_ms: Option<f64>,
) -> ChatCompletionResponse {
let usage = if prompt_tokens.is_none() && completion_tokens.is_none() {
None
} else {
let (completion_per_token_ms, completion_tps) = if let Some(prompt_tokens) = prompt_tokens
&& let Some(prompt_ms) = prompt_ms
{
let per_token_ms = prompt_ms / prompt_tokens as f64;
let tps = prompt_tokens as f64 / (prompt_ms / 1000.0);
(Some(per_token_ms), Some(tps))
} else {
(None, None)
};
Some(Usage {
prompt_tokens,
prompt_ms,
completion_tokens,
completion_ms,
completion_per_token_ms,
completion_tps,
total_tokens: prompt_tokens.unwrap_or(0) + completion_tokens.unwrap_or(0),
prompt_tokens_details: None,
completion_tokens_details: None,
})
};
build_response(res, model_name, usage)
}
pub fn build_completion_chunk_response(
res: String,
model_name: &str,