app.py

"""Stem-0 · missing-stem generation on ZeroGPU.

Design notes (ZeroGPU + credit awareness):
  * Weights are downloaded at STARTUP (CPU only) — downloads cost no GPU quota.
  * Per ZeroGPU docs the model is placed on cuda at MODULE level — a PyTorch CUDA emulation mode is
    active outside @spaces.GPU, and startup placement is what their transfer path is optimised for.
    (Lazy-loading inside the decorated function is explicitly discouraged as much slower.)
  * Duration is DYNAMIC: quota is billed per GPU-second, so we request only what the step count needs
    instead of a flat 120s.
  * Generation parameters come from OUR validated `run_ace_task_baseline.build_params` (imported,
    not reimplemented — a second hand-typed copy of that call already drifted in six places once).
"""
from __future__ import annotations

import json
import os
import shutil
import tempfile
import time
from pathlib import Path

import gradio as gr
import numpy as np
import soundfile as sf
import spaces
from huggingface_hub import snapshot_download

# ---- vendor the acestep SOURCE (not the pip package: its gradio/torch pins conflict with ZeroGPU)
ACE_SRC = Path(os.environ.get("STEM0_ACE_SRC", "/tmp/ace_src"))


def _fetch_ace_source() -> str:
    import io, tarfile, urllib.request
    if ACE_SRC.exists() and (ACE_SRC / "acestep").is_dir():
        return str(ACE_SRC)
    ACE_SRC.parent.mkdir(parents=True, exist_ok=True)
    url = "https://github.com/ace-step/ACE-Step-1.5/archive/refs/heads/main.tar.gz"
    raw = urllib.request.urlopen(url, timeout=180).read()
    tmp = ACE_SRC.with_suffix(".x")
    with tarfile.open(fileobj=io.BytesIO(raw), mode="r:gz") as tf:
        tf.extractall(tmp)
    inner = next(p for p in tmp.iterdir() if p.is_dir())
    if ACE_SRC.exists():
        shutil.rmtree(ACE_SRC, ignore_errors=True)
    inner.rename(ACE_SRC)
    shutil.rmtree(tmp, ignore_errors=True)
    return str(ACE_SRC)


CKPT = Path(os.environ.get("STEM0_CKPT", "/tmp/checkpoints"))
LORA_REPO = os.environ.get("STEM0_LORA_REPO", "fcolooo/stem-0-r128")
LORA_REV = os.environ.get("STEM0_LORA_REV", "main")   # step_040000 — content-verified sha256 e0fe86e7…
# FULL-FINETUNE MODE. Set STEM0_FULLFT_REPO and this Space serves a fully fine-tuned decoder instead
# of base+adapter; leave it empty and nothing below changes. The two are mutually exclusive.
#
# The weights are overlaid onto the LIVE model rather than swapped in as files, because the base
# checkpoint is the whole 4.987B model in 4 shards while a full-finetune checkpoint is the DECODER
# ONLY (4.169B, 2 shards, decoder-relative keys) -- a file swap would load without error and silently
# discard the 818M non-decoder parameters. COVER_STRENGTH below is unaffected and still applies.
FULLFT_REPO = os.environ.get("STEM0_FULLFT_REPO", "")
FULLFT_REV = os.environ.get("STEM0_FULLFT_REV", "main")
# Optional subfolder inside FULLFT_REPO holding the decoder (e.g. "step_065000/decoder" in a
# checkpoint-mirror repo). Empty = decoder files at the repo root, the rc.0 layout.
FULLFT_SUBDIR = os.environ.get("STEM0_FULLFT_SUBDIR", "").strip("/")
# How tightly the surrounding context binds the generation. run_take defaults this to 0.45, which was
# never a decision for this task -- it is the default of whichever copy of run_ace_task_baseline gets
# imported (another copy in the same codebase defaults to 1.0). Measured on 2026-08-09 at matched
# n=10 against the same weights at 0.45: 5 wins, 0 losses, six metrics at p=0.002 with unanimous
# 10/0 track splits. melody onset correlation 0.166 -> 0.375, bass onset 0.108 -> 0.206. The sweep is
# monotone across 0.30/0.45/0.65/0.80/1.00, and 1.00 is the top of the tested range -- the curve had
# not turned over, so the true optimum may be higher and is untested. Degeneracy was ruled out
# independently: CLAP up on all four roles, centroid error down on three of four, RMS normal.
COVER_STRENGTH = float(os.environ.get("STEM0_COVER_STRENGTH", "1.0"))
HF_TOK = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
MODEL = "acestep-v15-xl-base"
ROLES = {"drums": "drums", "bass": "bass", "guitar (melody)": "melody", "vocals": "vocals"}
# These four are EXACT full captions from the training corpus (captions.jsonl), verified
# present, not phrases composed to sound plausible. The previous defaults were v0.1 recipe
# prose: 0/1142 exact match per role, and most content words ("locked", "groove", "kit",
# "harmony", "phrasing") appear in ZERO captions of their role. "guitar" in the melody
# default was the melody/guitar mismatch again, arriving through the default path.
# They are deliberately genre-less (3.7% of the corpus, in distribution): a default fires
# when the user named no genre, so it must not assert one. "balanced" is the most neutral
# attested descriptor for each role.
DEFAULT_CAPTIONS = {
    "drums": "balanced drums.",
    "bass": "balanced bass.",
    "melody": "balanced lead melody.",
    "vocals": "balanced vocals.",
}
# The EXACT 24 keyscale strings the LoRA was conditioned on (from the training meta cache:
# sharps only, lowercase mode). Offering anything else would be a format the model never saw.
KEYSCALES = [f"{p} {m}" for m in ("major", "minor")
             for p in ("C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B")]
