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)