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", interactive=True) save_frequency_slider = gr.Slider(minimum=10, maximum=100, value=50, step=10, label="Save Frequency", interactive=True) batch_size_slider = gr.Slider(minimum=1, maximum=20, value=8, step=1, label="Batch Size", interactive=True) train_button = gr.Button("Train Model", interactive=True) 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()