app.py

"""Thin API wrapper around stabilityai/stable-audio-3 (Medium — general audio).

This Space owns no GPU and holds no credentials. It proxies to the upstream
ZeroGPU Space with **the caller's own token**, pinning
``variant_key="medium"`` (``stabilityai/stable-audio-3-medium``) and exposing one
clean, stable endpoint: ``/generate``.

Why a wrapper: upstream's ``/infer`` signature is positional and tracks its UI,
so it shifts whenever the demo's layout changes. ``/generate`` is keyword-only
and stays put.

Why the caller's token: ZeroGPU quota is charged to whoever's token is
presented. This Space deliberately has no ``HF_TOKEN`` secret, so a caller can
only ever spend their own quota — never the Space owner's.
"""

from __future__ import annotations

import hashlib
import os
import shutil
import tempfile
import threading
import time
from collections import OrderedDict
from typing import Optional

import gradio as gr
import httpx
from gradio_client import Client

UPSTREAM = os.environ.get("UPSTREAM_SPACE", "stabilityai/stable-audio-3")
VARIANT = "medium"  # stabilityai/stable-audio-3-medium — general audio (largest)

# Upstream medium: sample_size // sample_rate == 380s.
MAX_DURATION = 380
SAMPLERS = ["pingpong", "euler", "rk4", "dpmpp"]

# Wall-clock budget for one upstream job, including ZeroGPU queue wait.
JOB_TIMEOUT = float(os.environ.get("JOB_TIMEOUT", "900"))

# gradio_client's default httpx timeout kills the *download* of the result wav:
# a full 380s render is ~67 MB and reliably exceeds it. Budget generously.
HTTPX_TIMEOUT = httpx.Timeout(
    float(os.environ.get("HTTPX_TIMEOUT", "600")), connect=30.0
)

# Building a Client refetches the upstream config (seconds), so keep one per
# caller. Keyed by a digest, never the token itself, and bounded so a stream of
# distinct callers can't grow it without limit.
_MAX_CLIENTS = 32
_clients: "OrderedDict[str, Client]" = OrderedDict()
_clients_lock = threading.Lock()


def _key(token: Optional[str]) -> str:
    if not token:
        return "anonymous"
    return hashlib.sha256(token.encode()).hexdigest()


def _scrub(text: str, token: Optional[str]) -> str:
    """Never let a caller's token escape into an error message or log line."""
    if token and token in text:
        text = text.replace(token, "<token>")
    return text


def _get_client(token: Optional[str], force_new: bool = False) -> Client:
    key = _key(token)
    with _clients_lock:
        if force_new:
            _clients.pop(key, None)
        client = _clients.get(key)
        if client is not None:
            _clients.move_to_end(key)
            return client

    # Build outside the lock: this does network I/O and would otherwise stall
    # every other caller.
    client = Client(
        UPSTREAM,
        token=token or None,
        verbose=False,
        httpx_kwargs={"timeout": HTTPX_TIMEOUT},
    )

    with _clients_lock:
        _clients[key] = client
        _clients.move_to_end(key)
        while len(_clients) > _MAX_CLIENTS:
            _clients.popitem(last=False)
    return client


def _persist(path: str) -> str:
    """gradio_client writes downloads into a temp dir it may later reap. Copy
    into our own temp dir so the file outlives the call."""
    out_dir = tempfile.mkdtemp(prefix="sa3_")
    dest = os.path.join(out_dir, "stable_audio_3_medium.wav")
    shutil.copyfile(path, dest)
    return dest


def _call_upstream(args: tuple, api_name: str, token: Optional[str],
                   retry: bool = True):
    try:
        job = _get_client(token).submit(*args, api_name=api_name)
        return job.result(timeout=JOB_TIMEOUT)
    except Exception as exc:  # noqa: BLE001
        # A stale cached client (upstream restarted / slept) surfaces as a
        # connection or schema error. Rebuild once, then give up.
        if retry:
            _get_client(token, force_new=True)
            return _call_upstream(args, api_name, token, retry=False)
        raise exc


