diff --git a/download_and_run.sh b/download_and_run.sh new file mode 100755 index 0000000..867cb87 --- /dev/null +++ b/download_and_run.sh @@ -0,0 +1,160 @@ +#!/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 " glm-asr-nano-2512" + 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" + ;; + "glm-asr-nano-2512") + MODEL_ID="zai-org/GLM-ASR-Nano-2512" + ;; + *) + 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 a7dee6c..0071a81 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..a7b0a24 100644 --- a/tests/test_deepseek_ocr.rs +++ b/tests/test_deepseek_ocr.rs @@ -32,10 +32,12 @@ 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 +86,12 @@ 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..211f4e5 100644 --- a/tests/test_gelab_zero.rs +++ b/tests/test_gelab_zero.rs @@ -8,7 +8,9 @@ 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 +58,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_glm_asr_nano.rs b/tests/test_glm_asr_nano.rs index 7414cb5..b0cf26b 100644 --- a/tests/test_glm_asr_nano.rs +++ b/tests/test_glm_asr_nano.rs @@ -8,7 +8,9 @@ use rocket::futures::StreamExt; #[test] fn glm_asr_nano_generate() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda glm_asr_nano_generate -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/ZhipuAI/GLM-ASR-Nano-2512/"; + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/ZhipuAI/GLM-ASR-Nano-2512/", save_dir); let message = r#" { "model": "glm-asr-nano", @@ -34,7 +36,7 @@ fn glm_asr_nano_generate() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut glm_asr_model = GlmAsrNanoGenerateModel::init(model_path, None, None)?; + let mut glm_asr_model = GlmAsrNanoGenerateModel::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(); @@ -54,7 +56,9 @@ fn glm_asr_nano_generate() -> Result<()> { #[tokio::test] async fn glm_asr_nano_stream() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda glm_asr_nano_stream -r -- --nocapture - let model_path = "/home/jhq/huggingface_model/ZhipuAI/GLM-ASR-Nano-2512/"; + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/ZhipuAI/GLM-ASR-Nano-2512/", save_dir); let message = r#" { "model": "glm-asr-nano", @@ -80,7 +84,7 @@ async fn glm_asr_nano_stream() -> Result<()> { "#; let mes: ChatCompletionParameters = serde_json::from_str(message)?; let i_start = Instant::now(); - let mut glm_asr_model = GlmAsrNanoGenerateModel::init(model_path, None, None)?; + let mut glm_asr_model = GlmAsrNanoGenerateModel::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(); diff --git a/tests/test_hunyuan_ocr.rs b/tests/test_hunyuan_ocr.rs index adf887e..fa42c4d 100644 --- a/tests/test_hunyuan_ocr.rs +++ b/tests/test_hunyuan_ocr.rs @@ -31,10 +31,12 @@ 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 +81,12 @@ 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..72636b3 100644 --- a/tests/test_minicpm4.rs +++ b/tests/test_minicpm4.rs @@ -11,7 +11,9 @@ 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 +29,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 +52,9 @@ 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 +69,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..682afb3 100644 --- a/tests/test_paddleocr_vl.rs +++ b/tests/test_paddleocr_vl.rs @@ -32,10 +32,12 @@ 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 +81,12 @@ 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..886f03c 100644 --- a/tests/test_qwen2_5vl.rs +++ b/tests/test_qwen2_5vl.rs @@ -13,7 +13,9 @@ 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 +42,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 +67,9 @@ 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 +96,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..f52d662 100644 --- a/tests/test_qwen3vl.rs +++ b/tests/test_qwen3vl.rs @@ -9,7 +9,9 @@ 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 +38,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 +60,9 @@ 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 +89,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..346d8ab 100644 --- a/tests/test_rmbg2_0.rs +++ b/tests/test_rmbg2_0.rs @@ -8,7 +8,9 @@ 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 +33,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..ecf1cde 100644 --- a/tests/test_robo_brain.rs +++ b/tests/test_robo_brain.rs @@ -8,7 +8,9 @@ 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 +30,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..c45464b 100644 --- a/tests/test_voxcpm.rs +++ b/tests/test_voxcpm.rs @@ -13,7 +13,9 @@ use anyhow::{Ok, Result}; #[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,12 @@ 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 +110,10 @@ 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..4c7f452 100644 --- a/tests/test_voxcpm1_5.rs +++ b/tests/test_voxcpm1_5.rs @@ -13,7 +13,9 @@ 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 +42,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 +60,12 @@ 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 +110,10 @@ 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 bbfa740..730d1f3 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -7,8 +7,10 @@ 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 +25,10 @@ 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 +52,10 @@ 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 +74,10 @@ 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 +92,10 @@ 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 +113,10 @@ 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 {