Blame for 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 tempfile
from pathlib import Path

import gradio as gr
import torch
from pyharp import ModelCard, build_endpoint
from pyharp.tags import Subcategory

from muscriptor import TranscriptionModel
from muscriptor.tokenizer.mt3 import MT3_FULL_PLUS_GROUP_NAMES, resolve_instrument_names

model_card = ModelCard(
    name="MuScriptor",
    description=(
        "Multi-instrument audio-to-MIDI transcription, trained on 170k songs "
        "from classical music to heavy metal. The code is MIT, but the "
        "published weights are CC-BY-NC-4.0, so results cannot be used "
        "commercially. Repository: https://github.com/muscriptor/muscriptor "
        "Paper: https://arxiv.org/abs/2607.08168"
    ),
    author="Simon Rouard, Michael Krause, Axel Roebel, Carl-Johann Simon-Gabriel, Alexandre Défossez (Kyutai x Mirelo)",
    tags=[Subcategory.MUSIC_ANALYSIS, 'multi-instrument'],
)

_VALID_INSTRUMENTS = ", ".join(MT3_FULL_PLUS_GROUP_NAMES)

_models: dict[str, TranscriptionModel] = {}


def _get_model(variant: str) -> TranscriptionModel:
    """Load and cache a TranscriptionModel for a given size, one per variant.

    small/medium run on CPU; large runs on GPU. Built directly on its target
    device inside this GPU-decorated call, per ZeroGPU rules.
    """
    if variant not in _models:
        device = "cuda" if variant == "large" else "cpu"
        _models[variant] = TranscriptionModel.load_model(variant, device=device)
    return _models[variant]


def _resolve_instruments(text: str) -> tuple[list[str] | None, str]:
    """Resolve the comma-separated instruments box into exact group names.

    Names that don't resolve are dropped rather than raised: HARP's client
    shows generic error message, so an error here would not reach the user. 
    Second value returns what happened as a .txt output
    """
    tokens = [t for t in text.split(",") if t.strip()]
    if not tokens:
        return None, "No instrument restriction requested."

    resolved: list[str] = []
    problems: list[str] = []
    for token in tokens:
        try:
            resolved.extend(resolve_instrument_names([token]))
        except ValueError as e:
            problems.append(str(e))

    if not problems:
        return resolved, f"Matched: {', '.join(resolved)}."

    note = "; ".join(problems)
    if resolved:
        note = f"Matched: {', '.join(resolved)}. Ignored: {note}"
    else:
        note = f"No valid instrument names found, transcribing unrestricted. {note}"
    return (resolved or None), note


@spaces.GPU
@torch.inference_mode()
def process_fn(
    input_audio_path: str,
    variant: str,
    instruments_text: str,
    use_sampling: bool,
    temperature: float,
) -> tuple[str, str]:
    """Transcribe the input audio to MIDI, plus a note on the Instruments field."""
    model = _get_model(variant)
    instruments, instrument_note = _resolve_instruments(instruments_text)

    midi_bytes = model.transcribe_to_midi(
        input_audio_path,
        use_sampling=use_sampling,
        temperature=temperature,
        instruments=instruments,
    )

    with tempfile.NamedTemporaryFile(suffix=".mid", delete=False) as f:
        f.write(midi_bytes)
        output_midi_path = f.name

    notes_text = f"{Path(input_audio_path).name}\n\nInstrument Matching\n{instrument_note}\n"
    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
        f.write(notes_text)
        output_notes_path = f.name

    return output_midi_path, output_notes_path


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
                "Music to transcribe. Trained on a wide range of material, from "
                "classical to metal, and it names the instruments it hears."
            )),
            gr.Dropdown(
                choices=["small", "medium", "large"],
                value="medium",
                label="Model Size",
                info="small = fastest, least accurate. medium = balanced (default, per repo). large = most accurate, slower.",
            ),
            gr.Textbox(
                value="",
                label="Instruments",
                info=f"Separate names with commas. Leave blank to let the model detect instruments on its own. Valid names: {_VALID_INSTRUMENTS}.",
            ),
            gr.Checkbox(
                label="Use Sampling",
                info="Temperature sampling instead of greedy decoding (default: False, per repo).",
            ),
            gr.Slider(
                minimum=0.1,
                maximum=2.0,
                step=0.1,
                value=1.0,
                label="Temperature",
                info="Only used when Use Sampling is on. Higher values = note choices are more varied and unpredictable. Lower values = values are closer to model's most confident guesses. (default: 1.0, per repo).",
            ),
        ]
        output_components = [
            gr.File(type="filepath", file_types=[".mid", ".midi"], label="Transcription").set_info(
                "Transcribed MIDI notes."
            ),
            gr.File(type="filepath", file_types=[".txt"], label="Instrument Matching").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,
        )

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