Spaces:
Running
Running
File size: 3,981 Bytes
e976963 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import os
import shutil
import urllib.request
import zipfile
import gdown
import gradio as gr
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
rvc_models_dir = os.path.join(BASE_DIR, 'rvc_models')
def ignore_files(models_dir):
models_list = os.listdir(models_dir)
items_to_remove = ['hubert_base.pt', 'MODELS.txt', 'rmvpe.pt', 'fcpe.pt']
return [item for item in models_list if item not in items_to_remove]
def update_models_list():
models_l = ignore_files(rvc_models_dir)
return gr.update(choices=models_l)
def extract_zip(extraction_folder, zip_name):
os.makedirs(extraction_folder)
with zipfile.ZipFile(zip_name, 'r') as zip_ref:
zip_ref.extractall(extraction_folder)
os.remove(zip_name)
index_filepath, model_filepath = None, None
for root, dirs, files in os.walk(extraction_folder):
for name in files:
if name.endswith('.index') and os.stat(os.path.join(root, name)).st_size > 1024 * 100:
index_filepath = os.path.join(root, name)
if name.endswith('.pth') and os.stat(os.path.join(root, name)).st_size > 1024 * 1024 * 40:
model_filepath = os.path.join(root, name)
if not model_filepath:
raise gr.Error(f'Не найден файл модели .pth в распакованном zip-файле. Пожалуйста, проверьте {extraction_folder}.')
os.rename(model_filepath, os.path.join(extraction_folder, os.path.basename(model_filepath)))
if index_filepath:
os.rename(index_filepath, os.path.join(extraction_folder, os.path.basename(index_filepath)))
for filepath in os.listdir(extraction_folder):
if os.path.isdir(os.path.join(extraction_folder, filepath)):
shutil.rmtree(os.path.join(extraction_folder, filepath))
def download_from_url(url, dir_name, progress=gr.Progress()):
try:
progress(0, desc=f'[~] Загрузка голосовой модели с именем {dir_name}...')
zip_name = url.split('/')[-1]
extraction_folder = os.path.join(rvc_models_dir, dir_name)
if os.path.exists(extraction_folder):
raise gr.Error(f'Директория голосовой модели {dir_name} уже существует! Выберите другое имя для вашей голосовой модели.')
if 'huggingface.co' in url:
urllib.request.urlretrieve(url, zip_name)
elif 'pixeldrain.com' in url:
zip_name = dir_name + '.zip'
url = f'https://pixeldrain.com/api/file/{zip_name}'
urllib.request.urlretrieve(url, zip_name)
elif 'drive.google.com' in url:
zip_name = dir_name + '.zip'
file_id = url.split('/')[-2]
output = os.path.join('.', f'{dir_name}.zip')
gdown.download(id=file_id, output=output, quiet=False)
progress(0.5, desc='[~] Распаковка zip-файла...')
extract_zip(extraction_folder, zip_name)
return f'[+] Модель {dir_name} успешно загружена!'
except Exception as e:
raise gr.Error(str(e))
def upload_zip_model(zip_path, dir_name, progress=gr.Progress()):
try:
extraction_folder = os.path.join(rvc_models_dir, dir_name)
if os.path.exists(extraction_folder):
raise gr.Error(f'Директория голосовой модели {dir_name} уже существует! Выберите другое имя для вашей голосовой модели.')
zip_name = zip_path.name
progress(0.5, desc='[~] Распаковка zip-файла...')
extract_zip(extraction_folder, zip_name)
return f'[+] Модель {dir_name} успешно загружена!'
except Exception as e:
raise gr.Error(str(e))
|