File size: 2,647 Bytes
72ae837 |
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 |
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 Exception(f'No .pth model file was found in the extracted zip. Please check {extraction_folder}.')
# move model and index file to 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)))
# remove any unnecessary nested folders
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_online_model(url, dir_name, models_dir='./rvc_models'):
try:
print(f'[~] Downloading voice model with name {dir_name}...')
zip_name = url.split('/')[-1]
extraction_folder = os.path.join(models_dir, dir_name)
if os.path.exists(extraction_folder):
return f'[!] Voice model directory {dir_name} already exists! Choose a different name for your voice model.'
# Download from pixeldrain
if 'pixeldrain.com' in url:
url = f'https://pixeldrain.com/api/file/{zip_name}'
urllib.request.urlretrieve(url, zip_name)
# Download from Google Drive
elif 'drive.google.com' in url:
zip_name = dir_name + ".zip"
gdown.download(url, output=zip_name, use_cookies=True, quiet=True)
else:
# General URL download
urllib.request.urlretrieve(url, zip_name)
print(f'[~] Extracting zip file...')
extract_zip(extraction_folder, zip_name)
print(f'[+] {dir_name} Model successfully downloaded!')
# Return success message after successful download
return f"[+] {dir_name} Model successfully downloaded!"
except Exception as e:
# Return the error message instead of raising an exception
return f'[!] Error: {str(e)}'
|