demo.py

import os
import shutil
from os import listdir
from colorama import Fore
import os
import shutil
import numpy as np
import faiss
from pathlib import Path
from sklearn.cluster import MiniBatchKMeans
import traceback
import gradio as gr
import pathlib
import json
from random import shuffle
from subprocess import Popen, PIPE, STDOUT

# Function to preprocess data
def preprocess_data(model_name, dataset_folder):
    logs_path = f'/content/RVC/logs/{model_name}'
    temp_DG_path = '/content/temp_DG'

    if os.path.exists(logs_path):
        print("Model already exists, This will be resume training.")
        os.makedirs(temp_DG_path, exist_ok=True)

        # Move files for resuming training
        for item in os.listdir(logs_path):
            item_path = os.path.join(logs_path, item)
            if os.path.isfile(item_path) and (item.startswith('D_') or item.startswith('G_')) and item.endswith('.pth'):
                shutil.copy(item_path, temp_DG_path)

        for item in os.listdir(logs_path):
            item_path = os.path.join(logs_path, item)
            if os.path.isfile(item_path):
                os.remove(item_path)
            elif os.path.isdir(item_path):
                shutil.rmtree(item_path)

        for file_name in os.listdir(temp_DG_path):
            shutil.move(os.path.join(temp_DG_path, file_name), logs_path)

        shutil.rmtree(temp_DG_path)

    if len(os.listdir(dataset_folder)) < 1:
        return "Error: Dataset folder is empty."

    os.makedirs(f'./logs/{model_name}', exist_ok=True)
    !python infer/modules/train/preprocess.py {dataset_folder} 32000 2 ./logs/{model_name} False 3.0

    with open(f'./logs/{model_name}/preprocess.log', 'r') as f:
        log_content = f.read()

    if 'end preprocess' in log_content:
        return "Success: Data preprocessing complete."
    else:
        return "Error preprocessing data. Check your dataset folder."

# Function to extract F0 feature
def extract_f0_feature(model_name, f0method):
    if f0method != "rmvpe_gpu":
        !python infer/modules/train/extract/extract_f0_print.py ./logs/{model_name} 2 {f0method}
    else:
        !python infer/modules/train/extract/extract_f0_rmvpe.py 1 0 0 ./logs/{model_name} True

    with open(f'./logs/{model_name}/extract_f0_feature.log', 'r') as f:
        log_content = f.read()

    if 'all-feature-done' in log_content:
        return "Success: F0 feature extraction complete."
    else:
        return "Error extracting F0 feature."

# Function to train index
def train_index(exp_dir1, version19):
    exp_dir = f"logs/{exp_dir1}"
    os.makedirs(exp_dir, exist_ok=True)
    feature_dir = f"{exp_dir}/3_feature768" if version19 == "v2" else f"{exp_dir}/3_feature256"

    if not os.path.exists(feature_dir) or len(os.listdir(feature_dir)) == 0:
        return "Please run feature extraction first."

    npys = [np.load(f"{feature_dir}/{name}") for name in sorted(os.listdir(feature_dir))]
    big_npy = np.concatenate(npys, axis=0)
    big_npy_idx = np.arange(big_npy.shape[0])
    np.random.shuffle(big_npy_idx)
    big_npy = big_npy[big_npy_idx]

    if big_npy.shape[0] > 2e5:
        big_npy = MiniBatchKMeans(n_clusters=10000, batch_size=256, init="random").fit(big_npy).cluster_centers_

    np.save(f"{exp_dir}/total_fea.npy", big_npy)
    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)

    index = faiss.index_factory(768 if version19 == "v2" else 256, f"IVF{n_ivf},Flat")
    index.train(big_npy)
    faiss.write_index(index, f"{exp_dir}/trained_IVF{n_ivf}_Flat_nprobe_1_{exp_dir1}_{version19}.index")

    batch_size_add = 8192
    for i in range(0, big_npy.shape[0], batch_size_add):
        index.add(big_npy[i:i + batch_size_add])

    faiss.write_index(index, f"{exp_dir}/added_IVF{n_ivf}_Flat_nprobe_1_{exp_dir1}_{version19}.index")
    return f"Indexing completed. Index saved to {exp_dir}/added_IVF{n_ivf}_Flat_nprobe_1_{exp_dir1}_{version19}.index"





now_dir = os.getcwd()

