"""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 tempfile import requests 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 # --------------------------------------------------------------------------- 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): """Poll a job until done/error/timeout. Returns (status, elapsed).""" t0 = time.time() while time.time() - t0 < timeout: try: r = requests.get(f"{ACE_SERVER}/job", params={"id": job_id}, timeout=10) 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(2) 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 _run_pipeline(caption, lyrics, bpm, duration, seed, steps, output_format, adapter=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]" if bpm and int(bpm) > 0: req["bpm"] = int(bpm) if duration and float(duration) > 0: req["duration"] = min(float(duration), 300) if seed is not None and int(seed) >= 0: req["seed"] = int(seed) if steps and int(steps) > 0: req["inference_steps"] = int(steps) if adapter: req["adapter"] = adapter fmt = output_format if output_format in ("wav", "mp3") else "wav" 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=300) 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 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 # --------------------------------------------------------------------------- # 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: import shutil 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 # -- Generate tab handler -- def generate_music(caption, lyrics, instrumental, bpm, duration, seed, steps, output_format, progress=gr.Progress(track_tqdm=True)): 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) 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=output_format, 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 -- 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) 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_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 = """ .compact-row { gap: 8px !important; } .status-box textarea { font-family: monospace; font-size: 13px; } """ with gr.Blocks(title="ACE-Step 1.5 XL (CPU)", css=CSS) as demo: with gr.Tabs(): # ============================================================ # Tab 1: Generate Music # ============================================================ with gr.Tab("Generate Music"): gr.Markdown( "**[ACE-Step 1.5 XL (CPU)](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=2): 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=1): audio_out = gr.Audio(label="Output", type="filepath") status = gr.Textbox( label="Status", interactive=False, lines=2, elem_classes="status-box", ) with gr.Row(elem_classes="compact-row"): instrumental = gr.Checkbox(label="Instrumental", value=True, scale=1) bpm = gr.Number(label="BPM", value=120, minimum=0, maximum=300, scale=1) duration = gr.Slider( label="Duration (s)", minimum=10, maximum=120, value=10, step=5, scale=1, ) steps = gr.Slider( label="Steps", minimum=1, maximum=32, value=8, step=1, scale=1, ) seed = gr.Number(label="Seed (-1=random)", value=-1, scale=1) output_format = gr.Radio( label="Format", choices=["wav", "mp3"], value="wav", scale=1, ) with gr.Row(elem_classes="compact-row"): gen_btn = gr.Button("Generate Music", variant="primary", scale=2) status_btn = gr.Button("Server Status", scale=1) gen_btn.click( fn=generate_music, inputs=[caption, lyrics, instrumental, bpm, duration, seed, steps, output_format], outputs=[audio_out, status], api_name="generate", ) status_btn.click( fn=get_server_status, inputs=[], outputs=[status], api_name="server_status", ) # ============================================================ # Tab 2: Train LoRA # ============================================================ with gr.Tab("Train LoRA"): gr.Markdown( "### LoRA Training\n" "Fine-tune ACE-Step on your own audio data. " "CPU training is very slow. Checkpoints downloaded on first run (~10GB)." ) with gr.Row(elem_classes="compact-row"): with gr.Column(scale=2): train_audio = gr.File( label="Training Audio Files", file_count="multiple", file_types=["audio"], ) with gr.Column(scale=1): lora_name = gr.Textbox(label="LoRA Name", value="my-lora") 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=64) train_btn = gr.Button("Train", variant="primary") train_log = gr.Textbox( label="Training Log", interactive=False, lines=10, elem_classes="status-box", ) train_btn.click( fn=train_lora, inputs=[train_audio, lora_name, epochs, lr, rank], outputs=[train_log], api_name="train_lora", ) demo.launch( server_name="0.0.0.0", server_port=7860, mcp_server=True, ) # --------------------------------------------------------------------------- # 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()