add LoRA training, fix css kwarg

Dockerfile

@@ -69,8 +69,13 @@ RUN curl -fL --retry 3 --retry-delay 5 -o /app/models/Qwen3-Embedding-0.6B-Q8_0.
 RUN curl -fL --retry 3 --retry-delay 5 -o /app/models/vae-BF16.gguf \
     "https://huggingface.co/Serveurperso/ACE-Step-1.5-GGUF/resolve/main/vae-BF16.gguf"
 
-# Install Python deps for Gradio UI
-RUN pip3 install --no-cache-dir gradio==5.29.0 requests
+# Install Python deps for Gradio UI + training
+RUN pip3 install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu \
+    gradio==5.29.0 requests torch safetensors \
+    transformers>=4.51.0 peft>=0.18.0 accelerate>=1.12.0
+
+# Clone ACE-Step repo for training module
+RUN git clone --depth 1 https://github.com/ace-step/ACE-Step-1.5 /app/ace-step-source
 
 # Copy application files
 COPY app.py /app/app.py

app.py

@@ -12,6 +12,11 @@ ACE_SERVER = os.environ.get("ACE_SERVER", "http://127.0.0.1:8085")
 OUTPUT_DIR = os.environ.get("ACE_OUTPUT_DIR", "/app/outputs")
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+ACE_CHECKPOINT_DIR = os.environ.get("ACE_CHECKPOINT_DIR", "/app/checkpoints")
+ACE_SOURCE_DIR = "/app/ace-step-source"
+ACE_HF_MODEL = "ACE-Step/Ace-Step1.5"
+ADAPTER_DIR = os.environ.get("ACE_ADAPTER_DIR", "/app/adapters")
+
 # ---------------------------------------------------------------------------
 # ace-server helpers
 # ---------------------------------------------------------------------------
@@ -280,14 +285,143 @@ def gradio_main():
             lines.append(json.dumps(props, indent=2))
         return "\n".join(lines)
 
-    # -- Training placeholder --
-    def train_lora_placeholder(*args):
-        return ("Training requires PyTorch and the ACE-Step Python package.\n\n"
-                "To enable training, install dependencies:\n"
-                "  pip install torch torchaudio safetensors transformers "
-                "diffusers peft accelerate einops\n\n"
-                "Then restart the app. Training is not available on the "
-                "CPU-only HF Space — use a local GPU machine or a GPU Space.")
+    # -- Training --
+    def train_lora(audio_files, lora_name, epochs, lr, rank,
+                   progress=gr.Progress(track_tqdm=True)):
+        import shutil
+        import gc
+
+        if not audio_files:
+            return "No audio files uploaded."
+
+        lora_name = (lora_name or "").strip() or "my-lora"
+        epochs = max(1, min(int(epochs), 10))
+        lr = float(lr)
+        rank = max(1, min(int(rank), 64))
+
+        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}",
+            "",
+        ]
+
+        try:
+            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...")
+                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_lines.append("  Checkpoints downloaded.")
+
+            if ACE_SOURCE_DIR not in sys.path:
+                sys.path.insert(0, ACE_SOURCE_DIR)
+
+            log_lines.append("[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="float32",
+            )
+
+            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})")
+
+            if processed == 0:
+                log_lines.append("ERROR: No files preprocessed successfully.")
+                return "\n".join(log_lines)
+
+            log_lines.append("[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
+            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="float32",
+            )
+            model = model.float()
+
+            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="float32",
+                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_lines.append(f"Training complete! Final: step {step_count}, loss={last_loss:.4f}")
+            log_lines.append(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}")
+            import traceback
+            log_lines.append(traceback.format_exc())
+        except Exception as e:
+            import traceback
+            log_lines.append(f"ERROR: {e}")
+            log_lines.append(traceback.format_exc())
+
+        return "\n".join(log_lines)
 
     # -- Build UI --
     CSS = """
@@ -295,7 +429,7 @@ def gradio_main():
     .status-box textarea { font-family: monospace; font-size: 13px; }
     """
 
-    with gr.Blocks(title="ACE-Step 1.5 XL (CPU)") as demo:
+    with gr.Blocks(title="ACE-Step 1.5 XL (CPU)", css=CSS) as demo:
 
         with gr.Tabs():
             # ============================================================
@@ -373,7 +507,7 @@ def gradio_main():
                 gr.Markdown(
                     "### LoRA Training\n"
                     "Fine-tune ACE-Step on your own audio data. "
-                    "Requires PyTorch + GPU (not available on CPU Spaces)."
+                    "CPU training is very slow. Checkpoints downloaded on first run (~10GB)."
                 )
 
                 with gr.Row(elem_classes="compact-row"):
@@ -385,9 +519,9 @@ def gradio_main():
                         )
                     with gr.Column(scale=1):
                         lora_name = gr.Textbox(label="LoRA Name", value="my-lora")
-                        epochs = gr.Number(label="Epochs", value=100, minimum=1, maximum=10000)
+                        epochs = gr.Number(label="Epochs", value=5, minimum=1, maximum=10)
                         lr = gr.Number(label="Learning Rate", value=1e-4)
-                        rank = gr.Number(label="Rank (r)", value=16, minimum=1, maximum=256)
+                        rank = gr.Number(label="Rank (r)", value=16, minimum=1, maximum=64)
 
                 train_btn = gr.Button("Train", variant="primary")
                 train_log = gr.Textbox(
@@ -398,7 +532,7 @@ def gradio_main():
                 )
 
                 train_btn.click(
-                    fn=train_lora_placeholder,
+                    fn=train_lora,
                     inputs=[train_audio, lora_name, epochs, lr, rank],
                     outputs=[train_log],
                     api_name="train_lora",
@@ -408,7 +542,6 @@ def gradio_main():
             server_name="0.0.0.0",
             server_port=7860,
             mcp_server=True,
-            css=CSS,
         )