BPM_MIN, BPM_MAX = 40, 240   # outside this, treat the value as noise and fall back to N/A

_SESSION = None
_BOOT: dict[str, str] = {}
# measured locally on a 3090: ~0.22 s/step at 30s. Blackwell is faster; pad for VAE decode + IO.
STEP_SECONDS = 0.35
BASE_SECONDS = 25.0


# Output length. seconds=0 means "match the uploaded context", which is the default: the natural
# request is a stem for the track you handed over, not a stem for the first N seconds of it.
#
# MIN/MAX are not taste. 10 s is the decoder's floor. 120 s is the longest length ever measured end to
# end here (2026-08-11: 30 s -> 12.33 s GPU-held, 120 s -> 27.82 s, correct duration returned both
# times), and it is also the point where get_duration below still fits inside ZeroGPU's 300 s ceiling
# at the maximum 64 steps: 120 + 25 + 0.35*64*(120/20) = 279. Raising MAX_SECONDS without re-deriving
# that arithmetic would let the Space under-request GPU time and get reclaimed mid-generation.
#
# Quality is a separate matter from mechanics and does NOT hold flat across the range: the model is
# trained on 30 s, and measured over a 120 s vocal the voiced fraction fell 52.2% (first 30 s) to
# 30.4% after, with the spectral centroid drifting up. Long output works; it does not sound as good.
AUTO_SECONDS = 0.0
MIN_SECONDS = 10.0
MAX_SECONDS = 120.0


def _audio_seconds(path: str) -> float:
    """Duration of an upload, from the header. No decode, no GPU, no quota."""
    info = sf.info(path)
    return float(info.frames) / float(info.samplerate)


def resolve_seconds(context_path, requested) -> tuple[float, str]:
    """Turn a requested length into the one actually rendered. Returns (seconds, note).

    `requested` of 0 (AUTO_SECONDS) means match the upload. Anything else is an explicit override and
    is honoured, still clamped -- the research scripts pass a fixed 30 s on purpose so probe items
    stay comparable, and that must keep working.
    """
    try:
        source = _audio_seconds(context_path) if context_path else 0.0
    except Exception as exc:                                        # noqa: BLE001
        # A header we cannot read is not a reason to fail the request: fall back to the old default.
        return 20.0, f"length 20.00s (could not read upload duration: {type(exc).__name__})"

    if float(requested or 0) <= 0:
        seconds = min(max(source, MIN_SECONDS), MAX_SECONDS)
        note = f"length {seconds:.2f}s (matched to upload {source:.2f}s)"
        if source > MAX_SECONDS:
            note = (f"length {seconds:.2f}s — upload is {source:.2f}s, capped at {MAX_SECONDS:.0f}s "
                    f"(the longest length measured end to end)")
        elif source < MIN_SECONDS:
            note = (f"length {seconds:.2f}s — upload is only {source:.2f}s, raised to the "
                    f"{MIN_SECONDS:.0f}s decoder floor; the tail will be silence")
        return seconds, note

    seconds = min(max(float(requested), MIN_SECONDS), MAX_SECONDS)
    return seconds, f"length {seconds:.2f}s (explicit; upload is {source:.2f}s)"


def get_duration(item_str=None, steps=24, seed=0, role_label=None, lyrics=None, caption=None,
                 seconds=20, *a, **k):
    """Dynamic @spaces.GPU duration: ask for what we need, so quota isn't over-reserved.

    `seconds` here is ALREADY resolved -- generate() calls resolve_seconds before _gpu_generate, and
    ZeroGPU passes that same argument list to this function. If it ever received the raw slider value
    again, a 0 would fall back to 20 via `or 20` and under-request the GPU for a long render.
    """
    # A full finetune reads ~8.3 GiB of shards into the live decoder; an adapter attach is 1.2 GiB.
    extra = 0 if _LORA_READY else (120 if FULLFT_REPO else 45)
    return int(min(300, extra + BASE_SECONDS + STEP_SECONDS * float(steps or 24) * max(1.0, float(seconds or 20) / 20.0)))


# ----------------------------------------------------------------------------- startup (no GPU)
def _prepare_source() -> str:
    import sys
    root = _fetch_ace_source()
    if root not in sys.path:
        sys.path.insert(0, root)
    return root


