import argparse
import json
import os
import sys
import time
from pathlib import Path

os.environ.setdefault("TQDM_DISABLE", "1")
os.environ["PATH"] = "/opt/homebrew/bin:/usr/local/bin:" + os.environ.get("PATH", "")

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

prompt = (p.get("prompt") or "").strip()
if not prompt:
    print("ERROR: prompt (style tags) is required")
    sys.exit(1)

from acestep.pipeline_ace_step import ACEStepPipeline  # noqa: E402

seed = int(p.get("seed", 42))
out = Path(a.outdir) / f"music_{seed}.wav"
t0 = time.time()
pipe = ACEStepPipeline(dtype="float32")  # bfloat16 unsupported on MPS
print(f"pipeline ready in {time.time() - t0:.1f}s", flush=True)

t1 = time.time()
pipe(
    prompt=prompt,
    lyrics=(p.get("lyrics") or "[instrumental]").strip(),
    audio_duration=float(p.get("duration", 30)),
    infer_step=int(p.get("steps", 27)),
    guidance_scale=15.0,
    omega_scale=10.0,
    manual_seeds=[seed],
    save_path=str(out),
)
if not out.exists():
    # some versions append suffixes; register whatever wav landed in outdir
    wavs = list(Path(a.outdir).glob("*.wav"))
    if not wavs:
        print("ERROR: no wav produced")
        sys.exit(1)
    out = wavs[0]
print(f"done in {time.time() - t1:.1f}s: {out.name}", flush=True)