diff --git a/.github/scripts/runpod_api.py b/.github/scripts/runpod_api.py new file mode 100644 index 000000000..edc73c6b9 --- /dev/null +++ b/.github/scripts/runpod_api.py @@ -0,0 +1,238 @@ +import os +import json +import time +import requests +import sys +import argparse +import subprocess + + +def parse_arguments(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description='Run tests on RunPod GPU') + parser.add_argument('--gpu-type', type=str, help='GPU type to use') + parser.add_argument('--gpu-count', + type=int, + help='Number of GPUs to use', + default=1) + parser.add_argument('--test-command', type=str, help='Test command to run') + parser.add_argument('--disk-size', + type=int, + default=20, + help='Container disk size in GB (default: 20)') + parser.add_argument('--volume-size', + type=int, + default=20, + help='Persistent volume size in GB (default: 20)') + parser.add_argument( + '--image', + type=str, + default='runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04', + help='Docker image to use') + return parser.parse_args() + + +args = parse_arguments() +API_KEY = os.environ['RUNPOD_API_KEY'] +GITHUB_SHA = os.environ['GITHUB_SHA'] +GITHUB_REF = os.environ.get('GITHUB_REF', 'unknown') +GITHUB_REPOSITORY = os.environ['GITHUB_REPOSITORY'] +RUN_ID = os.environ['GITHUB_RUN_ID'] + +PODS_API = "https://rest.runpod.io/v1/pods" +HEADERS = { + "Content-Type": "application/json", + "Authorization": f"Bearer {API_KEY}" +} + + +def create_pod(): + """Create a RunPod instance""" + print(f"Creating RunPod instance with GPU: {args.gpu_type}...") + payload = { + "name": f"fastvideo-github-test-{RUN_ID}", + "containerDiskInGb": args.disk_size, + "volumeInGb": args.volume_size, + "env": { + "GITHUB_SHA": GITHUB_SHA, + "GITHUB_REF": GITHUB_REF + }, + "gpuTypeIds": [args.gpu_type], + "gpuCount": args.gpu_count, + "imageName": args.image + } + + response = requests.post(PODS_API, headers=HEADERS, json=payload) + response_data = response.json() + print(f"Response: {json.dumps(response_data, indent=2)}") + + return response_data["id"] + + +def wait_for_pod(pod_id): + """Wait for pod to be in RUNNING state and fully ready with SSH access""" + print("Waiting for RunPod to be ready...") + + # First wait for RUNNING status + max_attempts = 10 + attempts = 0 + while attempts < max_attempts: + response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS) + pod_data = response.json() + status = pod_data["desiredStatus"] + + if status == "RUNNING": + print("RunPod is running! Now waiting for ports to be assigned...") + break + + print(f"Current status: {status}, waiting... (attempt {attempts+1}/{max_attempts})") + time.sleep(2) + attempts += 1 + + if attempts >= max_attempts: + raise TimeoutError("Timed out waiting for RunPod to reach RUNNING state") + + # Wait for ports to be assigned + max_attempts = 6 + attempts = 0 + while attempts < max_attempts: + response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS) + pod_data = response.json() + port_mappings = pod_data.get("portMappings") + + if (port_mappings is not None and "22" in port_mappings + and pod_data.get("publicIp", "") != ""): + print("RunPod is ready with SSH access!") + print(f"SSH IP: {pod_data['publicIp']}") + print(f"SSH Port: {port_mappings['22']}") + break + + print(f"Waiting for SSH port and public IP to be available... (attempt {attempts+1}/{max_attempts})") + time.sleep(10) + attempts += 1 + + if attempts >= max_attempts: + raise TimeoutError("Timed out waiting for RunPod SSH access") + + +def execute_command(pod_id): + """Execute command on the pod via SSH using system SSH client""" + print(f"Running command: {args.test_command}") + + response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS) + pod_data = response.json() + ssh_ip = pod_data["publicIp"] + ssh_port = pod_data["portMappings"]["22"] + + # Copy the repository to the pod using scp + repo_dir = os.path.abspath(os.getcwd()) + repo_name = os.path.basename(repo_dir) + + print(f"Copying repository from {repo_dir} to RunPod...") + + tar_command = [ + "tar", "-czf", "/tmp/repo.tar.gz", "-C", + os.path.dirname(repo_dir), repo_name + ] + subprocess.run(tar_command, check=True) + + # Copy the tarball to the pod + scp_command = [ + "scp", "-o", "StrictHostKeyChecking=no", "-o", + "UserKnownHostsFile=/dev/null", "-P", + str(ssh_port), "/tmp/repo.tar.gz", f"root@{ssh_ip}:/tmp/" + ] + subprocess.run(scp_command, check=True) + + setup_steps = [ + "cd /workspace", + "wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh", + "bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda3", + "source $HOME/miniconda3/bin/activate", + "conda create --name venv python=3.10.0 -y", "conda activate venv", + "mkdir -p /workspace/repo", + "tar -xzf /tmp/repo.tar.gz --no-same-owner -C /workspace/", + f"cd /workspace/{repo_name}", args.test_command + ] + remote_command = " && ".join(setup_steps) + + ssh_command = [ + "ssh", "-o", "StrictHostKeyChecking=no", "-o", + "UserKnownHostsFile=/dev/null", "-p", + str(ssh_port), f"root@{ssh_ip}", remote_command + ] + + print(f"Connecting to {ssh_ip}:{ssh_port}...") + + try: + process = subprocess.Popen(ssh_command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + bufsize=1) + + stdout_lines = [] + + print("Command output:") + + for line in iter(process.stdout.readline, ''): + print(line.strip()) + stdout_lines.append(line) + + process.wait() + + return_code = process.returncode + success = return_code == 0 + + stdout_str = "".join(stdout_lines) + + if success: + print("Command executed successfully") + else: + print(f"Command failed with exit code {return_code}") + + result = { + "success": success, + "return_code": return_code, + "stdout": stdout_str, + "stderr": "" + } + return result + + except Exception as e: + print(f"Error executing SSH command: {str(e)}") + result = {"success": False, "error": str(e), "stdout": "", "stderr": ""} + return result + + +def terminate_pod(pod_id): + """Terminate the pod""" + print("Terminating RunPod...") + requests.delete(f"{PODS_API}/{pod_id}", headers=HEADERS) + print("RunPod terminated") + + +def main(): + pod_id = None + try: + pod_id = create_pod() + wait_for_pod(pod_id) + result = execute_command(pod_id) + + if result.get("error") is not None: + print(f"Error executing command: {result['error']}") + sys.exit(1) + + if not result.get("success", False): + print( + "Tests failed - check the output above for details on which tests failed" + ) + sys.exit(1) + + finally: + if pod_id: + terminate_pod(pod_id) + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml new file mode 100644 index 000000000..ca41d3206 --- /dev/null +++ b/.github/workflows/pr-test.yml @@ -0,0 +1,53 @@ +name: PR Test + +on: + push: + branches: [main] + pull_request: + branches: [main] + types: [opened, ready_for_review, synchronize, reopened] + +concurrency: + group: pr-test-${{ github.ref }} + cancel-in-progress: true + +jobs: + unit-test: + if: github.event.pull_request.draft == false + runs-on: ubuntu-latest + environment: runpod-runners + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + + - name: Set up SSH key + run: | + mkdir -p ~/.ssh + echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa + chmod 600 ~/.ssh/id_rsa + ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub + + - name: Install dependencies + run: pip install requests + + - name: Run tests on RunPod + env: + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + GITHUB_SHA: ${{ github.sha }} + GITHUB_REF: ${{ github.ref }} + GITHUB_REPOSITORY: ${{ github.repository }} + GITHUB_RUN_ID: ${{ github.run_id }} + run: >- + python .github/scripts/runpod_api.py + --gpu-type "NVIDIA A40" + --gpu-count 1 + --volume-size 100 + --test-command "./env_setup.sh && + pip install vllm && + pip install pytest && + pytest ./fastvideo/v1/tests/ --ignore=./fastvideo/v1/tests/old_tests/ -s" diff --git a/fastvideo/v1/tests/conftest.py b/fastvideo/v1/tests/conftest.py new file mode 100644 index 000000000..6f1e88b03 --- /dev/null +++ b/fastvideo/v1/tests/conftest.py @@ -0,0 +1,24 @@ +import pytest +import torch.distributed as dist +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, destroy_model_parallel + +@pytest.fixture(scope="function") +def distributed_setup(): + """ + Fixture to set up and tear down the distributed environment for tests. + + This ensures proper cleanup even if tests fail. + """ + init_distributed_environment(world_size=1, + rank=0, + distributed_init_method="env://", + local_rank=0, + backend="nccl") + initialize_model_parallel(tensor_model_parallel_size=1, + sequence_model_parallel_size=1, + backend="nccl") + yield + + if dist.is_initialized(): + destroy_model_parallel() + dist.destroy_process_group() diff --git a/fastvideo/v1/tests/distributed_example.py b/fastvideo/v1/tests/old_tests/distributed_example.py similarity index 100% rename from fastvideo/v1/tests/distributed_example.py rename to fastvideo/v1/tests/old_tests/distributed_example.py diff --git a/fastvideo/v1/tests/distributed_run.sh b/fastvideo/v1/tests/old_tests/distributed_run.sh similarity index 100% rename from fastvideo/v1/tests/distributed_run.sh rename to fastvideo/v1/tests/old_tests/distributed_run.sh diff --git a/fastvideo/v1/tests/test_hunyuan_vae.py b/fastvideo/v1/tests/old_tests/test_hunyuan_vae.py similarity index 100% rename from fastvideo/v1/tests/test_hunyuan_vae.py rename to fastvideo/v1/tests/old_tests/test_hunyuan_vae.py diff --git a/fastvideo/v1/tests/test_hunyuanvideo.py b/fastvideo/v1/tests/old_tests/test_hunyuanvideo.py similarity index 100% rename from fastvideo/v1/tests/test_hunyuanvideo.py rename to fastvideo/v1/tests/old_tests/test_hunyuanvideo.py diff --git a/fastvideo/v1/tests/test_hunyuanvideo_load.py b/fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py similarity index 100% rename from fastvideo/v1/tests/test_hunyuanvideo_load.py rename to fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py diff --git a/fastvideo/v1/tests/test_t5_encoder.py b/fastvideo/v1/tests/old_tests/test_t5_encoder.py similarity index 100% rename from fastvideo/v1/tests/test_t5_encoder.py rename to fastvideo/v1/tests/old_tests/test_t5_encoder.py diff --git a/fastvideo/v1/tests/test_wan_vae.py b/fastvideo/v1/tests/old_tests/test_wan_vae.py similarity index 100% rename from fastvideo/v1/tests/test_wan_vae.py rename to fastvideo/v1/tests/old_tests/test_wan_vae.py diff --git a/fastvideo/v1/tests/tp_example.py b/fastvideo/v1/tests/old_tests/tp_example.py similarity index 100% rename from fastvideo/v1/tests/tp_example.py rename to fastvideo/v1/tests/old_tests/tp_example.py diff --git a/fastvideo/v1/tests/test_clip_encoder.py b/fastvideo/v1/tests/test_clip_encoder.py index 29f9d7c55..73b07579c 100644 --- a/fastvideo/v1/tests/test_clip_encoder.py +++ b/fastvideo/v1/tests/test_clip_encoder.py @@ -1,108 +1,64 @@ # TODO: check if correct +import pytest from fastvideo.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer import os import torch import torch.nn as nn -import argparse import numpy as np from fastvideo.v1.logger import init_logger from transformers import AutoConfig from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel # from fastvideo.v1.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.utils import maybe_download_model +from fastvideo.v1.inference_args import InferenceArgs + logger = init_logger(__name__) +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" -def initialize_identical_weights(model1, model2, seed=42): - """Initialize both models with identical weights using a fixed seed for reproducibility.""" - # Get all parameters from both models - params1 = dict(model1.named_parameters()) - params2 = dict(model2.named_parameters()) - - # Initialize each layer with identical values - with torch.no_grad(): - # Initialize weights - for name1, param1 in params1.items(): - if 'weight' in name1: - # Set seed before each weight initialization - torch.manual_seed(seed) - nn.init.normal_(param1, mean=0.0, std=0.05) - - for name2, param2 in params2.items(): - if 'weight' in name2: - # Reset seed to get same initialization - torch.manual_seed(seed) - nn.init.normal_(param2, mean=0.0, std=0.05) - - # Initialize biases - for name1, param1 in params1.items(): - if 'bias' in name1: - torch.manual_seed(seed) - nn.init.normal_(param1, mean=0.0, std=0.05) - param1.data = param1.data.to(torch.bfloat16) - - for name2, param2 in params2.items(): - if 'bias' in name2: - torch.manual_seed(seed) - nn.init.normal_(param2, mean=0.0, std=0.05) - param2.data = param2.data.to(torch.bfloat16) - - logger.info("Both models initialized with identical weights in bfloat16") - return model1, model2 - - -def setup_args(): - parser = argparse.ArgumentParser(description='CLIP Text Encoder Test') - parser.add_argument('--model-path', - type=str, - default="openai/clip-vit-large-patch14", - help='Path to the CLIP model') - parser.add_argument( - '--precision', - type=str, - default="float16", - help='Precision to use for the model (float32, float16, bfloat16)') - return parser.parse_args() +# Set fixed random seed for reproducibility +torch.manual_seed(42) +np.random.seed(42) +BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH)) +TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder_2") +@pytest.mark.usefixtures("distributed_setup") def test_clip_encoder(): - init_distributed_environment(world_size=1, - rank=0, - distributed_init_method="env://", - local_rank=0, - backend="nccl") - initialize_model_parallel(tensor_model_parallel_size=1, - sequence_model_parallel_size=1, - backend="nccl") - args = setup_args() - - # Set fixed random seed for reproducibility - torch.manual_seed(42) - np.random.seed(42) - + """ + Tests compatibility between two different implementations for loading text encoders: + 1. load_text_encoder from fastvideo.models.hunyuan.text_encoder + 2. TextEncoderLoader from fastvideo.v1.models.loader + + The test verifies that both implementations: + - Load models with the same weights and parameters + - Produce nearly identical outputs for the same input prompts + """ + args = InferenceArgs(model_path="openai/clip-vit-large-patch14", precision="float16") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - # Initialize the two model implementations logger.info(f"Loading models from {args.model_path}") - model_path = "data/hunyuanvideo-community/HunyuanVideo/text_encoder_2" # config = json.load(open(os.path.join(model_path, "config.json"))) - hf_config = AutoConfig.from_pretrained(model_path) + hf_config = AutoConfig.from_pretrained(TEXT_ENCODER_PATH) print(hf_config) print(hf_config.use_return_dict) # Load our implementation using the loader from text_encoder/__init__.py model1, _ = load_text_encoder(text_encoder_type="clipL", text_encoder_precision='fp16', - text_encoder_path=model_path, + text_encoder_path=TEXT_ENCODER_PATH, logger=logger, device=device) from fastvideo.v1.models.loader.component_loader import TextEncoderLoader loader = TextEncoderLoader() args.device_str = "cuda:0" - model2 = loader.load_model(model_path, hf_config, device) + model2 = loader.load_model(TEXT_ENCODER_PATH, hf_config, device) # Load the HuggingFace implementation directly # model2 = CLIPTextModel(hf_config) @@ -210,18 +166,13 @@ def test_clip_encoder(): ) logger.info( f"Mean difference in pooler outputs: {mean_diff_pooler.item()}") - - # Check if outputs are similar (allowing for small numerical differences) - assert max_diff_hidden < 1e-4, \ + + # Check if outputs are similar (allowing for small numerical differences) + assert mean_diff_hidden < 1e-2, \ + f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" + assert mean_diff_pooler < 1e-2, \ + f"Pooler outputs differ significantly: mean diff = {mean_diff_pooler.item()}" + assert max_diff_hidden < 1e-1, \ f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" - assert max_diff_pooler < 1e-4, \ + assert max_diff_pooler < 1e-2, \ f"Pooler outputs differ significantly: max diff = {max_diff_pooler.item()}" - - logger.info( - "Test passed! Both CLIP text encoder implementations produce similar outputs." - ) - logger.info("Test completed successfully") - - -if __name__ == "__main__": - test_clip_encoder() diff --git a/fastvideo/v1/tests/test_llama_encoder.py b/fastvideo/v1/tests/test_llama_encoder.py index a85d3ae1a..62dfc819e 100644 --- a/fastvideo/v1/tests/test_llama_encoder.py +++ b/fastvideo/v1/tests/test_llama_encoder.py @@ -1,107 +1,62 @@ -from fastvideo.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer +import pytest +from fastvideo.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer import os import torch import torch.nn as nn -import argparse import numpy as np from fastvideo.v1.logger import init_logger from transformers import AutoConfig from fastvideo.v1.models.loader.component_loader import TextEncoderLoader - -from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel -from fastvideo.v1.pipelines.stages import DenoisingStage from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.utils import maybe_download_model +from fastvideo.v1.inference_args import InferenceArgs + logger = init_logger(__name__) +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29503" -def initialize_identical_weights(model1, model2, seed=42): - """Initialize both models with identical weights using a fixed seed for reproducibility.""" - # Get all parameters from both models - params1 = dict(model1.named_parameters()) - params2 = dict(model2.named_parameters()) - - # Initialize each layer with identical values - with torch.no_grad(): - # Initialize weights - for name1, param1 in params1.items(): - if 'weight' in name1: - # Set seed before each weight initialization - torch.manual_seed(seed) - nn.init.normal_(param1, mean=0.0, std=0.05) - - for name2, param2 in params2.items(): - if 'weight' in name2: - # Reset seed to get same initialization - torch.manual_seed(seed) - nn.init.normal_(param2, mean=0.0, std=0.05) - - # Initialize biases - for name1, param1 in params1.items(): - if 'bias' in name1: - torch.manual_seed(seed) - nn.init.normal_(param1, mean=0.0, std=0.05) - param1.data = param1.data.to(torch.bfloat16) - - for name2, param2 in params2.items(): - if 'bias' in name2: - torch.manual_seed(seed) - nn.init.normal_(param2, mean=0.0, std=0.05) - param2.data = param2.data.to(torch.bfloat16) - - logger.info("Both models initialized with identical weights in bfloat16") - return model1, model2 - - -def setup_args(): - parser = argparse.ArgumentParser(description='LLaMA Encoder Test') - parser.add_argument('--model-path', - type=str, - default="meta-llama/Llama-2-7b-hf", - help='Path to the LLaMA model') - parser.add_argument( - '--precision', - type=str, - default="float16", - help='Precision to use for the model (float32, float16, bfloat16)') - return parser.parse_args() +# Set fixed random seed for reproducibility +torch.manual_seed(42) +np.random.seed(42) +BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join('data', BASE_MODEL_PATH)) +TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder") +TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer") +@pytest.mark.usefixtures("distributed_setup") def test_llama_encoder(): - init_distributed_environment(world_size=1, - rank=0, - distributed_init_method="env://", - local_rank=0, - backend="nccl") - initialize_model_parallel(tensor_model_parallel_size=1, - sequence_model_parallel_size=1, - backend="nccl") - args = setup_args() - - # Set fixed random seed for reproducibility - torch.manual_seed(42) - np.random.seed(42) + """ + Tests compatibility between two different implementations for loading text encoders: + 1. load_text_encoder from fastvideo.models.hunyuan.text_encoder + 2. TextEncoderLoader from fastvideo.v1.models.loader + + The test verifies that both implementations: + - Load models with the same weights and parameters + - Produce nearly identical outputs for the same input prompts + """ + args = InferenceArgs(model_path="meta-llama/Llama-2-7b-hf", precision="float16") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Initialize the two model implementations logger.info(f"Loading models from {args.model_path}") - model_path = "data/hunyuanvideo-community/HunyuanVideo/text_encoder" - - hf_config = AutoConfig.from_pretrained(model_path) + hf_config = AutoConfig.from_pretrained(TEXT_ENCODER_PATH) print(hf_config) # Load our implementation using the loader from text_encoder/__init__.py model1, _ = load_text_encoder( text_encoder_type="llm", text_encoder_precision='fp16', - text_encoder_path=model_path, + text_encoder_path=TEXT_ENCODER_PATH, logger=logger, device=device ) loader = TextEncoderLoader() args.device_str = "cuda:0" device = torch.device(args.device_str) - model2 = loader.load_model(model_path, hf_config, device) + model2 = loader.load_model(TEXT_ENCODER_PATH, hf_config, device) # Convert to float16 and move to device model2 = model2.to(torch.float16) @@ -149,10 +104,8 @@ def test_llama_encoder(): except Exception as e: logger.info(f"Error comparing {name1} and {name2}: {e}") - tokenizer_path = "data/hunyuanvideo-community/HunyuanVideo/tokenizer" - # Load tokenizer tokenizer, _ = load_tokenizer(tokenizer_type="llm", - tokenizer_path=tokenizer_path, + tokenizer_path=TOKENIZER_PATH, logger=logger) # Test with some sample prompts @@ -214,11 +167,8 @@ def test_llama_encoder(): f"Mean difference in last hidden states: {mean_diff_hidden.item()}" ) - logger.info( - "Test passed! Both LLaMA encoder implementations produce similar outputs." - ) - logger.info("Test completed successfully") - - -if __name__ == "__main__": - test_llama_encoder() + # Check if outputs are similar (allowing for small numerical differences) + assert mean_diff_hidden < 1e-2, \ + f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" + assert max_diff_hidden < 1e-1, \ + f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" diff --git a/fastvideo/v1/utils.py b/fastvideo/v1/utils.py index 20e14ad49..0c4dea81d 100644 --- a/fastvideo/v1/utils.py +++ b/fastvideo/v1/utils.py @@ -380,7 +380,7 @@ def import_pynvml(): return pynvml -def maybe_download_model(model_path: str) -> str: +def maybe_download_model(model_path: str, local_dir: str = None) -> str: """ Check if the model path is a Hugging Face Hub model ID and download it if needed. @@ -404,6 +404,7 @@ def maybe_download_model(model_path: str) -> str: local_path = snapshot_download( repo_id=model_path, ignore_patterns=["*.onnx", "*.msgpack"], + local_dir=local_dir ) logger.info("Downloaded model to %s", local_path) return str(local_path)