def _download() -> str:
    """Pull the base model, VAE + text encoder, and the LoRA. ~23 GB, one time, no GPU quota."""
    CKPT.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    # base decoder -> checkpoints/acestep-v15-xl-base
    snapshot_download("ACE-Step/acestep-v15-xl-base", local_dir=str(CKPT / MODEL),
                      max_workers=8, tqdm_class=None)
    # vae/ + Qwen3-Embedding-0.6B/ + top-level config.json live in the umbrella repo
    snapshot_download("ACE-Step/Ace-Step1.5", local_dir=str(CKPT), max_workers=8, tqdm_class=None,
                      allow_patterns=["config.json", "vae/*", "Qwen3-Embedding-0.6B/*"])
    if FULLFT_REPO:
        patterns = ["*.safetensors", "model.safetensors.index.json"]
        if FULLFT_SUBDIR:  # scope the download: a mirror repo holds several 8.3 GiB checkpoints
            patterns = [f"{FULLFT_SUBDIR}/{p}" for p in patterns]
        full = snapshot_download(FULLFT_REPO, revision=FULLFT_REV, local_dir="/tmp/fullft",
                                 max_workers=4, tqdm_class=None, token=HF_TOK,
                                 allow_patterns=patterns)
        _BOOT["fullft_dir"] = str(Path(full) / FULLFT_SUBDIR) if FULLFT_SUBDIR else full
        return f"downloaded in {time.time() - t0:.0f}s (full finetune {FULLFT_REPO}@{FULLFT_REV})"
    # our adapter repo is PRIVATE -> needs the token from the Space secret
    lora = snapshot_download(LORA_REPO, revision=LORA_REV, local_dir="/tmp/lora",
                             max_workers=4, tqdm_class=None, token=HF_TOK,
                             allow_patterns=["adapter_config.json", "adapter_model.safetensors"])
    _BOOT["lora_dir"] = lora
    return f"downloaded in {time.time() - t0:.0f}s"


try:
    _BOOT["ace_src"] = _prepare_source()
    _BOOT["status"] = _download()
except Exception as exc:                                    # surface in the UI instead of a blank page
    _BOOT["status"] = f"DOWNLOAD FAILED: {type(exc).__name__}: {exc}"
_BOOT.setdefault("lora_dir", "/tmp/lora")


def _warm_analysis() -> None:
    """Run the bpm/key estimator once on synthetic audio at boot.

    Two reasons: librosa/numba JIT makes a cold call ~20s of (unbilled but user-visible) prep, and
    a version mismatch in the estimator should surface in the boot log rather than silently
    degrading the first real render to N/A. Off-GPU, so it costs no quota."""
    try:
        t0 = time.time()
        p = Path(tempfile.mkdtemp(prefix="warm_")) / "w.wav"
        n = _ANALYSIS_SR * 4
        t = np.arange(n) / _ANALYSIS_SR
        click = (np.sin(2 * np.pi * 220 * t) * (np.sin(2 * np.pi * 2 * t) > 0.9)).astype("float32")
        sf.write(p, click, _ANALYSIS_SR)
        bpm, key = _detect_bpm_key(str(p))
        _BOOT["analysis"] = (f"estimator ready in {time.time()-t0:.0f}s (warmup -> bpm={bpm} key={key or 'N/A'})"
                             if bpm or key else "ESTIMATOR BROKEN — bpm/key will fall back to N/A")
    except Exception as exc:
        _BOOT["analysis"] = f"ESTIMATOR BROKEN: {type(exc).__name__}: {exc}"
    print(f"[analysis] {_BOOT['analysis']}")



# --------------------------------------------- model on cuda at module level (ZeroGPU requirement)
def _get_session():
    global _SESSION
    if _SESSION is None:
        from run_ace_task_baseline import init_ace
        # NOTE: lora_path deliberately omitted here — PEFT injection requires a real GPU, which
        # only exists inside @spaces.GPU. The adapter is attached by _ensure_lora() on first call.
        _SESSION = init_ace(
            ace_root=_BOOT.get("ace_src", str(ACE_SRC)), checkpoints=str(CKPT), model=MODEL, device="cuda",
            use_lm=False, no_thinking=True,               # LM not needed: we pass explicit prompts
        )
    return _SESSION


_LORA_READY = False
_AOTI = {"state": "not compiled"}


def _demo_ctx(seconds: float = 20.0, sr: int = 48000) -> str:
    """Cheap stand-in context used only to capture real decoder inputs for torch.export."""
    import tempfile as _tf
    rng = np.random.default_rng(0)
    y = (rng.standard_normal((int(seconds * sr), 2)) * 0.05).astype("float32")
    p = Path(_tf.mkdtemp()) / "ctx.wav"
    sf.write(p, y, sr)
    return str(p)


