Update hevrvc.py

hevrvc.py

@@ -7,6 +7,119 @@ import faiss
 from sklearn.cluster import MiniBatchKMeans
 import traceback
 
+import gradio as gr
+import os
+from random import shuffle
+import json
+import pathlib
+from subprocess import Popen, PIPE, STDOUT
+
+# Define the function for training
+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
@@ -107,22 +220,52 @@ def train_index(exp_dir1, version19):
     return "\n".join(infos)
 
 with gr.Blocks() as demo:
-    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("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()