Music_Transcription_with_Transformers.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LHt2xZUSqIzV"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/magenta/mt3/blob/main/mt3/colab/music_transcription_with_transformers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n",
        "# Music Transcription with Transformers\n",
        "\n",
        "This notebook is an interactive demo of a few [music transcription models](g.co/magenta/mt3) created by Google's [Magenta](g.co/magenta) team.  You can upload audio and have one of our models automatically transcribe it.\n",
        "\n",
        "<img src=\"https://magenta.tensorflow.org/assets/transcription-with-transformers/architecture_diagram.png\" alt=\"Transformer-based transcription architecture\">\n",
        "\n",
        "The notebook supports two pre-trained models:\n",
        "1. the piano transcription model from [our ISMIR 2021 paper](https://archives.ismir.net/ismir2021/paper/000030.pdf)\n",
        "1. the multi-instrument transcription model from [our ICLR 2022 paper](https://openreview.net/pdf?id=iMSjopcOn0p)\n",
        "\n",
        "**Caveat**: neither model is trained on singing.  If you upload audio with vocals, you will likely get weird results.  Multi-instrument transcription is still not a completely-solved problem and so you may get weird results regardless.\n",
        "\n",
        "In any case, we hope you have fun transcribing!  Feel free to tweet any interesting output at [@GoogleMagenta](https://twitter.com/googlemagenta)...\n",
        "\n",
        "### Instructions for running:\n",
        "\n",
        "* Make sure to use a GPU runtime, click:  __Runtime >> Change Runtime Type >> GPU__\n",
        "* Press ▶️ on the left of each cell to execute the cell\n",
        "* In the __Load Model__ cell, choose either `ismir2021` for piano transcription or `mt3` for multi-instrument transcription\n",
        "* In the __Upload Audio__ cell, choose an MP3 or WAV file from your computer when prompted\n",
        "* Transcribe the audio using the __Transcribe Audio__ cell (it may take a few minutes depending on the length of the audio)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ibSG_uu0QXgc"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 Google LLC. All Rights Reserved.\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# ==============================================================================\n",
        "\n",
        "#@title Setup Environment\n",
        "#@markdown Install MT3 and its dependencies (may take a few minutes).\n",
        "\n",
        "!apt-get update -qq && apt-get install -qq libfluidsynth1 build-essential libasound2-dev libjack-dev\n",
        "\n",
        "!pip install --upgrade pip\n",
        "!pip install nest-asyncio\n",
        "!pip install pyfluidsynth\n",
        "\n",
        "# install t5x\n",
        "!git clone --branch=main https://github.com/google-research/t5x\n",
        "!mv t5x t5x_tmp; mv t5x_tmp/* .; rm -r t5x_tmp\n",
        "!sed -i 's:jax\\[tpu\\]:jax:' setup.py\n",
        "!python3 -m pip install -e .\n",
        "\n",
        "# install mt3\n",
        "!git clone --branch=main https://github.com/magenta/mt3\n",
        "!mv mt3 mt3_tmp; mv mt3_tmp/* .; rm -r mt3_tmp\n",
        "!python3 -m pip install -e .\n",
        "\n",
        "# copy checkpoints\n",
        "!gsutil -q -m cp -r gs://mt3/checkpoints .\n",
        "\n",
        "# copy soundfont (originally from https://sites.google.com/site/soundfonts4u)\n",
        "!gsutil -q -m cp gs://magentadata/soundfonts/SGM-v2.01-Sal-Guit-Bass-V1.3.sf2 .\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TSMSWDxxWmTS"
      },
      "outputs": [],
      "source": [
        "#@title Imports and Definitions\n",
        "\n",
        "import functools\n",
        "import os\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow.compat.v2 as tf\n",
        "\n",
        "import functools\n",
        "import gin\n",
        "import jax\n",
        "import librosa\n",
        "import note_seq\n",
        "import seqio\n",
        "import t5\n",
        "import t5x\n",
        "\n",
        "from mt3 import metrics_utils\n",
        "from mt3 import models\n",
        "from mt3 import network\n",
        "from mt3 import note_sequences\n",
        "from mt3 import preprocessors\n",
        "from mt3 import spectrograms\n",
        "from mt3 import vocabularies\n",
        "\n",
        "from google.colab import files\n",
        "\n",
        "import nest_asyncio\n",
        "nest_asyncio.apply()\n",
        "\n",
        "SAMPLE_RATE = 16000\n",
        "SF2_PATH = 'SGM-v2.01-Sal-Guit-Bass-V1.3.sf2'\n",
        "\n",
        "def upload_audio(sample_rate):\n",
        "  data = list(files.upload().values())\n",
        "  if len(data) > 1:\n",
        "    print('Multiple files uploaded; using only one.')\n",
        "  return note_seq.audio_io.wav_data_to_samples_librosa(\n",
        "    data[0], sample_rate=sample_rate)\n",
        "\n",
        "\n",
        "\n",
        "class InferenceModel(object):\n",
        "  \"\"\"Wrapper of T5X model for music transcription.\"\"\"\n",
        "\n",
        "  def __init__(self, checkpoint_path, model_type='mt3'):\n",
        "\n",
        "    # Model Constants.\n",
        "    if model_type == 'ismir2021':\n",
        "      num_velocity_bins = 127\n",
        "      self.encoding_spec = note_sequences.NoteEncodingSpec\n",
        "      self.inputs_length = 512\n",
        "    elif model_type == 'mt3':\n",
        "      num_velocity_bins = 1\n",
        "      self.encoding_spec = note_sequences.NoteEncodingWithTiesSpec\n",
        "      self.inputs_length = 256\n",
        "    else:\n",
        "      raise ValueError('unknown model_type: %s' % model_type)\n",
        "\n",
        "    gin_files = ['/content/mt3/gin/model.gin',\n",
        "                 f'/content/mt3/gin/{model_type}.gin']\n",
        "\n",
        "    self.batch_size = 8\n",
        "    self.outputs_length = 1024\n",
        "    self.sequence_length = {'inputs': self.inputs_length, \n",
        "                            'targets': self.outputs_length}\n",
        "\n",
        "    self.partitioner = t5x.partitioning.ModelBasedPjitPartitioner(\n",
        "        model_parallel_submesh=(1, 1, 1, 1), num_partitions=1)\n",
        "\n",
        "    # Build Codecs and Vocabularies.\n",
        "    self.spectrogram_config = spectrograms.SpectrogramConfig()\n",
        "    self.codec = vocabularies.build_codec(\n",
        "        vocab_config=vocabularies.VocabularyConfig(\n",
        "            num_velocity_bins=num_velocity_bins))\n",
        "    self.vocabulary = vocabularies.vocabulary_from_codec(self.codec)\n",
        "    self.output_features = {\n",
        "        'inputs': seqio.ContinuousFeature(dtype=tf.float32, rank=2),\n",
        "        'targets': seqio.Feature(vocabulary=self.vocabulary),\n",
        "    }\n",
        "\n",
        "    # Create a T5X model.\n",
        "    self._parse_gin(gin_files)\n",
        "    self.model = self._load_model()\n",
        "\n",
        "    # Restore from checkpoint.\n",
        "    self.restore_from_checkpoint(checkpoint_path)\n",
        "\n",
        "  @property\n",
        "  def input_shapes(self):\n",
        "    return {\n",
        "          'encoder_input_tokens': (self.batch_size, self.inputs_length),\n",
        "          'decoder_input_tokens': (self.batch_size, self.outputs_length)\n",
        "    }\n",
        "\n",
        "  def _parse_gin(self, gin_files):\n",
        "    \"\"\"Parse gin files used to train the model.\"\"\"\n",
        "    gin_bindings = [\n",
        "        'from __gin__ import dynamic_registration',\n",
        "        'from mt3 import vocabularies',\n",
        "        'VOCAB_CONFIG=@vocabularies.VocabularyConfig()',\n",
        "        'vocabularies.VocabularyConfig.num_velocity_bins=%NUM_VELOCITY_BINS'\n",
        "    ]\n",
        "    with gin.unlock_config():\n",
        "      gin.parse_config_files_and_bindings(\n",
        "          gin_files, gin_bindings, finalize_config=False)\n",
        "\n",
        "  def _load_model(self):\n",
        "    \"\"\"Load up a T5X `Model` after parsing training gin config.\"\"\"\n",
        "    model_config = gin.get_configurable(network.T5Config)()\n",
        "    module = network.Transformer(config=model_config)\n",
        "    return models.ContinuousInputsEncoderDecoderModel(\n",
        "        module=module,\n",
        "        input_vocabulary=self.output_features['inputs'].vocabulary,\n",
        "        output_vocabulary=self.output_features['targets'].vocabulary,\n",
        "        optimizer_def=t5x.adafactor.Adafactor(decay_rate=0.8, step_offset=0),\n",
        "        input_depth=spectrograms.input_depth(self.spectrogram_config))\n",
        "\n",
        "\n",
        "  def restore_from_checkpoint(self, checkpoint_path):\n",
        "    \"\"\"Restore training state from checkpoint, resets self._predict_fn().\"\"\"\n",
        "    train_state_initializer = t5x.utils.TrainStateInitializer(\n",
        "      optimizer_def=self.model.optimizer_def,\n",
        "      init_fn=self.model.get_initial_variables,\n",
        "      input_shapes=self.input_shapes,\n",
        "      partitioner=self.partitioner)\n",
        "\n",
        "    restore_checkpoint_cfg = t5x.utils.RestoreCheckpointConfig(\n",
        "        path=checkpoint_path, mode='specific', dtype='float32')\n",
        "\n",
        "    train_state_axes = train_state_initializer.train_state_axes\n",
        "    self._predict_fn = self._get_predict_fn(train_state_axes)\n",
        "    self._train_state = train_state_initializer.from_checkpoint_or_scratch(\n",
        "        [restore_checkpoint_cfg], init_rng=jax.random.PRNGKey(0))\n",
        "\n",
        "  @functools.lru_cache()\n",
        "  def _get_predict_fn(self, train_state_axes):\n",
        "    \"\"\"Generate a partitioned prediction function for decoding.\"\"\"\n",
        "    def partial_predict_fn(params, batch, decode_rng):\n",
        "      return self.model.predict_batch_with_aux(\n",
        "          params, batch, decoder_params={'decode_rng': None})\n",
        "    return self.partitioner.partition(\n",
        "        partial_predict_fn,\n",
        "        in_axis_resources=(\n",
        "            train_state_axes.params,\n",
        "            t5x.partitioning.PartitionSpec('data',), None),\n",
        "        out_axis_resources=t5x.partitioning.PartitionSpec('data',)\n",
        "    )\n",
        "\n",
        "  def predict_tokens(self, batch, seed=0):\n",
        "    \"\"\"Predict tokens from preprocessed dataset batch.\"\"\"\n",
        "    prediction, _ = self._predict_fn(\n",
        "        self._train_state.params, batch, jax.random.PRNGKey(seed))\n",
        "    return self.vocabulary.decode_tf(prediction).numpy()\n",
        "\n",
        "  def __call__(self, audio):\n",
        "    \"\"\"Infer note sequence from audio samples.\n",
        "    \n",
        "    Args:\n",
        "      audio: 1-d numpy array of audio samples (16kHz) for a single example.\n",
        "\n",
        "    Returns:\n",
        "      A note_sequence of the transcribed audio.\n",
        "    \"\"\"\n",
        "    ds = self.audio_to_dataset(audio)\n",
        "    ds = self.preprocess(ds)\n",
        "\n",
        "    model_ds = self.model.FEATURE_CONVERTER_CLS(pack=False)(\n",
        "        ds, task_feature_lengths=self.sequence_length)\n",
        "    model_ds = model_ds.batch(self.batch_size)\n",
        "\n",
        "    inferences = (tokens for batch in model_ds.as_numpy_iterator()\n",
        "                  for tokens in self.predict_tokens(batch))\n",
        "\n",
        "    predictions = []\n",
        "    for example, tokens in zip(ds.as_numpy_iterator(), inferences):\n",
        "      predictions.append(self.postprocess(tokens, example))\n",
        "\n",
        "    result = metrics_utils.event_predictions_to_ns(\n",
        "        predictions, codec=self.codec, encoding_spec=self.encoding_spec)\n",
        "    return result['est_ns']\n",
        "\n",
        "  def audio_to_dataset(self, audio):\n",
        "    \"\"\"Create a TF Dataset of spectrograms from input audio.\"\"\"\n",
        "    frames, frame_times = self._audio_to_frames(audio)\n",
        "    return tf.data.Dataset.from_tensors({\n",
        "        'inputs': frames,\n",
        "        'input_times': frame_times,\n",
        "    })\n",
        "\n",
        "  def _audio_to_frames(self, audio):\n",
        "    \"\"\"Compute spectrogram frames from audio.\"\"\"\n",
        "    frame_size = self.spectrogram_config.hop_width\n",
        "    padding = [0, frame_size - len(audio) % frame_size]\n",
        "    audio = np.pad(audio, padding, mode='constant')\n",
        "    frames = spectrograms.split_audio(audio, self.spectrogram_config)\n",
        "    num_frames = len(audio) // frame_size\n",
        "    times = np.arange(num_frames) / self.spectrogram_config.frames_per_second\n",
        "    return frames, times\n",
        "\n",
        "  def preprocess(self, ds):\n",
        "    pp_chain = [\n",
        "        functools.partial(\n",
        "            t5.data.preprocessors.split_tokens_to_inputs_length,\n",
        "            sequence_length=self.sequence_length,\n",
        "            output_features=self.output_features,\n",
        "            feature_key='inputs',\n",
        "            additional_feature_keys=['input_times']),\n",
        "        # Cache occurs here during training.\n",
        "        preprocessors.add_dummy_targets,\n",
        "        functools.partial(\n",
        "            preprocessors.compute_spectrograms,\n",
        "            spectrogram_config=self.spectrogram_config)\n",
        "    ]\n",
        "    for pp in pp_chain:\n",
        "      ds = pp(ds)\n",
        "    return ds\n",
        "\n",
        "  def postprocess(self, tokens, example):\n",
        "    tokens = self._trim_eos(tokens)\n",
        "    start_time = example['input_times'][0]\n",
        "    # Round down to nearest symbolic token step.\n",
        "    start_time -= start_time % (1 / self.codec.steps_per_second)\n",
        "    return {\n",
        "        'est_tokens': tokens,\n",
        "        'start_time': start_time,\n",
        "        # Internal MT3 code expects raw inputs, not used here.\n",
        "        'raw_inputs': []\n",
        "    }\n",
        "\n",
        "  @staticmethod\n",
        "  def _trim_eos(tokens):\n",
        "    tokens = np.array(tokens, np.int32)\n",
        "    if vocabularies.DECODED_EOS_ID in tokens:\n",
        "      tokens = tokens[:np.argmax(tokens == vocabularies.DECODED_EOS_ID)]\n",
        "    return tokens\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "yGQ-zpgy3raf"
      },
      "outputs": [],
      "source": [
        "#@title Load Model\n",
        "#@markdown The `ismir2021` model transcribes piano only, with note velocities.\n",
        "#@markdown The `mt3` model transcribes multiple simultaneous instruments,\n",
        "#@markdown but without velocities.\n",
        "\n",
        "MODEL = \"mt3\" #@param[\"ismir2021\", \"mt3\"]\n",
        "\n",
        "checkpoint_path = f'/content/checkpoints/{MODEL}/'\n",
        "\n",
        "inference_model = InferenceModel(checkpoint_path, MODEL)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "s2s3JAvBAgFB"
      },
      "outputs": [],
      "source": [
        "#@title Upload Audio\n",
        "\n",
        "audio = upload_audio(sample_rate=SAMPLE_RATE)\n",
        "note_seq.notebook_utils.colab_play(audio, sample_rate=SAMPLE_RATE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HSKNjUYYv1kV"
      },
      "outputs": [],
      "source": [
        "#@title Transcribe Audio\n",
        "#@markdown This may take a few minutes depending on the length of the WAV file\n",
        "#@markdown you uploaded.\n",
        "\n",
        "est_ns = inference_model(audio)\n",
        "\n",
        "note_seq.play_sequence(est_ns, synth=note_seq.fluidsynth, \n",
        "                       sample_rate=SAMPLE_RATE, sf2_path=SF2_PATH)\n",
        "note_seq.plot_sequence(est_ns)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-DiCjtDpyUMh"
      },
      "outputs": [],
      "source": [
        "#@title Download MIDI Transcription\n",
        "\n",
        "note_seq.sequence_proto_to_midi_file(est_ns, '/tmp/transcribed.mid')\n",
        "files.download('/tmp/transcribed.mid')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Music Transcription with Transformers",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}