weights.gg RVC model installer

model_installer.py

import zipfile
import os
import os.path as path
import json
import time
import shutil
from colorama import Fore, Back, Style

###########################################################
# Installs RVC models from Weights.gg zip files basically #
###########################################################

# Configure paths here
class FINAL:
    pth_dir   = "D:/AI/RVC/assets/weights/"
    index_dir = "D:/AI/RVC/logs/"

# Scripting stuff OoOooOOOoOo
location = os.getcwd()
print(f"Looking for model zips inside {location}")
files = os.listdir(location)
for i in range(len(files)):
    filename = files[i]
    if not filename.endswith(".zip"):
        continue
    
    # Extracting
    zipfolder_path: str = filename.replace(".zip", "")
    with zipfile.ZipFile(filename, 'r') as zip_ref:
        zip_ref.extractall(path.join(location, zipfolder_path))

    # Getting metadata
    metadata: dict = None
    with open(path.join(zipfolder_path, "metadata.json"), 'r') as f:
        metadata = json.loads("\n".join(f.readlines()))

    # Finding index and stuff
    print()
    name = input(
            f"What prefix + name should this file have? "
            f"[{Back.BLACK}{Fore.YELLOW} {zipfolder_path} {Style.RESET_ALL}]\n"
            f"{Back.BLACK}{Fore.GREEN}-> {Style.RESET_ALL}"
        ) \
        .replace(".pth", "") \
        .replace(".index", "")
    author = metadata["author"]["name"]
    ogfilelist: list[str] = metadata["files"]
    class initial_paths:
        pth   = path.join(location, zipfolder_path, [elem["name"] for elem in ogfilelist if elem["name"].endswith(".pth")][0])
        index = path.join(location, zipfolder_path, [elem["name"] for elem in ogfilelist if elem["name"].endswith(".index")][0])
    
    # Installing the files
    final_index_dir = path.join(FINAL.index_dir, author)
    if not path.exists(final_index_dir):
        os.mkdir(final_index_dir)
    shutil.move(initial_paths.pth,   path.join(FINAL.pth_dir, f"{name}.pth"))
    shutil.move(initial_paths.index, path.join(final_index_dir, f"{name}.index"))
    
    # Cleanup
    shutil.rmtree(zipfolder_path, True, lambda: print(f"{Back.BLACK}{Fore.RED}Error cleaning up \"{zipfolder_path}\"{Style.RESET_ALL}"))
    os.remove(path.join(location, filename))

print(f"Finished installing all models! (hopefully)")
time.sleep(1.0)