From fd65cf47f63b4041ec5e15f03729bc95677c97c3 Mon Sep 17 00:00:00 2001
From: FredWuCZ <fredwucz@outlook.com>
Date: Wed, 2 Oct 2024 12:19:01 +0800
Subject: [PATCH] Update Windows URL

---
 desktop_env/desktop_env.py              |  2 +-
 desktop_env/providers/docker/manager.py | 15 ++++++++++-----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/desktop_env/desktop_env.py b/desktop_env/desktop_env.py
index 9b24d05..7d7b409 100644
--- a/desktop_env/desktop_env.py
+++ b/desktop_env/desktop_env.py
@@ -36,7 +36,7 @@ class DesktopEnv(gym.Env):
             headless: bool = False,
             require_a11y_tree: bool = True,
             require_terminal: bool = False,
-            os_type: str = "Ubuntu",
+            os_type: str = "Windows",
     ):
         """
         Args:
diff --git a/desktop_env/providers/docker/manager.py b/desktop_env/providers/docker/manager.py
index e4b09a1..f3cb2d3 100644
--- a/desktop_env/providers/docker/manager.py
+++ b/desktop_env/providers/docker/manager.py
@@ -27,6 +27,7 @@ MAX_RETRY_TIMES = 10
 RETRY_INTERVAL = 5
 
 UBUNTU_X86_URL = "https://huggingface.co/datasets/xlangai/ubuntu_osworld/resolve/main/Ubuntu.qcow2"
+WINDOWS_X86_URL = "https://huggingface.co/datasets/xlangai/windows_osworld/resolve/main/Windows%2010%20x64%20-%20copy.qcow2"
 VMS_DIR = "./docker_vm_data"
 
 # Determine the platform and CPU architecture to decide the correct VM image to download
@@ -39,21 +40,20 @@ VMS_DIR = "./docker_vm_data"
 #     URL = UBUNTU_X86_URL
 # else:
 #     raise Exception("Unsupported platform or architecture")
-URL = UBUNTU_X86_URL
 
+URL = UBUNTU_X86_URL
 DOWNLOADED_FILE_NAME = URL.split('/')[-1]
 
 if platform.system() == 'Windows':
     docker_path = r"C:\Program Files\Docker\Docker"
     os.environ["PATH"] += os.pathsep + docker_path
 
-def _download_vm(vms_dir: str):
+def _download_vm(vms_dir: str, os_type: str):
+    global URL, DOWNLOADED_FILE_NAME
     # Download the virtual machine image
     logger.info("Downloading the virtual machine image...")
     downloaded_size = 0
-
-    URL = UBUNTU_X86_URL
-    DOWNLOADED_FILE_NAME = URL.split('/')[-1]
+    
     downloaded_file_name = DOWNLOADED_FILE_NAME
 
     os.makedirs(vms_dir, exist_ok=True)
@@ -118,6 +118,11 @@ class DockerVMManager(VMManager):
         pass
 
     def get_vm_path(self, os_type, region):
+        if os_type == "Ubuntu":
+            URL = WINDOWS_X86_URL
+        elif os_type == "Windows":
+            URL = WINDOWS_X86_URL
+        DOWNLOADED_FILE_NAME = URL.split('/')[-1]
         if not os.path.exists(os.path.join(VMS_DIR, DOWNLOADED_FILE_NAME)):
             _download_vm(VMS_DIR)
         return os.path.join(VMS_DIR, DOWNLOADED_FILE_NAME)
\ No newline at end of file