def click_train(exp_dir1, sr2, if_f0_3, spk_id5, save_epoch10, total_epoch11, batch_size12, 
                if_save_latest13, pretrained_G14, pretrained_D15, gpus16, if_cache_gpu17, 
                if_save_every_weights18, version19):
    exp_dir = "%s/logs/%s" % (now_dir, exp_dir1)
    os.makedirs(exp_dir, exist_ok=True)
    gt_wavs_dir = "%s/0_gt_wavs" % (exp_dir)
    feature_dir = (
        "%s/3_feature256" % (exp_dir)
        if version19 == "v1"
        else "%s/3_feature768" % (exp_dir)
    )
    
    if if_f0_3:
        f0_dir = "%s/2a_f0" % (exp_dir)
        f0nsf_dir = "%s/2b-f0nsf" % (exp_dir)
        names = (
            set([name.split(".")[0] for name in os.listdir(gt_wavs_dir)])
            & set([name.split(".")[0] for name in os.listdir(feature_dir)])
            & set([name.split(".")[0] for name in os.listdir(f0_dir)])
            & set([name.split(".")[0] for name in os.listdir(f0nsf_dir)])
        )
    else:
        names = set([name.split(".")[0] for name in os.listdir(gt_wavs_dir)]) & set(
            [name.split(".")[0] for name in os.listdir(feature_dir)]
        )
    
    opt = []
    for name in names:
        if if_f0_3:
            opt.append(
                "%s/%s.wav|%s/%s.npy|%s/%s.wav.npy|%s/%s.wav.npy|%s"
                % (
                    gt_wavs_dir.replace("\\", "\\\\"),
                    name,
                    feature_dir.replace("\\", "\\\\"),
                    name,
                    f0_dir.replace("\\", "\\\\"),
                    name,
                    f0nsf_dir.replace("\\", "\\\\"),
                    name,
                    spk_id5,
                )
            )
        else:
            opt.append(
                "%s/%s.wav|%s/%s.npy|%s"
                % (
                    gt_wavs_dir.replace("\\", "\\\\"),
                    name,
                    feature_dir.replace("\\", "\\\\"),
                    name,
                    spk_id5,
                )
            )
    
    fea_dim = 256 if version19 == "v1" else 768
    if if_f0_3:
        for _ in range(2):
            opt.append(
                "%s/logs/mute/0_gt_wavs/mute%s.wav|%s/logs/mute/3_feature%s/mute.npy|%s/logs/mute/2a_f0/mute.wav.npy|%s/logs/mute/2b-f0nsf/mute.wav.npy|%s"
                % (now_dir, sr2, now_dir, fea_dim, now_dir, now_dir, spk_id5)
            )
    else:
        for _ in range(2):
            opt.append(
                "%s/logs/mute/0_gt_wavs/mute%s.wav|%s/logs/mute/3_feature%s/mute.npy|%s"
                % (now_dir, sr2, now_dir, fea_dim, spk_id5)
            )
    
    shuffle(opt)
    with open("%s/filelist.txt" % exp_dir, "w") as f:
        f.write("\n".join(opt))
    
    print("Filelist generated")
    print("Using gpus:", gpus16)
    
    if pretrained_G14 == "":
        print("No pretrained Generator")
    if pretrained_D15 == "":
        print("No pretrained Discriminator")
    
    if version19 == "v1" or sr2 == "40k":
        config_path = "configs/v1/%s.json" % sr2
    else:
        config_path = "configs/v2/%s.json" % sr2
    config_save_path = os.path.join(exp_dir, "config.json")
    if not pathlib.Path(config_save_path).exists():
        with open(config_save_path, "w", encoding="utf-8") as f:
            with open(config_path, "r") as config_file:
                config_data = json.load(config_file)
                json.dump(
                    config_data,
                    f,
                    ensure_ascii=False,
                    indent=4,
                    sort_keys=True,
                )
    
    cmd = (
        'python infer/modules/train/train.py -e "%s" -sr %s -f0 %s -bs %s -g %s -te %s -se %s %s %s -l %s -c %s -sw %s -v %s'
        % (
            exp_dir1,
            sr2,
            1 if if_f0_3 else 0,
            batch_size12,
            gpus16,
            total_epoch11,
            save_epoch10,
            "-pg %s" % pretrained_G14 if pretrained_G14 != "" else "",
            "-pd %s" % pretrained_D15 if pretrained_D15 != "" else "",
            1 if if_save_latest13 == True else 0,
            1 if if_cache_gpu17 == True else 0,
            1 if if_save_every_weights18 == True else 0,
            version19,
        )
    )
    
    # Capture output
    p = Popen(cmd, shell=True, cwd=now_dir, stdout=PIPE, stderr=STDOUT, bufsize=1, universal_newlines=True)
    
    # Print output
    output_log = ""
    for line in p.stdout:
        print(line.strip())
        output_log += line.strip() + "\n"
    
    p.wait()
    return output_log


def launch_training(model_name, epochs, save_frequency, batch_size):
    sample_rate = '32k'
    OV2 = True
    G_file = f'assets/pretrained_v2/f0Ov2Super{sample_rate}G.pth' if OV2 else f'assets/pretrained_v2/f0G{sample_rate}.pth'
    D_file = f'assets/pretrained_v2/f0Ov2Super{sample_rate}D.pth' if OV2 else f'assets/pretrained_v2/f0D{sample_rate}.pth'
    
    # Call the training function
    training_log = click_train(
        model_name, 
        sample_rate, 
        True, 0, save_frequency, 
        epochs, batch_size, True, 
        G_file, D_file, 0, False, 
        True, 'v2'
    )
    
    return training_log




