dev space: r2c 65k via FULLFT_SUBDIR

README.md

@@ -1,13 +1,23 @@
 ---
-title: Stem 0 Web Dev
-emoji: 🚀
+title: Stem-0 DEV (r2c checkpoints)
+emoji: 🎚️
 colorFrom: blue
-colorTo: green
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.25.0
-python_version: '3.13'
+python_version: "3.12"
+sdk_version: 6.2.0
 app_file: app.py
+startup_duration_timeout: 1h
 pinned: false
+short_description: Generate the missing stem that fits your audio
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Stem-0 · missing-stem generation
+
+Give it an audio **context** (e.g. drums+bass, or a full instrumental) and it generates the
+**missing stem** so that it fits — tempo, key, arrangement. For vocals you supply the lyrics.
+
+Base model: `ACE-Step/acestep-v15-xl-base` · LoRA: `fcolooo/stem-0-r128` (r128, context-contrastive recipe).
+
+**ZeroGPU note:** the first generation loads ~20 GB of weights onto the GPU and is slow; later ones are fast.
+Keep `steps` low while testing — GPU quota is per-day.

app.py

@@ -0,0 +1,743 @@
+"""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()

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,20 @@
+# NOTE: deliberately NOT installing ace-step as a package, and NOT installing torch/gradio.
+# ace-step pins gradio==6.2.0 and exact torch==2.10.0+cu128 wheels; reinstalling torch inside a
+# ZeroGPU Space fights the torch build that ZeroGPU patches. We vendor the acestep SOURCE at
+# startup instead (app.py) and rely on the Space's own torch + gradio.
+torchaudio
+transformers>=4.51.0,<4.58.0
+diffusers>=0.37.0
+accelerate>=1.0.0
+peft>=0.13
+safetensors
+soundfile
+librosa
+scipy
+loguru
+einops
+diskcache
+vector-quantize-pytorch
+numba
+huggingface_hub>=0.25
+spaces

run_ace_task_baseline.py

@@ -0,0 +1,420 @@
+#!/usr/bin/env python3
+"""Run ACE-Step for one prepared pack item — the proven reference implementation.
+
+This file is the ONLY place GenerationParams is constructed. That is not a style
+preference, it is the lesson from 2026-07-16/17: apps/stack/infer_worker.py was a
+second, hand-typed copy of this call, it drifted in six places (global_caption='',
+audio_cover_strength=0.0, hardcoded bpm, wrong captions, timesignature='', steps=16),
+and it cost Francesco a night of "still crappy" output. Two hand-typed copies of the
+same call drift forever and you can never prove you have found the last divergence.
+One implementation is falsifiable; two are not.
+
+Three entry points, one code path:
+
+  init_ace()    load the model ONCE. Expensive (~20-40s: 4B decoder).
+  build_params() metadata.json -> GenerationParams. THE single construction.
+  run_take()    one generation against an already-loaded session. ~7s warm.
+  main()        the CLI, unchanged. scripts/stemgen_webapp.py and apps/stems shell
+                out to it and must keep working — it is the reference for
+                docs/INFERENCE_RECIPE.md.
+
+The CLI does init_ace() + run_take() and exits, so it pays the load every time. A
+resident caller (apps/stack/infer_worker.py) does init_ace() once and then run_take()
+per request — the same functions, so it cannot drift from the CLI by construction.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+# =========================================================================
+# THE MAY 7 RECIPE — the defaults, living in the ONE implementation.
+#
+# Provenance: artifacts/generated_batches/batch_10_actual_lego_20260507/*/
+# ace_lego_corrected/result.json — the last batch Francesco confirmed sounds good.
+# Copied, not invented. docs/INFERENCE_RECIPE.md is the writeup.
+#
+# WHY THESE LIVE HERE AND NOT IN A CALLER (found 2026-07-17, live):
+# apps/stack/infer_worker.py held these as its own constants. When serving switched
+# to shelling out to this script, the worker stopped being the path — and the recipe
+# went with it, because this script reads caption/global_caption from metadata.json
+# and the server's metadata.json has no `global_caption` key. So serving silently
+# reverted to global_caption="" — the single biggest prompt-side defect — with the
+# recipe still sitting "restored" in a file nothing called.
+#
+# A default in a caller is a default that gets lost. These are the defaults now, so
+# every caller (CLI, studio, resident worker) gets the proven recipe unless its
+# metadata deliberately overrides it.
+# =========================================================================
+
+# The instruction that makes the model emit a STEM rather than a mix.
+RECIPE_GLOBAL_CAPTION = (
+    "Generate only the requested missing isolated stem so that it fits the provided "
+    "audio context. Preserve timing, style, tempo, harmony, and arrangement. Do not "
+    "generate a full mix."
+)
+
+# Captions in ACE-Step's own register: short, natural descriptions of the target
+# stem — the way ACE was trained, not verbose instructions with negation lists. The
+# "isolated stem, no full mix" intent is carried by RECIPE_GLOBAL_CAPTION, so the
+# per-role caption just names the instrument and how it sits. (Was: long prose with
+# "Do not generate bass/guitars/vocals…", out of ACE's training distribution —
+# Francesco 2026-07-17. The proven long form is preserved in git if we A/B back.)
+RECIPE_CAPTIONS = {
+    "drums": "a tight drum kit locked to the groove and tempo",
+    "bass": "a groovy bass line locked to the drums and harmony",
+    "melody": "a melodic lead line locked to the harmony",
+    "vocals": "a lead vocal locked to the melody and phrasing",
+}
+RECIPE_TIMESIGNATURE = "4"
+
+
+def _meta_float(meta: dict, key: str):
+    value = meta.get(key)
+    if value in (None, "", "N/A"):
+        return None
+    try:
+        return float(value)
+    except Exception:
+        return None
+
+
+def _meta_int(meta: dict, key: str):
+    value = _meta_float(meta, key)
+    return int(round(value)) if value else None
+
+
+@dataclass
+class AceSession:
+    """A loaded ACE model. Hold one of these and run_take() is ~7s instead of ~45s."""
+    dit_handler: Any
+    llm_handler: Any
+    device: str
+    lora_path: str | None = None
+    lora_scale: float = 1.0
+    full_ft_checkpoint: str | None = None
+    load_seconds: float = 0.0
+
+
+def init_ace(
+    *,
+    ace_root: str = "/home/fcolo/ace-step-1.5-xl",
+    checkpoints: str = "/home/fcolo/ace-step/checkpoints",
+    model: str = "acestep-v15-xl-base",
+    device: str = "cuda",
+    lm_model: str = "acestep-5Hz-lm-1.7B",
+    lm_backend: str = "pt",
+    no_thinking: bool = True,
+    use_lm: bool = False,
+    use_cot: bool = False,
+    full_ft_checkpoint: str | None = None,
+    lora_path: str | None = None,
+    adapter_name: str = "stemgen",
+    lora_scale: float = 1.0,
+    log=print,
+) -> AceSession:
+    """Everything expensive, once. Safe to call from a long-lived process."""
+    if full_ft_checkpoint and lora_path:
+        raise ValueError("--full-ft-checkpoint and --lora-path are mutually exclusive")
+
+    t0 = time.time()
+    # Hard CPU mode: ACE/PEFT sometimes tries to stage LoRA weights on CUDA even when
+    # the requested generation device is CPU. Hide CUDA before importing ACE/torch so
+    # CPU jobs don't fight long-running GPU work.
+    if str(device).lower().split(":", 1)[0] == "cpu":
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""
+        os.environ.setdefault("ACESTEP_VAE_ON_CPU", "1")
+        os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+        log("[INFO] CPU mode: CUDA_VISIBLE_DEVICES cleared for this process")
+
+    ace_root_p = Path(ace_root).resolve()
+    if str(ace_root_p) not in sys.path:
+        sys.path.insert(0, str(ace_root_p))
+    os.environ["ACESTEP_CHECKPOINTS_DIR"] = checkpoints
+
+    from acestep.handler import AceStepHandler
+    from acestep.llm_inference import LLMHandler
+
+    dit_handler = AceStepHandler()
+    status, success = dit_handler.initialize_service(
+        project_root=str(ace_root_p),
+        config_path=model,
+        device=device,
+        prefer_source="huggingface",
+    )
+    if not success:
+        raise RuntimeError(f"ACE init failed: {status}")
+    log(status)
+
+    if full_ft_checkpoint:
+        # Full fine-tune: swap the whole DiT decoder. Traced 2026-07-17 --
+        # initialize_service() -> init_service_loader.py:175 sets
+        # self.model = AutoModel.from_pretrained(...), the SAME construction as
+        # training_v2/model_loader.py:load_decoder_for_training(); and
+        # AceStepConditionGenerationModel.__init__ (modeling_acestep_v15_base.py:1609)
+        # sets self.decoder = AceStepDiTModel(config). So handler.model.decoder is
+        # exactly the submodule train_full.py checkpoints.
+        from safetensors.torch import load_file
+
+        ckpt = Path(full_ft_checkpoint) / "decoder_model.safetensors"
+        if not ckpt.is_file():
+            raise RuntimeError(f"no decoder_model.safetensors under {full_ft_checkpoint}")
+        target = getattr(getattr(dit_handler, "model", None), "decoder", None)
+        if target is None:
+            raise RuntimeError("handler.model.decoder not found after initialize_service — "
+                               "ACE internals changed; re-trace before trusting this path.")
+        # Fingerprint before/after: a load that silently no-ops (wrong keys, empty
+        # dict) is otherwise indistinguishable from success.
+        probe = next(k for k, _ in target.named_parameters())
+        before = float(dict(target.named_parameters())[probe].detach().float().sum().item())
+        sd = load_file(str(ckpt))
+        target.load_state_dict(sd, strict=True)  # strict: any mismatch raises
+        after = float(dict(target.named_parameters())[probe].detach().float().sum().item())
+        if before == after:
+            raise RuntimeError(f"decoder weights UNCHANGED after load_state_dict (probe {probe} "
+                               f"sum {before}); refusing to run a checkpoint that did not apply.")
+        target.eval()
+        log(json.dumps({"full_ft_loaded": str(ckpt), "probe": probe,
+                        "sum_before": before, "sum_after": after, "tensors": len(sd)}))
+
+    if lora_path:
+        lora_status = dit_handler.add_lora(lora_path, adapter_name=adapter_name)
+        log(lora_status)
+        if not str(lora_status).startswith("✅"):
+            raise RuntimeError(f"ACE LoRA load failed: {lora_status}")
+        log(dit_handler.set_lora_scale(adapter_name, lora_scale))
+        log(dit_handler.set_use_lora(True))
+
+    llm_handler = None
+    if use_lm or use_cot or not no_thinking:
+        llm_handler = LLMHandler()
+        lm_status, lm_success = llm_handler.initialize(
+            checkpoint_dir=checkpoints,
+            lm_model_path=lm_model,
+            backend=lm_backend,
+            device=device,
+            offload_to_cpu=False,
+        )
+        if not lm_success:
+            raise RuntimeError(f"ACE LM init failed: {lm_status}")
+        log(lm_status)
+
+    return AceSession(
+        dit_handler=dit_handler,
+        llm_handler=llm_handler,
+        device=device,
+        lora_path=lora_path,
+        lora_scale=lora_scale,
+        full_ft_checkpoint=full_ft_checkpoint,
+        load_seconds=round(time.time() - t0, 2),
+    )
+
+
+def build_params(
+    meta: dict,
+    item_dir: Path,
+    *,
+    task: str = "lego",
+    steps: int = 64,
+    seed: int = 1234,
+    guidance_scale: float = 7.0,
+    cover_strength: float = 0.45,
+    no_thinking: bool = True,
+    use_cot: bool = False,
+    retake_seed: "int | None" = None,
+    retake_variance: float = 0.0,
+):
+    """metadata.json -> (GenerationParams, GenerationConfig).
+
+    ⚠️ THE SINGLE CONSTRUCTION. Every caller — CLI, resident worker, studio — comes
+    through here. Do not copy this into another file; import it. See the module
+    docstring for what a second copy cost.
+
+    Values default to docs/INFERENCE_RECIPE.md (the May 7 recipe); metadata.json
+    overrides where it carries a value. Note what the meta deliberately may omit:
+    bpm absent means lego locks tempo from src_audio itself, which is correct — a
+    GUESSED bpm is worse than none (the app sent a hardcoded 98 against audio at
+    119/170 and the model dutifully played out of tempo).
+    """
+    from acestep.inference import GenerationParams, GenerationConfig
+
+    role = meta.get("role", "")
+    ace_role = "guitar" if role == "melody" else role
+    # Recipe defaults, overridable by metadata. `or` not `.get(k, default)`: an empty
+    # string in the metadata means "absent", not "deliberately empty" — and empty is
+    # exactly the failure this guards.
+    caption = (meta.get("ace_caption") or meta.get("prompt") or RECIPE_CAPTIONS.get(role)
+               or f"Add {role} for this song.")
+    global_caption = meta.get("global_caption") or RECIPE_GLOBAL_CAPTION
+    source = item_dir / meta.get("source_audio", "context_mix_minus_target.wav")
+
+    if task == "lego":
+        instruction = f"Generate the {ace_role.upper()} track based on the audio context:"
+    elif task == "complete":
+        classes = meta.get("complete_track_classes") or [ace_role]
+        instruction = "Complete the input track with " + " | ".join(str(c).upper() for c in classes) + ":"
+    else:
+        instruction = "Generate audio semantic tokens based on the given conditions:"
+
+    params = GenerationParams(
+        task_type=task,
+        src_audio=str(source),
+        instruction=instruction,
+        caption=caption,
+        global_caption=global_caption,
+        lyrics=meta.get("lyrics", "[Instrumental]"),
+        instrumental=bool(meta.get("instrumental", True)),
+        vocal_language=meta.get("vocal_language", "unknown"),
+        bpm=_meta_int(meta, "bpm"),
+        keyscale=meta.get("keyscale") or meta.get("key") or "",
+        timesignature=str(meta.get("timesignature") or meta.get("time_signature") or RECIPE_TIMESIGNATURE),
+        duration=_meta_float(meta, "duration_seconds") or -1.0,
+        repainting_start=0.0,
+        repainting_end=-1,
+        # Retake: variance-preserving variation. With a FIXED base `seed` per part and a
+        # small `retake_variance`, each new sample is the SAME part subtly evolved — not
+        # an unrelated diffusion draw. This is how a part 'continues from what it was'.
+        retake_seed=retake_seed,
+        retake_variance=float(retake_variance or 0.0),
+        inference_steps=steps,
+        seed=seed,
+        thinking=not no_thinking,
+        guidance_scale=guidance_scale,
+        audio_cover_strength=cover_strength,
+        use_cot_metas=use_cot,
+        use_cot_caption=False,   # production engine.py uses metas + language, NOT caption
+        use_cot_language=use_cot,
+        shift=3.0,               # match production engine.py (was unset -> default)
+        dcw_enabled=False,       # #1255 NOISE FIX (production default) — the un-denoised culprit
+        cover_noise_strength=0.0,
+        use_adg=False,
+        use_constrained_decoding=True,
+        enable_normalization=True,
+        normalization_db=-1.0,
+        cfg_interval_start=0.0,
+        cfg_interval_end=1.0,
+    )
+    config = GenerationConfig(batch_size=1, use_random_seed=False, seeds=[seed], audio_format="wav")
+    return params, config
+
+
+def run_take(
+    session: AceSession,
+    item_dir: Path,
+    *,
+    task: str = "lego",
+    steps: int = 64,
+    seed: int = 1234,
+    guidance_scale: float = 7.0,
+    cover_strength: float = 0.45,
+    no_thinking: bool = True,
+    use_cot: bool = False,
+    out_dir: Path | None = None,
+    generated_name: str | None = None,
+    retake_seed: "int | None" = None,
+    retake_variance: float = 0.0,
+    log=print,
+) -> dict:
+    """One generation against an already-loaded session. Warm: ~7s.
+
+    Identical to what the CLI does per item — because the CLI calls this.
+    """
+    from acestep.inference import generate_music
+
+    item_dir = Path(item_dir).resolve()
+    meta = json.loads((item_dir / "metadata.json").read_text())
+    out_dir = Path(out_dir).resolve() if out_dir else item_dir / f"ace_{task}_corrected"
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    params, config = build_params(
+        meta, item_dir, task=task, steps=steps, seed=seed,
+        guidance_scale=guidance_scale, cover_strength=cover_strength,
+        no_thinking=no_thinking, use_cot=use_cot,
+        retake_seed=retake_seed, retake_variance=retake_variance,
+    )
+    t0 = time.time()
+    result = generate_music(session.dit_handler, session.llm_handler, params, config, save_dir=str(out_dir))
+    (out_dir / "result.json").write_text(
+        json.dumps(result.to_dict() if hasattr(result, "to_dict") else result.__dict__, indent=2, default=str) + "\n"
+    )
+    if not result.success:
+        raise RuntimeError(result.error or result.status_message)
+
+    generated = result.audios[0]["path"]
+    target_name = generated_name or ("generated_full_ace.wav" if task == "cover" else "generated_stem_ace.wav")
+    copied_to = item_dir / target_name
+    shutil.copy2(generated, copied_to)
+    return {
+        "generated": generated,
+        "copied_to": str(copied_to),
+        "task": task,
+        "lora_path": session.lora_path,
+        "lora_scale": session.lora_scale,
+        "seed": seed,
+        "steps": steps,
+        "guidance_scale": guidance_scale,
+        "cover_strength": cover_strength,
+        "generate_seconds": round(time.time() - t0, 2),
+    }
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="Run ACE-Step corrected baseline task for one prepared pack item.")
+    p.add_argument("item_dir", type=Path)
+    p.add_argument("--task", choices=["lego", "complete", "cover"], default="lego")
+    p.add_argument("--ace-root", default="/home/fcolo/ace-step-1.5-xl")
+    p.add_argument("--checkpoints", default="/home/fcolo/ace-step/checkpoints")
+    p.add_argument("--model", default="acestep-v15-xl-base")
+    p.add_argument("--lm-model", default="acestep-5Hz-lm-1.7B")
+    p.add_argument("--lm-backend", default="pt", choices=["pt", "vllm"])
+    p.add_argument("--device", default="cuda")
+    p.add_argument("--steps", type=int, default=64)
+    p.add_argument("--seed", type=int, default=1234)
+    p.add_argument("--guidance-scale", type=float, default=7.0)
+    p.add_argument("--cover-strength", type=float, default=0.45)
+    p.add_argument("--no-thinking", action="store_true")
+    p.add_argument("--use-lm", action="store_true", help="Initialize/use ACE 5Hz LM. Off by default for source-conditioned lego because CoT can rewrite the stem prompt.")
+    p.add_argument("--use-cot", action="store_true", help="Allow LM CoT to rewrite/fill caption/language/metas. Off by default to preserve explicit stem prompts.")
+    p.add_argument("--out-dir", type=Path, default=None, help="ACE raw output directory. Defaults to item_dir/ace_<task>_corrected.")
+    p.add_argument("--generated-name", default=None, help="Name to copy the generated wav to in item_dir. Defaults to generated_stem_ace.wav or generated_full_ace.wav.")
+    p.add_argument("--full-ft-checkpoint", default=None,
+                   help="Directory holding decoder_model.safetensors from a FULL fine-tune "
+                        "(e.g. artifacts/train/full_finetune_v4_.../best). Replaces the whole DiT "
+                        "decoder. Mutually exclusive with --lora-path: an adapter trained against "
+                        "the ORIGINAL decoder stacked on replaced weights is silently wrong, so "
+                        "passing both is refused rather than combined.")
+    p.add_argument("--lora-path", default=None, help="Optional PEFT LoRA adapter directory to load after base ACE init.")
+    p.add_argument("--adapter-name", default="stemgen", help="Adapter name used when loading --lora-path.")
+    p.add_argument("--lora-scale", type=float, default=1.0)
+    args = p.parse_args()
+
+    if args.full_ft_checkpoint and args.lora_path:
+        raise SystemExit("--full-ft-checkpoint and --lora-path are mutually exclusive")
+
+    session = init_ace(
+        ace_root=args.ace_root, checkpoints=args.checkpoints, model=args.model,
+        device=args.device, lm_model=args.lm_model, lm_backend=args.lm_backend,
+        no_thinking=args.no_thinking, use_lm=args.use_lm, use_cot=args.use_cot,
+        full_ft_checkpoint=args.full_ft_checkpoint, lora_path=args.lora_path,
+        adapter_name=args.adapter_name, lora_scale=args.lora_scale,
+    )
+    out = run_take(
+        session, args.item_dir, task=args.task, steps=args.steps, seed=args.seed,
+        guidance_scale=args.guidance_scale, cover_strength=args.cover_strength,
+        no_thinking=args.no_thinking, use_cot=args.use_cot,
+        out_dir=args.out_dir, generated_name=args.generated_name,
+    )
+    # Same shape the CLI has always printed. Callers parse this.
+    print(json.dumps({k: out[k] for k in ("generated", "copied_to", "task", "lora_path", "lora_scale")}, indent=2))
+
+
+if __name__ == "__main__":
+    main()