| import os |
| import re |
| import sys |
| import shutil |
| import zipfile |
| import requests |
| from bs4 import BeautifulSoup |
| from urllib.parse import unquote |
| from tqdm import tqdm |
|
|
| now_dir = os.getcwd() |
| sys.path.append(now_dir) |
|
|
| from rvc.lib.utils import format_title |
| from rvc.lib.tools import gdown |
|
|
|
|
| file_path = os.path.join(now_dir, "logs") |
| zips_path = os.path.join(file_path, "zips") |
| os.makedirs(zips_path, exist_ok=True) |
|
|
|
|
| def search_pth_index(folder): |
| pth_paths = [ |
| os.path.join(folder, file) |
| for file in os.listdir(folder) |
| if os.path.isfile(os.path.join(folder, file)) and file.endswith(".pth") |
| ] |
| index_paths = [ |
| os.path.join(folder, file) |
| for file in os.listdir(folder) |
| if os.path.isfile(os.path.join(folder, file)) and file.endswith(".index") |
| ] |
| return pth_paths, index_paths |
|
|
|
|
| def download_from_url(url): |
| os.chdir(zips_path) |
|
|
| try: |
| if "drive.google.com" in url: |
| file_id = extract_google_drive_id(url) |
| if file_id: |
| gdown.download( |
| url=f"https://drive.google.com/uc?id={file_id}", |
| quiet=False, |
| fuzzy=True, |
| ) |
| elif "/blob/" in url or "/resolve/" in url: |
| download_blob_or_resolve(url) |
| elif "/tree/main" in url: |
| download_from_huggingface(url) |
| else: |
| download_file(url) |
|
|
| rename_downloaded_files() |
| return "downloaded" |
| except Exception as error: |
| print(f"An error occurred downloading the file: {error}") |
| return None |
| finally: |
| os.chdir(now_dir) |
|
|
|
|
| def extract_google_drive_id(url): |
| if "file/d/" in url: |
| return url.split("file/d/")[1].split("/")[0] |
| if "id=" in url: |
| return url.split("id=")[1].split("&")[0] |
| return None |
|
|
|
|
| def download_blob_or_resolve(url): |
| if "/blob/" in url: |
| url = url.replace("/blob/", "/resolve/") |
| response = requests.get(url, stream=True) |
| if response.status_code == 200: |
| save_response_content(response) |
| else: |
| raise ValueError( |
| "Download failed with status code: " + str(response.status_code) |
| ) |
|
|
|
|
| def save_response_content(response): |
| content_disposition = unquote(response.headers.get("Content-Disposition", "")) |
| file_name = ( |
| re.search(r'filename="([^"]+)"', content_disposition) |
| .groups()[0] |
| .replace(os.path.sep, "_") |
| if content_disposition |
| else "downloaded_file" |
| ) |
|
|
| total_size = int(response.headers.get("Content-Length", 0)) |
| chunk_size = 1024 |
|
|
| with open(os.path.join(zips_path, file_name), "wb") as file, tqdm( |
| total=total_size, unit="B", unit_scale=True, desc=file_name |
| ) as progress_bar: |
| for data in response.iter_content(chunk_size): |
| file.write(data) |
| progress_bar.update(len(data)) |
|
|
|
|
| def download_from_huggingface(url): |
| response = requests.get(url) |
| soup = BeautifulSoup(response.content, "html.parser") |
| temp_url = next( |
| ( |
| link["href"] |
| for link in soup.find_all("a", href=True) |
| if link["href"].endswith(".zip") |
| ), |
| None, |
| ) |
| if temp_url: |
| url = temp_url.replace("blob", "resolve") |
| if "huggingface.co" not in url: |
| url = "https://huggingface.co" + url |
| download_file(url) |
| else: |
| raise ValueError("No zip file found in Huggingface URL") |
|
|
|
|
| def download_file(url): |
| response = requests.get(url, stream=True) |
| if response.status_code == 200: |
| save_response_content(response) |
| else: |
| raise ValueError( |
| "Download failed with status code: " + str(response.status_code) |
| ) |
|
|
|
|
| def rename_downloaded_files(): |
| for currentPath, _, zipFiles in os.walk(zips_path): |
| for file in zipFiles: |
| file_name, extension = os.path.splitext(file) |
| real_path = os.path.join(currentPath, file) |
| os.rename(real_path, file_name.replace(os.path.sep, "_") + extension) |
|
|
|
|
| def extract(zipfile_path, unzips_path): |
| try: |
| with zipfile.ZipFile(zipfile_path, "r") as zip_ref: |
| zip_ref.extractall(unzips_path) |
| os.remove(zipfile_path) |
| return True |
| except Exception as error: |
| print(f"An error occurred extracting the zip file: {error}") |
| return False |
|
|
|
|
| def unzip_file(zip_path, zip_file_name): |
| zip_file_path = os.path.join(zip_path, zip_file_name + ".zip") |
| extract_path = os.path.join(file_path, zip_file_name) |
| with zipfile.ZipFile(zip_file_path, "r") as zip_ref: |
| zip_ref.extractall(extract_path) |
| os.remove(zip_file_path) |
|
|
|
|
| def model_download_pipeline(url: str): |
| try: |
| result = download_from_url(url) |
| if result == "downloaded": |
| return handle_extraction_process() |
| else: |
| return "Error" |
| except Exception as error: |
| print(f"An unexpected error occurred: {error}") |
| return "Error" |
|
|
|
|
| def handle_extraction_process(): |
| extract_folder_path = "" |
| for filename in os.listdir(zips_path): |
| if filename.endswith(".zip"): |
| zipfile_path = os.path.join(zips_path, filename) |
| model_name = format_title(os.path.basename(zipfile_path).split(".zip")[0]) |
| extract_folder_path = os.path.join("logs", os.path.normpath(model_name)) |
| success = extract(zipfile_path, extract_folder_path) |
| clean_extracted_files(extract_folder_path, model_name) |
|
|
| if success: |
| print(f"Model {model_name} downloaded!") |
| else: |
| print(f"Error downloading {model_name}") |
| return "Error" |
| if not extract_folder_path: |
| print("Zip file was not found.") |
| return "Error" |
| return search_pth_index(extract_folder_path) |
|
|
|
|
| def clean_extracted_files(extract_folder_path, model_name): |
| macosx_path = os.path.join(extract_folder_path, "__MACOSX") |
| if os.path.exists(macosx_path): |
| shutil.rmtree(macosx_path) |
|
|
| subfolders = [ |
| f |
| for f in os.listdir(extract_folder_path) |
| if os.path.isdir(os.path.join(extract_folder_path, f)) |
| ] |
| if len(subfolders) == 1: |
| subfolder_path = os.path.join(extract_folder_path, subfolders[0]) |
| for item in os.listdir(subfolder_path): |
| shutil.move( |
| os.path.join(subfolder_path, item), |
| os.path.join(extract_folder_path, item), |
| ) |
| os.rmdir(subfolder_path) |
|
|
| for item in os.listdir(extract_folder_path): |
| source_path = os.path.join(extract_folder_path, item) |
| if ".pth" in item: |
| new_file_name = model_name + ".pth" |
| elif ".index" in item: |
| new_file_name = model_name + ".index" |
| else: |
| continue |
|
|
| destination_path = os.path.join(extract_folder_path, new_file_name) |
| if not os.path.exists(destination_path): |
| os.rename(source_path, destination_path) |
|
|
