| import os |
| import gradio as gr |
| from pytube import YouTube |
| from pydub import AudioSegment |
| import numpy as np |
| import faiss |
| from sklearn.cluster import MiniBatchKMeans |
| import traceback |
| from random import shuffle |
| import json |
| import pathlib |
| from subprocess import Popen, PIPE, STDOUT |
|
|
| |
| 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, |
| ): |
| now_dir = os.getcwd() |
| exp_dir = f"{now_dir}/logs/{exp_dir1}" |
| os.makedirs(exp_dir, exist_ok=True) |
| gt_wavs_dir = f"{exp_dir}/0_gt_wavs" |
| feature_dir = ( |
| f"{exp_dir}/3_feature256" if version19 == "v1" else f"{exp_dir}/3_feature768" |
| ) |
| |
| if if_f0_3: |
| f0_dir = f"{exp_dir}/2a_f0" |
| f0nsf_dir = f"{exp_dir}/2b-f0nsf" |
| 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( |
| f"{gt_wavs_dir.replace('\\', '\\\\')}/{name}.wav|{feature_dir.replace('\\', '\\\\')}/{name}.npy|{f0_dir.replace('\\', '\\\\')}/{name}.wav.npy|{f0nsf_dir.replace('\\', '\\\\')}/{name}.wav.npy|{spk_id5}" |
| ) |
| else: |
| opt.append( |
| f"{gt_wavs_dir.replace('\\', '\\\\')}/{name}.wav|{feature_dir.replace('\\', '\\\\')}/{name}.npy|{spk_id5}" |
| ) |
| |
| fea_dim = 256 if version19 == "v1" else 768 |
| if if_f0_3: |
| for _ in range(2): |
| opt.append( |
| f"{now_dir}/logs/mute/0_gt_wavs/mute{sr2}.wav|{now_dir}/logs/mute/3_feature{fea_dim}/mute.npy|{now_dir}/logs/mute/2a_f0/mute.wav.npy|{now_dir}/logs/mute/2b-f0nsf/mute.wav.npy|{spk_id5}" |
| ) |
| else: |
| for _ in range(2): |
| opt.append( |
| f"{now_dir}/logs/mute/0_gt_wavs/mute{sr2}.wav|{now_dir}/logs/mute/3_feature{fea_dim}/mute.npy|{spk_id5}" |
| ) |
| |
| shuffle(opt) |
| with open(f"{exp_dir}/filelist.txt", "w") as f: |
| f.write("\n".join(opt)) |
|
|
| print("Write filelist done") |
| print("Use gpus:", str(gpus16)) |
| if pretrained_G14 == "": |
| print("No pretrained Generator") |
| if pretrained_D15 == "": |
| print("No pretrained Discriminator") |
| |
| if version19 == "v1" or sr2 == "40k": |
| config_path = f"configs/v1/{sr2}.json" |
| else: |
| config_path = f"configs/v2/{sr2}.json" |
| |
| 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, |
| ) |
| f.write("\n") |
|
|
| cmd = ( |
| f'python infer/modules/train/train.py -e "{exp_dir1}" -sr {sr2} -f0 {1 if if_f0_3 else 0} -bs {batch_size12} -g {gpus16} -te {total_epoch11} -se {save_epoch10} {"-pg " + pretrained_G14 if pretrained_G14 != "" else ""} {"-pd " + pretrained_D15 if pretrained_D15 != "" else ""} -l {1 if if_save_latest13 else 0} -c {1 if if_cache_gpu17 else 0} -sw {1 if if_save_every_weights18 else 0} -v {version19}' |
| ) |
|
|
| p = Popen(cmd, shell=True, cwd=now_dir, stdout=PIPE, stderr=STDOUT, bufsize=1, universal_newlines=True) |
| |
| for line in p.stdout: |
| print(line.strip()) |
| |
| p.wait() |
| return "After the training is completed, you can view the console training log or train.log under the experiment folder" |
|
|
|
|
|
|
| def calculate_audio_duration(file_path): |
| duration_seconds = len(AudioSegment.from_file(file_path)) / 1000.0 |
| return duration_seconds |
|
|
| def youtube_to_wav(url, dataset_folder): |
| try: |
| yt = YouTube(url).streams.get_audio_only().download(output_path=dataset_folder) |
| mp4_path = os.path.join(dataset_folder, 'audio.mp4') |
| wav_path = os.path.join(dataset_folder, 'audio.wav') |
| os.rename(yt, mp4_path) |
| os.system(f'ffmpeg -i {mp4_path} -acodec pcm_s16le -ar 44100 {wav_path}') |
| os.remove(mp4_path) |
| return f'Audio downloaded and converted to WAV: {wav_path}' |
| except Exception as e: |
| return f"Error: {e}" |
|
|
| def create_training_files(model_name, dataset_folder, youtube_link): |
| if youtube_link: |
| youtube_to_wav(youtube_link, dataset_folder) |
|
|
| if not os.listdir(dataset_folder): |
| return "Your dataset folder is empty." |
|
|
| os.makedirs(f'./logs/{model_name}', exist_ok=True) |
| |
| os.system(f'python infer/modules/train/preprocess.py {dataset_folder} 32000 2 ./logs/{model_name} False 3.0 > /dev/null 2>&1') |
| |
| with open(f'./logs/{model_name}/preprocess.log', 'r') as f: |
| if 'end preprocess' in f.read(): |
| return "Preprocessing Success" |
| else: |
| return "Error preprocessing data... Make sure your dataset folder is correct." |
|
|
| def extract_features(model_name, f0method): |
| os.system(f'python infer/modules/train/extract/extract_f0_rmvpe.py 1 0 0 ./logs/{model_name} True' if f0method == "rmvpe_gpu" else |
| f'python infer/modules/train/extract/extract_f0_print.py ./logs/{model_name} 2 {f0method}') |
| os.system(f'python infer/modules/train/extract_feature_print.py cuda:0 1 0 ./logs/{model_name} v2 True') |
| |
| with open(f'./logs/{model_name}/extract_f0_feature.log', 'r') as f: |
| if 'all-feature-done' in f.read(): |
| return "Feature Extraction Success" |
| else: |
| return "Error in feature extraction... Make sure your data was preprocessed." |
|
|
| 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_feature256" if version19 == "v1" else f"{exp_dir}/3_feature768" |
| if not os.path.exists(feature_dir): |
| return "Please perform feature extraction first!" |
|
|
| listdir_res = list(os.listdir(feature_dir)) |
| if len(listdir_res) == 0: |
| return "Please perform feature extraction first!" |
|
|
| infos = [] |
| npys = [] |
| for name in sorted(listdir_res): |
| phone = np.load(f"{feature_dir}/{name}") |
| npys.append(phone) |
| big_npy = np.concatenate(npys, 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: |
| infos.append(f"Trying k-means with {big_npy.shape[0]} to 10k centers.") |
| try: |
| big_npy = MiniBatchKMeans( |
| n_clusters=10000, |
| verbose=True, |
| batch_size=256, |
| compute_labels=False, |
| init="random", |
| ).fit(big_npy).cluster_centers_ |
| except: |
| info = traceback.format_exc() |
| infos.append(info) |
| return "\n".join(infos) |
|
|
| 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) |
| infos.append(f"{big_npy.shape},{n_ivf}") |
|
|
| index = faiss.index_factory(256 if version19 == "v1" else 768, f"IVF{n_ivf},Flat") |
| infos.append("Training index") |
| index_ivf = faiss.extract_index_ivf(index) |
| index_ivf.nprobe = 1 |
| index.train(big_npy) |
| faiss.write_index(index, f"{exp_dir}/trained_IVF{n_ivf}_Flat_nprobe_{index_ivf.nprobe}_{exp_dir1}_{version19}.index") |
|
|
| infos.append("Adding to 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_{index_ivf.nprobe}_{exp_dir1}_{version19}.index") |
|
|
| infos.append(f"Successfully built index: added_IVF{n_ivf}_Flat_nprobe_{index_ivf.nprobe}_{exp_dir1}_{version19}.index") |
| return "\n".join(infos) |
|
|
| with gr.Blocks() as demo: |
| with gr.Tab("Training"): |
| with gr.Tab("CREATE TRANING FILES - This will process the data, extract the features and create your index file for you!"): |
| with gr.Row(): |
| model_name = gr.Textbox(label="Model Name", value="My-Voice") |
| dataset_folder = gr.Textbox(label="Dataset Folder", value="/content/dataset") |
| youtube_link = gr.Textbox(label="YouTube Link (optional)") |
| with gr.Row(): |
| start_button = gr.Button("Create Training Files") |
| f0method = gr.Dropdown(["pm", "harvest", "rmvpe", "rmvpe_gpu"], label="F0 Method", value="rmvpe_gpu") |
| extract_button = gr.Button("Extract Features") |
| train_button = gr.Button("Train Index") |
|
|
| output = gr.Textbox(label="Output") |
|
|
| start_button.click(create_training_files, inputs=[model_name, dataset_folder, youtube_link], outputs=output) |
| extract_button.click(extract_features, inputs=[model_name, f0method], outputs=output) |
| train_button.click(train_index, inputs=[model_name, "v2"], outputs=output) |
| with gr.Tab("train"): |
| exp_dir1 = gr.Textbox(label="Experiment Directory", value="mymodel") |
| sr2 = gr.Dropdown(choices=["32k", "40k", "48k"], label="Sample Rate", value="32k") |
| if_f0_3 = gr.Checkbox(label="Use F0", value=True) |
| spk_id5 = gr.Number(label="Speaker ID", value=0) |
| save_epoch10 = gr.Slider(label="Save Frequency", minimum=5, maximum=50, step=5, value=25) |
| total_epoch11 = gr.Slider(label="Total Epochs", minimum=10, maximum=2000, step=10, value=500) |
| batch_size12 = gr.Slider(label="Batch Size", minimum=1, maximum=20, step=1, value=8) |
| if_save_latest13 = gr.Checkbox(label="Save Latest", value=True) |
| pretrained_G14 = gr.Textbox(label="Pretrained Generator File", value="/content/pre/assets/pretrained_v2/f0Ov2Super32kG.pth") |
| pretrained_D15 = gr.Textbox(label="Pretrained Discriminator File", value="/content/pre/assets/pretrained_v2/f0Ov2Super32kD.pth") |
| gpus16 = gr.Number(label="GPUs", value=0) |
| if_cache_gpu17 = gr.Checkbox(label="Cache GPU", value=False) |
| if_save_every_weights18 = gr.Checkbox(label="Save Every Weights", value=True) |
| version19 = gr.Textbox(label="Version", value="v2") |
| training_log = gr.Textbox(label="Training Log", interactive=False) |
| train_button = gr.Button("Start Training") |
| |
| train_button.click( |
| fn=click_train, |
| inputs=[ |
| 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 |
| ], |
| outputs=training_log |
| ) |
|
|
|
|
|
|
| demo.launch() |
|
|
|
|
|
|
|
|
|
|
|
|
| |
|
|
