Aitron Emper commited on
Commit
9dea1d3
1 Parent(s): 037701f

Upload model_download.py

Browse files
Files changed (1) hide show
  1. rvc/lib/tools/model_download.py +47 -3
rvc/lib/tools/model_download.py CHANGED
@@ -6,6 +6,7 @@ from bs4 import BeautifulSoup
6
  import requests
7
  from urllib.parse import unquote
8
  import re
 
9
 
10
  def find_folder_parent(search_dir, folder_name):
11
  for dirpath, dirnames, _ in os.walk(search_dir):
@@ -16,6 +17,8 @@ def find_folder_parent(search_dir, folder_name):
16
  now_dir = os.getcwd()
17
  sys.path.append(now_dir)
18
 
 
 
19
  import rvc.lib.tools.gdown as gdown
20
 
21
  file_path = find_folder_parent(now_dir, "logs")
@@ -195,6 +198,10 @@ def unzip_file(zip_path, zip_file_name):
195
 
196
 
197
  url = sys.argv[1]
 
 
 
 
198
  verify = download_from_url(url)
199
 
200
  if verify == "downloaded":
@@ -204,20 +211,57 @@ if verify == "downloaded":
204
  zipfile_path = os.path.join(zips_path, filename)
205
  print("Proceeding with the extraction...")
206
 
207
- model_name = os.path.basename(zipfile_path)
 
208
  extract_folder_path = os.path.join(
209
  "logs",
210
- os.path.normpath(str(model_name).replace(".zip", "")),
211
  )
212
 
213
  success = extract_and_show_progress(zipfile_path, extract_folder_path)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
214
  if success:
215
  print(f"Model {model_name} downloaded!")
216
  else:
217
  print(f"Error downloading {model_name}")
218
  sys.exit()
219
  if extract_folder_path == "":
220
- print("No zip founded...")
221
  sys.exit()
222
  result = search_pth_index(extract_folder_path)
223
  else:
 
6
  import requests
7
  from urllib.parse import unquote
8
  import re
9
+ import shutil
10
 
11
  def find_folder_parent(search_dir, folder_name):
12
  for dirpath, dirnames, _ in os.walk(search_dir):
 
17
  now_dir = os.getcwd()
18
  sys.path.append(now_dir)
19
 
20
+ from rvc.lib.utils import format_title
21
+
22
  import rvc.lib.tools.gdown as gdown
23
 
24
  file_path = find_folder_parent(now_dir, "logs")
 
198
 
199
 
200
  url = sys.argv[1]
201
+
202
+ if "?download=true" in url:
203
+ url = url.replace("?download=true", "")
204
+
205
  verify = download_from_url(url)
206
 
207
  if verify == "downloaded":
 
211
  zipfile_path = os.path.join(zips_path, filename)
212
  print("Proceeding with the extraction...")
213
 
214
+ model_zip = os.path.basename(zipfile_path)
215
+ model_name = format_title(model_zip.split(".zip")[0])
216
  extract_folder_path = os.path.join(
217
  "logs",
218
+ os.path.normpath(model_name),
219
  )
220
 
221
  success = extract_and_show_progress(zipfile_path, extract_folder_path)
222
+
223
+ subfolders = [f for f in os.listdir(extract_folder_path) if os.path.isdir(os.path.join(extract_folder_path, f))]
224
+ if len(subfolders) == 1:
225
+ subfolder_path = os.path.join(extract_folder_path, subfolders[0])
226
+ for item in os.listdir(subfolder_path):
227
+ s = os.path.join(subfolder_path, item)
228
+ d = os.path.join(extract_folder_path, item)
229
+ shutil.move(s, d)
230
+ os.rmdir(subfolder_path)
231
+
232
+ for item in os.listdir(extract_folder_path):
233
+ if ".pth" in item:
234
+ file_name = item.split(".pth")[0]
235
+ if file_name != model_name:
236
+ os.rename(
237
+ os.path.join(extract_folder_path, item),
238
+ os.path.join(extract_folder_path, model_name + ".pth"),
239
+ )
240
+ else:
241
+ if "v2" not in item:
242
+ file_name = item.split("_nprobe_1_")[1].split("_v1")[0]
243
+ if file_name != model_name:
244
+ new_file_name = item.split("_nprobe_1_")[0] + "_nprobe_1_" + model_name + "_v1"
245
+ os.rename(
246
+ os.path.join(extract_folder_path, item),
247
+ os.path.join(extract_folder_path, new_file_name + ".index"),
248
+ )
249
+ else:
250
+ file_name = item.split("_nprobe_1_")[1].split("_v2")[0]
251
+ if file_name != model_name:
252
+ new_file_name = item.split("_nprobe_1_")[0] + "_nprobe_1_" + model_name + "_v2"
253
+ os.rename(
254
+ os.path.join(extract_folder_path, item),
255
+ os.path.join(extract_folder_path, new_file_name + ".index"),
256
+ )
257
+
258
  if success:
259
  print(f"Model {model_name} downloaded!")
260
  else:
261
  print(f"Error downloading {model_name}")
262
  sys.exit()
263
  if extract_folder_path == "":
264
+ print("Zip file was not found.")
265
  sys.exit()
266
  result = search_pth_index(extract_folder_path)
267
  else: