run training as detached subprocess to survive Gradio session timeout

app.py

@@ -285,15 +285,21 @@ def gradio_main():
             lines.append(json.dumps(props, indent=2))
         return "\n".join(lines)
 
-    # -- Training --
+    # -- Training (runs as detached subprocess to survive Gradio session timeout) --
+    TRAIN_LOG = "/app/outputs/train.log"
+
     def train_lora(audio_files, lora_name, epochs, lr, rank,
                    progress=gr.Progress(track_tqdm=True)):
-        import shutil
-        import gc
+        import shutil, subprocess
 
         if not audio_files:
             return "No audio files uploaded."
 
+        if os.path.exists(TRAIN_LOG):
+            last_line = open(TRAIN_LOG).readlines()[-1] if os.path.getsize(TRAIN_LOG) > 0 else ""
+            if "DONE" not in last_line and "ERROR" not in last_line and last_line.strip():
+                return f"Training already in progress. Click 'Check Log' to monitor.\n\nLast: {last_line.strip()}"
+
         lora_name = (lora_name or "").strip() or "my-lora"
         epochs = max(1, min(int(epochs), 10))
         lr = float(lr)
@@ -301,152 +307,126 @@ def gradio_main():
 
         output_dir = os.path.join(ADAPTER_DIR, lora_name)
         os.makedirs(output_dir, exist_ok=True)
-
         audio_dir = os.path.join(output_dir, "audio_input")
         os.makedirs(audio_dir, exist_ok=True)
         for f in audio_files:
             src = f.name if hasattr(f, "name") else str(f)
             shutil.copy2(src, os.path.join(audio_dir, os.path.basename(src)))
 
-        log_lines = [
-            f"LoRA Training: '{lora_name}'",
-            f"Audio files: {len(audio_files)}",
-            f"Epochs: {epochs}, LR: {lr}, Rank: {rank}",
-            f"Output: {output_dir}",
-            "",
-        ]
-
-        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("[Step 0] Downloading model checkpoints...")
-                progress(0.02, desc="Downloading checkpoints...")
-                from huggingface_hub import snapshot_download
-                snapshot_download(
-                    ACE_HF_MODEL,
-                    local_dir=ACE_CHECKPOINT_DIR,
-                    ignore_patterns=["*.md", "*.txt", ".gitattributes"],
-                )
-                _log("  Checkpoints downloaded.")
-
-            if ACE_SOURCE_DIR not in sys.path:
-                sys.path.insert(0, ACE_SOURCE_DIR)
-
-            import torchaudio
-            _orig_load = torchaudio.load
-            def _load_soundfile(filepath, *args, **kwargs):
-                kwargs.setdefault('backend', 'soundfile')
-                return _orig_load(filepath, *args, **kwargs)
-            torchaudio.load = _load_soundfile
-
-            _log("[Step 1/2] Preprocessing audio files...")
-            progress(0.10, desc="Preprocessing audio...")
-
-            tensor_dir = os.path.join(output_dir, "preprocessed_tensors")
-            os.makedirs(tensor_dir, exist_ok=True)
-
-            from acestep.training_v2.preprocess import preprocess_audio_files
-            result = preprocess_audio_files(
-                audio_dir=audio_dir,
-                output_dir=tensor_dir,
-                checkpoint_dir=ACE_CHECKPOINT_DIR,
-                variant="turbo",
-                max_duration=60.0,
-                device="cpu",
-                precision="bfloat16",
-            )
-
-            processed = result.get("processed", 0)
-            total_files = result.get("total", 0)
-            failed = result.get("failed", 0)
-            _log(f"  Preprocessed: {processed}/{total_files} (failed: {failed})")
-
-            if processed == 0:
-                _log("ERROR: No files preprocessed successfully.")
-                return "\n".join(log_lines)
-
-            _log("[Step 2/2] Training LoRA adapter (CPU, this will be slow)...")
-            progress(0.30, desc="Loading model for training...")
+        train_script = f"""
+import os, sys, time, gc
+sys.path.insert(0, "{ACE_SOURCE_DIR}")
+os.environ["TORCHAUDIO_USE_BACKEND_DISPATCHER"] = "1"
+
+LOG = "{TRAIN_LOG}"
+def log(msg):
+    print(f"[train] {{msg}}", flush=True)
+    with open(LOG, "a") as f:
+        f.write(msg + "\\n")
+        f.flush()
+
+open(LOG, "w").close()
+log("LoRA Training: '{lora_name}' | files={len(audio_files)} | epochs={epochs} lr={lr} rank={rank}")
+
+import subprocess
+log("Stopping ace-server...")
+subprocess.run(["pkill", "-f", "ace-server"], stderr=subprocess.DEVNULL)
+time.sleep(2)
+gc.collect()
+
+try:
+    import torchaudio
+    _orig = torchaudio.load
+    def _sf(p, *a, **kw):
+        kw.setdefault("backend", "soundfile")
+        return _orig(p, *a, **kw)
+    torchaudio.load = _sf
+
+    log("[Step 1/2] Preprocessing audio...")
+    from acestep.training_v2.preprocess import preprocess_audio_files
+    result = preprocess_audio_files(
+        audio_dir="{audio_dir}",
+        output_dir="{output_dir}/preprocessed_tensors",
+        checkpoint_dir="{ACE_CHECKPOINT_DIR}",
+        variant="turbo", max_duration=60.0,
+        device="cpu", precision="bfloat16",
+    )
+    processed = result.get("processed", 0)
+    failed = result.get("failed", 0)
+    log(f"  Preprocessed: {{processed}}/{{result.get('total',0)}} (failed: {{failed}})")
+    if processed == 0:
+        log("ERROR: No files preprocessed. DONE")
+        raise SystemExit(1)
+
+    gc.collect()
+    log("[Step 2/2] Training LoRA...")
+    from acestep.training_v2.model_loader import load_decoder_for_training
+    from acestep.training_v2.trainer_fixed import FixedLoRATrainer
+    from acestep.training_v2.configs import TrainingConfigV2, LoRAConfigV2
+
+    model = load_decoder_for_training(
+        checkpoint_dir="{ACE_CHECKPOINT_DIR}", variant="turbo",
+        device="cpu", precision="bfloat16",
+    ).bfloat16()
+
+    trainer = FixedLoRATrainer(model,
+        LoRAConfigV2(r={rank}, alpha={rank}, dropout=0.0),
+        TrainingConfigV2(
+            checkpoint_dir="{ACE_CHECKPOINT_DIR}", model_variant="turbo",
+            dataset_dir="{output_dir}/preprocessed_tensors",
+            output_dir="{output_dir}",
+            max_epochs={epochs}, batch_size=1, learning_rate={lr},
+            device="cpu", precision="bfloat16", seed=42,
+            num_workers=0, pin_memory=False,
+        ))
+
+    step_count, last_loss = 0, 0.0
+    for update in trainer.train():
+        if hasattr(update, "step"):
+            step_count, last_loss = update.step, update.loss
+        elif isinstance(update, tuple) and len(update) >= 2:
+            step_count, last_loss = update[0], update[1]
+        if step_count % 5 == 0:
+            log(f"  Step {{step_count}}: loss={{last_loss:.4f}}")
+
+    log(f"Training complete! step={{step_count}} loss={{last_loss:.4f}}")
+    log(f"LoRA saved to: {output_dir}")
+    del model, trainer
+    gc.collect()
+    log("DONE")
+
+except Exception as e:
+    import traceback
+    log(f"ERROR: {{e}}")
+    log(traceback.format_exc())
+    log("DONE")
+finally:
+    log("Restarting ace-server...")
+    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)
+"""
+        script_path = os.path.join(output_dir, "_train.py")
+        with open(script_path, "w") as f:
+            f.write(train_script)
+
+        subprocess.Popen(
+            ["python3", script_path],
+            stdout=open("/dev/null", "w"),
+            stderr=open("/dev/null", "w"),
+            start_new_session=True,
+        )
 