def generate(
    prompt: str,
    duration: float = 60,
    steps: float = 8,
    cfg_scale: float = 1.0,
    sampler_type: str = "pingpong",
    seed: int = 0,
    negative_prompt: str = "",
    hf_token: str = "",
):
    """Generate audio with Stable Audio 3 Medium (general audio).

    ``hf_token`` is your own Hugging Face token; upstream ZeroGPU quota is
    charged to it. Returns ``(wav_path, status)``.
    """
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    token = (hf_token or "").strip() or None
    if token is None:
        raise gr.Error(
            "An `hf_token` is required. This Space forwards your own token to "
            f"{UPSTREAM} so the ZeroGPU quota is charged to you, not the Space "
            "owner. Create one at https://huggingface.co/settings/tokens "
            "(a read token is enough) and pass it as `hf_token`."
        )

    duration = int(max(1, min(float(duration), MAX_DURATION)))
    steps = int(max(1, min(float(steps), 500)))
    cfg_scale = float(cfg_scale)
    sampler_type = sampler_type if sampler_type in SAMPLERS else "pingpong"
    seed = int(seed)
    negative_prompt = (negative_prompt or "").strip()

    t0 = time.time()
    try:
        if negative_prompt:
            # Only the advanced endpoint accepts a negative prompt. Its audio
            # inputs are optional in practice despite the schema marking them
            # required; if that ever stops holding, fall back to /infer.
            try:
                result = _call_upstream(
                    (
                        VARIANT, prompt, negative_prompt, duration, steps,
                        cfg_scale, sampler_type, seed,
                        1.0,    # sigma_max
                        1.0,    # apg_scale
                        6.0,    # duration_padding_sec
                        True,   # cut_to_seconds_total
                        None,   # init_audio
                        0.9,    # init_noise_level
                        None,   # inpaint_audio
                        0.0,    # mask_start_sec
                        0.0,    # mask_end_sec
                        0,      # preview_every
                    ),
                    api_name="/infer_advanced",
                    token=token,
                )
                audio_path = result[0] if isinstance(result, (list, tuple)) else result
            except Exception as exc:  # noqa: BLE001
                print("[generate] /infer_advanced failed "
                      f"({_scrub(str(exc), token)}); retrying without negative "
                      "prompt", flush=True)
                audio_path = _call_upstream(
                    (VARIANT, prompt, duration, steps, cfg_scale, sampler_type, seed),
                    api_name="/infer",
                    token=token,
                )
        else:
            audio_path = _call_upstream(
                (VARIANT, prompt, duration, steps, cfg_scale, sampler_type, seed),
                api_name="/infer",
                token=token,
            )
    except Exception as exc:  # noqa: BLE001
        raise gr.Error(
            f"Upstream {UPSTREAM} failed: {_scrub(str(exc), token)}"
        ) from None

    if isinstance(audio_path, (list, tuple)):
        audio_path = audio_path[0]
    if isinstance(audio_path, dict):
        audio_path = audio_path.get("path") or audio_path.get("name")
    if not audio_path or not os.path.exists(audio_path):
        raise gr.Error(f"Upstream returned no audio file (got {audio_path!r}).")

    elapsed = time.time() - t0
    status = (
        f"stable-audio-3-medium · {duration}s · {steps} steps · "
        f"cfg {cfg_scale} · {sampler_type} · seed {seed} · {elapsed:.1f}s"
    )
    return _persist(audio_path), status


DESCRIPTION = f"""
# 🎵 Stable Audio 3 Medium — Music API

Text-to-music via **`stabilityai/stable-audio-3-medium`** (general audio, the largest
SA3 variant), proxied from [`{UPSTREAM}`](https://huggingface.co/spaces/{UPSTREAM}).

**Bring your own token.** This Space stores no credentials — it forwards the
`hf_token` you pass straight to upstream, so the ZeroGPU quota is charged to
*you*. A plain [read token](https://huggingface.co/settings/tokens) is enough.

```python
from gradio_client import Client

client = Client("sleeper371/stable-audio-3-medium-api")
wav_path, status = client.predict(
    prompt="Warm analog house groove, deep sub bass, shuffled hats 122 BPM",
    duration=30,
    hf_token="hf_...",      # your token — your quota
    api_name="/generate",
)
```
"""

with gr.Blocks(title="Stable Audio 3 Medium API") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="A dream-like Synthpop instrumental that would accompany "
                            "a dream-sequence in a surrealist movie 120 BPM",
                lines=3,
            )
            hf_token = gr.Textbox(
                value="",
                label="Your Hugging Face token",
                info="Forwarded to the upstream ZeroGPU Space so the quota is "
                     "charged to you. Never stored.",
                placeholder="hf_...",
                type="password",
            )
            negative_prompt = gr.Textbox(
                value="",
                label="Negative prompt (optional)",
                placeholder="lo-fi, distorted, vocals",
                lines=1,
            )
            duration = gr.Slider(
                1, MAX_DURATION, value=60, step=1,
                label=f"Duration (s) · model max {MAX_DURATION}s",
            )
            with gr.Accordion("Advanced settings", open=False):
                steps = gr.Slider(1, 100, value=8, step=1, label="Steps")
                cfg_scale = gr.Slider(0.0, 25.0, value=1.0, step=0.1, label="CFG scale")
                sampler_type = gr.Dropdown(SAMPLERS, value="pingpong", label="Sampler")
                seed = gr.Number(value=0, precision=0, label="Seed (0 = random)")
            run_btn = gr.Button("🎼 Generate", variant="primary", size="lg")

        with gr.Column(scale=1):
            audio_out = gr.Audio(label="Output", type="filepath", autoplay=True)
            status_out = gr.Textbox(label="Status", lines=2, interactive=False)

    gr.Examples(
        examples=[
            ["House music that encapsulates the feeling of being at a festival in "
             "the sunny weather with all your friends 124 BPM", 30],
            ["Rainy night lo-fi hip-hop beat with vinyl crackle, mellow piano chords, "
             "soft kick and snare 80 BPM", 30],
            ["Sweeping orchestral score, soaring strings and brass, heroic and "
             "cinematic, 90 BPM", 45],
        ],
        inputs=[prompt, duration],
        label="Examples",
    )

    run_btn.click(
        fn=generate,
        inputs=[prompt, duration, steps, cfg_scale, sampler_type, seed,
                negative_prompt, hf_token],
        outputs=[audio_out, status_out],
        api_name="generate",
    )

if __name__ == "__main__":
    # Gradio 6 moved `theme` to launch() and dropped `show_api`.
    demo.queue(max_size=20).launch(theme=gr.themes.Citrus())