Update demo.py

demo.py

@@ -10,6 +10,10 @@ 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):
@@ -102,6 +106,160 @@ def train_index(exp_dir1, version19):
 
 
 
+
+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'
@@ -148,6 +306,9 @@ def run_inference(model_name, pitch, input_path, f0_method, save_as, index_rate,
     
     return f"Inference completed, output saved at {save_as}.", save_as
 
+
+
+
 # Gradio Interface
 
 with gr.Blocks() as demo:
@@ -173,22 +334,32 @@ with gr.Blocks() as demo:
             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)
+            #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"):
-            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("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()