add LoRA download button after training (gr.File output, like rvc-beatrice)

app.py

@@ -460,12 +460,12 @@ def gradio_main():
         # -- Validation --
         if not audio_files:
             _log("[FAIL] No audio files uploaded.")
-            yield _log_text(), gr.update(visible=True), gr.update(visible=False)
+            yield _log_text(), gr.update(visible=True), gr.update(visible=False), gr.update()
             return
 
         if len(audio_files) > MAX_AUDIO_FILES:
             _log(f"[FAIL] Too many files ({len(audio_files)}). Max: {MAX_AUDIO_FILES}")
-            yield _log_text(), gr.update(visible=True), gr.update(visible=False)
+            yield _log_text(), gr.update(visible=True), gr.update(visible=False), gr.update()
             return
 
         lora_name = (lora_name or "").strip() or "my-lora"
@@ -485,7 +485,7 @@ def gradio_main():
 
         # Copy uploaded audio files
         _log(f"[INFO] Preparing {len(audio_files)} audio files...")
-        yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+        yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
 
         for f in audio_files:
             src = f.name if hasattr(f, "name") else str(f)
@@ -493,18 +493,18 @@ def gradio_main():
 
         _log(f"[INFO] LoRA: '{lora_name}' | Files: {len(audio_files)} | "
              f"Epochs: {epochs} | LR: {lr} | Rank: {rank}")
-        yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+        yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
 
         # Stop ace-server before training (frees memory)
         _log("[INFO] Stopping ace-server for training...")
-        yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+        yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
         _stop_ace_server()
         _gc.collect()
 
         try:
             # -- Phase 1: Preprocessing --
             _log("[Step 1/2] Preprocessing audio...")
-            yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+            yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
 
             preprocessed_dir = os.path.join(work_dir, "preprocessed_tensors")
 
@@ -521,24 +521,24 @@ def gradio_main():
                 progress_callback=preprocess_progress,
                 cancel_check=lambda: False,
             )
-            yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+            yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
 
             processed = result.get("processed", 0)
             failed = result.get("failed", 0)
             total = result.get("total", 0)
             _log(f"[OK] Preprocessed: {processed}/{total} (failed: {failed})")
-            yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+            yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
 
             if processed == 0:
                 _log("[FAIL] No files preprocessed successfully. Cannot train.")
-                yield _log_text(), gr.update(visible=True), gr.update(visible=False)
+                yield _log_text(), gr.update(visible=True), gr.update(visible=False), gr.update()
                 return
 
             _gc.collect()
 
             # -- Phase 2: Training --
             _log("[Step 2/2] Training LoRA...")
-            yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+            yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
 
             for msg in train_lora_generator(
                 dataset_dir=preprocessed_dir,
@@ -568,31 +568,35 @@ def gradio_main():
                     break
 
                 _log(msg)
-                yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+                yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
 
                 if msg.strip() == "[DONE]":
                     break
 
             _log(f"[INFO] Total time: {time.time() - train_start:.0f}s")
-            yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+            yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
 
         except Exception as exc:
             _log(f"[FAIL] Training error: {exc}")
             import traceback
             _log(traceback.format_exc())
-            yield _log_text(), gr.update(visible=True), gr.update(visible=False)
+            yield _log_text(), gr.update(visible=True), gr.update(visible=False), gr.update()
 
         finally:
             # Always restart ace-server
             _log("[INFO] Restarting ace-server...")
-            yield _log_text(), gr.update(visible=False), gr.update(visible=True)
+            yield _log_text(), gr.update(visible=False), gr.update(visible=True), gr.update()
             _gc.collect()
             ok = _start_ace_server()
             if ok:
                 _log("[OK] ace-server restarted successfully")
             else:
                 _log("[WARN] ace-server may not have restarted -- check logs")
-            yield _log_text(), gr.update(visible=True), gr.update(visible=False)
+            adapter_safetensors = os.path.join(adapter_out, "adapter_model.safetensors")
+            if os.path.isfile(adapter_safetensors):
+                yield _log_text(), gr.update(visible=True), gr.update(visible=False), gr.update(value=adapter_safetensors, visible=True)
+            else:
+                yield _log_text(), gr.update(visible=True), gr.update(visible=False), gr.update()
 
     # -- Cancel handler --
     def _on_cancel():
@@ -728,18 +732,19 @@ def gradio_main():
                     cancel_btn = gr.Button("Cancel Training", variant="stop", visible=False, scale=1)
                     log_btn = gr.Button("Check Log", scale=1)
 
+                train_output_file = gr.File(label="Trained LoRA (download)", visible=False)
                 train_log = gr.Textbox(
                     label="Training Log",
                     interactive=False,
-                    lines=12,
+                    lines=10,
                     elem_classes="status-box",
                 )
 
-                # Training generator -- yields (log, train_btn, cancel_btn)
+                # Training generator -- yields (log, train_btn, cancel_btn, output_file)
                 train_event = train_btn.click(
                     train_lora_ui,
                     inputs=[train_audio, lora_name, train_epochs, train_lr, train_rank],
-                    outputs=[train_log, train_btn, cancel_btn],
+                    outputs=[train_log, train_btn, cancel_btn, train_output_file],
                     api_name="train_lora",
                     concurrency_limit=1,
                 )