def _attach_fullft(session):
    """Overlay a fully fine-tuned decoder onto the live model. Inside @spaces.GPU, like the adapter.

    strict=True is the safety argument: a partial load would leave most of the decoder at base weights
    and still generate plausible audio. The handler wraps the model in torch.compile when compiling is
    enabled, so unwrap _orig_mod first -- otherwise the overlay targets the wrapper and every request
    fails, which is exactly what happened on the dev Space.
    """
    from safetensors.torch import load_file

    shards = sorted(Path(_BOOT["fullft_dir"]).glob("*.safetensors"))
    if not shards:
        raise RuntimeError(f"no .safetensors in {_BOOT.get('fullft_dir')}")
    state = {}
    for shard in shards:
        state.update(load_file(str(shard)))
    model = getattr(session.dit_handler, "model", None)
    model = getattr(model, "_orig_mod", model)
    decoder = getattr(model, "decoder", None)
    _BOOT["overlay_target"] = (f"decoder={type(decoder).__name__} "
                               f"decoder_tensors={len(decoder.state_dict()) if decoder else 0} "
                               f"file_tensors={len(state)}")
    print(f"[fullft] {_BOOT['overlay_target']}", flush=True)
    if decoder is None:
        raise RuntimeError("dit_handler.model.decoder not reachable; cannot overlay a full finetune")
    decoder.load_state_dict(state, strict=True)
    print(f"[fullft] overlaid {len(state)} tensors ({FULLFT_REPO}@{FULLFT_REV})", flush=True)


def _ensure_lora(session):
    """Attach the trained weights — must run inside @spaces.GPU (needs a real GPU)."""
    global _LORA_READY
    if _LORA_READY:
        return
    if FULLFT_REPO:
        try:
            _attach_fullft(session)
        except Exception as exc:
            import traceback
            _BOOT["weights_error"] = traceback.format_exc()[-4000:]
            raise
        _BOOT["weights_error"] = "none — full finetune overlaid"
        _LORA_READY = True
        return
    h = session.dit_handler
    st = h.add_lora(_BOOT["lora_dir"], adapter_name="stem0")
    print("[lora]", st, flush=True)
    if not str(st).startswith("✅"):
        raise RuntimeError(f"LoRA load failed: {st}")
    print("[lora]", h.set_lora_scale("stem0", 1.0), flush=True)
    print("[lora]", h.set_use_lora(True), flush=True)
    session.lora_path = _BOOT["lora_dir"]
    _LORA_READY = True


def _to_flac(wav_path: Path) -> str:
    """Return FLAC: lossless, ~2x smaller than WAV. Runs off-GPU, so it costs latency only — never quota.

    Was MP3 (-q:a 2) for an 8x download saving. Replaced 2026-08-09: the generations are the product, and
    re-encoding them lossily to save 2 MB is the wrong trade when FLAC is free of artifacts and still halves
    the transfer. The measured alternative -- moving the VAE client-side and shipping 94 KiB of latents, 61x
    smaller than WAV -- was prototyped and shelved: VAE encode+decode costs 61 s on a CPU client, which is
    slower than the 19 s generation it was meant to accelerate. That design is worth revisiting only for
    clients with a GPU.
    """
    out = wav_path.with_suffix(".flac")
    try:
        import subprocess
        subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(wav_path),
                        "-codec:a", "flac", "-compression_level", "5", str(out)], check=True)
        return str(out)
    except Exception as exc:                       # never fail a good generation over encoding
        print("[flac] falling back to wav:", exc, flush=True)
        return str(wav_path)


_CTX_CACHE: dict[tuple, str] = {}


def _prepared_context(context_path: str, seconds: float) -> str:
    """Decode/resample/trim once per (file, length). Pure CPU — must stay OUT of @spaces.GPU,
    because ZeroGPU bills wall-clock while the GPU is held."""
    import hashlib
    h = hashlib.md5()
    with open(context_path, "rb") as f:
        for b in iter(lambda: f.read(1 << 20), b""):
            h.update(b)
    key = (h.hexdigest(), round(float(seconds), 2))
    hit = _CTX_CACHE.get(key)
    if hit and os.path.exists(hit):
        return hit
    y, sr = sf.read(context_path, always_2d=True)
    if sr != 48000:
        import librosa
        y = librosa.resample(y.T.astype("float32"), orig_sr=sr, target_sr=48000).T
        sr = 48000
    if y.shape[1] == 1:
        y = np.repeat(y, 2, axis=1)
    n = int(seconds * sr)
    y = y[:n] if len(y) > n else np.pad(y, ((0, n - len(y)), (0, 0)))
    out = Path(tempfile.mkdtemp(prefix="ctx_")) / "context_mix_minus_target.wav"
    sf.write(out, y.astype("float32"), sr)
    _CTX_CACHE[key] = str(out)
    return str(out)


# Krumhansl key-profile weights + analysis SR, COPIED VERBATIM from the script that produced our
# training labels (scripts/data_recipe/build_track_metas.py). Do not "improve" these: the point is
# to reproduce the exact label-generating process the LoRA was conditioned on, biases included.
_ANALYSIS_SR = 22050
_KMAJ = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88])
_KMIN = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17])
_NOTES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
_ANALYSIS_CACHE: dict = {}