def run_inference(model_name, pitch, input_path, f0_method, save_as, index_rate, volume_normalization, consonant_protection):
    # Setting paths for model and index files
    model_filename = model_name + '.pth'
    index_temp = 'Index_Temp'
    
    # Ensure Index_Temp exists
    if not os.path.exists(index_temp):
        os.mkdir(index_temp)
        print("Index_Temp folder created.")
    else:
        print("Index_Temp folder found.")
    
    # Copy .index file to Index_Temp
    index_file_path = os.path.join('logs/', model_name, '')
    for file_name in listdir(index_file_path):
        if file_name.startswith('added') and file_name.endswith('.index'):
            shutil.copy(index_file_path + file_name, os.path.join(index_temp, file_name))
            print('Index file copied successfully.')
    
    # Get the .index file
    indexfile_directory = os.getcwd() + '/' + index_temp
    files = os.listdir(indexfile_directory)
    index_filename = files[0] if files else None
    if index_filename is None:
        raise ValueError("Index file not found.")
    
    shutil.rmtree(index_temp)
    
    model_path = "assets/weights/" + model_filename
    index_path = os.path.join('logs', model_name, index_filename)
    
    if not os.path.exists(input_path):
        raise ValueError(f"{input_path} was not found.")
    
    os.environ['index_root'] = os.path.dirname(index_path)
    index_path = os.path.basename(index_path)
    
    os.environ['weight_root'] = os.path.dirname(model_path)
    
    # Run the command
    cmd = f"python tools/cmd/infer_cli.py --f0up_key {pitch} --input_path {input_path} --index_path {index_path} --f0method {f0_method} --opt_path {save_as} --model_name {os.path.basename(model_path)} --index_rate {index_rate} --device 'cuda:0' --is_half True --filter_radius 3 --resample_sr 0 --rms_mix_rate {volume_normalization} --protect {consonant_protection}"
    os.system(f"rm -f {save_as}")
    os.system(cmd)
    
    return f"Inference completed, output saved at {save_as}.", save_as




# Gradio Interface

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# RVC V2 - EASY GUI")
    with gr.Row():
        with gr.Tab("Inference"):
            with gr.Row():
                model_name = gr.Textbox(label="Model Name For Inference")
            with gr.Row():
                input_path = gr.Audio(label="Input Audio Path", type="filepath")
            with gr.Row():
                with gr.Accordion("Inference Settings"):
                    pitch = gr.Slider(minimum=-12, maximum=12, step=1, label="Pitch", value=0)
                    f0_method = gr.Dropdown(choices=["rmvpe", "pm", "harvest"], label="f0 Method", value="rmvpe")
                    index_rate = gr.Slider(minimum=0, maximum=1, step=0.01, label="Index Rate", value=0.5)
                    volume_normalization = gr.Slider(minimum=0, maximum=1, step=0.01, label="Volume Normalization", value=0)
                    consonant_protection = gr.Slider(minimum=0, maximum=1, step=0.01, label="Consonant Protection", value=0.5)
            with gr.Row():
                save_as = gr.Textbox(value="/content/RVC/audios/output_audio.wav", label="Output Audio Path")
                
            run_btn = gr.Button("Run Inference")
            with gr.Row():
                output_message = gr.Textbox(label="Output Message",interactive=False)
                output_audio = gr.Audio(label="Output Audio",interactive=False)
            #run_btn.click(run_inference, [model_name, pitch, input_path, f0_method, save_as, index_rate, volume_normalization, consonant_protection], output_message)

        with gr.Tab("Training"):
            with gr.TabItem("Create Index and stuff"):
                model_name = gr.Textbox(label="Model Name (No spaces or symbols)")
                dataset_folder = gr.Textbox(label="Dataset Folder", value="/content/dataset")
                f0method = gr.Dropdown(["pm", "harvest", "rmvpe", "rmvpe_gpu"], label="F0 Method", value="rmvpe_gpu")
                preprocess_btn = gr.Button("Start Preprocessing")
                f0_btn = gr.Button("Extract F0 Feature")
                train_btn = gr.Button("Train Index")
                preprocess_output = gr.Textbox(label="Preprocessing Log")
                f0_output = gr.Textbox(label="F0 Feature Extraction Log")
                train_output = gr.Textbox(label="Training Log")

                #preprocess_btn.click(preprocess_data, inputs=[model_name, dataset_folder], outputs=preprocess_output)
                #f0_btn.click(extract_f0_feature, inputs=[model_name, f0method], outputs=f0_output)
                #train_btn.click(train_index, inputs=[model_name, "v2"], outputs=train_output)
            with gr.TabItem("Train Your Model"):
                model_name_input = gr.Textbox(label="Model Name", placeholder="Enter the model name", interactive=True)
                epochs_slider = gr.Slider(minimum=50, maximum=2000, value=200, step=10, label="Epochs")
                save_frequency_slider = gr.Slider(minimum=10, maximum=100, value=50, step=10, label="Save Frequency")
                batch_size_slider = gr.Slider(minimum=1, maximum=20, value=8, step=1, label="Batch Size")
    
                train_button = gr.Button("Train Model")
                training_output = gr.Textbox(label="Training Log", interactive=False)
    
                #train_button.click(launch_training, inputs=[model_name_input, epochs_slider, save_frequency_slider, batch_size_slider], outputs=training_output)
    
demo.launch()