.

.gitignore

@@ -0,0 +1,2 @@
+amt/
+examples/

app.py

@@ -0,0 +1,56 @@
+import glob
+import gradio as gr
+
+from gradio_helper import *
+
+AUDIO_EXAMPLES = glob.glob('/content/examples/*.*', recursive=True)
+YOUTUBE_EXAMPLES = ["https://www.youtube.com/watch?v=vMboypSkj3c"]
+
+theme = 'gradio/dracula_revamped' #'Insuz/Mocha' #gr.themes.Soft()
+with gr.Blocks(theme=theme) as demo:
+
+    with gr.Row():
+        with gr.Column(scale=10):
+            gr.Markdown(
+            """
+            # YourMT3+: Bridging the Gap in Multi-instrument Music Transcription with Advanced Model Architectures and Cross-dataset Stem Augmentation
+            """)
+
+    with gr.Group():
+        with gr.Tab("Upload audio"):
+            # Input
+            audio_input = gr.Audio(label="Record Audio", type="filepath",
+                                show_share_button=True, show_download_button=True)
+            # Display examples
+            gr.Examples(examples=AUDIO_EXAMPLES, inputs=audio_input)
+            # Submit button
+            transcribe_audio_button = gr.Button("Transcribe", variant="primary")
+            # Transcribe
+            output_tab1 = gr.HTML()
+            # audio_output = gr.Text(label="Audio Info")
+            # transcribe_audio_button.click(process_audio, inputs=audio_input, outputs=output_tab1)
+            transcribe_audio_button.click(process_audio, inputs=audio_input, outputs=output_tab1)
+
+        with gr.Tab("From YouTube"):
+            with gr.Row():
+                # Input URL
+                youtube_url = gr.Textbox(label="YouTube Link URL",
+                        placeholder="https://youtu.be/...")
+                # Play youtube
+                youtube_player = gr.HTML(render=True)
+            with gr.Row():
+                # Play button
+                play_video_button = gr.Button("Play", variant="primary")
+                # Submit button
+                transcribe_video_button = gr.Button("Transcribe", variant="primary")
+            # Transcribe
+            output_tab2 = gr.HTML(render=True)
+            # video_output = gr.Text(label="Video Info")
+            transcribe_video_button.click(process_video, inputs=youtube_url, outputs=output_tab2)
+            # Play
+            play_video_button.click(play_video, inputs=youtube_url, outputs=youtube_player)
+
+            # Display examples
+            gr.Examples(examples=YOUTUBE_EXAMPLES, inputs=youtube_url)
+
+demo.launch(debug=True)

gradio_helper.py

@@ -0,0 +1,80 @@
+# @title GradIO helper
+import os
+import subprocess
+import glob
+from typing import Tuple, Dict, Literal
+from ctypes import ArgumentError
+# from google.colab import output
+
+from model_helper import *
+from html_helper import *
+
+from pytube import YouTube
+import gradio as gr
+import torchaudio
+
+def prepare_media(source_path_or_url: os.PathLike,
+                  source_type: Literal['audio_filepath', 'youtube_url'],
+                  delete_video: bool = True) -> Dict:
+    """prepare media from source path or youtube, and return audio info"""
+    # Get audio_file
+    if source_type == 'audio_filepath':
+        audio_file = source_path_or_url
+    elif source_type == 'youtube_url':
+        # Download from youtube
+        try:
+            # Try PyTube first
+            yt = YouTube(source_path_or_url)
+            audio_stream = min(yt.streams.filter(only_audio=True), key=lambda s: s.bitrate)
+            mp4_file = audio_stream.download(output_path='downloaded') # ./downloaded
+            audio_file = mp4_file[:-3] + 'mp3'
+            subprocess.run(['ffmpeg', '-i', mp4_file, '-ac', '1', audio_file])
+            os.remove(mp4_file)
+        except Exception as e:
+            try:
+                # Try alternative
+                print(f"Failed with PyTube, error: {e}. Trying yt-dlp...")
+                audio_file = './downloaded/yt_audio'
+                subprocess.run(['yt-dlp', '-x', source_path_or_url, '-f', 'bestaudio',
+                    '-o', audio_file, '--audio-format', 'mp3', '--restrict-filenames',
+                    '--force-overwrites'])
+                audio_file += '.mp3'
+            except Exception as e:
+                print(f"Alternative downloader failed, error: {e}. Please try again later!")
+                return None
+    else:
+        raise ValueError(source_type)
+
+    # Create info
+    info = torchaudio.info(audio_file)
+    return {
+        "filepath": audio_file,
+        "track_name": os.path.basename(audio_file).split('.')[0],
+        "sample_rate": int(info.sample_rate),
+        "bits_per_sample": int(info.bits_per_sample),
+        "num_channels": int(info.num_channels),
+        "num_frames": int(info.num_frames),
+        "duration": int(info.num_frames / info.sample_rate),
+        "encoding": str.lower(info.encoding),
+        }
+
+def process_audio(audio_filepath):
+    if audio_filepath is None:
+        return None
+    audio_info = prepare_media(audio_filepath, source_type='audio_filepath')
+    midifile = transcribe(model, audio_info)
+    midifile = to_data_url(midifile)
+    return create_html_from_midi(midifile) # html midiplayer
+
+def process_video(youtube_url):
+    if 'youtu' not in youtube_url:
+        return None
+    audio_info = prepare_media(youtube_url, source_type='youtube_url')
+    midifile = transcribe(model, audio_info)
+    midifile = to_data_url(midifile)
+    return create_html_from_midi(midifile) # html midiplayer
+
+def play_video(youtube_url):
+    if 'youtu' not in youtube_url:
+        return None
+    return create_html_youtube_player(youtube_url)

