import spaces import random import argparse import glob import json import os import time from concurrent.futures import ThreadPoolExecutor import gradio as gr import numpy as np import onnxruntime as rt import tqdm from huggingface_hub import hf_hub_download import MIDI from midi_synthesizer import MidiSynthesizer from midi_tokenizer import MIDITokenizer MAX_SEED = np.iinfo(np.int32).max in_space = os.getenv("SYSTEM") == "spaces" def softmax(x, axis): x_max = np.amax(x, axis=axis, keepdims=True) exp_x_shifted = np.exp(x - x_max) return exp_x_shifted / np.sum(exp_x_shifted, axis=axis, keepdims=True) def sample_top_p_k(probs, p, k, generator=None): if generator is None: generator = np.random probs_idx = np.argsort(-probs, axis=-1) probs_sort = np.take_along_axis(probs, probs_idx, -1) probs_sum = np.cumsum(probs_sort, axis=-1) mask = probs_sum - probs_sort > p probs_sort[mask] = 0.0 mask = np.zeros(probs_sort.shape[-1]) mask[:k] = 1 probs_sort = probs_sort * mask probs_sort /= np.sum(probs_sort, axis=-1, keepdims=True) shape = probs_sort.shape probs_sort_flat = probs_sort.reshape(-1, shape[-1]) probs_idx_flat = probs_idx.reshape(-1, shape[-1]) next_token = np.stack([generator.choice(idxs, p=pvals) for pvals, idxs in zip(probs_sort_flat, probs_idx_flat)]) next_token = next_token.reshape(*shape[:-1]) return next_token def apply_io_binding(model: rt.InferenceSession, inputs, outputs, batch_size, past_len, cur_len): io_binding = model.io_binding() for input_ in model.get_inputs(): name = input_.name if name.startswith("past_key_values"): present_name = name.replace("past_key_values", "present") if present_name in outputs: v = outputs[present_name] else: v = rt.OrtValue.ortvalue_from_shape_and_type( (batch_size, input_.shape[1], past_len, input_.shape[3]), element_type=np.float32, device_type=device) inputs[name] = v else: v = inputs[name] io_binding.bind_ortvalue_input(name, v) for output in model.get_outputs(): name = output.name if name.startswith("present"): v = rt.OrtValue.ortvalue_from_shape_and_type( (batch_size, output.shape[1], cur_len, output.shape[3]), element_type=np.float32, device_type=device) outputs[name] = v else: v = outputs[name] io_binding.bind_ortvalue_output(name, v) return io_binding def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98, top_k=20, disable_patch_change=False, disable_control_change=False, disable_channels=None, generator=None): tokenizer = model[2] if disable_channels is not None: disable_channels = [tokenizer.parameter_ids["channel"][c] for c in disable_channels] else: disable_channels = [] if generator is None: generator = np.random max_token_seq = tokenizer.max_token_seq if prompt is None: input_tensor = np.full((1, max_token_seq), tokenizer.pad_id, dtype=np.int64) input_tensor[0, 0] = tokenizer.bos_id # bos input_tensor = input_tensor[None, :, :] input_tensor = np.repeat(input_tensor, repeats=batch_size, axis=0) else: if len(prompt.shape) == 2: prompt = prompt[None, :] prompt = np.repeat(prompt, repeats=batch_size, axis=0) elif prompt.shape[0] == 1: prompt = np.repeat(prompt, repeats=batch_size, axis=0) elif len(prompt.shape) != 3 or prompt.shape[0] != batch_size: raise ValueError(f"invalid shape for prompt, {prompt.shape}") prompt = prompt[..., :max_token_seq] if prompt.shape[-1] < max_token_seq: prompt = np.pad(prompt, ((0, 0), (0, 0), (0, max_token_seq - prompt.shape[-1])), mode="constant", constant_values=tokenizer.pad_id) input_tensor = prompt cur_len = input_tensor.shape[1] bar = tqdm.tqdm(desc="generating", total=max_len - cur_len, disable=in_space) model0_inputs = {} model0_outputs = {} emb_size = 1024 for output in model[0].get_outputs(): if output.name == "hidden": emb_size = output.shape[2] past_len = 0 with bar: while cur_len < max_len: end = [False] * batch_size model0_inputs["x"] = rt.OrtValue.ortvalue_from_numpy(input_tensor[:, past_len:], device_type=device) model0_outputs["hidden"] = rt.OrtValue.ortvalue_from_shape_and_type( (batch_size, cur_len - past_len, emb_size), element_type=np.float32, device_type=device) io_binding = apply_io_binding(model[0], model0_inputs, model0_outputs, batch_size, past_len, cur_len) io_binding.synchronize_inputs() model[0].run_with_iobinding(io_binding) io_binding.synchronize_outputs() hidden = model0_outputs["hidden"].numpy()[:, -1:] next_token_seq = np.zeros((batch_size, 0), dtype=np.int64) event_names = [""] * batch_size model1_inputs = {"hidden": rt.OrtValue.ortvalue_from_numpy(hidden, device_type=device)} model1_outputs = {} for i in range(max_token_seq): mask = np.zeros((batch_size, tokenizer.vocab_size), dtype=np.int64) for b in range(batch_size): if end[b]: mask[b, tokenizer.pad_id] = 1 continue if i == 0: mask_ids = list(tokenizer.event_ids.values()) + [tokenizer.eos_id] if disable_patch_change: mask_ids.remove(tokenizer.event_ids["patch_change"]) if disable_control_change: mask_ids.remove(tokenizer.event_ids["control_change"]) mask[b, mask_ids] = 1 else: param_names = tokenizer.events[event_names[b]] if i > len(param_names): mask[b, tokenizer.pad_id] = 1 continue param_name = param_names[i - 1] mask_ids = tokenizer.parameter_ids[param_name] if param_name == "channel": mask_ids = [i for i in mask_ids if i not in disable_channels] mask[b, mask_ids] = 1 mask = mask[:, None, :] x = next_token_seq if i != 0: # cached if i == 1: hidden = np.zeros((batch_size, 0, emb_size), dtype=np.float32) model1_inputs["hidden"] = rt.OrtValue.ortvalue_from_numpy(hidden, device_type=device) x = x[:, -1:] model1_inputs["x"] = rt.OrtValue.ortvalue_from_numpy(x, device_type=device) model1_outputs["y"] = rt.OrtValue.ortvalue_from_shape_and_type( (batch_size, 1, tokenizer.vocab_size), element_type=np.float32, device_type=device ) io_binding = apply_io_binding(model[1], model1_inputs, model1_outputs, batch_size, i, i+1) io_binding.synchronize_inputs() model[1].run_with_iobinding(io_binding) io_binding.synchronize_outputs() logits = model1_outputs["y"].numpy() scores = softmax(logits / temp, -1) * mask samples = sample_top_p_k(scores, top_p, top_k, generator) if i == 0: next_token_seq = samples for b in range(batch_size): if end[b]: continue eid = samples[b].item() if eid == tokenizer.eos_id: end[b] = True else: event_names[b] = tokenizer.id_events[eid] else: next_token_seq = np.concatenate([next_token_seq, samples], axis=1) if all([len(tokenizer.events[event_names[b]]) == i for b in range(batch_size) if not end[b]]): break if next_token_seq.shape[1] < max_token_seq: next_token_seq = np.pad(next_token_seq, ((0, 0), (0, max_token_seq - next_token_seq.shape[-1])), mode="constant", constant_values=tokenizer.pad_id) next_token_seq = next_token_seq[:, None, :] input_tensor = np.concatenate([input_tensor, next_token_seq], axis=1) past_len = cur_len cur_len += 1 bar.update(1) yield next_token_seq[:, 0] if all(end): break def create_msg(name, data): return {"name": name, "data": data} def send_msgs(msgs): return json.dumps(msgs) def get_duration(model_name, tab, mid_seq, continuation_state, continuation_select, instruments, drum_kit, bpm, time_sig, key_sig, mid, midi_events, reduce_cc_st, remap_track_channel, add_default_instr, remove_empty_channels, seed, seed_rand, gen_events, temp, top_p, top_k, allow_cc): t = gen_events // 28 if "large" in model_name: t = gen_events // 20 return t + 10 @spaces.GPU(duration=get_duration) def run(model_name, tab, mid_seq, continuation_state, continuation_select, instruments, drum_kit, bpm, time_sig, key_sig, mid, midi_events, reduce_cc_st, remap_track_channel, add_default_instr, remove_empty_channels, seed, seed_rand, gen_events, temp, top_p, top_k, allow_cc): model = models[model_name] model_base = rt.InferenceSession(model[0], providers=providers) model_token = rt.InferenceSession(model[1], providers=providers) tokenizer = model[2] model = [model_base, model_token, tokenizer] bpm = int(bpm) if time_sig == "auto": time_sig = None time_sig_nn = 4 time_sig_dd = 2 else: time_sig_nn, time_sig_dd = time_sig.split('/') time_sig_nn = int(time_sig_nn) time_sig_dd = {2: 1, 4: 2, 8: 3}[int(time_sig_dd)] if key_sig == 0: key_sig = None key_sig_sf = 0 key_sig_mi = 0 else: key_sig = (key_sig - 1) key_sig_sf = key_sig // 2 - 7 key_sig_mi = key_sig % 2 gen_events = int(gen_events) max_len = gen_events if seed_rand: seed = random.randint(0, MAX_SEED) generator = np.random.RandomState(seed) disable_patch_change = False disable_channels = None if tab == 0: i = 0 mid = [[tokenizer.bos_id] + [tokenizer.pad_id] * (tokenizer.max_token_seq - 1)] if tokenizer.version == "v2": if time_sig is not None: mid.append(tokenizer.event2tokens(["time_signature", 0, 0, 0, time_sig_nn - 1, time_sig_dd - 1])) if key_sig is not None: mid.append(tokenizer.event2tokens(["key_signature", 0, 0, 0, key_sig_sf + 7, key_sig_mi])) if bpm != 0: mid.append(tokenizer.event2tokens(["set_tempo", 0, 0, 0, bpm])) patches = {} if instruments is None: instruments = [] for instr in instruments: patches[i] = patch2number[instr] i = (i + 1) if i != 8 else 10 if drum_kit != "None": patches[9] = drum_kits2number[drum_kit] for i, (c, p) in enumerate(patches.items()): mid.append(tokenizer.event2tokens(["patch_change", 0, 0, i + 1, c, p])) mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64) mid_seq = mid.tolist() if len(instruments) > 0: disable_patch_change = True disable_channels = [i for i in range(16) if i not in patches] elif tab == 1 and mid is not None: eps = 4 if reduce_cc_st else 0 mid = tokenizer.tokenize(MIDI.midi2score(mid), cc_eps=eps, tempo_eps=eps, remap_track_channel=remap_track_channel, add_default_instr=add_default_instr, remove_empty_channels=remove_empty_channels) mid = mid[:int(midi_events)] mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64) mid_seq = mid.tolist() elif tab == 2 and mid_seq is not None: mid = np.asarray(mid_seq, dtype=np.int64) if continuation_select > 0: continuation_state.append(mid_seq) mid = np.repeat(mid[continuation_select - 1:continuation_select], repeats=OUTPUT_BATCH_SIZE, axis=0) mid_seq = mid.tolist() else: continuation_state.append(mid.shape[1]) else: continuation_state = [0] mid = [[tokenizer.bos_id] + [tokenizer.pad_id] * (tokenizer.max_token_seq - 1)] mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64) mid_seq = mid.tolist() if mid is not None: max_len += mid.shape[1] init_msgs = [create_msg("progress", [0, gen_events])] if not (tab == 2 and continuation_select == 0): for i in range(OUTPUT_BATCH_SIZE): events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]] init_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]), create_msg("visualizer_append", [i, events])] yield mid_seq, continuation_state, seed, send_msgs(init_msgs) midi_generator = generate(model, mid, batch_size=OUTPUT_BATCH_SIZE, max_len=max_len, temp=temp, top_p=top_p, top_k=top_k, disable_patch_change=disable_patch_change, disable_control_change=not allow_cc, disable_channels=disable_channels, generator=generator) events = [list() for i in range(OUTPUT_BATCH_SIZE)] t = time.time() + 1 for i, token_seqs in enumerate(midi_generator): token_seqs = token_seqs.tolist() for j in range(OUTPUT_BATCH_SIZE): token_seq = token_seqs[j] mid_seq[j].append(token_seq) events[j].append(tokenizer.tokens2event(token_seq)) if time.time() - t > 0.5: msgs = [create_msg("progress", [i + 1, gen_events])] for j in range(OUTPUT_BATCH_SIZE): msgs += [create_msg("visualizer_append", [j, events[j]])] events[j] = list() yield mid_seq, continuation_state, seed, send_msgs(msgs) t = time.time() yield mid_seq, continuation_state, seed, send_msgs([]) def finish_run(model_name, mid_seq): if mid_seq is None: outputs = [None] * OUTPUT_BATCH_SIZE return *outputs, [] tokenizer = models[model_name][2] outputs = [] end_msgs = [create_msg("progress", [0, 0])] if not os.path.exists("outputs"): os.mkdir("outputs") for i in range(OUTPUT_BATCH_SIZE): events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]] mid = tokenizer.detokenize(mid_seq[i]) with open(f"outputs/output{i + 1}.mid", 'wb') as f: f.write(MIDI.score2midi(mid)) outputs.append(f"outputs/output{i + 1}.mid") end_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]), create_msg("visualizer_append", [i, events]), create_msg("visualizer_end", i)] return *outputs, send_msgs(end_msgs) def synthesis_task(mid): return synthesizer.synthesis(MIDI.score2opus(mid)) def render_audio(model_name, mid_seq, should_render_audio): if (not should_render_audio) or mid_seq is None: outputs = [None] * OUTPUT_BATCH_SIZE return tuple(outputs) tokenizer = models[model_name][2] outputs = [] if not os.path.exists("outputs"): os.mkdir("outputs") audio_futures = [] for i in range(OUTPUT_BATCH_SIZE): mid = tokenizer.detokenize(mid_seq[i]) audio_future = thread_pool.submit(synthesis_task, mid) audio_futures.append(audio_future) for future in audio_futures: outputs.append((44100, future.result())) if OUTPUT_BATCH_SIZE == 1: return outputs[0] return tuple(outputs) def undo_continuation(model_name, mid_seq, continuation_state): if mid_seq is None or len(continuation_state) < 2: return mid_seq, continuation_state, send_msgs([]) tokenizer = models[model_name][2] if isinstance(continuation_state[-1], list): mid_seq = continuation_state[-1] else: mid_seq = [ms[:continuation_state[-1]] for ms in mid_seq] continuation_state = continuation_state[:-1] end_msgs = [create_msg("progress", [0, 0])] for i in range(OUTPUT_BATCH_SIZE): events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]] end_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]), create_msg("visualizer_append", [i, events]), create_msg("visualizer_end", i)] return mid_seq, continuation_state, send_msgs(end_msgs) def load_javascript(dir="javascript"): scripts_list = glob.glob(f"{dir}/*.js") javascript = "" for path in scripts_list: with open(path, "r", encoding="utf8") as jsfile: js_content = jsfile.read() js_content = js_content.replace("const MIDI_OUTPUT_BATCH_SIZE=4;", f"const MIDI_OUTPUT_BATCH_SIZE={OUTPUT_BATCH_SIZE};") javascript += f"\n" template_response_ori = gr.routes.templates.TemplateResponse def template_response(*args, **kwargs): res = template_response_ori(*args, **kwargs) res.body = res.body.replace( b'', f'{javascript}'.encode("utf8")) res.init_headers() return res gr.routes.templates.TemplateResponse = template_response def hf_hub_download_retry(repo_id, filename): print(f"downloading {repo_id} {filename}") retry = 0 err = None while retry < 30: try: return hf_hub_download(repo_id=repo_id, filename=filename) except Exception as e: err = e retry += 1 if err: raise err def get_tokenizer(repo_id): config_path = hf_hub_download_retry(repo_id=repo_id, filename=f"config.json") with open(config_path, "r") as f: config = json.load(f) tokenizer = MIDITokenizer(config["tokenizer"]["version"]) tokenizer.set_optimise_midi(config["tokenizer"]["optimise_midi"]) return tokenizer number2drum_kits = {-1: "None", 0: "Standard", 8: "Room", 16: "Power", 24: "Electric", 25: "TR-808", 32: "Jazz", 40: "Blush", 48: "Orchestra"} patch2number = {v: k for k, v in MIDI.Number2patch.items()} drum_kits2number = {v: k for k, v in number2drum_kits.items()} key_signatures = ['C♭', 'A♭m', 'G♭', 'E♭m', 'D♭', 'B♭m', 'A♭', 'Fm', 'E♭', 'Cm', 'B♭', 'Gm', 'F', 'Dm', 'C', 'Am', 'G', 'Em', 'D', 'Bm', 'A', 'F♯m', 'E', 'C♯m', 'B', 'G♯m', 'F♯', 'D♯m', 'C♯', 'A♯m'] if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--share", action="store_true", default=False, help="share gradio app") parser.add_argument("--port", type=int, default=7860, help="gradio server port") parser.add_argument("--device", type=str, default="cuda", help="device to run model") parser.add_argument("--batch", type=int, default=8, help="batch size") parser.add_argument("--max-gen", type=int, default=1024, help="max") opt = parser.parse_args() OUTPUT_BATCH_SIZE = opt.batch soundfont_path = hf_hub_download_retry(repo_id="skytnt/midi-model", filename="soundfont.sf2") thread_pool = ThreadPoolExecutor(max_workers=OUTPUT_BATCH_SIZE) synthesizer = MidiSynthesizer(soundfont_path) models_info = { "generic pretrain model (tv2o-medium) by breadlicker45": [ "breadlicker45/bread-tv2o-medium", "", {} ], "generic pretrain model (tv2o-medium) by skytnt": [ "skytnt/midi-model-tv2o-medium", "", { "jpop": "skytnt/midi-model-tv2om-jpop-lora", "touhou": "skytnt/midi-model-tv2om-touhou-lora" } ], "generic pretrain model (tv2o-large) by asigalov61": [ "asigalov61/Music-Llama", "", {} ], "generic pretrain model (tv2o-medium) by asigalov61": [ "asigalov61/Music-Llama-Medium", "", {} ], "generic pretrain model (tv1-medium) by skytnt": [ "skytnt/midi-model", "", {} ] } models = {} providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] device = "cuda" for name, (repo_id, path, loras) in models_info.items(): model_base_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}onnx/model_base.onnx") model_token_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}onnx/model_token.onnx") tokenizer = get_tokenizer(repo_id) models[name] = [model_base_path, model_token_path, tokenizer] for lora_name, lora_repo in loras.items(): model_base_path = hf_hub_download_retry(repo_id=lora_repo, filename=f"onnx/model_base.onnx") model_token_path = hf_hub_download_retry(repo_id=lora_repo, filename=f"onnx/model_token.onnx") models[f"{name} with {lora_name} lora"] = [model_base_path, model_token_path, tokenizer] load_javascript() app = gr.Blocks(theme=gr.themes.Soft()) with app: gr.Markdown("