app.py

from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import uuid
from pathlib import Path
from tempfile import gettempdir

import gradio as gr
import soundfile as sf

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

from musicgen_runtime import generate, warmup

MAX_INPUT_SECONDS = 30
OUTPUT_ROOT = Path(gettempdir()) / "musicgen_outputs"

model_card = ModelCard(
    name="MusicGen",
    description=(
        "Generate music from a text prompt with melody conditioning, as 32 kHz mono audio. "
        "The code "
        "is MIT, but Meta's MusicGen weights are CC-BY-NC-4.0, so results "
        "cannot be used commercially. Repository: "
        "https://github.com/facebookresearch/audiocraft Paper: "
        "https://arxiv.org/abs/2306.05284"
    ),
    author="Jade Copet, Felix Kreuk, Itai Gat, et al.",
    tags=[Subcategory.HOLISTIC, SampleRate(32000), Channels(1), 'melody-conditioning'],
)


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server. The model is downloaded, built and placed there alone, where ZeroGPU's emulation records
# the CUDA placement, so that no request moves it inside its GPU reservation.
if __name__ != "__main__":
    warmup()


def _validate_audio(path: str | None) -> str | None:
    if not path:
        return None
    try:
        duration = sf.info(path).duration
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if duration == 0:
        raise gr.Error("The melody file is empty.")
    if duration > MAX_INPUT_SECONDS:
        raise gr.Error(
            f"Melody audio must be no longer than {MAX_INPUT_SECONDS} seconds. "
            f"Received {duration:.1f} seconds."
        )
    return path


def process_fn(
    audio_path: str | None,
    prompt: str,
    duration: str,
) -> str:
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a text description.")

    return compose(_validate_audio(audio_path), prompt, int(duration))


@spaces.GPU(duration=120)
def compose(melody_path: str | None, prompt: str, duration_seconds: int) -> str:
    try:
        audio, sample_rate = generate(
            prompt=prompt,
            duration_seconds=duration_seconds,
            audio_path=melody_path,
        )
    except Exception as exc:
        raise gr.Error(f"MusicGen generation failed: {exc}") from exc

    # Generation can overshoot full scale, which writing would flatten, so the result is scaled
    # down to fit when it exceeds it
    peak = audio.abs().max()
    if peak > 1.0:
        audio = audio / peak

    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "musicgen_output.wav"
    sf.write(str(output_path), audio.numpy(), sample_rate)

    return str(output_path)


if __name__ == "__main__":
    with gr.Blocks(title="MusicGen") as demo:
        input_components = [
            gr.Audio(type="filepath", label="Melody (optional)").harp_required(False)
            .set_info("Melody audio to condition on, up to 30 seconds."),
            gr.Textbox(
                label="Text Description",
                placeholder="e.g. An upbeat electronic dance track with heavy bass",
                info=(
                    "Describe the music to generate: genre, instrumentation, mood, "
                    "tempo."
                ),
            ),
            gr.Dropdown(
                choices=["5", "10", "15", "20"],
                value="10",
                label="Duration (seconds)",
                info=(
                    "Length of the result. Longer generations drift further from the "
                    "prompt."
                ),
            ),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Generated Music").set_info((
                "The generated music, at 32 kHz mono."
            )),
        ]
        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)