def _detect_bpm_key(wav_path: str) -> tuple[int | None, str]:
    """Estimate (bpm, keyscale) with the SAME librosa procedure that labelled the training set.

    Our training bpm/keyscale were never ground truth — build_track_metas.py estimated them from
    audio. So the closest thing to the conditioning the model actually learned is to re-run that
    identical estimator here, rather than to send N/A (a branch only ~0.7% of training saw).
    CPU only, cached per file — must stay outside @spaces.GPU so it is never billed.
    """
    hit = _ANALYSIS_CACHE.get(wav_path)
    if hit is not None:
        return hit
    def _scalar(x) -> float:
        """librosa's tempo return shape moves between versions (scalar / (1,) / (n,)) and numpy 2
        refuses float() on ndim>0. Take the first element whatever the shape."""
        a = np.asarray(x, dtype="float64").ravel()
        return float(a[0]) if a.size else 0.0

    try:
        import librosa
        y, _ = librosa.load(wav_path, sr=_ANALYSIS_SR, mono=True)
        try:
            tempo = _scalar(librosa.feature.rhythm.tempo(y=y, sr=_ANALYSIS_SR))
        except Exception:
            tempo = _scalar(librosa.beat.beat_track(y=y, sr=_ANALYSIS_SR)[0])
        chroma = librosa.feature.chroma_cqt(y=y, sr=_ANALYSIS_SR).mean(axis=1)
        mode, root, _ = max(
            ((m, i, np.corrcoef(np.roll(k, i), chroma)[0, 1])
             for m, k in (("maj", _KMAJ), ("min", _KMIN)) for i in range(12)),
            key=lambda x: (x[2] if np.isfinite(x[2]) else -9))
        out = (int(round(tempo)), f"{_NOTES[root]} {'major' if mode == 'maj' else 'minor'}")
    except Exception as e:                      # never fail a render over metadata
        print(f"[analysis] bpm/key detection failed: {e}")
        out = (None, "")
    _ANALYSIS_CACHE[wav_path] = out
    return out


def _prep_item(context_path: str, role: str, caption: str, lyrics: str, seconds: float,
               bpm: float = 0, keyscale: str = "") -> Path:
    item = Path(tempfile.mkdtemp(prefix="stem0_"))
    prepared = _prepared_context(context_path, seconds)
    shutil.copy2(prepared, item / "context_mix_minus_target.wav")
    y, sr = sf.read(item / "context_mix_minus_target.wav", always_2d=True)
    has_lyrics = role == "vocals" and bool(lyrics.strip())
    meta = {
        "role": role,
        "prompt": caption.strip() or DEFAULT_CAPTIONS[role],
        "ace_task_type": "lego",
        "source_audio": "context_mix_minus_target.wav",
        # NOTE: this drives the actual generation LENGTH (run_ace_task_baseline: duration=
        # _meta_float(meta,"duration_seconds")), not just the metas line — never round it to 30.
        "duration_seconds": float(len(y) / sr),
        "timesignature": "4",
        "lyrics": lyrics.strip() if has_lyrics else "[Instrumental]",
        "instrumental": not has_lyrics,
        "vocal_language": "en" if has_lyrics else "unknown",
    }
    # bpm / keyscale: pass through ONLY when the caller actually knows them.
    #
    # These are not cosmetic. Both metas builders emit the field unconditionally — training
    # (preprocess_utils.build_metas_str) and inference (metadata_utils._dict_to_meta_string) —
    # so leaving them out does not remove the line, it ships "- bpm: N/A". 1164/1172 (99%) of
    # our training tracks carried a REAL bpm and keyscale, so N/A is the ~0.7% branch of the
    # conditioning distribution, not a neutral default. A caller with ground truth (the live
    # app knows its own click tempo) should send it. A caller that would be GUESSING must not:
    # a wrong tempo makes the model play out of time, which is worse than the rare-token hit.
    # A caller-supplied value is honoured as-is. Absent/garbage/"auto" -> field stays out unless
    # "auto" opts into estimation below. float() is guarded because bpm may arrive as "auto".
    try:
        _b = float(bpm)
    except (TypeError, ValueError):
        _b = 0.0
    if _b and BPM_MIN <= _b <= BPM_MAX:
        meta["bpm"] = int(round(_b))
    if keyscale and str(keyscale).strip() in KEYSCALES:
        meta["keyscale"] = str(keyscale).strip()
    #
    # Measured on 40 tracks, estimator-on-partial-context vs the full-mix training label:
    #                        key exact   bpm exact   bpm within 5%
    #   drums in context        88%         82%          85%
    #   drums absent            82%         52%          65%
    #
    # Key holds up either way -> always estimate it. Tempo collapses without drums to lock onto,
    # and a wrong tempo makes the model play out of time (worse than N/A), so we only estimate bpm
    # when the context plausibly contains drums — i.e. every role EXCEPT drums itself.
    # ESTIMATION IS OPT-IN ("auto"), NOT the default. Measured on a labelled track, 64 steps,
    # fixed seed, generated melody's chroma fit to the context's true key (D# major, ctx = +0.398):
    #
    #   keyscale sent      generated fit to the true key
    #   ---------------    -----------------------------
    #   N/A (absent)             +0.406   <- best; matches the context almost exactly
    #   "D# major" (right)       +0.192   <- WORSE than sending nothing
    #   "A major"  (wrong)       -0.347   <- actively clashing
    #
    # And supplying a correct bpm did not improve tempo at all (117.5 generated either way against
    # a 117 context). So the model already reads tempo AND harmony out of the context audio, and an
    # explicit label competes with that evidence — a coarse 30s key summary is strictly less
    # informative than the audio itself. The original "omit these" recipe was right; estimating
    # them by default made output worse. Kept available for callers who genuinely want to force a
    # key, plus "auto" for experiments.
    # bpm rides a gr.Number, which rejects the string "auto" — so -1 is the numeric opt-in sentinel.
    want_bpm_auto = str(bpm).strip().lower() == "auto" or _b == -1
    want_key_auto = str(keyscale).strip().lower() == "auto"
    if want_bpm_auto or want_key_auto:
        det_bpm, det_key = _detect_bpm_key(str(item / "context_mix_minus_target.wav"))
        est = []
        if want_bpm_auto and role != "drums" and det_bpm and BPM_MIN <= det_bpm <= BPM_MAX:
            meta["bpm"] = det_bpm
            est.append("bpm")
        if want_key_auto and det_key in KEYSCALES:
            meta["keyscale"] = det_key
            est.append("keyscale")
        if est:
            meta["_estimated"] = est
    (item / "metadata.json").write_text(json.dumps(meta, indent=2))
    return item