html_helper.py

@@ -0,0 +1,100 @@
+# @title HTML helper
+import re
+import base64
+def to_data_url(midi_filename):
+    """ This is crucial for Colab/WandB support. Thanks to Scott Hawley!!
+        https://github.com/drscotthawley/midi-player/blob/main/midi_player/midi_player.py
+
+    """
+    with open(midi_filename, "rb") as f:
+        encoded_string = base64.b64encode(f.read())
+    return 'data:audio/midi;base64,'+encoded_string.decode('utf-8')
+
+
+def to_youtube_embed_url(video_url):
+    regex = r"(?:https:\/\/)?(?:www\.)?(?:youtube\.com|youtu\.be)\/(?:watch\?v=)?(.+)"
+    return re.sub(regex, r"https://www.youtube.com/embed/\1",video_url)
+
+
+def create_html_from_midi(midifile):
+    html_template = """
+<!DOCTYPE html>
+<html>
+<head>
+  <title>Awesome MIDI Player</title>
+  <script src="https://cdn.jsdelivr.net/combine/npm/tone@14.7.58,npm/@magenta/music@1.23.1/es6/core.js,npm/focus-visible@5,npm/html-midi-player@1.5.0">
+  </script>
+  <style>
+    /* Background color for the section */
+    #proll {{background-color:transparent}}
+
+    /* Custom player style */
+    #proll midi-player {{
+      display: block;
+      width: inherit;
+      margin: 4px;
+      margin-bottom: 0;
+    }}
+
+    #proll midi-player::part(control-panel) {{
+      background: #D8DAE8;
+      border-radius: 8px 8px 0 0;
+      border: 1px solid #A0A0A0;
+    }}
+
+    /* Custom visualizer style */
+    #proll midi-visualizer .piano-roll-visualizer {{
+      background: #F7FAFA;
+      border-radius: 0 0 8px 8px;
+      border: 1px solid #A0A0A0;
+      margin: 4px;
+      margin-top: 2;
+      overflow: auto;
+    }}
+
+    #proll midi-visualizer svg rect.note {{
+      opacity: 0.6;
+      stroke-width: 2;
+    }}
+
+    #proll midi-visualizer svg rect.note[data-instrument="0"] {{
+      fill: #e22;
+      stroke: #055;
+    }}
+
+    #proll midi-visualizer svg rect.note[data-instrument="2"] {{
+      fill: #2ee;
+      stroke: #055;
+    }}
+
+    #proll midi-visualizer svg rect.note[data-is-drum="true"] {{
+      fill: #888;
+      stroke: #888;
+    }}
+
+    #proll midi-visualizer svg rect.note.active {{
+      opacity: 0.9;
+      stroke: #34384F;
+    }}
+  </style>
+</head>
+<body>
+  <div>
+    <a href="{midifile}" target="_blank">Download MIDI</a> <br>
+    <section id="proll">
+      <midi-player src="{midifile}" sound-font="https://storage.googleapis.com/magentadata/js/soundfonts/sgm_plus" visualizer="#proll midi-visualizer">
+      </midi-player>
+      <midi-visualizer src="{midifile}">
+      </midi-visualizer>
+    </section>
+  </div>
+</body>
+</html>
+""".format(midifile=midifile)
+    html = f"""<iframe style="width: 100%; height: 400px; overflow:auto" srcdoc='{html_template}'></iframe>"""
+    return html
+
+def create_html_youtube_player(youtube_url):
+    youtube_url = to_youtube_embed_url(youtube_url)
+    html = f"""<iframe width="560" height="315" src='{youtube_url}' title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>"""
+    return html

