tabs/train/train.py

import os
import shutil
import sys

import gradio as gr

from assets.i18n.i18n import I18nAuto

i18n = I18nAuto()
now_dir = os.getcwd()
sys.path.append(now_dir)

# Custom Pretraineds
pretraineds_custom_path = os.path.join(now_dir, "rvc", "models", "pretraineds", "pretraineds_custom")

pretraineds_custom_path_relative = os.path.relpath(pretraineds_custom_path, now_dir)

custom_embedder_root = os.path.join(now_dir, "rvc", "models", "embedders", "embedders_custom")

os.makedirs(custom_embedder_root, exist_ok=True)
os.makedirs(pretraineds_custom_path_relative, exist_ok=True)


def get_pretrained_list(suffix):
    return [
        os.path.join(dirpath, filename)
        for dirpath, _, filenames in os.walk(pretraineds_custom_path_relative)
        for filename in filenames
        if filename.endswith(".pth") and suffix in filename
    ]


# Dataset Creator
datasets_path = os.path.join(now_dir, "assets", "datasets")

if not os.path.exists(datasets_path):
    os.makedirs(datasets_path)


# Drop Model
def save_drop_model(dropbox):
    if ".pth" not in dropbox:
        gr.Info(i18n("The file you dropped is not a valid pretrained file. Please try again."))
    else:
        file_name = os.path.basename(dropbox)
        pretrained_path = os.path.join(pretraineds_custom_path_relative, file_name)
        if os.path.exists(pretrained_path):
            os.remove(pretrained_path)
        shutil.copy(dropbox, pretrained_path)
        gr.Info(i18n("Click the refresh button to see the pretrained file in the dropdown menu."))
    return None