"""engine.py — Stable Audio 3 Small Music continuation core. CODA's job is one thing done well: take a short, unfinished clip and continue it into a finished-sounding track in the same key, tempo and feel. SA3 does that in a SINGLE call. Its `generate_diffusion_cond_inpaint` is a native audio-inpainting diffusion sampler: place the user's clip at the front of the buffer, mask the region after it, and the model fills the masked region conditioned on the kept audio — true long-form continuation, 44.1 kHz stereo, no multi-pass chaining, no energy guards, no re-roll logic. This module is the whole generation core. It returns ONLY the newly generated tail plus the source length in seconds; `stitch.py` joins that tail onto the user's *pristine* original so the real recording (and any vocals) plays untouched up to the seam. Bounded lead-in (the deployed-bug fix): SA3 only needs a short run-up to know where the song is going. We therefore condition on at most MAX_LEAD_SECONDS of the clip's TAIL, not the whole clip. Feeding a long clip (e.g. 100 s) into the buffer and masking only a few seconds makes the 8-step distilled sampler collapse to near-silence in that tiny window — the bug that shipped. A bounded lead keeps the masked (generated) region substantial and healthy, and because stitch rejoins the tail onto the full pristine original, the listener still hears their entire clip before the seam. Mask convention (verified against the installed library source): inpaint_mask = ones(buffer); inpaint_mask[start:end] = 0 -> 1 = keep the input audio, 0 = generate. We place `lead` seconds of source at the front and mask [lead, lead+new], so SA3 keeps the lead and generates a fresh `new`-second tail that continues from the clip's end. """ import numpy as np import torch MODEL_ID = "stabilityai/stable-audio-3-small-music" SR = 44100 # SA3 native sample rate (model_config: sample_rate) STEPS = 8 # SA3 Small is an 8-step adversarially-distilled model SAMPLER = "pingpong" # the sampler the distilled model was tuned for DEFAULT_CFG = 1.0 # distilled-model guidance; the prompt still conditions # at 1.0 (CFG amplification off, conditional path on) MAX_TOTAL_SECONDS = 120 # SA3 Small duration cap (sample_size / sample_rate) MIN_NEW_SECONDS = 5 # below this a "continuation" isn't worth a GPU call MAX_LEAD_SECONDS = 30 # how much of the clip's TAIL to feed SA3 as run-up. # SA3 generates a healthy continuation from a bounded # lead-in; keeping a very long source in the buffer and # masking only a few seconds makes the distilled sampler # produce near-silence. 30 s is inside the model's # healthy range (the verified lab takes used ~30 s leads) # and the splice restores the full clip anyway. _model = None _model_config = None _sample_size = None _on_device = None # which device the weights currently live on def _device(): return "cuda" if torch.cuda.is_available() else "cpu" def preload(): """Load model + autoencoder + T5Gemma conditioner into CPU RAM at process start. On ZeroGPU the per-call GPU window is the scarce resource, so weights must come off disk at boot, not inside the window. The CUDA placement + fp16 cast is deferred to the first `continue_audio` call (i.e. the @spaces.GPU window), matching how Stability's own Space defers it.""" global _model, _model_config, _sample_size if _model is None: from stable_audio_tools import get_pretrained_model _model, _model_config = get_pretrained_model(MODEL_ID) _sample_size = int(_model_config["sample_size"]) _model.eval() print(f"[coda] preload: SA3 resident " f"(sr={_model_config['sample_rate']}, " f"sample_size={_sample_size} = " f"{_sample_size / int(_model_config['sample_rate']):.0f}s)", flush=True) return _model, _model_config def _ensure_on_device(): """Ensure weights are on the GPU in fp16. Called inside the @spaces.GPU window on every generation. `.to()` is a cheap no-op when the model is already placed, so we re-ensure each call rather than caching device state — that stays correct even if ZeroGPU detaches the GPU between calls. fp16 is only valid on CUDA; on CPU the model stays fp32.""" global _model, _on_device dev = _device() _model = _model.to(dev) if dev == "cuda": _model = _model.to(torch.float16) _on_device = dev return _model def _load_source(clip_path): """Load the clip as stereo float32 @44.1k as a (2, N) tensor. SA3's autoencoder is stereo; `prepare_audio` inside the sampler will pad/crop to the buffer length and place this at the FRONT (PadCrop, randomize=False).""" import librosa y, _ = librosa.load(clip_path, sr=SR, mono=False) y = np.asarray(y, dtype=np.float32) if y.ndim == 1: y = np.stack([y, y]) # mono -> stereo elif y.shape[0] > 2: y = y[:2] return torch.from_numpy(np.ascontiguousarray(y)) #: longest clip that still leaves room for MIN_NEW of continuation under the cap MAX_SOURCE_SECONDS = MAX_TOTAL_SECONDS - MIN_NEW_SECONDS def plan_continuation(source_seconds, total_seconds): """Pure helper (unit-testable, no model): turn a (source, requested-total) pair into the SA3 generation buffer and return (lead, new_seconds, buffer). - `lead` : seconds of the clip's TAIL used as run-up context, capped at MAX_LEAD_SECONDS so a long clip can't drown the masked region. - `new` : seconds of fresh audio to generate. We extend to the requested finished length (`total - source`), floored at MIN_NEW_SECONDS so every call earns its GPU time, and bounded so the buffer (lead + new) never exceeds SA3's MAX_TOTAL_SECONDS cap. - `buffer` : lead + new, i.e. the full generation buffer. The mask runs [lead, buffer]; buffer > lead always, so it never inverts. Raises ValueError only for a clip longer than MAX_SOURCE_SECONDS — at that point it's a full track, not an unfinished clip to continue. """ source_seconds = float(source_seconds) total_seconds = float(total_seconds) if source_seconds > MAX_SOURCE_SECONDS: raise ValueError( f"clip is {source_seconds:.0f}s — that's a finished-length track, " f"not an unfinished clip. CODA continues clips up to " f"{MAX_SOURCE_SECONDS:.0f}s; trim it shorter and re-upload.") total_seconds = min(total_seconds, MAX_TOTAL_SECONDS) lead = min(source_seconds, MAX_LEAD_SECONDS) # extend to the requested finished length; floor at MIN_NEW, and never let # lead + new exceed the buffer cap. new_seconds = max(total_seconds - source_seconds, MIN_NEW_SECONDS) new_seconds = min(new_seconds, MAX_TOTAL_SECONDS - lead) buffer_seconds = lead + new_seconds return lead, new_seconds, buffer_seconds def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG, seed=-1, progress=None): """Continue `clip_path` up to `total_seconds` in one SA3 inpaint call. Returns (new_tail, source_seconds, SR) where: new_tail : (2, M) float32 @44.1k — ONLY the generated region [source_end, total]. Peak-normalized to <= 1.0. source_seconds : the clip's true length (the splice boundary, in seconds) SR : 44100 `progress(stage_name)` is called (best-effort) at each stage so the UI can paint a live status. SA3 is one diffusion call, so progress is stage-based. """ from einops import rearrange from stable_audio_tools.inference.generation import ( generate_diffusion_cond_inpaint) def _notify(stage): if progress is not None: try: progress(stage) except Exception as e: print(f"[coda] progress callback failed ({e})", flush=True) preload() model = _ensure_on_device() dev = _device() # The library does `np.random.randint(0, 2**32-1)` when seed == -1, which # overflows int32 on Windows/numpy<2. Draw a safe seed ourselves so the # default path works everywhere, not just on the Linux Space. if seed is None or seed < 0: seed = int(np.random.randint(0, 2 ** 31 - 1)) _notify("reading") source = _load_source(clip_path) source_seconds = source.shape[-1] / SR lead, new_seconds, buffer_seconds = plan_continuation( source_seconds, total_seconds) # condition on only the TAIL `lead` seconds of the clip. This is the bug fix: # a long source no longer fills the buffer and starves the masked region. lead_samples = min(int(round(lead * SR)), source.shape[-1]) lead_audio = source[:, -lead_samples:] # the autoencoder runs in the model's dtype (fp16 on CUDA); the conditioning # audio must match it or the encoder's conv1d rejects the input dtype. model_dtype = next(model.model.parameters()).dtype lead_audio = lead_audio.to(model_dtype) mask_start, mask_end = lead, buffer_seconds prompt = (prompt or "").strip() print(f"[coda] continuation: source={source_seconds:.1f}s, " f"lead={lead:.1f}s -> buffer={buffer_seconds:.1f}s " f"(+{new_seconds:.1f}s new), mask=[{mask_start:.1f}s, {mask_end:.1f}s], " f"steps={STEPS}, cfg={cfg_scale}, prompt={prompt!r}", flush=True) _notify("composing") with torch.no_grad(): output = generate_diffusion_cond_inpaint( model, steps=STEPS, cfg_scale=cfg_scale, conditioning=[{"prompt": prompt, "seconds_total": buffer_seconds}], sample_size=_sample_size, sampler_type=SAMPLER, inpaint_audio=(SR, lead_audio), inpaint_mask_start_seconds=mask_start, inpaint_mask_end_seconds=mask_end, seed=seed, device=dev, ) _notify("finalizing") # (b, d, n) -> (d, b*n) output = rearrange(output, "b d n -> d (b n)") audio = output.to(torch.float32).cpu().numpy() if audio.shape[0] == 1: # safety: ensure stereo audio = np.repeat(audio, 2, axis=0) # the generated region is [lead, buffer]; slice it out first, THEN normalize # by the tail's OWN peak. Normalizing the whole buffer (as before) let a loud # lead transient divide the tail down toward silence; per-tail normalization # returns the continuation at a healthy standalone level and stitch re-levels # it to the seam. start = int(round(lead * SR)) end = min(int(round(buffer_seconds * SR)), audio.shape[-1]) new_tail = np.ascontiguousarray(audio[:, start:end].astype(np.float32)) peak = float(np.abs(new_tail).max()) if peak > 1e-9: new_tail = new_tail / peak print(f"[coda] generated tail: shape={new_tail.shape} " f"({new_tail.shape[-1] / SR:.1f}s), peak after norm " f"{float(np.abs(new_tail).max()):.3f}, " f"rms {float(np.sqrt(np.mean(new_tail ** 2))):.3f}", flush=True) return new_tail, source_seconds, SR