From 4e35866a591eb21f68acf8e5205e123fbc1fb55d Mon Sep 17 00:00:00 2001 From: KingBright Date: Mon, 5 Jan 2026 15:21:15 +0800 Subject: [PATCH] dynamic home dir & model donwload script --- Cargo.lock | 138 +++++++++++++++++++++++++++++- download_and_run.sh | 156 ++++++++++++++++++++++++++++++++++ src/models/voxcpm/generate.rs | 4 + tests/test_deepseek_ocr.rs | 10 ++- tests/test_gelab_zero.rs | 5 +- tests/test_hunyuan_ocr.rs | 10 ++- tests/test_minicpm4.rs | 10 ++- tests/test_paddleocr_vl.rs | 10 ++- tests/test_qwen2_5vl.rs | 10 ++- tests/test_qwen3vl.rs | 10 ++- tests/test_rmbg2_0.rs | 5 +- tests/test_robo_brain.rs | 5 +- tests/test_voxcpm.rs | 16 ++-- tests/test_voxcpm1_5.rs | 15 ++-- tests/weight_test.rs | 30 ++++--- 15 files changed, 378 insertions(+), 56 deletions(-) create mode 100755 download_and_run.sh diff --git a/Cargo.lock b/Cargo.lock index fdb881a..6f9f9c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -377,6 +377,12 @@ version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6099cdc01846bc367c4e7dd630dc5966dccf36b652fae7a74e17b640411a91b2" +[[package]] +name = "block" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" + [[package]] name = "built" version = "0.7.7" @@ -438,10 +444,12 @@ checksum = "a9f51e2ecf6efe9737af8f993433c839f956d2b6ed4fd2dd4a7c6d8b0fa667ff" dependencies = [ "byteorder", "candle-kernels", + "candle-metal-kernels", "cudarc", "gemm 0.17.1", "half", "memmap2", + "metal 0.27.0", "num-traits", "num_cpus", "rand 0.9.2", @@ -451,6 +459,7 @@ dependencies = [ "thiserror 1.0.69", "ug", "ug-cuda", + "ug-metal", "yoke 0.7.5", "zip", ] @@ -476,6 +485,19 @@ dependencies = [ "bindgen_cuda", ] +[[package]] +name = "candle-metal-kernels" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a323ee9c813707f73b6e59300661b354a70410f31fe4135170c4eda8a061534" +dependencies = [ + "half", + "metal 0.27.0", + "once_cell", + "thiserror 1.0.69", + "tracing", +] + [[package]] name = "candle-nn" version = "0.9.1" @@ -483,7 +505,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1980d53280c8f9e2c6cbe1785855d7ff8010208b46e21252b978badf13ad69d" dependencies = [ "candle-core", + "candle-metal-kernels", "half", + "metal 0.27.0", "num-traits", "rayon", "safetensors", @@ -700,6 +724,17 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core-graphics-types" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf" +dependencies = [ + "bitflags 1.3.2", + "core-foundation", + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -1153,7 +1188,28 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" dependencies = [ - "foreign-types-shared", + "foreign-types-shared 0.1.1", +] + +[[package]] +name = "foreign-types" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" +dependencies = [ + "foreign-types-macros", + "foreign-types-shared 0.3.1", +] + +[[package]] +name = "foreign-types-macros" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -1162,6 +1218,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" +[[package]] +name = "foreign-types-shared" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -2249,6 +2311,15 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" +[[package]] +name = "malloc_buf" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb" +dependencies = [ + "libc", +] + [[package]] name = "matchers" version = "0.2.0" @@ -2284,6 +2355,36 @@ dependencies = [ "stable_deref_trait", ] +[[package]] +name = "metal" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c43f73953f8cbe511f021b58f18c3ce1c3d1ae13fe953293e13345bf83217f25" +dependencies = [ + "bitflags 2.10.0", + "block", + "core-graphics-types", + "foreign-types 0.5.0", + "log", + "objc", + "paste", +] + +[[package]] +name = "metal" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ecfd3296f8c56b7c1f6fbac3c71cefa9d78ce009850c45000015f206dc7fa21" +dependencies = [ + "bitflags 2.10.0", + "block", + "core-graphics-types", + "foreign-types 0.5.0", + "log", + "objc", + "paste", +] + [[package]] name = "mime" version = "0.3.17" @@ -2590,6 +2691,25 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "objc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1" +dependencies = [ + "malloc_buf", + "objc_exception", +] + +[[package]] +name = "objc_exception" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad970fb455818ad6cba4c122ad012fae53ae8b4795f86378bce65e4f6bab2ca4" +dependencies = [ + "cc", +] + [[package]] name = "object" version = "0.37.3" @@ -2641,7 +2761,7 @@ checksum = "24ad14dd45412269e1a30f52ad8f0664f0f4f4a89ee8fe28c3b3527021ebb654" dependencies = [ "bitflags 2.10.0", "cfg-if", - "foreign-types", + "foreign-types 0.3.2", "libc", "once_cell", "openssl-macros", @@ -4236,6 +4356,20 @@ dependencies = [ "ug", ] +[[package]] +name = "ug-metal" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76daec3c7a32a1b4a0e3307b6b057fa067aa64e750713987410a2c402e5cd731" +dependencies = [ + "half", + "metal 0.29.0", + "objc", + "serde", + "thiserror 1.0.69", + "ug", +] + [[package]] name = "uncased" version = "0.9.10" diff --git a/download_and_run.sh b/download_and_run.sh new file mode 100755 index 0000000..9350c41 --- /dev/null +++ b/download_and_run.sh @@ -0,0 +1,156 @@ +#!/bin/bash +set -e + +# Default values +# MIRROR_URL="https://hf-mirror.com" +DEFAULT_SAVE_DIR="$HOME/.aha" +MODEL_ALIAS="" + +# Help function +show_help() { + echo "Usage: $0 [model_alias]" + echo "" + echo "Arguments:" + echo " model_alias The model alias to download (e.g., voxcpm, qwen2.5vl-3b)" + echo "" + echo "Available models:" + echo " minicpm4-0.5b" + echo " qwen2.5vl-3b" + echo " qwen2.5vl-7b" + echo " qwen3vl-2b" + echo " qwen3vl-4b" + echo " qwen3vl-8b" + echo " qwen3vl-32b" + echo " deepseek-ocr" + echo " hunyuan-ocr" + echo " paddleocr-vl" + echo " RMBG2.0" + echo " voxcpm" + echo " voxcpm1.5" + echo "" + exit 1 +} + +# Check if model alias is provided +if [ -z "$1" ]; then + show_help +fi + +MODEL_ALIAS=$1 + +# Map alias to Repo ID +MODEL_ID="" +case $MODEL_ALIAS in + "minicpm4-0.5b") + MODEL_ID="OpenBMB/MiniCPM4-0.5B" + ;; + "qwen2.5vl-3b") + MODEL_ID="Qwen/Qwen2.5-VL-3B-Instruct" + ;; + "qwen2.5vl-7b") + MODEL_ID="Qwen/Qwen2.5-VL-7B-Instruct" + ;; + "qwen3vl-2b") + MODEL_ID="Qwen/Qwen3-VL-2B-Instruct" + ;; + "qwen3vl-4b") + MODEL_ID="Qwen/Qwen3-VL-4B-Instruct" + ;; + "qwen3vl-8b") + MODEL_ID="Qwen/Qwen3-VL-8B-Instruct" + ;; + "qwen3vl-32b") + MODEL_ID="Qwen/Qwen3-VL-32B-Instruct" + ;; + "deepseek-ocr") + MODEL_ID="deepseek-ai/DeepSeek-OCR" + ;; + "hunyuan-ocr") + MODEL_ID="Tencent-Hunyuan/HunyuanOCR" + ;; + "paddleocr-vl") + MODEL_ID="PaddlePaddle/PaddleOCR-VL" + ;; + "RMBG2.0") + MODEL_ID="AI-ModelScope/RMBG-2.0" + ;; + "voxcpm") + MODEL_ID="OpenBMB/VoxCPM-0.5B" + ;; + "voxcpm1.5") + MODEL_ID="OpenBMB/VoxCPM1.5" + ;; + *) + echo "Error: Unknown model alias '$MODEL_ALIAS'" + show_help + ;; +esac + +echo "Selected Model: $MODEL_ALIAS" +echo "Repo ID: $MODEL_ID" +echo "Target Directory: $DEFAULT_SAVE_DIR/$MODEL_ID" + +# Prepare environment for acceleration (Default to mirror if not set) +if [ -z "$HF_ENDPOINT" ]; then + export HF_ENDPOINT=$MIRROR_URL + echo "Using default HF Mirror: $HF_ENDPOINT" +else + echo "Using custom HF Endpoint: $HF_ENDPOINT" +fi + +# Check if huggingface-cli and hf_transfer are installed +CLI_CMD="" +if command -v huggingface-cli &> /dev/null; then + CLI_CMD="huggingface-cli" +elif command -v hf &> /dev/null; then + CLI_CMD="hf" +else + echo "huggingface-cli not found. Installing via pip..." + if command -v pip3 &> /dev/null; then + pip3 install -U "huggingface_hub[cli]" hf_transfer + elif command -v pip &> /dev/null; then + pip install -U "huggingface_hub[cli]" hf_transfer + else + echo "Error: pip is not available. Please install python and pip." + exit 1 + fi + CLI_CMD="huggingface-cli" +fi + +# Try to install hf_transfer if simple python check fails (optional but recommended for speed) +if ! python3 -c "import hf_transfer" &> /dev/null; then + echo "Installing hf_transfer for faster downloads..." + pip3 install -U hf_transfer || echo "Warning: hf_transfer install failed, falling back to standard download." +fi + +# Enable HF Transfer (Rust-based downloader) - Default to 1 (On) unless explicitly disabled +export HF_HUB_ENABLE_HF_TRANSFER=${HF_HUB_ENABLE_HF_TRANSFER:-1} +if [ "$HF_HUB_ENABLE_HF_TRANSFER" == "1" ]; then + echo "Enabled HF_HUB_ENABLE_HF_TRANSFER for high-speed download" +else + echo "HF_HUB_ENABLE_HF_TRANSFER disabled. Using standard Python downloader." +fi + +# Download model +echo "Downloading model with acceleration via $CLI_CMD..." +# Note: --resume-download is deprecated/implicit in newer versions, removing it. +# --local-dir-use-symlinks matches standard CLI if version is recent. +if ! $CLI_CMD download "$MODEL_ID" --local-dir "$DEFAULT_SAVE_DIR/$MODEL_ID" --local-dir-use-symlinks False $TOKEN_ARG; then + echo "" + echo "Error: Download failed." + echo "If you received a 429 Rate Limit error, please provide a Hugging Face Token." + echo "Usage: ./download_and_run.sh $MODEL_ALIAS [hf_token]" + echo "Or set the HF_TOKEN environment variable." + exit 1 +fi + +# Run cargo +echo "Starting application with cargo..." +# Detect if on Mac to add 'metal' feature +if [[ "$OSTYPE" == "darwin"* ]]; then + echo "Detected MacOS. Running with 'metal' feature..." + cargo run -r --features metal -- --model "$MODEL_ALIAS" --weight-path "$DEFAULT_SAVE_DIR/$MODEL_ID" +else + echo "Running with default features (cuda)..." + cargo run -r --features cuda -- --model "$MODEL_ALIAS" --weight-path "$DEFAULT_SAVE_DIR/$MODEL_ID" +fi diff --git a/src/models/voxcpm/generate.rs b/src/models/voxcpm/generate.rs index 97268b2..50478dd 100644 --- a/src/models/voxcpm/generate.rs +++ b/src/models/voxcpm/generate.rs @@ -199,6 +199,10 @@ impl VoxCPMGenerate { )?; Ok(audio) } + + pub fn sample_rate(&self) -> usize { + self.sample_rate + } } impl GenerateModel for VoxCPMGenerate { diff --git a/tests/test_deepseek_ocr.rs b/tests/test_deepseek_ocr.rs index e73dee1..0d632dd 100644 --- a/tests/test_deepseek_ocr.rs +++ b/tests/test_deepseek_ocr.rs @@ -32,10 +32,11 @@ fn deepseek_ocr_generate() -> Result<()> { "metadata": {"base_size": "640", "image_size": "640", "crop_mode": "false"} } "#; - let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/deepseek-ai/DeepSeek-OCR/", save_dir); let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = DeepseekOCRGenerateModel::init(model_path, None, None)?; + let mut model = DeepseekOCRGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); let i_start = Instant::now(); @@ -84,10 +85,11 @@ async fn deepseek_ocr_stream() -> Result<()> { "metadata": {"base_size": "640", "image_size": "640", "crop_mode": "false"} } "#; - let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/deepseek-ai/DeepSeek-OCR/", save_dir); let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = DeepseekOCRGenerateModel::init(model_path, None, None)?; + let mut model = DeepseekOCRGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); let mut stream = pin!(model.generate_stream(mes)?); diff --git a/tests/test_gelab_zero.rs b/tests/test_gelab_zero.rs index 4aff3e8..a906a3b 100644 --- a/tests/test_gelab_zero.rs +++ b/tests/test_gelab_zero.rs @@ -8,7 +8,8 @@ use anyhow::Result; fn gelab_zero_generate() -> Result<()> { // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda gelab_zero_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/stepfun-ai/GELab-Zero-4B-preview"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/stepfun-ai/GELab-Zero-4B-preview", save_dir); let message = r#" { @@ -56,7 +57,7 @@ fn gelab_zero_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut qwen3vl = Qwen3VLGenerateModel::init(model_path, None, None)?; + let mut qwen3vl = Qwen3VLGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); diff --git a/tests/test_hunyuan_ocr.rs b/tests/test_hunyuan_ocr.rs index adf887e..104d092 100644 --- a/tests/test_hunyuan_ocr.rs +++ b/tests/test_hunyuan_ocr.rs @@ -31,10 +31,11 @@ fn hunyuan_ocr_generate() -> Result<()> { ] } "#; - let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Tencent-Hunyuan/HunyuanOCR/", save_dir); let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = HunyuanOCRGenerateModel::init(model_path, None, None)?; + let mut model = HunyuanOCRGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); let i_start = Instant::now(); @@ -79,10 +80,11 @@ async fn hunyuan_ocr_stream() -> Result<()> { ] } "#; - let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Tencent-Hunyuan/HunyuanOCR/", save_dir); let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = HunyuanOCRGenerateModel::init(model_path, None, None)?; + let mut model = HunyuanOCRGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); let mut stream = pin!(model.generate_stream(mes)?); diff --git a/tests/test_minicpm4.rs b/tests/test_minicpm4.rs index 5477e17..a79a985 100644 --- a/tests/test_minicpm4.rs +++ b/tests/test_minicpm4.rs @@ -11,7 +11,8 @@ fn minicpm_generate() -> Result<()> { // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda minicpm_generate -r -- --nocapture // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/OpenBMB/MiniCPM4-0.5B/", save_dir); let message = r#" { "temperature": 0.3, @@ -27,7 +28,7 @@ fn minicpm_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = MiniCPMGenerateModel::init(model_path, None, None)?; + let mut model = MiniCPMGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); @@ -50,7 +51,8 @@ fn minicpm_generate() -> Result<()> { async fn minicpm_stream() -> Result<()> { // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_stream -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/OpenBMB/MiniCPM4-0.5B/", save_dir); let message = r#" { @@ -65,7 +67,7 @@ async fn minicpm_stream() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = MiniCPMGenerateModel::init(model_path, None, None)?; + let mut model = MiniCPMGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); diff --git a/tests/test_paddleocr_vl.rs b/tests/test_paddleocr_vl.rs index 834a396..4901ee5 100644 --- a/tests/test_paddleocr_vl.rs +++ b/tests/test_paddleocr_vl.rs @@ -32,10 +32,11 @@ fn paddleocr_vl_generate() -> Result<()> { "stream": false } "#; - let model_path = "/home/jhq/huggingface_model/PaddlePaddle/PaddleOCR-VL/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/PaddlePaddle/PaddleOCR-VL/", save_dir); let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = PaddleOCRVLGenerateModel::init(model_path, None, None)?; + let mut model = PaddleOCRVLGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); let i_start = Instant::now(); @@ -79,10 +80,11 @@ async fn paddleocr_vl_stream() -> Result<()> { ] } "#; - let model_path = "/home/jhq/huggingface_model/PaddlePaddle/PaddleOCR-VL/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/PaddlePaddle/PaddleOCR-VL/", save_dir); let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = PaddleOCRVLGenerateModel::init(model_path, None, None)?; + let mut model = PaddleOCRVLGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); let mut stream = pin!(model.generate_stream(mes)?); diff --git a/tests/test_qwen2_5vl.rs b/tests/test_qwen2_5vl.rs index d0423be..ffb4079 100644 --- a/tests/test_qwen2_5vl.rs +++ b/tests/test_qwen2_5vl.rs @@ -13,7 +13,8 @@ fn qwen2_5vl_generate() -> Result<()> { // let device = Device::cuda_if_available(0)?; // let dtype = DType::BF16; - let model_path = "/home/jhq/huggingface_model/Qwen/Qwen2.5-VL-3B-Instruct/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Qwen/Qwen2.5-VL-3B-Instruct/", save_dir); let message = r#" { @@ -40,7 +41,7 @@ fn qwen2_5vl_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = Qwen2_5VLGenerateModel::init(model_path, None, None)?; + let mut model = Qwen2_5VLGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); @@ -65,7 +66,8 @@ async fn qwen2_5vl_stream() -> Result<()> { // let device = Device::cuda_if_available(0)?; // let dtype = DType::BF16; - let model_path = "/home/jhq/huggingface_model/Qwen/Qwen2.5-VL-3B-Instruct/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Qwen/Qwen2.5-VL-3B-Instruct/", save_dir); let message = r#" { @@ -92,7 +94,7 @@ async fn qwen2_5vl_stream() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = Qwen2_5VLGenerateModel::init(model_path, None, None)?; + let mut model = Qwen2_5VLGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); diff --git a/tests/test_qwen3vl.rs b/tests/test_qwen3vl.rs index f7a5ce7..67d23e2 100644 --- a/tests/test_qwen3vl.rs +++ b/tests/test_qwen3vl.rs @@ -9,7 +9,8 @@ use rocket::futures::StreamExt; fn qwen3vl_generate() -> Result<()> { // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda,ffmpeg qwen3vl_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-2B-Instruct/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Qwen/Qwen3-VL-2B-Instruct/", save_dir); let message = r#" { @@ -36,7 +37,7 @@ fn qwen3vl_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut qwen3vl = Qwen3VLGenerateModel::init(model_path, None, None)?; + let mut qwen3vl = Qwen3VLGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); @@ -58,7 +59,8 @@ fn qwen3vl_generate() -> Result<()> { async fn qwen3vl_stream() -> Result<()> { // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda,ffmpeg qwen3vl_stream -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-2B-Instruct/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Qwen/Qwen3-VL-2B-Instruct/", save_dir); let message = r#" { @@ -85,7 +87,7 @@ async fn qwen3vl_stream() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut qwen3vl = Qwen3VLGenerateModel::init(model_path, None, None)?; + let mut qwen3vl = Qwen3VLGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); diff --git a/tests/test_rmbg2_0.rs b/tests/test_rmbg2_0.rs index ef75833..9b1f1e5 100644 --- a/tests/test_rmbg2_0.rs +++ b/tests/test_rmbg2_0.rs @@ -8,7 +8,8 @@ use anyhow::Result; fn rmbg2_0_generate() -> Result<()> { // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda rmbg2_0_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/AI-ModelScope/RMBG-2.0/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/AI-ModelScope/RMBG-2.0/", save_dir); let message = r#" { @@ -31,7 +32,7 @@ fn rmbg2_0_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let model = RMBG2_0Model::init(model_path, None, None)?; + let model = RMBG2_0Model::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); diff --git a/tests/test_robo_brain.rs b/tests/test_robo_brain.rs index 2352c14..874fa3b 100644 --- a/tests/test_robo_brain.rs +++ b/tests/test_robo_brain.rs @@ -8,7 +8,8 @@ use anyhow::Result; fn robo_brain_generate() -> Result<()> { // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda robo_brain_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/BAAI/RoboBrain2.0-3B/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/BAAI/RoboBrain2.0-3B/", save_dir); let message = r#" { @@ -28,7 +29,7 @@ fn robo_brain_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut model = Qwen2_5VLGenerateModel::init(model_path, None, None)?; + let mut model = Qwen2_5VLGenerateModel::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); diff --git a/tests/test_voxcpm.rs b/tests/test_voxcpm.rs index 783fb7c..7efdcf4 100644 --- a/tests/test_voxcpm.rs +++ b/tests/test_voxcpm.rs @@ -10,10 +10,12 @@ use aha::{ use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; use anyhow::{Ok, Result}; +#[test] #[test] fn voxcpm_use_message_generate() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda voxcpm_use_message_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/OpenBMB/VoxCPM-0.5B/", save_dir); let message = r#" { "model": "voxcpm", @@ -40,7 +42,7 @@ fn voxcpm_use_message_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut voxcpm_generate = VoxCPMGenerate::init(model_path, None, None)?; + let mut voxcpm_generate = VoxCPMGenerate::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); @@ -59,10 +61,11 @@ fn voxcpm_use_message_generate() -> Result<()> { #[test] fn voxcpm_generate() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda voxcpm_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/OpenBMB/VoxCPM-0.5B/", save_dir); let i_start = Instant::now(); - let mut voxcpm_generate = VoxCPMGenerate::init(model_path, None, None)?; + let mut voxcpm_generate = VoxCPMGenerate::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); @@ -106,8 +109,9 @@ fn voxcpm_generate() -> Result<()> { #[test] fn voxcpm_tokenizer() -> Result<()> { - let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; - let tokenizer = SingleChineseTokenizer::new(model_path)?; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/OpenBMB/VoxCPM-0.5B/", save_dir); + let tokenizer = SingleChineseTokenizer::new(&model_path)?; let ids = tokenizer.encode("你好啊，你吃饭了吗".to_string())?; println!("ids: {:?}", ids); Ok(()) diff --git a/tests/test_voxcpm1_5.rs b/tests/test_voxcpm1_5.rs index da1111a..9e39b50 100644 --- a/tests/test_voxcpm1_5.rs +++ b/tests/test_voxcpm1_5.rs @@ -13,7 +13,8 @@ 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 - let model_path = "/home/jhq/huggingface_model/OpenBMB/VoxCPM1.5/"; + 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); let message = r#" { "model": "voxcpm1.5", @@ -40,7 +41,7 @@ fn voxcpm1_5_use_message_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut voxcpm_generate = VoxCPMGenerate::init(model_path, None, None)?; + let mut voxcpm_generate = VoxCPMGenerate::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); @@ -58,10 +59,11 @@ fn voxcpm1_5_use_message_generate() -> Result<()> { #[test] fn voxcpm1_5_generate() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda voxcpm1_5_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/OpenBMB/VoxCPM1.5/"; + 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); let i_start = Instant::now(); - let mut voxcpm_generate = VoxCPMGenerate::init(model_path, None, None)?; + let mut voxcpm_generate = VoxCPMGenerate::init(&model_path, None, None)?; let i_duration = i_start.elapsed(); println!("Time elapsed in load model is: {:?}", i_duration); @@ -106,8 +108,9 @@ fn voxcpm1_5_generate() -> Result<()> { #[test] fn voxcpm1_5_tokenizer() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda voxcpm1_5_tokenizer -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/OpenBMB/VoxCPM1.5/"; - let tokenizer = SingleChineseTokenizer::new(model_path)?; + 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); + let tokenizer = SingleChineseTokenizer::new(&model_path)?; let ids = tokenizer.encode("你好啊，你吃饭了吗".to_string())?; println!("ids: {:?}", ids); Ok(()) diff --git a/tests/weight_test.rs b/tests/weight_test.rs index ab68db6..e91de49 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -7,8 +7,9 @@ use candle_nn::VarBuilder; #[test] fn minicpm4_weight() -> Result<()> { - let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/"; - let model_list = find_type_files(model_path, "safetensors")?; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/OpenBMB/MiniCPM4-0.5B/", save_dir); + let model_list = find_type_files(&model_path, "safetensors")?; let device = Device::Cpu; for m in model_list { let weights = safetensors::load(m, &device)?; @@ -23,8 +24,9 @@ fn minicpm4_weight() -> Result<()> { #[test] fn voxcpm_weight() -> Result<()> { - let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; - let model_list = find_type_files(model_path, "pth")?; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/OpenBMB/VoxCPM-0.5B/", save_dir); + let model_list = find_type_files(&model_path, "pth")?; println!("model_list: {:?}", model_list); let dev = get_device(None); let mut dict_to_hashmap = HashMap::new(); @@ -48,8 +50,9 @@ fn voxcpm_weight() -> Result<()> { #[test] fn voxcpm1_5_weight() -> Result<()> { - let model_path = "/home/jhq/huggingface_model/OpenBMB/VoxCPM1.5/"; - let model_list = find_type_files(model_path, "pth")?; + 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); + let model_list = find_type_files(&model_path, "pth")?; println!("model_list: {:?}", model_list); // let dev = get_device(None); let mut dict_to_hashmap = HashMap::new(); @@ -68,8 +71,9 @@ fn voxcpm1_5_weight() -> Result<()> { #[test] fn qwen3vl_weight() -> Result<()> { - let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-4B-Instruct/"; - let model_list = find_type_files(model_path, "safetensors")?; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Qwen/Qwen3-VL-4B-Instruct/", save_dir); + let model_list = find_type_files(&model_path, "safetensors")?; let device = Device::Cpu; for m in &model_list { @@ -84,8 +88,9 @@ fn qwen3vl_weight() -> Result<()> { #[test] fn deepseekocr_weight() -> Result<()> { - let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; - let model_list = find_type_files(model_path, "safetensors")?; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/deepseek-ai/DeepSeek-OCR/", save_dir); + let model_list = find_type_files(&model_path, "safetensors")?; let device = Device::Cpu; for m in &model_list { @@ -103,8 +108,9 @@ fn deepseekocr_weight() -> Result<()> { #[test] fn hunyuanocr_weight() -> Result<()> { - let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/"; - let model_list = find_type_files(model_path, "safetensors")?; + let save_dir = aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/Tencent-Hunyuan/HunyuanOCR/", save_dir); + let model_list = find_type_files(&model_path, "safetensors")?; let device = Device::Cpu; for m in &model_list {