Politrees commited on
Commit
8226d13
1 Parent(s): 3b09304

Update src/modules/model_management.py

Browse files
Files changed (1) hide show
  1. src/modules/model_management.py +14 -14
src/modules/model_management.py CHANGED
@@ -34,7 +34,7 @@ def extract_zip(extraction_folder, zip_name):
34
  model_filepath = os.path.join(root, name)
35
 
36
  if not model_filepath:
37
- raise gr.Error(f'Не найден файл модели .pth в распакованном zip-файле. Пожалуйста, проверьте {extraction_folder}.')
38
 
39
  os.rename(model_filepath, os.path.join(extraction_folder, os.path.basename(model_filepath)))
40
  if index_filepath:
@@ -46,47 +46,47 @@ def extract_zip(extraction_folder, zip_name):
46
 
47
  def download_from_url(url, dir_name, progress=gr.Progress()):
48
  try:
49
- progress(0, desc=f'[~] Загрузка голосовой модели с именем {dir_name}...')
50
  zip_name = os.path.join(rvc_models_dir, dir_name + '.zip')
51
  extraction_folder = os.path.join(rvc_models_dir, dir_name)
52
  if os.path.exists(extraction_folder):
53
- raise gr.Error(f'Директория голосовой модели {dir_name} уже существует! Выберите другое имя для вашей голосовой модели.')
54
 
55
  if 'drive.google.com' in url:
56
- progress(0.5, desc='[~] Загрузка модели с Google Grive...')
57
  file_id = url.split("file/d/")[1].split("/")[0] if "file/d/" in url else url.split("id=")[1].split("&")[0]
58
  output = zip_name
59
  gdown.download(id=file_id, output=output, quiet=False)
60
 
61
  elif 'huggingface.co' in url:
62
- progress(0.5, desc='[~] Загрузка модели с HuggingFace...')
63
  urllib.request.urlretrieve(url, zip_name)
64
 
65
  elif 'pixeldrain.com' in url:
66
- progress(0.5, desc='[~] Загрузка модели с Pixeldrain...')
67
  file_id = url.split("pixeldrain.com/u/")[1]
68
  response = requests.get(f"https://pixeldrain.com/api/file/{file_id}")
69
  with open(zip_name, 'wb') as f:
70
  f.write(response.content)
71
 
72
  elif 'mega.nz' in url:
73
- progress(0.5, desc='[~] Загрузка модели с Mega...')
74
  m = Mega()
75
  m.download_url(url, dest_filename=zip_name)
76
 
77
  elif 'yadi.sk' in url or 'disk.yandex.ru' in url:
78
- progress(0.5, desc='[~] Загрузка модели с Яндекс Диска...')
79
  yandex_api_url = "https://cloud-api.yandex.net/v1/disk/public/resources/download?public_key={}".format(url)
80
  response = requests.get(yandex_api_url)
81
  if response.status_code == 200:
82
  download_link = response.json().get('href')
83
  urllib.request.urlretrieve(download_link, zip_name)
84
  else:
85
- raise gr.Error(f"Ошибка при получении ссылки на скачивание с Яндекс Диск: {response.status_code}")
86
 
87
- progress(0.8, desc='[~] Распаковка zip-файла...')
88
  extract_zip(extraction_folder, zip_name)
89
- return f'[+] Модель {dir_name} успешно загружена!'
90
  except Exception as e:
91
  raise gr.Error(str(e))
92
 
@@ -94,12 +94,12 @@ def upload_zip_model(zip_path, dir_name, progress=gr.Progress()):
94
  try:
95
  extraction_folder = os.path.join(rvc_models_dir, dir_name)
96
  if os.path.exists(extraction_folder):
97
- raise gr.Error(f'Директория голосовой модели {dir_name} уже существует! Выберите другое имя для вашей голосовой модели.')
98
 
99
  zip_name = zip_path.name
100
- progress(0.8, desc='[~] Распаковка zip-файла...')
101
  extract_zip(extraction_folder, zip_name)
102
- return f'[+] Модель {dir_name} успешно загружена!'
103
 
104
  except Exception as e:
105
  raise gr.Error(str(e))
 
34
  model_filepath = os.path.join(root, name)
35
 
36
  if not model_filepath:
37
+ raise gr.Error(f'The .pth model file was not found in the unzipped zip file. Please check the {extraction_folder}.')
38
 
39
  os.rename(model_filepath, os.path.join(extraction_folder, os.path.basename(model_filepath)))
40
  if index_filepath:
 
46
 
47
  def download_from_url(url, dir_name, progress=gr.Progress()):
48
  try:
49
+ progress(0, desc=f'[~] Downloading voice model named {dir_name}...')
50
  zip_name = os.path.join(rvc_models_dir, dir_name + '.zip')
51
  extraction_folder = os.path.join(rvc_models_dir, dir_name)
52
  if os.path.exists(extraction_folder):
53
+ raise gr.Error(f'The voice model directory {dir_name} already exists! Please choose a different name for your voice model.')
54
 
55
  if 'drive.google.com' in url:
56
+ progress(0.5, desc='[~] Downloading model from Google Grive...')
57
  file_id = url.split("file/d/")[1].split("/")[0] if "file/d/" in url else url.split("id=")[1].split("&")[0]
58
  output = zip_name
59
  gdown.download(id=file_id, output=output, quiet=False)
60
 
61
  elif 'huggingface.co' in url:
62
+ progress(0.5, desc='[~] Downloading model from HuggingFace...')
63
  urllib.request.urlretrieve(url, zip_name)
64
 
65
  elif 'pixeldrain.com' in url:
66
+ progress(0.5, desc='[~] Downloading model from Pixeldrain...')
67
  file_id = url.split("pixeldrain.com/u/")[1]
68
  response = requests.get(f"https://pixeldrain.com/api/file/{file_id}")
69
  with open(zip_name, 'wb') as f:
70
  f.write(response.content)
71
 
72
  elif 'mega.nz' in url:
73
+ progress(0.5, desc='[~] Downloading model from Mega...')
74
  m = Mega()
75
  m.download_url(url, dest_filename=zip_name)
76
 
77
  elif 'yadi.sk' in url or 'disk.yandex.ru' in url:
78
+ progress(0.5, desc='[~] Downloading model from Yandex Disk...')
79
  yandex_api_url = "https://cloud-api.yandex.net/v1/disk/public/resources/download?public_key={}".format(url)
80
  response = requests.get(yandex_api_url)
81
  if response.status_code == 200:
82
  download_link = response.json().get('href')
83
  urllib.request.urlretrieve(download_link, zip_name)
84
  else:
85
+ raise gr.Error(f"Error when receiving a download link from Yandex Disk: {response.status_code}")
86
 
87
+ progress(0.8, desc='[~] Extracting zip file...')
88
  extract_zip(extraction_folder, zip_name)
89
+ return f'[+] The {dir_name} model was loaded successfully!'
90
  except Exception as e:
91
  raise gr.Error(str(e))
92
 
 
94
  try:
95
  extraction_folder = os.path.join(rvc_models_dir, dir_name)
96
  if os.path.exists(extraction_folder):
97
+ raise gr.Error(f'The voice model directory {dir_name} already exists! Choose a different name for your voice model.')
98
 
99
  zip_name = zip_path.name
100
+ progress(0.8, desc='[~] Extracting zip file...')
101
  extract_zip(extraction_folder, zip_name)
102
+ return f'[+] The {dir_name} model was loaded successfully!'
103
 
104
  except Exception as e:
105
  raise gr.Error(str(e))