app.py

import sys

sys.stdout.reconfigure(line_buffering=True)

try:
    import spaces
except ImportError:
    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=60):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                func = args[0]
                return func


import urllib.request
from pathlib import Path

# mmt/ keeps the original repo's bare intra-package imports (`import utils`,
# `import representation`), so it needs to be on sys.path directly rather
# than imported as a "mmt.*" package.
sys.path.insert(0, str(Path(__file__).parent / "mmt"))

import gradio as gr
import muspy
import numpy as np
import torch
from pyharp import ModelCard, build_endpoint, get_default_path

import music_x_transformers
import representation
import utils as mmt_utils

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CKPT_DIR = Path(__file__).parent / "checkpoint"

train_args = mmt_utils.load_json(CKPT_DIR / "train-args.json")
encoding = representation.get_encoding()

_VALID_INSTRUMENTS = ", ".join(sorted(n for n in encoding["instrument_code_map"] if isinstance(n, str) and n != "null"))


def load_model():
    """Build the model on CPU and load the sod/ape checkpoint weights. Stays on
    CPU here even if GPU hardware is available -- ZeroGPU only allows touching
    CUDA from inside an @spaces.GPU call, not at module load time."""
    model = music_x_transformers.MusicXTransformer(
        dim=train_args["dim"],
        encoding=encoding,
        depth=train_args["layers"],
        heads=train_args["heads"],
        max_seq_len=train_args["max_seq_len"],
        max_beat=train_args["max_beat"],
        rotary_pos_emb=train_args["rel_pos_emb"],
        use_abs_pos_emb=train_args["abs_pos_emb"],
        emb_dropout=train_args["dropout"],
        attn_dropout=train_args["dropout"],
        ff_dropout=train_args["dropout"],
    )
    state_dict = torch.load(
        CKPT_DIR / "checkpoints" / "best_model.pt", map_location="cpu"
    )
    model.load_state_dict(state_dict)
    model.eval()
    return model


# Checkpoint is small (~80MB) and lives in the Space repo, so load it
# synchronously at startup rather than in a background thread.
model = load_model()
model_ready = False  # has the model been moved onto DEVICE yet?


def download_musescore_soundfont():
    """Fetches the MuseScore General soundfont over HTTPS -- muspy's own
    downloader uses FTP, which HF Spaces blocks."""
    from muspy.external import (
        get_musescore_soundfont_dir,
        get_musescore_soundfont_path,
    )

    sf_path = get_musescore_soundfont_path()
    if sf_path.is_file():
        return
    get_musescore_soundfont_dir().mkdir(parents=True, exist_ok=True)
    prefix = "https://ftp.osuosl.org/pub/musescore/soundfont/MuseScore_General/"
    urllib.request.urlretrieve(prefix + "MuseScore_General.sf3", sf_path)


# Fetch upfront so failures show in logs, not after a request.
download_musescore_soundfont()


def encode_seed(midi_path, encoding, max_beat):
    """Turn an uploaded MIDI file into the token prefix the model continues
    from. Mirrors mmt/convert_sod.py's resolution-adjustment step and
    mmt/dataset.py's max_beat trim, since generate.py normally gets both for
    free from the preprocessed dataset it reads conditioning prefixes from."""
    music = muspy.read(midi_path)
    music.adjust_resolution(encoding["resolution"])
    for track in music:
        for note in track:
            if note.duration == 0:
                note.duration = 1
    music.remove_duplicate()

    notes = representation.extract_notes(music, encoding["resolution"])
    if len(notes) == 0:
        raise gr.Error("No usable notes found in the uploaded MIDI (check instruments are General MIDI, non-drum).")
    n_beats = notes[-1, 0] + 1
    if n_beats > max_beat:
        notes = notes[notes[:, 0] < max_beat]

    codes = representation.encode_notes(notes, encoding)
    return codes[:-1]  # drop the trailing EOS row so generation continues instead of stopping immediately


def resolve_instruments(text):
    """Matches a comma-separated instrument list against the model's fixed
    64-name vocabulary. Unmatched names are dropped rather than raised --
    HARP's client only shows a generic error banner, so per-name feedback
    has to go through the Instrument Matching output file instead."""
    valid = encoding["instrument_code_map"]
    tokens = [t.strip() for t in text.split(",") if t.strip()]
    if not tokens:
        return [], "No instruments requested -- fully unconditioned generation."

    matched, unmatched, seen = [], [], set()
    for token in tokens:
        key = token.lower().replace(" ", "-").replace("_", "-")
        if key != "null" and key in valid:
            if key not in seen:
                matched.append(key)
                seen.add(key)
        else:
            unmatched.append(token)

    if matched:
        note = f"Matched: {', '.join(matched)}."
        if unmatched:
            note += f" Not recognized: {', '.join(unmatched)}."
    else:
        note = f"No valid instrument names found, generating unconditioned. Not recognized: {', '.join(unmatched)}."
    return matched, note


def build_instrument_prefix(instrument_names):
    """Builds a (start-of-song, instrument..., start-of-notes) prefix with
    zero notes, so the model writes a whole piece for exactly this
    instrument list. Mirrors mmt/representation.py's encode_notes() up
    through its instrument block."""
    type_code_map = encoding["type_code_map"]
    instrument_code_map = encoding["instrument_code_map"]

    codes = [[type_code_map["start-of-song"], 0, 0, 0, 0, 0]]
    instrument_rows = sorted(
        [type_code_map["instrument"], 0, 0, 0, 0, instrument_code_map[name]]
        for name in instrument_names
    )
    codes.extend(instrument_rows)
    codes.append([type_code_map["start-of-notes"], 0, 0, 0, 0, 0])
    return np.array(codes)


