stop ace-server during training to free RAM, restart after, add log visibility

app.py

@@ -316,10 +316,20 @@ def gradio_main():
             "",
         ]
 
+        def _log(msg):
+            log_lines.append(msg)
+            print(f"[train] {msg}", flush=True)
+
         try:
+            import subprocess, signal
+            _log("Stopping ace-server to free RAM for training...")
+            subprocess.run(["pkill", "-f", "ace-server"], stderr=subprocess.DEVNULL)
+            time.sleep(2)
+            gc.collect()
+
             ckpt_files = os.listdir(ACE_CHECKPOINT_DIR) if os.path.isdir(ACE_CHECKPOINT_DIR) else []
             if len(ckpt_files) < 3:
-                log_lines.append("[Step 0] Downloading model checkpoints...")
+                _log("[Step 0] Downloading model checkpoints...")
                 progress(0.02, desc="Downloading checkpoints...")
                 from huggingface_hub import snapshot_download
                 snapshot_download(
@@ -327,7 +337,7 @@ def gradio_main():
                     local_dir=ACE_CHECKPOINT_DIR,
                     ignore_patterns=["*.md", "*.txt", ".gitattributes"],
                 )
-                log_lines.append("  Checkpoints downloaded.")
+                _log("  Checkpoints downloaded.")
 
             if ACE_SOURCE_DIR not in sys.path:
                 sys.path.insert(0, ACE_SOURCE_DIR)
@@ -339,7 +349,7 @@ def gradio_main():
                 return _orig_load(filepath, *args, **kwargs)
             torchaudio.load = _load_soundfile
 
-            log_lines.append("[Step 1/2] Preprocessing audio files...")
+            _log("[Step 1/2] Preprocessing audio files...")
             progress(0.10, desc="Preprocessing audio...")
 
             tensor_dir = os.path.join(output_dir, "preprocessed_tensors")
@@ -359,13 +369,13 @@ def gradio_main():
             processed = result.get("processed", 0)
             total_files = result.get("total", 0)
             failed = result.get("failed", 0)
-            log_lines.append(f"  Preprocessed: {processed}/{total_files} (failed: {failed})")
+            _log(f"  Preprocessed: {processed}/{total_files} (failed: {failed})")
 
             if processed == 0:
-                log_lines.append("ERROR: No files preprocessed successfully.")
+                _log("ERROR: No files preprocessed successfully.")
                 return "\n".join(log_lines)
 
-            log_lines.append("[Step 2/2] Training LoRA adapter (CPU, this will be slow)...")
+            _log("[Step 2/2] Training LoRA adapter (CPU, this will be slow)...")
             progress(0.30, desc="Loading model for training...")
 
             from acestep.training_v2.model_loader import load_decoder_for_training
@@ -412,21 +422,29 @@ def gradio_main():
                     pct = 0.30 + 0.65 * min(step_count / max(epochs * processed, 1), 1.0)
                     progress(pct, desc=f"Step {step_count}, loss={last_loss:.4f}")
 
-            log_lines.append(f"Training complete! Final: step {step_count}, loss={last_loss:.4f}")
-            log_lines.append(f"LoRA saved to: {output_dir}")
+            _log(f"Training complete! Final: step {step_count}, loss={last_loss:.4f}")
+            _log(f"LoRA saved to: {output_dir}")
 
             del model, trainer
             gc.collect()
 
         except ImportError as e:
-            log_lines.append(f"Import error: {e}")
-            log_lines.append(f"Check ACE-Step source at {ACE_SOURCE_DIR}")
+            _log(f"Import error: {e}")
+            _log(f"Check ACE-Step source at {ACE_SOURCE_DIR}")
             import traceback
             log_lines.append(traceback.format_exc())
         except Exception as e:
             import traceback
-            log_lines.append(f"ERROR: {e}")
+            _log(f"ERROR: {e}")
             log_lines.append(traceback.format_exc())
+        finally:
+            _log("Restarting ace-server...")
+            import subprocess
+            subprocess.Popen([
+                "/app/ace-server", "--host", "127.0.0.1", "--port", "8085",
+                "--models", "/app/models", "--adapters", "/app/adapters",
+                "--max-batch", "1",
+            ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
 
         return "\n".join(log_lines)