rvc/lib/tools/model_download.py

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)