Update demo.py

demo.py

@@ -173,8 +173,23 @@ 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)
     
 demo.launch()