"""ACE-Step 1.5 XL (CPU) - Gradio frontend + CLI for ace-server GGUF inference""" import os import sys import time import json import argparse import base64 import tempfile import subprocess import shutil import string import random import requests import logging import threading from train_engine import ( preprocess_audio, train_lora_generator, cancel_training, _training_cancel, get_trained_loras as _get_trained_loras_engine, MAX_TRAINING_TIME, ) logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Configurable limits (edit here, not buried in code) # --------------------------------------------------------------------------- MAX_TOTAL_AUDIO = 1800 # seconds total across all uploaded files (30 min) # MAX_TRAINING_TIME is imported from train_engine (single source of truth) MAX_AUDIO_FILES = 50 # max number of training audio files per run # --------------------------------------------------------------------------- # Paths & constants # --------------------------------------------------------------------------- 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) # Clean up old inference temp files (older than 1 hour) at startup _CLEANUP_MAX_AGE = 3600 # seconds try: _now = time.time() for _fname in os.listdir(OUTPUT_DIR): if _fname.lower().endswith((".wav", ".mp3")): _fpath = os.path.join(OUTPUT_DIR, _fname) try: if os.path.isfile(_fpath) and (_now - os.path.getmtime(_fpath)) > _CLEANUP_MAX_AGE: os.remove(_fpath) except OSError: pass except Exception: pass 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") MODELS_DIR = os.environ.get("ACE_MODELS_DIR", "/app/models") ACE_SERVER_BIN = "/app/ace-server" # Detect if running on HF Space (ace-server available) vs locally (PyTorch only) _is_space = os.path.isfile(ACE_SERVER_BIN) or os.environ.get("SPACE_ID") is not None _training_lock = threading.Lock() # HF repo for on-demand GGUF downloads GGUF_HF_REPO = "Serveurperso/ACE-Step-1.5-GGUF" # --------------------------------------------------------------------------- # ace-server helpers # --------------------------------------------------------------------------- def _server_ok(): try: return requests.get(f"{ACE_SERVER}/health", timeout=5).status_code == 200 except Exception: return False def _get_props(): """Fetch server properties (models, adapters).""" try: r = requests.get(f"{ACE_SERVER}/props", timeout=10) if r.status_code == 200: return r.json() except Exception: pass return {} def _poll_job(job_id, timeout=600, progress_cb=None, cancel_check=None): """Poll a job until done/error/timeout/cancelled. Returns (status, elapsed).""" t0 = time.time() while time.time() - t0 < timeout: if cancel_check and cancel_check(): return "cancelled", time.time() - t0 try: r = requests.get(f"{ACE_SERVER}/job", params={"id": job_id}, timeout=5) data = r.json() status = data.get("status", "unknown") if progress_cb: progress_cb(status, data) if status in ("done", "error"): return status, time.time() - t0 except Exception: pass time.sleep(1) return "timeout", time.time() - t0 def _fetch_result(job_id, timeout=60): """Fetch result bytes/json for a completed job.""" r = requests.get( f"{ACE_SERVER}/job", params={"id": job_id, "result": 1}, timeout=timeout, ) return r def _caption_via_understand(audio_path, timeout=600, cancel_check=None): """Call ace-server /understand for a rich caption. Returns dict or None.""" fname = os.path.basename(audio_path) try: with open(audio_path, "rb") as f: r = requests.post( f"{ACE_SERVER}/understand", files={"audio": (fname, f, "audio/mpeg")}, timeout=30, ) if r.status_code != 200: logger.warning("[Caption] %s: /understand %d: %s", fname, r.status_code, r.text[:200]) return None job_id = r.json().get("id") if not job_id: return None except Exception as exc: logger.warning("[Caption] %s: /understand submit failed: %s", fname, exc) return None status, elapsed = _poll_job(job_id, timeout=timeout, cancel_check=cancel_check) if status != "done": logger.warning("[Caption] %s: /understand -> %s (%.0fs)", fname, status, elapsed) return None try: r = _fetch_result(job_id, timeout=30) if r.status_code != 200: return None data = r.json() if isinstance(data, dict) and data.get("caption"): return data return None except Exception: return None def _run_pipeline(caption, lyrics, bpm, duration, seed, steps, output_format, adapter=None, lm_model=None, progress_cb=None): """Run full LM -> synth pipeline. Returns (audio_path, status_msg) or raises.""" t0 = time.time() # -- Build LM request -- req = {"caption": caption or "upbeat electronic dance music"} req["lyrics"] = lyrics if lyrics and lyrics.strip() else "[Instrumental]" try: if bpm and int(float(bpm)) > 0: req["bpm"] = int(float(bpm)) if duration and float(duration) > 0: req["duration"] = min(float(duration), 300) if seed is not None and int(float(seed)) >= 0: req["seed"] = int(float(seed)) if steps and int(float(steps)) > 0: req["inference_steps"] = int(float(steps)) except (ValueError, TypeError): pass if adapter: req["adapter"] = adapter if lm_model: req["model"] = lm_model fmt = output_format if output_format in ("wav", "mp3") else "mp3" synth_fmt = "wav16" if fmt == "wav" else "mp3" suffix = f".{fmt}" # -- LM phase -- if progress_cb: progress_cb("lm_submit", None) r = requests.post(f"{ACE_SERVER}/lm", json=req, timeout=30) if r.status_code != 200: raise RuntimeError(f"LM submit failed: {r.status_code} {r.text}") lm_job_id = r.json().get("id") if progress_cb: progress_cb("lm_poll", {"job_id": lm_job_id}) lm_status, lm_elapsed = _poll_job(lm_job_id, timeout=900) if lm_status != "done": raise RuntimeError(f"LM {lm_status} after {lm_elapsed:.0f}s") # Fetch LM result r = _fetch_result(lm_job_id) lm_results = r.json() if not isinstance(lm_results, list) or len(lm_results) == 0: raise RuntimeError(f"LM returned no results: {lm_results}") synth_request = lm_results[0] # -- Synth phase -- synth_request["output_format"] = synth_fmt if adapter: synth_request["adapter"] = adapter synth_request["synth_model"] = "acestep-v15-turbo-Q4_K_M.gguf" if progress_cb: progress_cb("synth_submit", None) r = requests.post(f"{ACE_SERVER}/synth", json=synth_request, timeout=30) if r.status_code != 200: raise RuntimeError(f"Synth submit failed: {r.status_code} {r.text}") synth_job_id = r.json().get("id") if progress_cb: progress_cb("synth_poll", {"job_id": synth_job_id}) synth_status, synth_elapsed = _poll_job(synth_job_id, timeout=600) if synth_status != "done": raise RuntimeError(f"Synth {synth_status} after {synth_elapsed:.0f}s") # Fetch audio if progress_cb: progress_cb("fetch", None) r = _fetch_result(synth_job_id, timeout=60) if r.status_code != 200: raise RuntimeError(f"Audio fetch failed: {r.status_code}") tmp = tempfile.NamedTemporaryFile(suffix=suffix, dir=OUTPUT_DIR, delete=False) tmp.write(r.content) tmp.close() elapsed = time.time() - t0 msg = f"Done in {elapsed:.0f}s | {duration}s audio, {steps} steps, {fmt}" return tmp.name, msg # --------------------------------------------------------------------------- # LM model scanning & on-demand download # --------------------------------------------------------------------------- DEFAULT_LM = "acestep-5Hz-lm-1.7B-Q8_0.gguf" AVAILABLE_LM_MODELS = [ "acestep-5Hz-lm-1.7B-Q8_0.gguf", "acestep-5Hz-lm-0.6B-Q8_0.gguf", "acestep-5Hz-lm-4B-Q5_K_M.gguf", ] def _scan_lm_models(): """Return LM model choices. Installed shown as-is, others need download.""" installed = set() if os.path.isdir(MODELS_DIR): for f in os.listdir(MODELS_DIR): if "-lm-" in f and f.endswith(".gguf"): installed.add(f) choices = [] for m in AVAILABLE_LM_MODELS: if m in installed: choices.append(m) else: choices.append(f"{m} [not installed]") return choices def _download_lm_model(filename): """Download a GGUF LM model from HF if not already present.""" dest = os.path.join(MODELS_DIR, filename) if os.path.isfile(dest): return dest try: from huggingface_hub import hf_hub_download path = hf_hub_download( repo_id=GGUF_HF_REPO, filename=filename, local_dir=MODELS_DIR, ) return path except Exception as exc: logger.error("Failed to download %s: %s", filename, exc) return None # --------------------------------------------------------------------------- # LoRA listing for UI dropdowns # --------------------------------------------------------------------------- def _list_lora_choices(): """Return list of LoRA choices for dropdown, including 'None'.""" choices = ["None (no LoRA)"] if os.path.isdir(ADAPTER_DIR): for d in os.listdir(ADAPTER_DIR): if os.path.isdir(os.path.join(ADAPTER_DIR, d)): choices.append(d) return choices # --------------------------------------------------------------------------- # ace-server stop/start helpers # --------------------------------------------------------------------------- _ace_proc = None def _stop_ace_server(): """Stop ace-server process.""" global _ace_proc logger.info("[ace-server] Stopping...") if _ace_proc and _ace_proc.poll() is None: _ace_proc.terminate() try: _ace_proc.wait(timeout=10) except subprocess.TimeoutExpired: _ace_proc.kill() _ace_proc = None logger.info("[ace-server] Stopped (tracked PID)") else: try: subprocess.run(["pkill", "ace-server"], stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL, timeout=10) logger.info("[ace-server] Stopped (pkill)") except Exception: pass time.sleep(1) def _start_ace_server(max_retries: int = 3, retry_delay: float = 5.0): """Start ace-server in background and wait for health. Retries up to max_retries times with retry_delay seconds between attempts. """ global _ace_proc for attempt in range(1, max_retries + 1): logger.info( "[ace-server] Starting (attempt %d/%d) with --adapters %s", attempt, max_retries, ADAPTER_DIR, ) try: _ace_proc = subprocess.Popen( [ACE_SERVER_BIN, "--host", "127.0.0.1", "--port", "8085", "--models", MODELS_DIR, "--adapters", ADAPTER_DIR, "--max-batch", "1"], ) except Exception as exc: logger.error("[ace-server] Failed to start: %s", exc) if attempt < max_retries: time.sleep(retry_delay) continue return False for _ in range(30): if _server_ok(): logger.info("[ace-server] Healthy") return True time.sleep(2) logger.warning("[ace-server] Health check timeout on attempt %d/%d", attempt, max_retries) # Kill the failed process before retrying if _ace_proc and _ace_proc.poll() is None: _ace_proc.kill() try: _ace_proc.wait(timeout=5) except subprocess.TimeoutExpired: pass if attempt < max_retries: time.sleep(retry_delay) logger.error("[ace-server] Failed to start after %d attempts", max_retries) return False # --------------------------------------------------------------------------- # CLI mode # --------------------------------------------------------------------------- def cli_main(): parser = argparse.ArgumentParser( description="ACE-Step 1.5 XL (CPU) - CLI inference via ace-server", ) parser.add_argument("caption", nargs="?", default="upbeat electronic dance music", help="Music description / caption") parser.add_argument("--lyrics", "-l", default="[Instrumental]", help="Lyrics text (use '[Instrumental]' for no vocals)") parser.add_argument("--bpm", type=int, default=120, help="Beats per minute") parser.add_argument("--duration", "-d", type=float, default=10, help="Duration in seconds (max 300)") parser.add_argument("--steps", "-s", type=int, default=8, help="Inference steps (1-32)") parser.add_argument("--seed", type=int, default=-1, help="Random seed (-1 for random)") parser.add_argument("--format", "-f", choices=["wav", "mp3"], default="wav", help="Output audio format") parser.add_argument("--adapter", "-a", default=None, help="LoRA adapter name") parser.add_argument("-o", "--output", default=None, help="Output file path (default: auto in outputs dir)") parser.add_argument("--server", default=None, help="ace-server URL (default: http://127.0.0.1:8085)") args = parser.parse_args() if args.server: global ACE_SERVER ACE_SERVER = args.server if not _server_ok(): print(f"ERROR: ace-server not reachable at {ACE_SERVER}", file=sys.stderr) sys.exit(1) seed = args.seed if args.seed >= 0 else None def cli_progress(phase, data): phases = { "lm_submit": "Submitting LM job...", "lm_poll": f"LM generating (job {data['job_id']})..." if data else "LM generating...", "synth_submit": "Submitting synth job...", "synth_poll": f"Synthesizing (job {data['job_id']})..." if data else "Synthesizing...", "fetch": "Fetching audio...", } msg = phases.get(phase, phase) print(f" [{phase}] {msg}") print(f"ACE-Step CLI | caption: {args.caption}") print(f" lyrics: {args.lyrics} | bpm: {args.bpm} | duration: {args.duration}s " f"| steps: {args.steps} | seed: {args.seed} | format: {args.format}") try: audio_path, status = _run_pipeline( caption=args.caption, lyrics=args.lyrics, bpm=args.bpm, duration=args.duration, seed=seed, steps=args.steps, output_format=args.format, adapter=args.adapter, progress_cb=cli_progress, ) except RuntimeError as e: print(f"ERROR: {e}", file=sys.stderr) sys.exit(1) # Move to requested output path if specified if args.output: out_dir = os.path.dirname(os.path.abspath(args.output)) os.makedirs(out_dir, exist_ok=True) shutil.move(audio_path, args.output) audio_path = args.output print(f" {status}") print(f" Output: {audio_path}") # --------------------------------------------------------------------------- # Gradio UI mode # --------------------------------------------------------------------------- def gradio_main(): import gradio as gr import gc # -- Persistent training log buffer (survives across yields) -- _train_log_lines = [] # -- Generate tab handler -- def generate_music(caption, lyrics, instrumental, bpm, duration, seed, steps, lora_select, lm_model_select, progress=gr.Progress(track_tqdm=True)): if not _training_lock.acquire(blocking=False): return None, "Training in progress. Inference unavailable until training completes. Press Cancel to stop training." _training_lock.release() if not _server_ok(): return None, "ace-server not running. Check logs." if instrumental or not lyrics or lyrics.strip() == "": lyrics = "[Instrumental]" actual_seed = None if seed is None or int(seed) < 0 else int(seed) adapter = None if lora_select == "None (no LoRA)" else lora_select lm_model_file = lm_model_select.replace(" [not installed]", "") if lm_model_select else None if lm_model_file and "[not installed]" in (lm_model_select or ""): _download_lm_model(lm_model_file) lm_model = lm_model_file progress_map = { "lm_submit": (0.05, "Submitting LM job..."), "lm_poll": (0.10, "LM generating..."), "synth_submit": (0.40, "Submitting synth job..."), "synth_poll": (0.50, "Synthesizing audio..."), "fetch": (0.90, "Fetching audio..."), } def gr_progress(phase, data): pct, desc = progress_map.get(phase, (0.5, phase)) if data and "job_id" in data: desc += f" (job {data['job_id']})" progress(pct, desc=desc) try: audio_path, status = _run_pipeline( caption=caption, lyrics=lyrics, bpm=bpm, duration=duration, seed=actual_seed, steps=steps, output_format="mp3", adapter=adapter, lm_model=lm_model, progress_cb=gr_progress, ) return audio_path, status except RuntimeError as e: return None, str(e) except Exception as e: return None, f"Unexpected error: {e}" # -- Server info helper -- def get_server_status(): if not _server_ok(): return "ace-server: OFFLINE" props = _get_props() lines = ["ace-server: ONLINE"] if props: lines.append(json.dumps(props, indent=2)) return "\n".join(lines) # -- Training generator (direct integration, no subprocess) -- def train_lora_ui(audio_files, lora_name, epochs, lr, rank): """Generator that yields (train_log, train_btn_update, cancel_btn_update).""" import gc as _gc _train_log_lines.clear() train_start = time.time() def _log(msg): elapsed = int(time.time() - train_start) m, s = divmod(elapsed, 60) h, m = divmod(m, 60) ts = f"+{h}:{m:02d}:{s:02d}" if h else f"+{m:02d}:{s:02d}" line = f"[{ts}] {msg}" _train_log_lines.append(line) logger.info(msg) if len(_train_log_lines) > 2000: _train_log_lines[:] = _train_log_lines[-1000:] def _log_text(): return "\n".join(_train_log_lines) # -- Validation -- if not audio_files: _log("[FAIL] No audio files uploaded.") yield _log_text(), gr.Button(visible=True), gr.Button(visible=False), gr.File() 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.Button(visible=True), gr.Button(visible=False), gr.File() return lora_name = (lora_name or "").strip() or "my-lora" lora_name = "".join(c if c.isalnum() or c in "-_" else "-" for c in lora_name) epochs = max(1, min(int(epochs), 1000)) lr = float(lr) rank = max(1, min(int(rank), 128)) work_dir = os.path.join(OUTPUT_DIR, "train_workspace", lora_name) os.makedirs(work_dir, exist_ok=True) audio_dir = os.path.join(work_dir, "audio_input") if os.path.exists(audio_dir): shutil.rmtree(audio_dir) os.makedirs(audio_dir) adapter_out = os.path.join(ADAPTER_DIR, lora_name) os.makedirs(adapter_out, exist_ok=True) # Copy uploaded audio files + check total duration _log(f"[INFO] Preparing {len(audio_files)} audio files...") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() import librosa as _lr total_dur = 0.0 accepted = 0 skipped_names = [] truncated_names = [] for f in audio_files: src = f.name if hasattr(f, "name") else str(f) fname = os.path.basename(src) # .txt/.json sidecars: copy as caption files, skip duration check if fname.lower().endswith((".txt", ".json")): shutil.copy2(src, os.path.join(audio_dir, fname)) continue try: dur = _lr.get_duration(path=src) except Exception: dur = 0.0 if dur <= 0: skipped_names.append(f"{fname} (invalid/empty)") continue remaining = MAX_TOTAL_AUDIO - total_dur if remaining <= 0: skipped_names.append(fname) continue if dur > remaining: # Truncate this file to fit import soundfile as _sf y, sr = _lr.load(src, sr=None, mono=False) max_samples = int(remaining * sr) if y.ndim == 1: y = y[:max_samples] else: y = y[:, :max_samples] dst = os.path.join(audio_dir, fname) _sf.write(dst, y.T if y.ndim > 1 else y, sr) truncated_names.append(f"{fname} ({dur:.0f}s -> {remaining:.0f}s)") total_dur += remaining accepted += 1 else: shutil.copy2(src, os.path.join(audio_dir, fname)) total_dur += dur accepted += 1 if truncated_names: _log(f"[WARN] Truncated: {', '.join(truncated_names)}") if skipped_names: _log(f"[WARN] Skipped (over {MAX_TOTAL_AUDIO/60:.0f} min cap): {', '.join(skipped_names)}") _log(f"[INFO] Total audio: {total_dur:.0f}s ({total_dur/60:.1f} min), {accepted} files") _log(f"[INFO] LoRA: '{lora_name}' | Files: {len(audio_files)} | " f"Epochs: {epochs} | LR: {lr} | Rank: {rank}") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() # Caption audio files: GGUF LM if ace-server running, else librosa use_understand = _server_ok() method = "GGUF LM (BPM, key, mood, lyrics)" if use_understand else "librosa (BPM only)" _log(f"[INFO] Auto-captioning via {method}...") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() for audio_fname in sorted(os.listdir(audio_dir)): if _training_cancel.is_set(): break full_path = os.path.join(audio_dir, audio_fname) if not os.path.isfile(full_path): continue ext = audio_fname.lower().rsplit(".", 1)[-1] if "." in audio_fname else "" if ext in ("json", "txt"): continue stem = audio_fname.rsplit(".", 1)[0] if "." in audio_fname else audio_fname sidecar_json = os.path.join(audio_dir, stem + ".json") sidecar_txt = os.path.join(audio_dir, stem + ".txt") if os.path.isfile(sidecar_json) or os.path.isfile(sidecar_txt): _log(f" {audio_fname}: using caption file") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() continue caption_data = None if use_understand: _log(f" {audio_fname}: GGUF LM captioning...") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() caption_data = _caption_via_understand( full_path, timeout=600, cancel_check=lambda: _training_cancel.is_set(), ) if not caption_data: use_understand = False _log(f" {audio_fname}: GGUF LM too slow, skipping (preprocessing will analyze)") if caption_data: bpm_s = caption_data.get("bpm", "?") key_s = caption_data.get("keyscale", caption_data.get("key", "?")) _log(f" {audio_fname}: OK (BPM={bpm_s}, key={key_s})") with open(sidecar_json, "w") as cj: json.dump(caption_data, cj) else: _log(f" {audio_fname}: will be analyzed in preprocessing (BPM + key + caption)") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() if _training_cancel.is_set(): _training_cancel.clear() _log("[CANCELLED] Stopped") yield _log_text(), gr.Button(visible=True), gr.Button(visible=False), gr.File() shutil.rmtree(work_dir, ignore_errors=True) return # Stop ace-server before training (frees memory) _training_lock.acquire() _log("[INFO] Stopping ace-server for training...") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() _stop_ace_server() _gc.collect() _cleanup_done = False try: # -- Phase 1: Preprocessing (runs in thread for live progress) -- preprocessed_dir = os.path.join(work_dir, "preprocessed_tensors") _preprocess_log_len = len(_train_log_lines) def preprocess_progress(current, total, desc): _log(f" {desc} ({current}/{total})") _preprocess_result = [None] _preprocess_error = [None] def _run_preprocess(): try: _preprocess_result[0] = preprocess_audio( audio_dir=audio_dir, output_dir=preprocessed_dir, checkpoint_dir=ACE_CHECKPOINT_DIR, device="cpu", variant="turbo", max_duration=float(MAX_TOTAL_AUDIO), progress_callback=preprocess_progress, cancel_check=lambda: _training_cancel.is_set(), ) except Exception as exc: _preprocess_error[0] = exc _log("[Step 1/2] Encoding audio → training data (VAE + text encoder)...") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() t = threading.Thread(target=_run_preprocess, daemon=True) t.start() while t.is_alive(): t.join(timeout=3) if len(_train_log_lines) > _preprocess_log_len: _preprocess_log_len = len(_train_log_lines) yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() if _preprocess_error[0]: raise _preprocess_error[0] result = _preprocess_result[0] if _training_cancel.is_set(): _training_cancel.clear() _log("[CANCELLED] Stopped during preprocessing") yield _log_text(), gr.Button(visible=True), gr.Button(visible=False), gr.File() return processed = result.get("processed", 0) failed = result.get("failed", 0) total = result.get("total", 0) _log(f"[OK] Preprocessed: {processed}/{total} files (failed: {failed})") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() if processed == 0: _log("[FAIL] No files preprocessed successfully. Cannot train.") yield _log_text(), gr.Button(visible=True), gr.Button(visible=False), gr.File() return _gc.collect() # -- Phase 2: Training (random 60s crops for speed + augmentation) -- _log("[Step 2/2] Training LoRA...") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() for msg in train_lora_generator( dataset_dir=preprocessed_dir, output_dir=adapter_out, checkpoint_dir=ACE_CHECKPOINT_DIR, epochs=epochs, lr=lr, rank=rank, alpha=rank * 2, dropout=0.0, batch_size=1, gradient_accumulation_steps=4, warmup_steps=100, weight_decay=0.01, max_grad_norm=1.0, save_every_n_epochs=0, seed=42, variant="turbo", device="cpu", chunk_duration=60, log_every=5, ): elapsed = time.time() - train_start if elapsed > MAX_TRAINING_TIME: _log(f"[WARN] Training timed out after {int(elapsed)}s") cancel_training() break _log(msg) yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() if msg.strip() == "[DONE]": break _log(f"[INFO] Total time: {time.time() - train_start:.0f}s") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() except GeneratorExit: _training_cancel.set() logger.info("Generator closed by Gradio, cleaning up") _cleanup_done = True _training_lock.release() _gc.collect() _start_ace_server() shutil.rmtree(work_dir, ignore_errors=True) return except Exception as exc: _log(f"[FAIL] Training error: {exc}") import traceback _log(traceback.format_exc()) yield _log_text(), gr.Button(visible=True), gr.Button(visible=False), gr.File() finally: if not _cleanup_done: _training_lock.release() _log("[INFO] Restarting ace-server...") yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File() _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") adapter_safetensors = os.path.join(adapter_out, "adapter_model.safetensors") if os.path.isfile(adapter_safetensors): tmp_out = tempfile.NamedTemporaryFile( suffix=".safetensors", prefix=f"{lora_name}_", delete=False, ) tmp_out.close() shutil.copy2(adapter_safetensors, tmp_out.name) _log(f"[OK] LoRA saved: {lora_name}") yield _log_text(), gr.Button(visible=True), gr.Button(visible=False), gr.File(value=tmp_out.name, visible=True) else: yield _log_text(), gr.Button(visible=True), gr.Button(visible=False), gr.File() shutil.rmtree(work_dir, ignore_errors=True) # -- Cancel handler -- def _on_cancel(): cancel_training() logger.info("Cancel requested by user") return "Cancelling..." # -- Build LM model choices -- def _lm_model_choices(): return _scan_lm_models() # -- Build UI -- CSS = """ .compact-row { gap: 8px !important; } .status-box textarea { font-family: monospace; font-size: 13px; } """ with gr.Blocks(title="ACE-Step 1.5 XL (CPU)") as demo: with gr.Tabs(): # ============================================================ # Tab 1: Generate Music # ============================================================ with gr.Tab("Generate Music"): gr.Markdown("**[ACE-Step 1.5 XL](https://github.com/ace-step/ACE-Step-1.5)** GGUF Q4_K_M via [acestep.cpp](https://github.com/ServeurpersoCom/acestep.cpp)") with gr.Row(elem_classes="compact-row"): with gr.Column(scale=3): audio_out = gr.Audio(label="Output", type="filepath") status = gr.Textbox( label="Status", interactive=False, lines=2, elem_classes="status-box", ) caption = gr.Textbox( label="Music Description", lines=2, value="upbeat electronic dance music, energetic synth leads", ) lyrics = gr.Textbox( label="Lyrics", lines=3, value="[Instrumental]", placeholder="Enter lyrics or [Instrumental] for no vocals", ) with gr.Column(scale=2): gen_btn = gr.Button("Generate Music", variant="primary") with gr.Row(elem_classes="compact-row"): instrumental = gr.Checkbox(label="Instrumental", value=True) bpm = gr.Number(label="BPM", value=120, minimum=0, maximum=300) with gr.Row(elem_classes="compact-row"): duration = gr.Slider( label="Duration (s)", minimum=10, maximum=120, value=10, step=5, ) steps = gr.Slider( label="Steps", minimum=1, maximum=32, value=8, step=1, ) seed = gr.Number(label="Seed (-1=random)", value=-1) lora_select = gr.Dropdown( label="LoRA", choices=_list_lora_choices(), value="None (no LoRA)", allow_custom_value=True, ) lm_model_select = gr.Dropdown( label="LM Model", choices=_lm_model_choices(), value=DEFAULT_LM, ) gen_btn.click( fn=generate_music, inputs=[caption, lyrics, instrumental, bpm, duration, seed, steps, lora_select, lm_model_select], outputs=[audio_out, status], api_name="generate", ) # ============================================================ # Tab 2: Train LoRA # ============================================================ with gr.Tab("Train LoRA"): gr.Markdown("LoRA training ported from [Side-Step](https://github.com/koda-dernet/Side-Step) | Model: [ACE-Step 1.5](https://github.com/ace-step/ACE-Step-1.5)") with gr.Row(elem_classes="compact-row"): with gr.Column(scale=3): train_log = gr.Textbox( label="Training Log", interactive=False, lines=12, elem_classes="status-box", ) train_output_file = gr.File(label="Trained LoRA (download)", visible=False) train_audio = gr.File( label="Training Audio (optional caption .txt)", file_count="multiple", file_types=["audio", ".txt", ".json"], height=120, ) with gr.Column(scale=2): with gr.Row(elem_classes="compact-row"): train_btn = gr.Button("Train", variant="primary", scale=2) cancel_btn = gr.Button("Cancel Training", variant="stop", visible=False, scale=1) lora_name = gr.Textbox(label="LoRA Name", value="my-lora") train_epochs = gr.Slider( label="Epochs", minimum=1, maximum=1000, value=3, step=1, ) train_lr = gr.Number(label="Learning Rate", value=3e-4) train_rank = gr.Slider( label="Rank (r)", minimum=1, maximum=128, value=16, step=1, ) # Button swap on click (separate handler, like rvc-beatrice) # This fires immediately so user sees Cancel even if training # queues behind concurrency_limit=1 train_btn.click( lambda: (gr.Button(visible=False), gr.Button(visible=True)), outputs=[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, train_output_file], api_name="train_lora", concurrency_limit=1, ) # After training completes, restore buttons and refresh LoRA dropdown # This ensures cleanup even if the user navigated away def _post_training(): return ( gr.Button(visible=True), gr.Button(visible=False), gr.Dropdown(choices=_list_lora_choices()), ) train_event.then( _post_training, outputs=[train_btn, cancel_btn, lora_select], ) # Cancel: set the flag, update status cancel_btn.click( _on_cancel, outputs=[train_log], ) demo.launch( server_name="0.0.0.0", server_port=7860, mcp_server=True, css=CSS, ) # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- if __name__ == "__main__": # If any CLI arguments besides the script name, run CLI mode # (Gradio sets no extra args; start.sh calls `python3 /app/app.py`) if len(sys.argv) > 1: cli_main() else: gradio_main()