Spaces:
Running
Running
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)) | |