From a194907d091dfa462afc42f303ec56c928df67d3 Mon Sep 17 00:00:00 2001
From: Politrees <143968312+Bebra777228@users.noreply.github.com>
Date: Fri, 6 Sep 2024 10:23:50 +0500
Subject: [PATCH] Update model_management.py

---
 rvc/modules/model_management.py | 42 ++++++++++++++++++++-------------
 1 file changed, 25 insertions(+), 17 deletions(-)

diff --git a/rvc/modules/model_management.py b/rvc/modules/model_management.py
index 937f540..8937fe6 100644
--- a/rvc/modules/model_management.py
+++ b/rvc/modules/model_management.py
@@ -110,37 +110,45 @@ def download_file(url, zip_name, progress):
 
         elif "onedrive.live.com" in url:
             progress(0.5, desc="[~] Загрузка модели с OneDrive...")
-            direct_url = url.replace("?download=1", "?download=0")
-            response = requests.get(direct_url, stream=True)
-            with open(zip_name, "wb") as f:
-                for chunk in response.iter_content(chunk_size=8192):
-                    f.write(chunk)
+            response = requests.get(url, allow_redirects=True)
+            if response.status_code == 200:
+                download_url = re.search(r'href="([^"]+)"', response.text)
+                if download_url:
+                    urllib.request.urlretrieve(download_url.group(1), zip_name)
+            else:
+                raise gr.Error("Ошибка загрузки с OneDrive.")
 
         elif "dropbox.com" in url:
             progress(0.5, desc="[~] Загрузка модели с Dropbox...")
-            direct_url = url.split("?")[0]
-            response = requests.get(direct_url, stream=True)
-            with open(zip_name, "wb") as f:
-                for chunk in response.iter_content(chunk_size=8192):
-                    f.write(chunk)
+            direct_url = url.split("?")[0] + "?dl=1"
+            urllib.request.urlretrieve(direct_url, zip_name)
 
         elif "box.com" in url:
             progress(0.5, desc="[~] Загрузка модели с Box...")
             response = requests.get(url)
-            direct_url = re.search(r'href="([^"]+)"\s+class="download-btn"', response.text).group(1)
-            urllib.request.urlretrieve(direct_url, zip_name)
+            direct_url = re.search(r'href="([^"]+)"\s+class="download-btn"', response.text)
+            if direct_url:
+                urllib.request.urlretrieve(direct_url.group(1), zip_name)
+            else:
+                raise gr.Error("Не удалось найти ссылку для скачивания с Box.")
 
         elif "mediafire.com" in url:
             progress(0.5, desc="[~] Загрузка модели с MediaFire...")
             response = requests.get(url)
-            direct_url = re.search(r'href="([^"]+)"\s+class="download_link"', response.text).group(1)
-            urllib.request.urlretrieve(direct_url, zip_name)
+            direct_url = re.search(r'href="([^"]+)"\s+class="download_link"', response.text)
+            if direct_url:
+                urllib.request.urlretrieve(direct_url.group(1), zip_name)
+            else:
+                raise gr.Error("Не удалось найти ссылку для скачивания с MediaFire.")
 
         elif "pcloud.com" in url:
             progress(0.5, desc="[~] Загрузка модели с pCloud...")
             response = requests.get(url)
-            direct_url = re.search(r'href="([^"]+)"\s+class="download-button"', response.text).group(1)
-            urllib.request.urlretrieve(direct_url, zip_name)
+            direct_url = re.search(r'href="([^"]+)"\s+class="download-button"', response.text)
+            if direct_url:
+                urllib.request.urlretrieve(direct_url.group(1), zip_name)
+            else:
+                raise gr.Error("Не удалось найти ссылку для скачивания с pCloud.")
 
     except Exception as e:
         raise gr.Error(f"Ошибка при загрузке файла: {str(e)}")
@@ -230,4 +238,4 @@ def main():
         sys.exit(1)
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()