app.py

from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import os
import uuid
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import gradio as gr
import numpy as np
import soundfile as sf
import torch
from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint
from pyharp.tags import Channels, SampleRate, Subcategory

from magenta_rt import paths
from magenta_rt.torch import MagentaRT2
from magenta_rt.torch.musiccoca import MusicCoCa


MODEL_REPO = "google/magenta-realtime-2"
MODEL_NAME = "mrt2_small"
CHECKPOINT = f"{MODEL_NAME}.safetensors"
AOTI_REPO = "magenta-torch/magenta-rt-aoti-small"
SAMPLE_RATE = 48_000
FRAMES_PER_SECOND = 25
OUTPUT_DIR = Path("/tmp/magenta_rt_outputs")

model_card = ModelCard(
    name="Magenta RealTime 2",
    description=(
        "Generates short instrumental music clips from text prompts with the small Magenta "
        "RealTime 2 model, as 48 kHz stereo audio. The code is Apache-2.0 and the weights are "
        "CC-BY-4.0, which requires attribution. Repository: "
        "https://github.com/magenta/magenta-realtime Paper: "
        "https://arxiv.org/abs/2508.04651"
    ),
    author="Antoine Caillon, Brian McWilliams, Cassie Tarakajian, et al.",
    tags=[Subcategory.HOLISTIC, SampleRate(SAMPLE_RATE), Channels(2), 'real-time-music', 'domain:music'],
)

OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the checkpoint is downloaded and the model built there alone
if __name__ != "__main__":
    model_root = Path("/data" if Path("/data").is_dir() else "/tmp/magenta")
    magenta_home = model_root / "magenta-rt-v2"
    magenta_home.mkdir(parents=True, exist_ok=True)

    hf_hub_download(
        repo_id=MODEL_REPO,
        filename=f"checkpoints/{CHECKPOINT}",
        local_dir=magenta_home,
        revision="010aa0dcb0dfd27b24f0ad07b4dad63e8f9521cc",
    )
    paths.set_magenta_home(magenta_home)

    style_model = MusicCoCa(device="cpu")
    model = MagentaRT2(
        size=MODEL_NAME,
        device="cuda",
        dtype=torch.bfloat16,
        style_model=style_model,
    )

    try:
        model.load_compiled(repo_id=AOTI_REPO)
    except Exception as exc:
        print(f"AOTI loading failed; using eager inference: {exc}", flush=True)


def process_fn(
    prompt: str,
    duration: str,
    temperature: float,
    top_k: float,
    seed: float,
) -> str:
    # Checked before the GPU is reserved, so a prompt the model cannot take spends none of the
    # caller's ZeroGPU allowance
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please describe the music you want to generate.")
    if len(prompt) > 300:
        raise gr.Error("The prompt must be 300 characters or fewer.")

    return generate(prompt, duration, temperature, top_k, seed)


@spaces.GPU(duration=45)
def generate(prompt: str, duration: str, temperature: float, top_k: float, seed: float) -> str:
    duration_seconds = int(duration)
    expected_samples = duration_seconds * SAMPLE_RATE
    frames = duration_seconds * FRAMES_PER_SECOND + 1

    try:
        if style_model.device != "cuda":
            style_model.to("cuda")
        style_tokens = style_model.embed_tokens(prompt)
        audio, _ = model.generate(
            style=style_tokens,
            temperature=float(temperature),
            top_k=int(top_k),
            frames=frames,
            seed=int(seed),
            flush=True,
        )
    except Exception as exc:
        raise gr.Error(f"Magenta RealTime 2 inference failed: {exc}") from exc

    audio = np.asarray(audio, dtype=np.float32)
    if audio.ndim != 2 or audio.shape[1] != 2:
        raise gr.Error("The model returned an unexpected audio shape.")
    if len(audio) < expected_samples:
        raise gr.Error("The model returned less audio than requested.")

    output_path = OUTPUT_DIR / f"{uuid.uuid4().hex}.wav"
    sf.write(
        output_path,
        audio[:expected_samples],
        SAMPLE_RATE,
        subtype="PCM_16",
    )
    return str(output_path)


if __name__ == "__main__":
    with gr.Blocks(title="Magenta RealTime 2") as demo:
        input_components = [
            gr.Textbox(
                value="warm analog synthesizer with a gentle rhythmic pulse",
                label="Music Prompt",
                info="Describe the instruments, texture, style, or mood.",
                lines=2,
                max_lines=4,
            ),
            gr.Dropdown(
                choices=["2", "4", "8"],
                value="4",
                label="Duration (seconds)",
                info="Length of the generated clip.",
            ),
            gr.Slider(
                minimum=0.1,
                maximum=2.0,
                step=0.1,
                value=1.1,
                label="Temperature",
                info="Higher values produce more variation.",
            ),
            gr.Slider(
                minimum=10,
                maximum=100,
                step=5,
                value=50,
                label="Top-k",
                info="Limits each sampling step to the most likely tokens.",
            ),
            gr.Number(
                value=0,
                minimum=0,
                maximum=2_147_483_647,
                precision=0,
                label="Seed",
                info="Use the same seed and controls to reproduce a result.",
            ),
        ]
        output_components = [
            gr.Audio(
                type="filepath",
                label="Generated Music",
            ).set_info("A 48 kHz stereo WAV file."),
        ]
        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)