|
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() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
