Update app.py

app.py

@@ -19,7 +19,7 @@ def upload_audio(audio_files, model_name):
     return f"✅ {len(audio_files)} audio files saved to '{save_dir}' folder!"
 
 # ---- Start training function ----
-def train_rvc(model_name, sample_rate, epochs):
+def train_rvc(model_name, sample_rate, epochs, batch_size):
     dataset_dir = os.path.join("dataset", model_name)
     if not os.path.exists(dataset_dir):
         return f"⚠️ {dataset_dir} folder not found. Upload audio files first."
@@ -30,31 +30,34 @@ def train_rvc(model_name, sample_rate, epochs):
             "--model_name", model_name,
             "--dataset", dataset_dir,
             "--sample_rate", str(sample_rate),
-            "--epochs", str(epochs)
+            "--epochs", str(epochs),
+            "--batch_size", str(batch_size)
         ]
         subprocess.Popen(cmd)
-        return f"🚀 Training started! Model name: {model_name}, Epochs: {epochs}, Sample rate: {sample_rate}"
+        return f"🚀 RVC v2 Training started!\nModel: {model_name}\nEpochs: {epochs}\nBatch: {batch_size}\nSample Rate: {sample_rate}Hz\n\nCheck console for progress!"
     except Exception as e:
         return f"❌ Training failed to start: {e}"
 
 # ---- Interface ----
-with gr.Blocks() as demo:
-    gr.Markdown("# 🎙️ RVC Voice Model Training Tool")
+with gr.Blocks(title="RVC v2 Training") as demo:
+    gr.Markdown("# 🎙️ RVC v2 Voice Model Training Tool")
     gr.Markdown("1️⃣ Upload audio → 2️⃣ Enter model name → 3️⃣ Start training")
 
     with gr.Tab("Upload Audio"):
         with gr.Row():
             audio_files = gr.File(file_count="multiple", label="🎧 Upload Audio Files (.wav)")
             model_name = gr.Textbox(label="Model Name", placeholder="e.g.: zeynep_rvc")
-        output_upload = gr.Textbox(label="Status")
-        upload_button = gr.Button("📦 Upload Audio")
+        output_upload = gr.Textbox(label="Status", lines=3)
+        upload_button = gr.Button("📦 Upload Audio", variant="primary")
         upload_button.click(upload_audio, inputs=[audio_files, model_name], outputs=output_upload)
 
     with gr.Tab("Start Training"):
         sample_rate = gr.Dropdown(choices=[32000, 40000, 48000], value=40000, label="Sample Rate (Hz)")
-        epochs = gr.Slider(50, 1000, value=750, step=50, label="Number of Epochs")
-        output_train = gr.Textbox(label="Status")
-        train_button = gr.Button("🚀 Start Training")
-        train_button.click(train_rvc, inputs=[model_name, sample_rate, epochs], outputs=output_train)
+        epochs = gr.Slider(50, 1000, value=200, step=50, label="Number of Epochs")
+        batch_size = gr.Slider(4, 16, value=8, step=4, label="Batch Size")
+        output_train = gr.Textbox(label="Training Status", lines=5)
+        train_button = gr.Button("🚀 Start RVC v2 Training", variant="primary")
+        train_button.click(train_rvc, inputs=[model_name, sample_rate, epochs, batch_size], outputs=output_train)
+
 
 demo.launch()