Spaces:
Running on Zero
Running on Zero
| """abv1 SA3 engine — shared Stable Audio 3 backend for the music-AI demo apps. | |
| Two surfaces on one Space: | |
| * **UI** — Text→audio and Audio→audio tabs for humans (visitors burn their own | |
| ZeroGPU quota, so these are never token-gated). | |
| * **Headless API** — ``/generate`` and ``/generate_a2a``, registered off hidden | |
| components so gradio_client / plain HTTP can drive batch jobs (the trend-radio | |
| cron) and the demo pages' server side. | |
| Model wiring follows the official ``stabilityai/stable-audio-3`` Space exactly: | |
| ``stable_audio_tools.get_pretrained_model`` → ``.to("cuda").to(float16)`` at | |
| module level, ``generate_diffusion_cond_inpaint`` under a bare ZeroGPU shim. | |
| Set ``SA3_STUB=1`` to run the whole app (UI + both endpoints) with no torch and | |
| no weights — gens return a synthesized placeholder WAV. That is the only way to | |
| exercise this file on a laptop. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import random | |
| import tempfile | |
| import time | |
| from dataclasses import dataclass | |
| from typing import Optional | |
| STUB = os.environ.get("SA3_STUB", "").strip() in ("1", "true", "yes") | |
| API_TOKEN = os.environ.get("SA3_API_TOKEN", "").strip() | |
| # ZeroGPU shim. Absent locally, so fall back to a decorator that ignores its | |
| # kwargs and returns the function untouched. | |
| try: | |
| import spaces # type: ignore | |
| except Exception: # pragma: no cover — local / stub runs | |
| class _NoSpaces: | |
| def GPU(*d_args, **d_kwargs): | |
| if d_args and callable(d_args[0]): | |
| return d_args[0] | |
| def _wrap(fn): | |
| return fn | |
| return _wrap | |
| spaces = _NoSpaces() # type: ignore | |
| import gradio as gr | |
| import numpy as np | |
| import soundfile as sf | |
| # --------------------------------------------------------------------------- | |
| # Variants | |
| # --------------------------------------------------------------------------- | |
| class Variant: | |
| key: str | |
| repo: str | |
| label: str | |
| # Hard cap we expose in the UI/API. The real ceiling is the model's | |
| # sample_size // sample_rate; we take the min of the two at load time. | |
| max_seconds: int | |
| default_seconds: int | |
| placeholder: str | |
| VARIANTS: list[Variant] = [ | |
| Variant( | |
| key="small-music", | |
| repo="stabilityai/stable-audio-3-small-music", | |
| label="Small Music — 0.6B, fast (seconds per gen)", | |
| max_seconds=120, | |
| default_seconds=30, | |
| placeholder="Cinematic neo-soul groove with electric piano, brushed drums, " | |
| "walking upright bass, smoky vibe 92 BPM", | |
| ), | |
| Variant( | |
| key="medium", | |
| repo="stabilityai/stable-audio-3-medium", | |
| label="Medium — general audio, higher quality", | |
| max_seconds=180, | |
| default_seconds=45, | |
| placeholder="A dream-like synthpop instrumental for a surrealist dream " | |
| "sequence, warm analog pads, 120 BPM", | |
| ), | |
| ] | |
| BY_KEY = {v.key: v for v in VARIANTS} | |
| VARIANT_CHOICES = [(v.label, v.key) for v in VARIANTS] | |
| DEFAULT_VARIANT = "small-music" | |
| STUB_SAMPLE_RATE = 44100 | |
| # --------------------------------------------------------------------------- | |
| # Model preload (skipped entirely in stub mode) | |
| # --------------------------------------------------------------------------- | |
| class Loaded: | |
| variant: Variant | |
| model: object | |
| sample_rate: int | |
| sample_size: int | |
| max_seconds: int | |
| LOADED: dict[str, Loaded] = {} | |
| def _ensure_stable_audio_tools() -> None: | |
| import subprocess | |
| import sys | |
| try: | |
| import stable_audio_tools # noqa: F401 | |
| return | |
| except ImportError: | |
| pass | |
| # stable-audio-tools strict-pins torch==2.7.1, which lacks sm_120 kernels. | |
| # Install --no-deps; its transitive deps are in requirements.txt. | |
| print("[startup] installing stable-audio-tools (--no-deps) …", flush=True) | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", "--quiet", "--no-deps", | |
| "stable-audio-tools"], | |
| ) | |
| if not STUB: | |
| _ensure_stable_audio_tools() | |
| import torch | |
| import torchaudio | |
| from einops import rearrange | |
| from stable_audio_tools import get_pretrained_model | |
| from stable_audio_tools.inference.generation import generate_diffusion_cond_inpaint | |
| for _v in VARIANTS: | |
| print(f"[startup] loading {_v.repo} …", flush=True) | |
| _t0 = time.time() | |
| _model, _config = get_pretrained_model(_v.repo) | |
| _sr = int(_config["sample_rate"]) | |
| _ss = int(_config["sample_size"]) | |
| _model = _model.to("cuda").to(torch.float16) | |
| LOADED[_v.key] = Loaded( | |
| variant=_v, | |
| model=_model, | |
| sample_rate=_sr, | |
| sample_size=_ss, | |
| max_seconds=min(_v.max_seconds, _ss // _sr), | |
| ) | |
| print(f"[startup] {_v.key} ready in {time.time() - _t0:.1f}s · " | |
| f"sr={_sr} · sample_size={_ss} (~{_ss // _sr}s model max)", flush=True) | |
| else: | |
| print("[startup] SA3_STUB=1 — no torch, no weights, placeholder audio only.", | |
| flush=True) | |
| def variant_max_seconds(key: str) -> int: | |
| if key in LOADED: | |
| return LOADED[key].max_seconds | |
| return BY_KEY[key].max_seconds if key in BY_KEY else 120 | |
| # --------------------------------------------------------------------------- | |
| # Input hygiene | |
| # --------------------------------------------------------------------------- | |
| def _clamp_inputs(variant_key: str, seconds, steps, cfg_scale, seed): | |
| key = (variant_key or DEFAULT_VARIANT).strip() | |
| if key not in BY_KEY: | |
| raise gr.Error(f"Unknown model_variant {variant_key!r}. " | |
| f"Use one of: {', '.join(BY_KEY)}") | |
| seconds = max(1, min(int(float(seconds or 30)), variant_max_seconds(key))) | |
| steps = max(1, min(int(float(steps or 8)), 100)) | |
| cfg_scale = max(0.0, min(float(cfg_scale if cfg_scale is not None else 1.0), 15.0)) | |
| seed = int(float(seed if seed is not None else -1)) | |
| if seed < 0: | |
| seed = random.randint(0, 2**31 - 1) | |
| return key, seconds, steps, cfg_scale, seed | |
| def _check_token(token: Optional[str]) -> None: | |
| """API-only gate. Unset secret == open (dev). UI paths never call this.""" | |
| if not API_TOKEN: | |
| return | |
| if (token or "").strip() != API_TOKEN: | |
| raise gr.Error("Bad or missing API token.") | |
| def _gpu_seconds(seconds: int) -> int: | |
| """ZeroGPU budget request. small-music does 120s of audio in <2s of GPU and | |
| medium a few seconds; 20s base + a quarter of the requested length is ample | |
| headroom without hogging a slot.""" | |
| return int(min(120, 20 + int(seconds) // 4)) | |
| # --------------------------------------------------------------------------- | |
| # Generation | |
| # --------------------------------------------------------------------------- | |
| def _stub_wav(seconds: int, seed: int) -> str: | |
| """Band-limited sweep + noise bed so the shape of the payload (stereo wav, | |
| right duration, real sample rate) matches a genuine gen.""" | |
| rng = np.random.default_rng(seed) | |
| sr = STUB_SAMPLE_RATE | |
| n = int(seconds * sr) | |
| t = np.arange(n, dtype=np.float32) / sr | |
| f0, f1 = 110.0, 1760.0 | |
| phase = 2 * np.pi * f0 * seconds / np.log(f1 / f0) * ( | |
| np.power(f1 / f0, t / max(seconds, 1e-6)) - 1.0 | |
| ) | |
| sweep = 0.35 * np.sin(phase).astype(np.float32) | |
| noise = 0.05 * rng.standard_normal(n).astype(np.float32) | |
| # cheap one-pole lowpass on the noise so it isn't harsh white | |
| for _ in range(2): | |
| noise = np.convolve(noise, np.ones(16, dtype=np.float32) / 16, mode="same") | |
| env = np.minimum(1.0, np.minimum(t * 4.0, (seconds - t) * 4.0)).astype(np.float32) | |
| mono = (sweep + noise) * env | |
| stereo = np.stack([mono, np.roll(mono, 64)], axis=1) | |
| out_path = os.path.join(tempfile.mkdtemp(), "sa3_stub.wav") | |
| sf.write(out_path, stereo, sr, subtype="PCM_16") | |
| return out_path | |
| def _load_init_audio(path: str, target_sr: int, dtype): | |
| """Filepath → (sample_rate, tensor[C,N]) at the model's rate and dtype. | |
| Decode via soundfile, not torchaudio.load — torchaudio 2.x delegates | |
| load() to torchcodec, which isn't in the ZeroGPU image (same trap as the | |
| audiofix Space). torchaudio.functional.resample is pure torch, so it's | |
| still fine to use. | |
| Pre-resampling in fp32 keeps prepare_audio's fp32 resample kernel a no-op, | |
| which otherwise trips on the fp16 model.""" | |
| data, sr = sf.read(path, dtype="float32", always_2d=True) # [N, C] | |
| wav = torch.from_numpy(data.T.copy()) # [C, N] | |
| if sr != target_sr: | |
| wav = torchaudio.functional.resample(wav, sr, target_sr) | |
| return target_sr, wav.to(dtype) | |
| def _generate(variant_key: str, | |
| prompt: str, | |
| negative_prompt: str, | |
| seconds: int, | |
| steps: int, | |
| cfg_scale: float, | |
| seed: int, | |
| init_audio_path: Optional[str] = None, | |
| init_noise_level: Optional[float] = None) -> tuple[str, dict]: | |
| """Returns (wav_path, meta). Inputs must already be clamped.""" | |
| prompt = (prompt or "").strip() | |
| if not prompt: | |
| raise gr.Error("Please enter a prompt.") | |
| t0 = time.time() | |
| if STUB: | |
| out_path = _stub_wav(seconds, seed) | |
| sample_rate = STUB_SAMPLE_RATE | |
| else: | |
| lv = LOADED[variant_key] | |
| conditioning = [{"prompt": prompt, "seconds_total": int(seconds)}] | |
| negative_conditioning = None | |
| neg = (negative_prompt or "").strip() | |
| if neg: | |
| negative_conditioning = [{"prompt": neg, "seconds_total": int(seconds)}] | |
| gen_kwargs: dict = dict( | |
| steps=steps, | |
| cfg_scale=cfg_scale, | |
| conditioning=conditioning, | |
| negative_conditioning=negative_conditioning, | |
| sample_size=lv.sample_size, | |
| sampler_type="pingpong", | |
| seed=seed, | |
| device="cuda", | |
| ) | |
| if init_audio_path: | |
| model_dtype = next(lv.model.parameters()).dtype | |
| gen_kwargs["init_audio"] = _load_init_audio( | |
| init_audio_path, lv.sample_rate, model_dtype) | |
| gen_kwargs["init_noise_level"] = float( | |
| init_noise_level if init_noise_level is not None else 0.4) | |
| output = generate_diffusion_cond_inpaint(lv.model, **gen_kwargs) | |
| output = rearrange(output, "b d n -> d (b n)") | |
| output = (output.to(torch.float32) | |
| .div(torch.max(torch.abs(output)).clamp(min=1e-9)) | |
| .clamp(-1, 1).mul(32767).to(torch.int16).cpu()) | |
| output = output[:, : seconds * lv.sample_rate] | |
| sample_rate = lv.sample_rate | |
| out_path = os.path.join(tempfile.mkdtemp(), "sa3.wav") | |
| sf.write(out_path, output.numpy().T, sample_rate, subtype="PCM_16") | |
| meta = { | |
| "ok": True, | |
| "seed": seed, | |
| "model_variant": variant_key, | |
| "seconds": seconds, | |
| "steps": steps, | |
| "cfg_scale": cfg_scale, | |
| "sample_rate": sample_rate, | |
| "gen_wall_s": round(time.time() - t0, 2), | |
| "stub": STUB, | |
| } | |
| if init_audio_path: | |
| meta["init_noise_level"] = float( | |
| init_noise_level if init_noise_level is not None else 0.4) | |
| print(f"[gen] {meta}", flush=True) | |
| return out_path, meta | |
| # --- UI handlers ------------------------------------------------------------ | |
| def _ui_duration_t2a(variant_key, prompt, negative_prompt, seconds, steps, cfg_scale, seed): | |
| return _gpu_seconds(seconds) | |
| def ui_text_to_audio(variant_key, prompt, negative_prompt, seconds, steps, cfg_scale, seed): | |
| key, seconds, steps, cfg_scale, seed = _clamp_inputs( | |
| variant_key, seconds, steps, cfg_scale, seed) | |
| path, meta = _generate(key, prompt, negative_prompt, seconds, steps, cfg_scale, seed) | |
| return path, f"seed **{meta['seed']}** · {meta['gen_wall_s']}s · {key}" | |
| def _ui_duration_a2a(variant_key, init_audio, prompt, negative_prompt, | |
| init_noise_level, seconds, steps, cfg_scale, seed): | |
| return _gpu_seconds(seconds) | |
| def ui_audio_to_audio(variant_key, init_audio, prompt, negative_prompt, | |
| init_noise_level, seconds, steps, cfg_scale, seed): | |
| if not init_audio: | |
| raise gr.Error("Upload an audio file to transform.") | |
| key, seconds, steps, cfg_scale, seed = _clamp_inputs( | |
| variant_key, seconds, steps, cfg_scale, seed) | |
| noise = max(0.0, min(float(init_noise_level if init_noise_level is not None else 0.4), 1.0)) | |
| path, meta = _generate(key, prompt, negative_prompt, seconds, steps, cfg_scale, | |
| seed, init_audio_path=init_audio, init_noise_level=noise) | |
| return path, f"seed **{meta['seed']}** · {meta['gen_wall_s']}s · {key} · noise {noise}" | |
| # --- API handlers ----------------------------------------------------------- | |
| def _api_duration_gen(token, variant_key, prompt, negative_prompt, | |
| seconds, steps, cfg_scale, seed): | |
| return _gpu_seconds(seconds) | |
| def api_generate(token, variant_key, prompt, negative_prompt, | |
| seconds, steps, cfg_scale, seed): | |
| _check_token(token) | |
| key, seconds, steps, cfg_scale, seed = _clamp_inputs( | |
| variant_key, seconds, steps, cfg_scale, seed) | |
| path, meta = _generate(key, prompt, negative_prompt, seconds, steps, cfg_scale, seed) | |
| return path, meta | |
| def _api_duration_a2a(token, variant_key, init_audio, prompt, negative_prompt, | |
| init_noise_level, seconds, steps, cfg_scale, seed): | |
| return _gpu_seconds(seconds) | |
| def api_generate_a2a(token, variant_key, init_audio, prompt, negative_prompt, | |
| init_noise_level, seconds, steps, cfg_scale, seed): | |
| _check_token(token) | |
| if not init_audio: | |
| raise gr.Error("init_audio is required for generate_a2a.") | |
| key, seconds, steps, cfg_scale, seed = _clamp_inputs( | |
| variant_key, seconds, steps, cfg_scale, seed) | |
| noise = max(0.0, min(float(init_noise_level if init_noise_level is not None else 0.4), 1.0)) | |
| path, meta = _generate(key, prompt, negative_prompt, seconds, steps, cfg_scale, | |
| seed, init_audio_path=init_audio, init_noise_level=noise) | |
| return path, meta | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| HEADER = """ | |
| # abv1 · SA3 engine | |
| Text→audio and audio→audio with **Stable Audio 3**. Shared backend for the abv1 | |
| music-AI demos; also exposes headless `/generate` and `/generate_a2a` endpoints. | |
| """ | |
| FOOTER = """ | |
| --- | |
| **Powered by Stability AI.** Models are used under the | |
| [Stability AI Community License](https://stability.ai/community-license-agreement); | |
| see `NOTICE.md` in this repo. All audio here is **AI-generated — not real | |
| releases** by any artist. | |
| """ | |
| NOISE_INFO = ("Low (0.1–0.3) = stay close to your upload · " | |
| "mid (0.4–0.6) = recognisable but reimagined · " | |
| "high (0.8–1.0) = barely related / go wild") | |
| def _on_variant_change(variant_key): | |
| mx = variant_max_seconds(variant_key) | |
| v = BY_KEY.get(variant_key, BY_KEY[DEFAULT_VARIANT]) | |
| return ( | |
| gr.update(maximum=mx, value=min(v.default_seconds, mx), | |
| label=f"Seconds · max {mx}"), | |
| gr.update(placeholder=v.placeholder), | |
| ) | |
| _start = BY_KEY[DEFAULT_VARIANT] | |
| _start_max = variant_max_seconds(DEFAULT_VARIANT) | |
| with gr.Blocks(theme=gr.themes.Citrus(), title="abv1 sa3 engine") as demo: | |
| gr.Markdown(HEADER) | |
| with gr.Tabs(): | |
| # ---------------- Text → audio ---------------- | |
| with gr.Tab("Text → audio"): | |
| t_variant = gr.Radio(VARIANT_CHOICES, value=DEFAULT_VARIANT, label="Model") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| t_prompt = gr.Textbox(label="Prompt", lines=3, | |
| placeholder=_start.placeholder) | |
| t_negative = gr.Textbox(label="Negative prompt (optional)", lines=1, | |
| placeholder="lo-fi, distorted, vocals") | |
| t_seconds = gr.Slider(1, _start_max, value=_start.default_seconds, | |
| step=1, label=f"Seconds · max {_start_max}") | |
| with gr.Accordion("Advanced", open=False): | |
| t_steps = gr.Slider(1, 100, value=8, step=1, label="Steps") | |
| t_cfg = gr.Slider(0.0, 15.0, value=1.0, step=0.1, label="CFG scale") | |
| t_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)") | |
| t_btn = gr.Button("Generate", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| t_audio = gr.Audio(label="Output", type="filepath") | |
| t_meta = gr.Markdown("") | |
| t_variant.change(_on_variant_change, [t_variant], [t_seconds, t_prompt]) | |
| t_btn.click( | |
| ui_text_to_audio, | |
| [t_variant, t_prompt, t_negative, t_seconds, t_steps, t_cfg, t_seed], | |
| [t_audio, t_meta], | |
| ) | |
| # ---------------- Audio → audio ---------------- | |
| with gr.Tab("Audio → audio"): | |
| a_variant = gr.Radio(VARIANT_CHOICES, value=DEFAULT_VARIANT, label="Model") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| a_init = gr.Audio(label="Your audio", type="filepath") | |
| a_noise = gr.Slider(0.0, 1.0, value=0.4, step=0.05, | |
| label="How far to travel from your audio", | |
| info=NOISE_INFO) | |
| a_prompt = gr.Textbox(label="Prompt", lines=3, | |
| placeholder=_start.placeholder) | |
| a_negative = gr.Textbox(label="Negative prompt (optional)", lines=1) | |
| a_seconds = gr.Slider(1, _start_max, value=_start.default_seconds, | |
| step=1, label=f"Seconds · max {_start_max}") | |
| with gr.Accordion("Advanced", open=False): | |
| a_steps = gr.Slider(1, 100, value=8, step=1, label="Steps") | |
| a_cfg = gr.Slider(0.0, 15.0, value=1.0, step=0.1, label="CFG scale") | |
| a_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)") | |
| a_btn = gr.Button("Transform", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| a_audio = gr.Audio(label="Output", type="filepath") | |
| a_meta = gr.Markdown("") | |
| a_variant.change(_on_variant_change, [a_variant], [a_seconds, a_prompt]) | |
| a_btn.click( | |
| ui_audio_to_audio, | |
| [a_variant, a_init, a_prompt, a_negative, a_noise, | |
| a_seconds, a_steps, a_cfg, a_seed], | |
| [a_audio, a_meta], | |
| ) | |
| gr.Markdown(FOOTER) | |
| # ── Headless API surface ──────────────────────────────────────────── | |
| # Hidden components + hidden buttons carry the api_name routes. Visibility | |
| # is UI-only; the endpoints are registered server-side regardless, and this | |
| # plays nicely with gradio_client.handle_file for the a2a upload. | |
| with gr.Group(visible=False): | |
| _g_token = gr.Textbox(value="", label="token") | |
| _g_variant = gr.Textbox(value=DEFAULT_VARIANT, label="model_variant") | |
| _g_prompt = gr.Textbox(value="", label="prompt") | |
| _g_negative = gr.Textbox(value="", label="negative_prompt") | |
| _g_seconds = gr.Number(value=30, label="seconds") | |
| _g_steps = gr.Number(value=8, label="steps") | |
| _g_cfg = gr.Number(value=1.0, label="cfg_scale") | |
| _g_seed = gr.Number(value=-1, label="seed") | |
| _g_file = gr.File(label="audio out") | |
| _g_json = gr.JSON(label="meta") | |
| _g_btn = gr.Button("generate") | |
| _g_btn.click( | |
| fn=api_generate, | |
| inputs=[_g_token, _g_variant, _g_prompt, _g_negative, | |
| _g_seconds, _g_steps, _g_cfg, _g_seed], | |
| outputs=[_g_file, _g_json], | |
| api_name="generate", | |
| show_progress="hidden", | |
| ) | |
| _a_token = gr.Textbox(value="", label="token") | |
| _a_variant = gr.Textbox(value=DEFAULT_VARIANT, label="model_variant") | |
| _a_init = gr.File(type="filepath", label="init_audio", | |
| file_types=["audio", ".wav", ".mp3", ".flac", ".ogg", | |
| ".m4a", ".aiff", ".aif", ".opus"]) | |
| _a_prompt = gr.Textbox(value="", label="prompt") | |
| _a_negative = gr.Textbox(value="", label="negative_prompt") | |
| _a_noise = gr.Number(value=0.4, label="init_noise_level") | |
| _a_seconds = gr.Number(value=30, label="seconds") | |
| _a_steps = gr.Number(value=8, label="steps") | |
| _a_cfg = gr.Number(value=1.0, label="cfg_scale") | |
| _a_seed = gr.Number(value=-1, label="seed") | |
| _a_file = gr.File(label="audio out") | |
| _a_json = gr.JSON(label="meta") | |
| _a_btn = gr.Button("generate_a2a") | |
| _a_btn.click( | |
| fn=api_generate_a2a, | |
| inputs=[_a_token, _a_variant, _a_init, _a_prompt, _a_negative, | |
| _a_noise, _a_seconds, _a_steps, _a_cfg, _a_seed], | |
| outputs=[_a_file, _a_json], | |
| api_name="generate_a2a", | |
| show_progress="hidden", | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=32).launch( | |
| server_name=os.environ.get("SA3_HOST", "0.0.0.0"), | |
| server_port=int(os.environ.get("SA3_PORT", "7860")), | |
| ) | |