-            from acestep.training_v2.model_loader import load_decoder_for_training
-            from acestep.training_v2.trainer_fixed import FixedLoRATrainer
-            from acestep.training_v2.configs import TrainingConfigV2, LoRAConfigV2
+        return (f"Training started in background for '{lora_name}'.\n"
+                f"Audio: {len(audio_files)} files, Epochs: {epochs}, Rank: {rank}\n\n"
+                f"Click 'Check Log' to monitor progress.\n"
+                f"Inference will be unavailable until training completes (ace-server stopped).")
 
-            model = load_decoder_for_training(
-                checkpoint_dir=ACE_CHECKPOINT_DIR,
-                variant="turbo",
-                device="cpu",
-                precision="bfloat16",
-            )
-            model = model.bfloat16()
-
-            adapter_cfg = LoRAConfigV2(r=rank, alpha=rank, dropout=0.0)
-            train_cfg = TrainingConfigV2(
-                checkpoint_dir=ACE_CHECKPOINT_DIR,
-                model_variant="turbo",
-                dataset_dir=tensor_dir,
-                output_dir=output_dir,
-                max_epochs=epochs,
-                batch_size=1,
-                learning_rate=lr,
-                device="cpu",
-                precision="bfloat16",
-                seed=42,
-                num_workers=0,
-                pin_memory=False,
-            )
-
-            trainer = FixedLoRATrainer(model, adapter_cfg, train_cfg)
-
-            step_count = 0
-            last_loss = 0.0
-            for update in trainer.train():
-                if hasattr(update, "step"):
-                    step_count = update.step
-                    last_loss = update.loss
-                elif isinstance(update, tuple) and len(update) >= 2:
-                    step_count = update[0]
-                    last_loss = update[1]
-                if step_count % 5 == 0:
-                    log_lines.append(f"  Step {step_count}: loss={last_loss:.4f}")
-                    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(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(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(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)
+    def check_train_log():
+        if not os.path.exists(TRAIN_LOG):
+            return "No training log found."
+        with open(TRAIN_LOG) as f:
+            return f.read() or "Log is empty."
 
     # -- Build UI --
     CSS = """
@@ -548,11 +528,13 @@ def gradio_main():
                         lr = gr.Number(label="Learning Rate", value=1e-4)
                         rank = gr.Number(label="Rank (r)", value=16, minimum=1, maximum=64)
 
-                train_btn = gr.Button("Train", variant="primary")
+                with gr.Row(elem_classes="compact-row"):
+                    train_btn = gr.Button("Train", variant="primary", scale=2)
+                    log_btn = gr.Button("Check Log", scale=1)
                 train_log = gr.Textbox(
                     label="Training Log",
                     interactive=False,
-                    lines=10,
+                    lines=12,
                     elem_classes="status-box",
                 )
 
@@ -562,6 +544,12 @@ def gradio_main():
                     outputs=[train_log],
                     api_name="train_lora",
                 )
+                log_btn.click(
+                    fn=check_train_log,
+                    inputs=[],
+                    outputs=[train_log],
+                    api_name="check_train_log",
+                )
 
         demo.launch(
             server_name="0.0.0.0",