Spaces:
Running on Zero
Running on Zero
| import spaces # imported first, as ZeroGPU requires | |
| import os | |
| import pickle | |
| import torch | |
| import torch.nn as nn | |
| import gradio as gr | |
| from transformers import T5Tokenizer | |
| from huggingface_hub import hf_hub_download | |
| from time import time_ns | |
| from uuid import uuid4 | |
| from transformer_model import Transformer | |
| from pyharp.core import ModelCard, build_endpoint | |
| from pyharp.tags import Subcategory | |
| # HARP Model Card | |
| model_card = ModelCard( | |
| name="text2midi", | |
| description=( | |
| "Generates a MIDI sequence from a text description of the music. Repository: " | |
| "https://github.com/amaai-lab/text2midi Paper: https://arxiv.org/abs/2412.16526" | |
| ), | |
| author="Keshav Bhandari, Abhinaba Roy, Kyra Wang, Geeta Puri, Simon Colton, Dorien Herremans", | |
| tags=[Subcategory.HOLISTIC, 'domain:music'] | |
| ) | |
| # Model/artifacts from HF Hub | |
| REPO_ID = "amaai-lab/text2midi" | |
| REVISION = "9bc91479c4a9b08a9a28a7dcfb3f1a56ee2ded36" | |
| # Helpers | |
| def _unique_path(ext: str) -> str: | |
| """Create a unique file path in /tmp to avoid naming collisions.""" | |
| return os.path.join("/tmp", f"t2m_{time_ns()}_{uuid4().hex[:8]}{ext}") | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = None | |
| r_tokenizer = None | |
| tokenizer = None | |
| def load_model(): | |
| """Builds the model on the CPU, then places it on DEVICE.""" | |
| global model, r_tokenizer, tokenizer | |
| # Load REMI vocab/tokenizer (pickle dict used by the provided model) | |
| with open(hf_hub_download(repo_id=REPO_ID, filename="vocab_remi.pkl", revision=REVISION), "rb") as f: | |
| r_tokenizer = pickle.load(f) | |
| vocab_size = len(r_tokenizer) | |
| model = Transformer( | |
| vocab_size, # vocab size | |
| 768, # d_model | |
| 8, # nhead | |
| 2048, # dim_feedforward | |
| 18, # nlayers | |
| 1024, # max_seq_len | |
| False, # use_rotary | |
| 8, # rotary_dim | |
| device="cpu" # device | |
| ) | |
| model_path = hf_hub_download(repo_id=REPO_ID, filename="pytorch_model.bin", revision=REVISION) | |
| model.load_state_dict(torch.load(model_path, map_location="cpu")) | |
| model.eval() | |
| model.to(DEVICE) | |
| tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base", revision="7bcac572ce56db69c1ea7c8af255c5d7c9672fc2") | |
| # Built in the worker, which runs process_fn, on its main thread, where ZeroGPU's emulation records | |
| # the CUDA placement, so that no request spends the caller's allowance on building it. The server | |
| # never runs process_fn. | |
| if __name__ != "__main__": | |
| load_model() | |
| # Core Text -> MIDI | |
| def generate_midi(prompt: str, temperature: float = 0.9, max_len: int = 500) -> str: | |
| inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True) | |
| input_ids = nn.utils.rnn.pad_sequence(inputs.input_ids, batch_first=True, padding_value=0).to(DEVICE) | |
| attention_mask = nn.utils.rnn.pad_sequence(inputs.attention_mask, batch_first=True, padding_value=0).to(DEVICE) | |
| with torch.no_grad(): | |
| output = model.generate(input_ids, attention_mask, max_len=max_len, temperature=temperature) | |
| output_list = output[0].tolist() | |
| generated_midi = r_tokenizer.decode(output_list) | |
| midi_path = _unique_path(".mid") | |
| generated_midi.dump_midi(midi_path) | |
| return midi_path | |
| def gpu_seconds(prompt: str, temperature: float, max_length: int) -> float: | |
| # Measured at about 42 s for the longest, 1500 tokens, padded | |
| return 10 + 0.05 * int(max_length) | |
| def process_fn(prompt: str, temperature: float, max_length: int): | |
| """ | |
| Generate a MIDI from a text description. | |
| Args: | |
| prompt (str): What the music should be. | |
| temperature (float): How freely the model departs from its likeliest | |
| continuation. | |
| max_length (int): Ceiling on the number of tokens generated. | |
| Returns: | |
| midi_path (str): The generated MIDI. | |
| """ | |
| # Checked before the GPU is reserved, so a request the model cannot take spends none of the | |
| # caller's ZeroGPU allowance | |
| if not (prompt or "").strip(): | |
| raise gr.Error("Describe the music to generate.") | |
| return compose(prompt, temperature, max_length) | |
| def compose(prompt: str, temperature: float, max_length: int) -> str: | |
| try: | |
| return generate_midi(prompt, float(temperature), int(max_length)) | |
| except Exception as error: | |
| raise gr.Error(str(error)) | |
| # Gradio + HARP UI | |
| if __name__ == "__main__": | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## 🎶 text2midi") | |
| # Inputs | |
| prompt_in = gr.Textbox( | |
| label="Describe Your Music", | |
| info="Type a short phrase like 'calm piano with flowing arpeggios' ", | |
| ).harp_required(True) | |
| temperature_in = gr.Slider(minimum=0.8, maximum=1.1, value=0.9, step=0.1, label="Creativity", info=( | |
| "Adjusts how much freedom the model takes while composing.\n" | |
| "Lower = safer and more predictable (structured), " | |
| "Higher = more varied and expressive." | |
| ), interactive=True) | |
| maxlen_in = gr.Slider(minimum=500, maximum=1500, step=100, value=500, label="Composition Length", info=( | |
| "Determines how long the generated piece is in musical tokens.\n" | |
| "Higher values produce longer phrases (roughly more measures of music).") | |
| ) | |
| # Outputs | |
| midi_out = gr.File(label="Generated MIDI", file_types=[".mid", ".midi"], type="filepath").set_info(( | |
| "The generated notes, as a MIDI file." | |
| )) | |
| # Build HARP endpoint | |
| _ = build_endpoint( | |
| model_card=model_card, | |
| input_components=[prompt_in, temperature_in, maxlen_in], | |
| output_components=[midi_out], | |
| process_fn=process_fn | |
| ) | |
| # Launch App | |
| demo.queue().launch(pwa=True, show_error=True, ssr_mode=False) | |