From 4a5d48000fe26711aa11947afb9e256a649e1fdb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E9=80=B8=E7=BE=A4?= <875977494@qq.com>
Date: Wed, 23 Jul 2025 03:40:35 +0800
Subject: [PATCH] feat: add HuggingFace mirror support for VM providers (#278)

Add support for HuggingFace mirror (hf-mirror.com) to improve download
speeds in regions where huggingface.co access is slow.

- Support HF_ENDPOINT environment variable detection
- Automatically switch to hf-mirror.com when HF_ENDPOINT is set
- Apply to Docker, VMware, and VirtualBox providers
- Maintain backward compatibility with default huggingface.co URLs

Users can now set HF_ENDPOINT=https://hf-mirror.com to use the mirror.
---
 desktop_env/providers/docker/manager.py     | 13 +++++++++++++
 desktop_env/providers/virtualbox/manager.py |  6 ++++++
 desktop_env/providers/vmware/manager.py     |  7 +++++++
 3 files changed, 26 insertions(+)

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