model_helper.py

@@ -0,0 +1,161 @@
+# @title Model helper
+%cd /content/amt/src
+from collections import Counter
+import argparse
+import torch
+import numpy as np
+
+from model.init_train import initialize_trainer, update_config
+from utils.task_manager import TaskManager
+from config.vocabulary import drum_vocab_presets
+from utils.utils import str2bool
+from utils.utils import Timer
+from utils.audio import slice_padded_array
+from utils.note2event import mix_notes
+from utils.event2note import merge_zipped_note_events_and_ties_to_notes
+from utils.utils import write_model_output_as_midi, write_err_cnt_as_json
+from model.ymt3 import YourMT3
+
+
+def load_model_checkpoint(args=None):
+    parser = argparse.ArgumentParser(description="YourMT3")
+    # General
+    parser.add_argument('exp_id', type=str, help='A unique identifier for the experiment is used to resume training. The "@" symbol can be used to load a specific checkpoint.')
+    parser.add_argument('-p', '--project', type=str, default='ymt3', help='project name')
+    parser.add_argument('-ac', '--audio-codec', type=str, default=None, help='audio codec (default=None). {"spec", "melspec"}. If None, default value defined in config.py will be used.')
+    parser.add_argument('-hop', '--hop-length', type=int, default=None, help='hop length in frames (default=None). {128, 300} 128 for MT3, 300 for PerceiverTFIf None, default value defined in config.py will be used.')
+    parser.add_argument('-nmel', '--n-mels', type=int, default=None, help='number of mel bins (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-if', '--input-frames', type=int, default=None, help='number of audio frames for input segment (default=None). If None, default value defined in config.py will be used.')
+    # Model configurations
+    parser.add_argument('-sqr', '--sca-use-query-residual', type=str2bool, default=None, help='sca use query residual flag. Default follows config.py')
+    parser.add_argument('-enc', '--encoder-type', type=str, default=None, help="Encoder type. 't5' or 'perceiver-tf' or 'conformer'. Default is 't5', following config.py.")
+    parser.add_argument('-dec', '--decoder-type', type=str, default=None, help="Decoder type. 't5' or 'multi-t5'. Default is 't5', following config.py.")
+    parser.add_argument('-preenc', '--pre-encoder-type', type=str, default='default', help="Pre-encoder type. None or 'conv' or 'default'. By default, t5_enc:None, perceiver_tf_enc:conv, conformer:None")
+    parser.add_argument('-predec', '--pre-decoder-type', type=str, default='default', help="Pre-decoder type. {None, 'linear', 'conv1', 'mlp', 'group_linear'} or 'default'. Default is {'t5': None, 'perceiver-tf': 'linear', 'conformer': None}.")
+    parser.add_argument('-cout', '--conv-out-channels', type=int, default=None, help='Number of filters for pre-encoder conv layer. Default follows "model_cfg" of config.py.')
+    parser.add_argument('-tenc', '--task-cond-encoder', type=str2bool, default=True, help='task conditional encoder (default=True). True or False')
+    parser.add_argument('-tdec', '--task-cond-decoder', type=str2bool, default=True, help='task conditional decoder (default=True). True or False')
+    parser.add_argument('-df', '--d-feat', type=int, default=None, help='Audio feature will be projected to this dimension for Q,K,V of T5 or K,V of Perceiver (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-pt', '--pretrained', type=str2bool, default=False, help='pretrained T5(default=False). True or False')
+    parser.add_argument('-b', '--base-name', type=str, default="google/t5-v1_1-small", help='base model name (default="google/t5-v1_1-small")')
+    parser.add_argument('-epe', '--encoder-position-encoding-type', type=str, default='default', help="Positional encoding type of encoder. By default, pre-defined PE for T5 or Perceiver-TF encoder in config.py. For T5: {'sinusoidal', 'trainable'}, conformer: {'rotary', 'trainable'}, Perceiver-TF: {'trainable', 'rope', 'alibi', 'alibit', 'None', '0', 'none', 'tkd', 'td', 'tk', 'kdt'}.")
+    parser.add_argument('-dpe', '--decoder-position-encoding-type', type=str, default='default', help="Positional encoding type of decoder. By default, pre-defined PE for T5 in config.py. {'sinusoidal', 'trainable'}.")
+    parser.add_argument('-twe', '--tie-word-embedding', type=str2bool, default=None, help='tie word embedding (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-el', '--event-length', type=int, default=None, help='event length (default=None). If None, default value defined in model cfg of config.py will be used.')
+    # Perceiver-TF configurations
+    parser.add_argument('-dl', '--d-latent', type=int, default=None, help='Latent dimension of Perceiver. On T5, this will be ignored (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-nl', '--num-latents', type=int, default=None, help='Number of latents of Perceiver. On T5, this will be ignored (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-dpm', '--perceiver-tf-d-model', type=int, default=None, help='Perceiver-TF d_model (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-npb', '--num-perceiver-tf-blocks', type=int, default=None, help='Number of blocks of Perceiver-TF. On T5, this will be ignored (default=None). If None, default value defined in config.py.')
+    parser.add_argument('-npl', '--num-perceiver-tf-local-transformers-per-block', type=int, default=None, help='Number of local layers per block of Perceiver-TF. On T5, this will be ignored (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-npt', '--num-perceiver-tf-temporal-transformers-per-block', type=int, default=None, help='Number of temporal layers per block of Perceiver-TF. On T5, this will be ignored (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-atc', '--attention-to-channel', type=str2bool, default=None, help='Attention to channel flag of Perceiver-TF. On T5, this will be ignored (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-ln', '--layer-norm-type', type=str, default=None, help='Layer normalization type (default=None). {"layer_norm", "rms_norm"}. If None, default value defined in config.py will be used.')
+    parser.add_argument('-ff', '--ff-layer-type', type=str, default=None, help='Feed forward layer type (default=None). {"mlp", "moe", "gmlp"}. If None, default value defined in config.py will be used.')
+    parser.add_argument('-wf', '--ff-widening-factor', type=int, default=None, help='Feed forward layer widening factor for MLP/MoE/gMLP (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-nmoe', '--moe-num-experts', type=int, default=None, help='Number of experts for MoE (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-kmoe', '--moe-topk', type=int, default=None, help='Top-k for MoE (default=None). If None, default value defined in config.py will be used.')
+    parser.add_argument('-act', '--hidden-act', type=str, default=None, help='Hidden activation function (default=None). {"gelu", "silu", "relu", "tanh"}. If None, default value defined in config.py will be used.')
+    parser.add_argument('-rt', '--rotary-type', type=str, default=None, help='Rotary embedding type expressed in three letters. e.g. ppl: "pixel" for SCA and latents, "lang" for temporal transformer. If None, use config.')
+    parser.add_argument('-rk', '--rope-apply-to-keys', type=str2bool, default=None, help='Apply rope to keys (default=None). If None, use config.')
+    parser.add_argument('-rp', '--rope-partial-pe', type=str2bool, default=None, help='Whether to apply RoPE to partial positions (default=None). If None, use config.')
+    # Decoder configurations
+    parser.add_argument('-dff', '--decoder-ff-layer-type', type=str, default=None, help='Feed forward layer type of decoder (default=None). {"mlp", "moe", "gmlp"}. If None, default value defined in config.py will be used.')
+    parser.add_argument('-dwf', '--decoder-ff-widening-factor', type=int, default=None, help='Feed forward layer widening factor for decoder MLP/MoE/gMLP (default=None). If None, default value defined in config.py will be used.')
+    # Task and Evaluation configurations
+    parser.add_argument('-tk', '--task', type=str, default='mt3_full_plus', help='tokenizer type (default=mt3_full_plus). See config/task.py for more options.')
+    parser.add_argument('-epv', '--eval-program-vocab', type=str, default=None, help='evaluation vocabulary (default=None). If None, default vocabulary of the data preset will be used.')
+    parser.add_argument('-edv', '--eval-drum-vocab', type=str, default=None, help='evaluation vocabulary for drum (default=None). If None, default vocabulary of the data preset will be used.')
+    parser.add_argument('-etk', '--eval-subtask-key', type=str, default='default', help='evaluation subtask key (default=default). See config/task.py for more options.')
+    parser.add_argument('-t', '--onset-tolerance', type=float, default=0.05, help='onset tolerance (default=0.05).')
+    parser.add_argument('-os', '--test-octave-shift', type=str2bool, default=False, help='test optimal octave shift (default=False). True or False')
+    parser.add_argument('-w', '--write-model-output', type=str2bool, default=True, help='write model test output to file (default=False). True or False')
+    # Trainer configurations
+    parser.add_argument('-pr','--precision', type=str, default="bf16-mixed", help='precision (default="bf16-mixed") {32, 16, bf16, bf16-mixed}')
+    parser.add_argument('-st', '--strategy', type=str, default='auto', help='strategy (default=auto). auto or deepspeed or ddp')
+    parser.add_argument('-n', '--num-nodes', type=int, default=1, help='number of nodes (default=1)')
+    parser.add_argument('-g', '--num-gpus', type=str, default='auto', help='number of gpus (default="auto")')
+    parser.add_argument('-wb', '--wandb-mode', type=str, default="disabled", help='wandb mode for logging (default=None). "disabled" or "online" or "offline". If None, default value defined in config.py will be used.')
+    # Debug
+    parser.add_argument('-debug', '--debug-mode', type=str2bool, default=False, help='debug mode (default=False). True or False')
+    parser.add_argument('-tps', '--test-pitch-shift', type=int, default=None, help='use pitch shift when testing. debug-purpose only. (default=None). semitone in int.')
+    args = parser.parse_args(args)
+    # yapf: enable
+    if torch.__version__ >= "1.13":
+        torch.set_float32_matmul_precision("high")
+    args.epochs = None
+
+    # Initialize and update config
+    _, _, dir_info, shared_cfg = initialize_trainer(args, stage='test')
+    shared_cfg, audio_cfg, model_cfg = update_config(args, shared_cfg, stage='test')
+
+    if args.eval_drum_vocab != None:  # override eval_drum_vocab
+        eval_drum_vocab = drum_vocab_presets[args.eval_drum_vocab]
+
+    # Initialize task manager
+    tm = TaskManager(task_name=args.task,
+                     max_shift_steps=int(shared_cfg["TOKENIZER"]["max_shift_steps"]),
+                     debug_mode=args.debug_mode)
+    print(f"Task: {tm.task_name}, Max Shift Steps: {tm.max_shift_steps}")
+
+    # Use GPU if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Model
+    model = YourMT3(
+        audio_cfg=audio_cfg,
+        model_cfg=model_cfg,
+        shared_cfg=shared_cfg,
+        optimizer=None,
+        task_manager=tm,  # tokenizer is a member of task_manager
+        eval_subtask_key=args.eval_subtask_key,
+        write_output_dir=dir_info["lightning_dir"] if args.write_model_output or args.test_octave_shift else None
+        ).to(device)
+    checkpoint = torch.load(dir_info["last_ckpt_path"])
+    state_dict = checkpoint['state_dict']
+    new_state_dict = {k: v for k, v in state_dict.items() if 'pitchshift' not in k}
+    model.load_state_dict(new_state_dict, strict=False)
+    return model.eval()
+
+
+def transcribe(model, audio_info):
+    t = Timer()
+
+    # Converting Audio
+    t.start()
+    audio, sr = torchaudio.load(uri=audio_info['filepath'])
+    audio = torch.mean(audio, dim=0).unsqueeze(0)
+    audio = torchaudio.functional.resample(audio, sr, model.audio_cfg['sample_rate'])
+    audio_segments = slice_padded_array(audio, model.audio_cfg['input_frames'], model.audio_cfg['input_frames'])
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    audio_segments = torch.from_numpy(audio_segments.astype('float32')).to(device).unsqueeze(1) # (n_seg, 1, seg_sz)
+    t.stop(); t.print_elapsed_time("converting audio");
+
+    # Inference
+    t.start()
+    pred_token_arr, _ = model.inference_file(bsz=8, audio_segments=audio_segments)
+    t.stop(); t.print_elapsed_time("model inference");
+
+    # Post-processing
+    t.start()
+    num_channels = model.task_manager.num_decoding_channels
+    n_items = audio_segments.shape[0]
+    start_secs_file = [model.audio_cfg['input_frames'] * i / model.audio_cfg['sample_rate'] for i in range(n_items)]
+    pred_notes_in_file = []
+    n_err_cnt = Counter()
+    for ch in range(num_channels):
+        pred_token_arr_ch = [arr[:, ch, :] for arr in pred_token_arr]  # (B, L)
+        zipped_note_events_and_tie, list_events, ne_err_cnt = model.task_manager.detokenize_list_batches(
+            pred_token_arr_ch, start_secs_file, return_events=True)
+        pred_notes_ch, n_err_cnt_ch = merge_zipped_note_events_and_ties_to_notes(zipped_note_events_and_tie)
+        pred_notes_in_file.append(pred_notes_ch)
+        n_err_cnt += n_err_cnt_ch
+    pred_notes = mix_notes(pred_notes_in_file)  # This is the mixed notes from all channels
+
+    # Write MIDI
+    write_model_output_as_midi(pred_notes, '/content/',
+                              audio_info['track_name'], model.midi_output_inverse_vocab)
+    t.stop(); t.print_elapsed_time("post processing");
+    midifile =  os.path.join('/content/model_output/', audio_info['track_name']  + '.mid')
+    assert os.path.exists(midifile)
+    return midifile