Blame for app.py

"""
SonicMaster on HARP: text-controlled music restoration and mastering.

This Space is a frontend for the authors' own Space, amaai-lab/SonicMaster, which runs the model on
ZeroGPU. The model decodes through Stable Audio Open's VAE, whose weights are gated, so serving it
from here would hand it to callers who never accepted their terms (see "Calling Another Space" in
HARP's deployment guidelines). Each request is forwarded with the caller's x-ip-token instead,
which bills it to the caller's own ZeroGPU allowance, as a direct call to the authors' Space would.
"""

import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
import soundfile as sf
import soxr
from gradio.context import LocalContext
from gradio_client import Client, handle_file
from gradio_client.exceptions import AppError
from pyharp import ModelCard, Subcategory, build_endpoint

BACKEND = "amaai-lab/SonicMaster"
# Connected to by URL. Given the repository id, gradio_client looks the host up through the Hub's
# API without a token, which the Hub rate-limits from the address a Space shares with others.
BACKEND_URL = "https://amaai-lab-sonicmaster.hf.space"
API_NAME = "/enhance_audio_ui"
# The parameters of API_NAME, in order, as the backend declared them when this frontend was written
BACKEND_PARAMETERS = ["audio_path", "prompt"]

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}


def check_backend():
    """
    Stops the Space from starting if the backend's API no longer matches what process_fn sends.

    An upstream change then shows up as a Space that fails to start, rather than as a puzzle on
    every request. A backend that cannot be reached is only reported, since a sleeping or
    restarting Space is no reason to keep this one down.
    """
    try:
        api = Client(BACKEND_URL, verbose=False).view_api(return_format="dict", print_info=False)
    except Exception as e:
        print(f"SonicMaster: could not reach {BACKEND} to check its API: {e}", flush=True)
        return

    endpoint = api["named_endpoints"].get(API_NAME)
    names = [p.get("parameter_name") for p in endpoint["parameters"]] if endpoint else None
    if names != BACKEND_PARAMETERS:
        raise RuntimeError(f"{BACKEND}{API_NAME} now takes {names}, but this frontend sends "
                           f"{BACKEND_PARAMETERS}. Update process_fn to match.")


model_card = ModelCard(
    name="SonicMaster",
    description=(
        "Text-controlled all-in-one music restoration and mastering: one instruction in plain "
        "language fixes reverb, clipping, tonal balance, dynamics, or the stereo image. The "
        "model is a flow-matching transformer over Stable Audio's latent space, conditioned on "
        "a FLAN-T5 encoding of the instruction, so the wording of the prompt is what selects the "
        "treatment. It runs on the authors' Space (amaai-lab/SonicMaster), to which this Space "
        "forwards each request, so a request spends the caller's own ZeroGPU allowance and "
        "waits in that Space's queue. The model works at 44.1 kHz stereo, and the result is "
        "returned at the sample rate of the input. SonicMaster itself is Apache 2.0, but it "
        "decodes through Stable Audio Open's VAE, so results are also governed by the Stability "
        "AI Community License: "
        "https://huggingface.co/stabilityai/stable-audio-open-1.0/blob/main/LICENSE.md "
        "Repository: https://github.com/AMAAI-Lab/SonicMaster Paper: "
        "https://arxiv.org/abs/2508.03448"
    ),
    author="Jan Melechovsky, Ambuj Mehrish, Dorien Herremans",
    tags=[Subcategory.POST_PROCESSING, "text-guided", "zerogpu", "domain:music"],
)


def _caller_client():
    """A client of the backend carrying the caller's x-ip-token, so that ZeroGPU bills the caller."""
    # The Hub adds x-ip-token to every request a Space receives, naming the caller. PyHARP
    # restores the request inside its worker, so it is read here.
    request = LocalContext.request.get(None)
    ip_token = request.headers.get("x-ip-token") if request is not None else None
    print(f"SonicMaster: forwarding to {BACKEND} "
          f"{'with' if ip_token else 'without'} the caller's x-ip-token", flush=True)
    return Client(BACKEND_URL, headers={"x-ip-token": ip_token} if ip_token else None,
                  verbose=False)


def _in_the_format_of(output_path, input_path, seconds):
    """
    The backend's audio at the rate and in the container of the input, cut to `seconds`, as
    16-bit within full scale. The channels stay as the model made them.
    """
    audio, rate = sf.read(output_path, dtype="float32", always_2d=True)
    input_rate = sf.info(input_path).samplerate
    if rate != input_rate:
        audio = soxr.resample(audio, rate, input_rate)
        rate = input_rate
    audio = audio[: int(round(seconds * rate))]
    peak = float(np.abs(audio).max()) if audio.size else 0.0
    if peak > 1.0:
        audio = audio / peak
    ext = Path(input_path).suffix.lower()
    path = Path(tempfile.mkdtemp()) / f"sonicmaster{ext if ext in LOSSLESS else '.wav'}"
    sf.write(path, audio, rate, subtype="PCM_16")
    return str(path)


def process_fn(input_audio_path, prompt):
    """
    Apply the instruction to the recording.

    Args:
        input_audio_path (str): The recording to edit.
        prompt (str): What to change about it.

    Returns:
        output_path (str): The edited audio.
    """

    # Checked before forwarding, so a request the model cannot take waits in no queue and spends
    # none of the caller's allowance
    if not input_audio_path:
        raise gr.Error("Provide the recording to treat.")
    try:
        info = sf.info(input_audio_path)
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if info.frames == 0:
        raise gr.Error("The audio file is empty.")
    if not (prompt or "").strip():
        raise gr.Error("Describe what to change about the audio.")

    try:
        output, status = _caller_client().predict(handle_file(input_audio_path), prompt.strip(),
                                                  api_name=API_NAME)
    except AppError as e:
        # The backend's own message, a quota error included, so the caller sees why
        raise gr.Error(str(e))

    if not output:
        raise gr.Error(status or "SonicMaster returned no audio.")
    return _in_the_format_of(output, input_audio_path, float(info.duration))


# The processing worker imports this file, so the app must not be built there
if __name__ == "__main__":
    check_backend()

    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input audio").harp_required(True).set_info((
                "The recording to treat. The model works at 44.1 kHz stereo, and a mono input "
                "is heard on both channels."
            )),
            gr.Textbox(label="Instruction", value="Remove the reverb from this recording.",
                       info=(
                           "What to change, in plain language, for example "
                           "\"reduce the reverb\", \"brighten the high end\", "
                           "\"even out the loudness\", or \"widen the stereo "
                           "image\". One instruction at a time works better "
                           "than several combined."
                       )),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Treated audio").set_info((
                "The result, at 16-bit, at the sample rate and in the container of the input."
            )),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

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