Merge pull request #75 from mobiusml/hf_transfer

Add hf-transfer for faster HF Hub ops
mobiusml · Mar 15, 2024 · d918ba9 · d918ba9
2 parents 11a1443 + a1db2c7
commit d918ba9
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 6 deletions.
diff --git a/.env b/.env
@@ -1,3 +1,4 @@
 CUDA_VISIBLE_DEVICES=""
 USE_DEPLOYMENT_CACHE = True
-SAVE_DEPLOYMENT_CACHE = True
+SAVE_DEPLOYMENT_CACHE = True
+HF_HUB_ENABLE_HF_TRANSFER = 1
diff --git a/README.md b/README.md
@@ -30,7 +30,7 @@ sh install.sh
 5. Run the SDK.
 
 ```bash
-CUDA_VISIBLE_DEVICES=0 poetry run aana --port 8000 --host 0.0.0.0 --target chat_with_video
+HF_HUB_ENABLE_HF_TRANSFER=1 CUDA_VISIBLE_DEVICES=0 poetry run aana --port 8000 --host 0.0.0.0 --target chat_with_video
 ```
 
 The target parameter specifies the set of endpoints to deploy.
@@ -211,4 +211,5 @@ Here are the environment variables that can be used to configure the Aaana SDK:
 - NUM_WORKERS: The number of request workers. Default: `2`.
 - DB_CONFIG: The database configuration in the format `{"datastore_type": "sqlite", "datastore_config": {"path": "/path/to/sqlite.db"}}`. Currently only SQLite and PostgreSQL are supported. Default: `{"datastore_type": "sqlite", "datastore_config": {"path": "/var/lib/aana_data"}}`.
 - USE_DEPLOYMENT_CACHE (testing only): If set to `true`, the tests will use the deployment cache to avoid downloading the models and running the deployments. Default: `false`.
-- SAVE_DEPLOYMENT_CACHE (testing only): If set to `true`, the tests will save the deployment cache after running the deployments. Default: `false`.
+- SAVE_DEPLOYMENT_CACHE (testing only): If set to `true`, the tests will save the deployment cache after running the deployments. Default: `false`.
+- HF_HUB_ENABLE_HF_TRANSFER: If set to `1`, the HuggingFace Transformers will use the HF Transfer library to download the models from HuggingFace Hub to speed up the process. Recommended to always set to it `1`. Default: `0`.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -19,6 +19,7 @@ deepdiff = "^6.7.0"
 diffusers = "^0.23.1"
 fastapi = "^0.104.0"
 faster-whisper = ">=0.10.0"
+hf-transfer = "^0.1.6"
 onnxruntime = "1.16.1"
 opencv-python = "^4.8.1.78"
 portpicker = "^1.6.0"

diff --git a/startup.sh b/startup.sh
@@ -1,3 +1,3 @@
 #!/bin/bash
 # TODO: pass arguments to the docker to set target instead of environment variable
-poetry run aana --port 8000 --host 0.0.0.0 --target $TARGET
+HF_HUB_ENABLE_HF_TRANSFER=1 poetry run aana --port 8000 --host 0.0.0.0 --target $TARGET