diff --git a/desktop_env/providers/docker/manager.py b/desktop_env/providers/docker/manager.py index fcc6a68..01fc096 100644 --- a/desktop_env/providers/docker/manager.py +++ b/desktop_env/providers/docker/manager.py @@ -34,6 +34,12 @@ def _download_vm(vms_dir: str): logger.info("Downloading the virtual machine image...") downloaded_size = 0 + # Check for HF_ENDPOINT environment variable and replace domain if set to hf-mirror.com + hf_endpoint = os.environ.get('HF_ENDPOINT') + if hf_endpoint and 'hf-mirror.com' in hf_endpoint: + URL = URL.replace('huggingface.co', 'hf-mirror.com') + logger.info(f"Using HF mirror: {URL}") + downloaded_file_name = DOWNLOADED_FILE_NAME os.makedirs(vms_dir, exist_ok=True) @@ -113,6 +119,13 @@ class DockerVMManager(VMManager): URL = UBUNTU_X86_URL elif os_type == "Windows": URL = WINDOWS_X86_URL + + # Check for HF_ENDPOINT environment variable and replace domain if set to hf-mirror.com + hf_endpoint = os.environ.get('HF_ENDPOINT') + if hf_endpoint and 'hf-mirror.com' in hf_endpoint: + URL = URL.replace('huggingface.co', 'hf-mirror.com') + logger.info(f"Using HF mirror: {URL}") + DOWNLOADED_FILE_NAME = URL.split('/')[-1] if DOWNLOADED_FILE_NAME.endswith(".zip"): diff --git a/desktop_env/providers/virtualbox/manager.py b/desktop_env/providers/virtualbox/manager.py index 3fa053e..66c6b48 100644 --- a/desktop_env/providers/virtualbox/manager.py +++ b/desktop_env/providers/virtualbox/manager.py @@ -57,6 +57,12 @@ def _install_vm(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu else: raise Exception("Unsupported platform or architecture.") + # Check for HF_ENDPOINT environment variable and replace domain if set to hf-mirror.com + hf_endpoint = os.environ.get('HF_ENDPOINT') + if hf_endpoint and 'hf-mirror.com' in hf_endpoint: + url = url.replace('huggingface.co', 'hf-mirror.com') + logger.info(f"Using HF mirror: {url}") + # Download the virtual machine image logger.info("Downloading the virtual machine image...") downloaded_size = 0 diff --git a/desktop_env/providers/vmware/manager.py b/desktop_env/providers/vmware/manager.py index 806232e..8f73520 100644 --- a/desktop_env/providers/vmware/manager.py +++ b/desktop_env/providers/vmware/manager.py @@ -134,6 +134,13 @@ def _install_vm(vm_name, vms_dir, downloaded_file_name, os_type, original_vm_nam elif os_type == "Windows": if platform.machine().lower() in ['amd64', 'x86_64']: URL = WINDOWS_X86_URL + + # Check for HF_ENDPOINT environment variable and replace domain if set to hf-mirror.com + hf_endpoint = os.environ.get('HF_ENDPOINT') + if hf_endpoint and 'hf-mirror.com' in hf_endpoint: + URL = URL.replace('huggingface.co', 'hf-mirror.com') + logger.info(f"Using HF mirror: {URL}") + DOWNLOADED_FILE_NAME = URL.split('/')[-1] downloaded_file_name = DOWNLOADED_FILE_NAME