@spaces.GPU
@torch.inference_mode()
def generate_tokens(prefix_codes, generation_length, temperature, focus):
    """Runs the model's autoregressive generation on GPU -- the only part of
    the pipeline that touches CUDA, so it's the only part wrapped in
    @spaces.GPU. Everything else (MIDI decoding, audio synthesis) is
    CPU-bound and shouldn't burn ZeroGPU quota."""
    global model, model_ready
    if not model_ready:
        model = model.to(DEVICE)
        model_ready = True

    sos = encoding["type_code_map"]["start-of-song"]
    eos = encoding["type_code_map"]["end-of-song"]

    if prefix_codes is not None:
        tgt_start = torch.tensor(prefix_codes, dtype=torch.long, device=DEVICE).unsqueeze(0)
    else:
        tgt_start = torch.zeros((1, 1, 6), dtype=torch.long, device=DEVICE)
        tgt_start[:, 0, 0] = sos

    generated = model.generate(
        tgt_start,
        int(generation_length),
        eos_token=eos,
        temperature=temperature,
        filter_logits_fn="top_k",
        filter_thres=focus,
        monotonicity_dim=("type", "beat"),
    )
    return torch.cat((tgt_start, generated), 1)[0].cpu().numpy()


def process_fn(input_midi_path, instruments_text, generation_length, temperature, focus, render_audio):
    """Generate a multi-instrument continuation of an uploaded MIDI seed, an
    instrument-informed piece for a requested instrument list, or a fully
    unconditioned sample -- in that priority order -- and return it as MIDI
    + (if requested) a rendered audio preview + a report of what happened
    with the requested instruments."""
    if input_midi_path:
        prefix_codes = encode_seed(input_midi_path, encoding, train_args["max_beat"])
        instrument_note = "Seed MIDI provided -- Instruments field ignored."
        report_name = Path(input_midi_path).name
    else:
        instrument_names, instrument_note = resolve_instruments(instruments_text)
        prefix_codes = build_instrument_prefix(instrument_names) if instrument_names else None
        report_name = "(no seed MIDI)"

    full_codes = generate_tokens(prefix_codes, generation_length, temperature, focus)
    music = representation.decode(full_codes, encoding)

    midi_path = get_default_path(ext=".mid")
    music.write(midi_path)

    audio_path = None
    if render_audio:
        audio_path = get_default_path(ext=".wav")
        # Uses the MuseScore General soundfont muspy fetches on first use. The original repo's
        # polyphony option isn't supported by muspy==0.5.0's write_audio() (no passthrough kwarg).
        music.write(audio_path)

    report_path = get_default_path(ext=".txt")
    Path(report_path).write_text(f"{report_name}\n\nInstrument Matching\n{instrument_note}\n")

    return midi_path, audio_path, report_path


model_card = ModelCard(
    name="Multitrack Music Transformer",
    description=(
        "Generates symbolic multi-instrument music. Upload a MIDI seed to continue it, "
        "list instruments to write a fresh piece for exactly those, or leave both empty "
        "for a fully unconditioned generation. Trained on the Symbolic Orchestral "
        "Database (SOD)."
    ),
    author="Hao-Wen Dong, Ke Chen, Shlomo Dubnov, Julian McAuley, Taylor Berg-Kirkpatrick",
    tags=["symbolic-music", "midi", "multitrack", "transformer"],
)


with gr.Blocks() as demo:
    input_components = [
        gr.File(
            type="filepath",
            label="Seed MIDI (optional)",
            file_types=[".mid", ".midi"],
        ).harp_required(False),
        gr.Textbox(
            value="",
            label="Instruments",
            info=(
                "Only used when no Seed MIDI is given. Comma-separated instrument names -- "
                "the model writes a fresh piece for exactly this set. Leave blank to let the "
                f"model pick its own instruments. Valid names: {_VALID_INSTRUMENTS}."
            ),
        ),
        gr.Slider(
            minimum=64,
            maximum=1024,
            step=64,
            value=1024,
            label="Generation Length",
            info="Number of musical events to generate (default: 1024, per repo config)",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            step=0.1,
            value=1.0,
            label="Randomness",
            info="Sampling temperature -- higher wanders more (default: 1.0, per repo config)",
        ),
        gr.Slider(
            minimum=0.5,
            maximum=0.99,
            step=0.01,
            value=0.9,
            label="Focus",
            info="Higher keeps only the most likely notes; lower allows more variety (default: 0.9, per repo config)",
        ),
        gr.Checkbox(
            value=True,
            label="Render Audio Preview",
            info="Turn off to skip audio synthesis and get MIDI only (faster).",
        ),
    ]
    output_components = [
        gr.File(
            label="Generated MIDI",
            file_types=[".mid", ".midi"],
        ).set_info("The generated multi-instrument MIDI."),
        gr.Audio(
            type="filepath",
            label="Audio Preview",
        ).set_info("A synthesized audio preview of the generated MIDI."),
        gr.File(
            type="filepath",
            label="Instrument Matching",
            file_types=[".txt"],
        ).set_info("Which requested instrument names were matched or ignored."),
    ]

    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

if __name__ == "__main__":
    demo.queue().launch(pwa=True)