app.py

from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import gradio as gr
import soundfile as sf

from pyharp import ModelCard, build_endpoint
from pyharp.tags import Subcategory

from stable_audio_runtime import edit_audio, load_model

MIN_AUDIO_SECONDS = 5
MAX_AUDIO_SECONDS = 30

model_card = ModelCard(
    name="Stable Audio 3",
    description=(
        "Restyle, inpaint, or continue music audio using a text prompt. "
        "Repository: https://github.com/Stability-AI/stable-audio-3 Paper: "
        "https://arxiv.org/abs/2605.17991"
    ),
    author="Zach Evans, Julian D. Parker, Matthew Rice, et al.",
    tags=[Subcategory.EDITING],
)


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server. stable-audio-tools is installed and the model downloaded, built and placed there, on the
# worker's main thread, so that no request does any of it inside its GPU reservation.
if __name__ != "__main__":
    load_model()


def _audio_duration(path: str | None) -> float:
    if not path:
        raise gr.Error("Please upload a music clip.")
    try:
        duration = float(sf.info(path).duration)
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if duration < MIN_AUDIO_SECONDS:
        raise gr.Error(f"Audio must be at least {MIN_AUDIO_SECONDS} seconds.")
    if duration > MAX_AUDIO_SECONDS:
        raise gr.Error(f"Audio must be no longer than {MAX_AUDIO_SECONDS} seconds.")
    return duration


def process_fn(
    audio_path: str | None,
    prompt: str,
    mode: str,
    edit_start: float,
    edit_end: float,
    continuation_length: float,
    strength: float,
    seed: int,
) -> str:
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    duration = _audio_duration(audio_path)
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter an editing prompt.")

    if mode == "Inpaint" and (
        edit_start < 0 or edit_end > duration or edit_end <= edit_start
    ):
        raise gr.Error(
            "The inpaint region must be inside the uploaded clip, "
            "with the end after the start."
        )

    return run_edit(audio_path, prompt, mode, edit_start, edit_end, continuation_length,
                    strength, seed)


# A 6 s restyle measured about 24 s, and inputs are capped at 30 s
@spaces.GPU(duration=120)
def run_edit(audio_path, prompt, mode, edit_start, edit_end, continuation_length, strength,
             seed) -> str:
    try:
        return edit_audio(
            audio_path=audio_path,
            prompt=prompt,
            mode=mode,
            edit_start=float(edit_start),
            edit_end=float(edit_end),
            continuation_length=float(continuation_length),
            strength=float(strength),
            seed=int(seed),
        )
    except gr.Error:
        raise
    except Exception as exc:
        raise gr.Error(f"Stable Audio 3 inference failed: {exc}") from exc


if __name__ == "__main__":
    with gr.Blocks(title="Stable Audio 3") as demo:
        input_components = [
            gr.Audio(type="filepath", label="Music Audio")
            .harp_required(True)
            .set_info("Music clip between 5 and 30 seconds."),
            gr.Textbox(
                label="Prompt",
                placeholder="A warm synthwave groove with punchy drums",
                info=(
                    "Describe what the audio should become. What this changes depends "
                    "on the Edit Mode below."
                ),
            ).harp_required(True),
            gr.Dropdown(
                choices=["Restyle", "Inpaint", "Continue"],
                value="Restyle",
                label="Edit Mode",
                info=(
                    "Restyle reworks the whole clip, Inpaint replaces the region "
                    "between the edit points, and Continue extends past the end."
                ),
            ),
            gr.Slider(
                minimum=0,
                maximum=30,
                value=4,
                step=0.1,
                label="Edit Start (seconds)",
            ).set_info("Used by Inpaint mode."),
            gr.Slider(
                minimum=0,
                maximum=30,
                value=8,
                step=0.1,
                label="Edit End (seconds)",
            ).set_info("Used by Inpaint mode."),
            gr.Slider(
                minimum=1,
                maximum=15,
                value=8,
                step=1,
                label="Continuation Length (seconds)",
            ).set_info("Used by Continue mode."),
            gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.75,
                step=0.05,
                label="Transformation Strength",
            ).set_info("Used by Restyle mode."),
            gr.Slider(
                minimum=0,
                maximum=99999,
                value=0,
                step=1,
                label="Seed",
            ).set_info("Use 0 for a random seed."),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Edited Audio").set_info((
                "The edited result, at 16-bit."
            )),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )


    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True, ssr_mode=False)