@spaces.GPU(duration=get_duration)
def _gpu_generate(item_str, steps, seed, *_):
    """ONLY the model call. Everything CPU-bound is kept outside so it isn't billed as GPU time."""
    from run_ace_task_baseline import run_take
    session = _get_session()
    _ensure_lora(session)                 # first call attaches the adapter on the real GPU
    item = Path(item_str)
    _t = time.time()
    run_take(session, item, task="lego", steps=int(steps), seed=int(seed),
             cover_strength=COVER_STRENGTH,
             out_dir=item / "out", generated_name="generated_stem.wav", log=print)
    held = time.time() - _t          # actual GPU-held time == what ZeroGPU bills
    return f"{item / 'generated_stem.wav'}|{held:.3f}"


def generate(context_audio, role_label, lyrics, caption, steps, seconds, seed, bpm=0, keyscale=""):
    """bpm/keyscale are TRAILING and OPTIONAL: existing 7-arg callers keep working unchanged
    (Gradio fills missing trailing inputs from component defaults)."""
    if context_audio is None:
        raise gr.Error("Upload a context audio first (e.g. drums+bass, or a full instrumental).")
    if _BOOT.get("status", "").startswith("DOWNLOAD FAILED"):
        raise gr.Error(_BOOT["status"])
    role = ROLES[role_label]
    if role == "vocals" and not lyrics.strip():
        raise gr.Error("Vocals need lyrics — type some, or pick a different stem.")

    # Resolve the output length BEFORE anything else touches `seconds`: _prep_item pads or trims the
    # context to exactly this value, and get_duration sizes the ZeroGPU request from it.
    seconds, length_note = resolve_seconds(context_audio, seconds)

    tp = time.time()
    item = _prep_item(context_audio, role, caption, lyrics, float(seconds), bpm, keyscale or "")
    _m = json.loads((item / "metadata.json").read_text())
    _est = _m.get("_estimated", [])             # '~' in the note marks an estimate, not ground truth
    prep_s = time.time() - tp
    t0 = time.time()
    _ret = _gpu_generate(str(item), int(steps), int(seed), role_label, lyrics, caption, seconds)
    _path, _held = _ret.rsplit("|", 1)
    gen = Path(_path)
    held_s = float(_held)                       # GPU actually held (billed)
    outer_s = time.time() - t0                  # held + ZeroGPU allocation/queue wait
    alloc_s = max(0.0, outer_s - held_s)
    tm = time.time()

    ctx, sr = sf.read(item / "context_mix_minus_target.wav", always_2d=True)
    g, _ = sf.read(gen, always_2d=True)
    n = min(len(ctx), len(g))
    mix = ctx[:n] + g[:n]
    pk = float(np.abs(mix).max()) or 1.0
    mix = (mix * (0.97 / pk)).astype("float32")
    mix_path = item / "mix.wav"
    sf.write(mix_path, mix, sr)
    t_enc = time.time()
    gen_out, mix_out = _to_flac(gen), _to_flac(mix_path)
    enc_s = time.time() - t_enc
    mix_s = time.time() - tm
    note = (f"role={role} · steps={int(steps)} · {length_note} · "
            f"bpm={_m.get('bpm', 'N/A')}{'~' if 'bpm' in _est else ''} · "
            f"key={_m.get('keyscale', 'N/A')}{'~' if 'keyscale' in _est else ''} · "
            f"GPU-held {held_s:.2f}s ({held_s/max(int(steps),1):.3f}s/step) · "
            f"alloc/queue {alloc_s:.2f}s · prep {prep_s:.2f}s · mix {mix_s:.2f}s · "
            f"flac {enc_s:.2f}s (all off-GPU)")
    return gen_out, mix_out, note


# Place the model on cuda NOW (startup), per ZeroGPU guidance. Safe outside @spaces.GPU thanks to
# their CUDA emulation mode; makes the first real request fast instead of paying a 20GB load then.
if not _BOOT.get("status", "").startswith("DOWNLOAD FAILED"):
    try:
        _t = time.time()
        _get_session()
        _BOOT["load"] = f"model on cuda in {time.time()-_t:.0f}s"
    except Exception as exc:
        _BOOT["load"] = f"MODEL LOAD FAILED: {type(exc).__name__}: {exc}"
    print("[boot]", _BOOT.get("load"), flush=True)


