Spaces:
Paused
Paused
| """ACE-Step 1.5 engine for Cowboy Beatz Song Lab: one clean API endpoint (/generate) on ZeroGPU. | |
| Model init mirrors the official ACE-Step/Ace-Step-v1.5 Space (MIT).""" | |
| import os, sys, json, time, tempfile | |
| here = os.path.dirname(os.path.abspath(__file__)) | |
| nano = os.path.join(here, "acestep", "third_parts", "nano-vllm") | |
| if os.path.exists(nano): | |
| sys.path.insert(0, nano) | |
| os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" | |
| import spaces # before torch | |
| import torch | |
| import gradio as gr | |
| from acestep.handler import AceStepHandler | |
| from acestep.llm_inference import LLMHandler | |
| from acestep.inference import GenerationParams, GenerationConfig, generate_music | |
| DIT_MODEL = os.environ.get("ENGINE_DIT_MODEL", "acestep-v15-xl-turbo") | |
| LM_MODEL = os.environ.get("ENGINE_LM_MODEL", "acestep-5Hz-lm-1.7B") | |
| store = os.path.join(here, "data"); os.makedirs(store, exist_ok=True) | |
| # ZeroGPU now runs on Blackwell (sm_120); the prebuilt flash-attn / flash-attn3 kernels have no sm_120 | |
| # image ("no kernel image is available"), so force PyTorch SDPA everywhere. | |
| LLMHandler.get_best_attn_implementation = lambda self: "sdpa" | |
| AceStepHandler.is_flash_attention_available = lambda self: False | |
| DIT = AceStepHandler(persistent_storage_path=store) | |
| LLM = LLMHandler(persistent_storage_path=store) | |
| st, ok = DIT.initialize_service(project_root=here, config_path=DIT_MODEL, device="auto", | |
| use_flash_attention=DIT.is_flash_attention_available(), compile_model=False, | |
| offload_to_cpu=False, offload_dit_to_cpu=False) | |
| print("DiT:", ok, st, flush=True) | |
| lst, lok = LLM.initialize(checkpoint_dir=DIT._get_checkpoint_dir(), lm_model_path=LM_MODEL, backend="pt", | |
| device="auto", offload_to_cpu=False, dtype=DIT.dtype) | |
| print("LM:", lok, lst, flush=True) | |
| def _dur(caption, lyrics, instrumental, bpm, keyscale, duration, seed, negative, thinking, steps): | |
| return int(min(170, 40 + float(duration) * 0.5)) | |
| def generate(caption, lyrics, instrumental, bpm, keyscale, duration, seed, negative, thinking, steps): | |
| t0 = time.time() | |
| seed = int(seed) if seed is not None else -1 | |
| instrumental = bool(instrumental) | |
| params = GenerationParams( | |
| caption=caption or "", lyrics="[Instrumental]" if instrumental else (lyrics or ""), | |
| instrumental=instrumental, vocal_language="unknown" if instrumental else "en", | |
| bpm=int(bpm) if bpm else None, keyscale=keyscale or "", timesignature="4", | |
| duration=float(duration), inference_steps=int(steps or 8), seed=seed, | |
| thinking=bool(thinking), lm_negative_prompt=negative or "NO USER INPUT", | |
| use_cot_metas=False, use_cot_caption=False, use_cot_language=False, | |
| ) | |
| cfg = GenerationConfig(batch_size=1, use_random_seed=seed < 0, seeds=[seed] if seed >= 0 else None, | |
| audio_format="wav") | |
| res = generate_music(DIT, LLM, params, cfg, save_dir=tempfile.mkdtemp()) | |
| if not res.success or not res.audios or not res.audios[0].get("path"): | |
| raise gr.Error("Generation failed: " + str(res.error or res.status_message)[:300]) | |
| a = res.audios[0] | |
| info = {"seed": a["params"].get("seed"), "gpu_seconds": round(time.time() - t0, 1), | |
| "time_costs": res.extra_outputs.get("time_costs", {})} | |
| return a["path"], json.dumps(info, default=str) | |
| with gr.Blocks(title="ACE-Step engine (Cowboy Beatz)") as demo: | |
| gr.Markdown("ACE-Step 1.5 engine for the Cowboy Beatz Song Lab. API only.") | |
| cap = gr.Textbox(label="caption"); lyr = gr.Textbox(label="lyrics"); ins = gr.Checkbox(True, label="instrumental") | |
| bpm = gr.Number(150, label="bpm"); key = gr.Textbox("F minor", label="key"); dur = gr.Number(60, label="duration") | |
| seed = gr.Number(-1, label="seed"); neg = gr.Textbox(label="negative"); think = gr.Checkbox(True, label="thinking") | |
| steps = gr.Number(8, label="steps") | |
| out = gr.Audio(type="filepath"); meta = gr.Textbox(label="info") | |
| gr.Button("Generate").click(generate, [cap, lyr, ins, bpm, key, dur, seed, neg, think, steps], [out, meta], api_name="generate") | |
| demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860, show_error=True) | |