"""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()