def boot_status() -> str:
    """CPU-only: verify download + cuda placement without spending any GPU quota."""
    info = {k: str(v)[:300] for k, v in _BOOT.items()} | {"aoti": _AOTI["state"]}
    # WHICH weights are actually loaded, answered by CONTENT not by branch name. A branch name has
    # silently pointed at the wrong checkpoint twice (a fresh branch inherits main's commit, so a
    # failed upload leaves a branch that looks right and isn't). Compare this to the local
    # checkpoint's sha256 before trusting any claim about which step is being served.
    info["lora_rev_requested"] = LORA_REV
    if "lora_sha256" not in _BOOT:
        try:
            import hashlib
            p = Path(_BOOT.get("lora_dir", "/tmp/lora")) / "adapter_model.safetensors"
            h = hashlib.sha256()
            with open(p, "rb") as f:
                for b in iter(lambda: f.read(1 << 22), b""):
                    h.update(b)
            _BOOT["lora_sha256"] = f"{h.hexdigest()} ({p.stat().st_size/1e6:.0f} MB)"
        except Exception as e:
            _BOOT["lora_sha256"] = f"UNAVAILABLE: {type(e).__name__}: {e}"
    info["lora_sha256"] = _BOOT["lora_sha256"]
    if FULLFT_REPO and "fullft_id" not in _BOOT:
        try:
            import hashlib
            from huggingface_hub import HfApi
            sha = HfApi(token=HF_TOK).model_info(FULLFT_REPO, revision=FULLFT_REV).sha
            shards = sorted(Path(_BOOT.get("fullft_dir", "/tmp/fullft")).glob("*.safetensors"))
            h2, total = hashlib.sha256(), 0
            for shard in shards:
                size = shard.stat().st_size
                total += size
                h2.update(f"{shard.name}:{size}".encode())
                with open(shard, "rb") as f:
                    h2.update(f.read(4 << 20))
                    if size > (8 << 20):
                        f.seek(-(4 << 20), 2)
                        h2.update(f.read(4 << 20))
            _BOOT["fullft_id"] = (f"commit {sha[:12]} · partial-digest {h2.hexdigest()[:16]} · "
                                  f"{len(shards)} shards · {total/1e9:.2f} GB")
        except Exception as e:
            _BOOT["fullft_id"] = f"UNAVAILABLE: {type(e).__name__}: {e}"
    if FULLFT_REPO:
        info["fullft_id"] = _BOOT["fullft_id"]
        info["fullft_rev_requested"] = FULLFT_REV
        # REPORT THE REPO, not just the commit. deploy.py verifies the served commit against
        # `status["fullft_repo"] or "fcolooo/stem-0-fullft"` -- and this dict never carried
        # fullft_repo, so every non-prod Space silently fell back to prod's repo and FAILED
        # verification while serving exactly what it was asked to. A deploy gate that cannot pass
        # on a dev Space is a gate people learn to ignore, which is worse than no gate.
        info["fullft_repo"] = FULLFT_REPO
    try:
        import torch
        info["torch"] = torch.__version__
        try:
            import torchaudio; info["torchaudio"] = torchaudio.__version__
        except Exception as e: info["torchaudio"] = f"MISSING: {e}"
        import gradio; info["gradio"] = gradio.__version__
        try:
            import librosa, numpy
            info["librosa"] = f"{librosa.__version__} (numpy {numpy.__version__})"
        except Exception as e: info["librosa"] = f"MISSING: {e}"
    except Exception as e:
        info["torch"] = f"ERR {e}"
    return json.dumps(info, indent=2)


@spaces.GPU(duration=1500)
def optimize_aoti(seconds: float = 20.0):
    """Ahead-of-time compile the DiT decoder (torch.compile is unsupported on ZeroGPU).
    Opt-in: compilation itself costs GPU quota, and the graph is specialised to the clip length
    it was captured at, so recompile if you change `seconds`."""
    import torch
    from run_ace_task_baseline import run_take
    if _AOTI["state"].startswith("applied"):
        return _AOTI["state"]
    session = _get_session()
    _ensure_lora(session)
    dec = session.dit_handler.model.decoder
    t0 = time.time()
    try:
        # capture the real args the decoder is called with, by running a tiny generation
        # aoti_capture aborts the call with a sentinel exception once it has the args, but ACE's
        # generate_music catches everything and re-raises it as "Generation failed" — so swallow it
        # here and check whether the args were recorded anyway.
        with spaces.aoti_capture(dec) as call:
            try:
                item = _prep_item(_demo_ctx(seconds), "drums", "", "", float(seconds))
                run_take(session, item, task="lego", steps=4, seed=1,
                         cover_strength=COVER_STRENGTH,
                         out_dir=item / "o", generated_name="g.wav", log=print)
            except Exception as cap_exc:
                print("[aoti] capture aborted as expected:", type(cap_exc).__name__, flush=True)
        if not getattr(call, "args", None) and not getattr(call, "kwargs", None):
            raise RuntimeError("aoti_capture recorded no decoder call")
        print(f"[aoti] captured {len(call.args)} args / {len(call.kwargs)} kwargs", flush=True)
        exported = torch.export.export(dec, args=call.args, kwargs=call.kwargs)
        compiled = spaces.aoti_compile(exported)
        spaces.aoti_apply(compiled, session.dit_handler.model.decoder)
        _AOTI["state"] = f"applied in {time.time() - t0:.0f}s (captured at {seconds:.0f}s clips)"
    except Exception as exc:
        # eager still works — never let a failed optimisation break generation
        _AOTI["state"] = f"AOTI failed, staying on eager: {type(exc).__name__}: {str(exc)[:300]}"
    print("[aoti]", _AOTI["state"], flush=True)
    return _AOTI["state"]


_warm_analysis()   # JIT + version check up front, so no user's first render pays for it

with gr.Blocks(title="Stem-0 · missing-stem generation") as demo:
    gr.Markdown(
        "## Stem-0 — generate the missing stem\n"
        "Upload an audio **context** (drums+bass, an instrumental, even a solo stem) and it writes the "
        "missing part to fit. For **vocals**, type your own lyrics.\n\n"
        f"*Boot: {_BOOT.get('status','?')} · {_BOOT.get('load','')}* · base `ACE-Step/acestep-v15-xl-base` + LoRA "
        f"`{LORA_REPO}@{LORA_REV}`"
    )
    gr.Markdown(
        "⚠️ **ZeroGPU quota is per-day.** The first run also loads ~20 GB onto the GPU and is slow; "
        "later runs are quick. Keep **steps** low (24–32) while experimenting — cost scales with steps × seconds."
    )
    with gr.Row():
        with gr.Column():
            ctx_in = gr.Audio(label="Context audio (the stems you already have)", type="filepath")
            role_in = gr.Dropdown(list(ROLES), value="vocals", label="Stem to generate")
            cap_in = gr.Textbox(label="Caption (optional — describes the part you want)",
                                placeholder="a raw, intimate indie lead vocal, breathy and close")
            lyr_in = gr.Textbox(label="Lyrics (vocals only)", lines=8,
                                placeholder="[Verse]\nyou're a cigarette in my head\nstill burning when i'm in bed")
            with gr.Row():
                steps_in = gr.Slider(16, 64, value=24, step=4, label="Steps (↑quality, ↑GPU cost)")
                # 0 is the default and means "as long as the file you uploaded". The slider stays in
                # the same position in the API's positional `data` array, so existing callers that
                # send an explicit length keep working unchanged.
                secs_in = gr.Slider(0, MAX_SECONDS, value=AUTO_SECONDS, step=5,
                                    label="Seconds (0 = match the uploaded audio)")
                seed_in = gr.Number(value=901, precision=0, label="Seed")
            with gr.Row():
                # Only fill these if you KNOW them (e.g. you played to a click). 99% of training
                # carried a real bpm/key, so a correct value matches the conditioning the LoRA
                # learned — but a guess is worse than leaving it blank.
                bpm_in = gr.Number(value=0, precision=0,
                                   label=f"BPM — 0 = omit, -1 = estimate ({BPM_MIN}-{BPM_MAX})")
                # allow_custom_value: without it Gradio hard-errors an API caller that sends a key
                # outside the list. We'd rather accept the string and let _prep_item fall back to
                # N/A, so an unrecognised key degrades instead of failing the whole render.
                key_in = gr.Dropdown([""] + KEYSCALES, value="", allow_custom_value=True,
                                     label="Key — blank = omit, \"auto\" = estimate")
            go = gr.Button("Generate", variant="primary")
        with gr.Column():
            stem_out = gr.Audio(label="Generated stem")
            mix_out = gr.Audio(label="Context + generated (mix)")
            info_out = gr.Markdown()
    with gr.Accordion("Speed: ahead-of-time compile (AOTI)", open=False):
        gr.Markdown(
            "One-off AOTI compile of the DiT decoder — HF report **1.3–1.8x** on ZeroGPU. "
            "Compilation itself consumes GPU quota and is specialised to the clip length used here, "
            "so run it once for the `seconds` you actually generate at. Falls back to eager on failure."
        )
        with gr.Row():
            aoti_secs = gr.Slider(10, 30, value=20, step=5, label="Compile for this clip length (s)")
            aoti_btn = gr.Button("Compile now")
        aoti_out = gr.Textbox(label="AOTI state", lines=2)
        aoti_btn.click(optimize_aoti, aoti_secs, aoti_out, api_name="optimize")
    with gr.Accordion("Boot status (free, no GPU)", open=False):
        st_btn = gr.Button("Check")
        st_out = gr.Textbox(label="status", lines=6)
        st_btn.click(boot_status, None, st_out, api_name="status")
    go.click(generate, [ctx_in, role_in, lyr_in, cap_in, steps_in, secs_in, seed_in, bpm_in, key_in],
             [stem_out, mix_out, info_out], api_name="generate")

if __name__ == "__main__":
    demo.queue(max_size=8).launch()