colab-notebooks/MusicVAE.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "MusicVAE.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "hYaJ6dvF0v7g",
        "R122bwRNbTus",
        "ZLfb2a_12wcj",
        "C_TD5psbv9Ax",
        "lEJptw-V4CEJ",
        "RTsrpipz4cFc",
        "moLOftFqBS-0",
        "wM6gOe6X3hWB",
        "8YxEHHI937Oa",
        "mDKI2rmOk0Dv",
        "l47dxtR82s0t",
        "_ss6V0582zpU",
        "0d4st_BlUBdl",
        "6PxW0_7Z2fvb",
        "LGIZPuZc2dIa"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bhOAxQyU0rhs",
        "colab_type": "text"
      },
      "source": [
        "Copyright 2017 Google LLC.\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",
        "https://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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hYaJ6dvF0v7g",
        "colab_type": "text"
      },
      "source": [
        "# MusicVAE: A Hierarchical Latent Vector Model for Learning Long-Term Structure in Music.\n",
        "### ___Adam Roberts, Jesse Engel, Colin Raffel, Curtis Hawthorne, and Douglas Eck___\n",
        "\n",
        "[MusicVAE](https://g.co/magenta/music-vae) learns a latent space of musical scores, providing different modes\n",
        "of interactive musical creation, including:\n",
        "\n",
        "* Random sampling from the prior distribution.\n",
        "* Interpolation between existing sequences.\n",
        "* Manipulation of existing sequences via attribute vectors.\n",
        "\n",
        "Examples of these interactions can be generated below, and selections can be heard in our\n",
        "[YouTube playlist](https://www.youtube.com/playlist?list=PLBUMAYA6kvGU8Cgqh709o5SUvo-zHGTxr).\n",
        "\n",
        "For short sequences (e.g., 2-bar \"loops\"), we use a bidirectional LSTM encoder\n",
        "and LSTM decoder. For longer sequences, we use a novel hierarchical LSTM\n",
        "decoder, which helps the model learn longer-term structures.\n",
        "\n",
        "We also model the interdependencies between instruments by training multiple\n",
        "decoders on the lowest-level embeddings of the hierarchical decoder.\n",
        "\n",
        "For additional details, check out our [blog post](https://g.co/magenta/music-vae) and [paper](https://goo.gl/magenta/musicvae-paper).\n",
        "___\n",
        "\n",
        "This colab notebook is self-contained and should run natively on google cloud. The [code](https://github.com/tensorflow/magenta/tree/master/magenta/models/music_vae) and [checkpoints](http://download.magenta.tensorflow.org/models/music_vae/checkpoints.tar.gz) can be downloaded separately and run locally, which is required if you want to train your own model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R122bwRNbTus",
        "colab_type": "text"
      },
      "source": [
        "# Basic Instructions\n",
        "\n",
        "1. Double click on the hidden cells to make them visible, or select \"View > Expand Sections\" in the menu at the top.\n",
        "2. Hover over the \"`[ ]`\" in the top-left corner of each cell and click on the \"Play\" button to run it, in order.\n",
        "3. Listen to the generated samples.\n",
        "4. Make it your own: copy the notebook, modify the code, train your own models, upload your own MIDI, etc.!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZLfb2a_12wcj",
        "colab_type": "text"
      },
      "source": [
        "# Environment Setup\n",
        "Includes package installation for sequence synthesis. Will take a few minutes.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PfRDVhNs3UFx",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Setup Environment\n",
        "#@test {\"output\": \"ignore\"}\n",
        "\n",
        "import glob\n",
        "\n",
        "BASE_DIR = \"gs://download.magenta.tensorflow.org/models/music_vae/colab2\"\n",
        "\n",
        "print('Installing dependencies...')\n",
        "!apt-get update -qq && apt-get install -qq libfluidsynth1 fluid-soundfont-gm build-essential libasound2-dev libjack-dev\n",
        "!pip install -q pyfluidsynth\n",
        "!pip install -qU magenta\n",
        "\n",
        "# Hack to allow python to pick up the newly-installed fluidsynth lib.\n",
        "# This is only needed for the hosted Colab environment.\n",
        "import ctypes.util\n",
        "orig_ctypes_util_find_library = ctypes.util.find_library\n",
        "def proxy_find_library(lib):\n",
        "  if lib == 'fluidsynth':\n",
        "    return 'libfluidsynth.so.1'\n",
        "  else:\n",
        "    return orig_ctypes_util_find_library(lib)\n",
        "ctypes.util.find_library = proxy_find_library\n",
        "\n",
        "\n",
        "print('Importing libraries and defining some helper functions...')\n",
        "from google.colab import files\n",
        "import magenta.music as mm\n",
        "from magenta.models.music_vae import configs\n",
        "from magenta.models.music_vae.trained_model import TrainedModel\n",
        "import numpy as np\n",
        "import os\n",
        "import tensorflow.compat.v1 as tf\n",
        "\n",
        "tf.disable_v2_behavior()\n",
        "\n",
        "# Necessary until pyfluidsynth is updated (>1.2.5).\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
        "\n",
        "def play(note_sequence):\n",
        "  mm.play_sequence(note_sequence, synth=mm.fluidsynth)\n",
        "\n",
        "def interpolate(model, start_seq, end_seq, num_steps, max_length=32,\n",
        "                assert_same_length=True, temperature=0.5,\n",
        "                individual_duration=4.0):\n",
        "  \"\"\"Interpolates between a start and end sequence.\"\"\"\n",
        "  note_sequences = model.interpolate(\n",
        "      start_seq, end_seq,num_steps=num_steps, length=max_length,\n",
        "      temperature=temperature,\n",
        "      assert_same_length=assert_same_length)\n",
        "\n",
        "  print('Start Seq Reconstruction')\n",
        "  play(note_sequences[0])\n",
        "  print('End Seq Reconstruction')\n",
        "  play(note_sequences[-1])\n",
        "  print('Mean Sequence')\n",
        "  play(note_sequences[num_steps // 2])\n",
        "  print('Start -> End Interpolation')\n",
        "  interp_seq = mm.sequences_lib.concatenate_sequences(\n",
        "      note_sequences, [individual_duration] * len(note_sequences))\n",
        "  play(interp_seq)\n",
        "  mm.plot_sequence(interp_seq)\n",
        "  return interp_seq if num_steps > 3 else note_sequences[num_steps // 2]\n",
        "\n",
        "def download(note_sequence, filename):\n",
        "  mm.sequence_proto_to_midi_file(note_sequence, filename)\n",
        "  files.download(filename)\n",
        "\n",
        "print('Done')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C_TD5psbv9Ax",
        "colab_type": "text"
      },
      "source": [
        "# 2-Bar Drums Model\n",
        "\n",
        "Below are 4 pre-trained models to experiment with. The first 3 map the 61 MIDI drum \"pitches\" to a reduced set of 9 classes (bass, snare, closed hi-hat, open hi-hat, low tom, mid tom, high tom, crash cymbal, ride cymbal) for a simplified but less expressive output space. The last model uses a [NADE](http://homepages.inf.ed.ac.uk/imurray2/pub/11nade/) to represent all possible MIDI drum \"pitches\".\n",
        "\n",
        "* **drums_2bar_oh_lokl**: This *low* KL model was trained for more *realistic* sampling. The output is a one-hot encoding of 2^9 combinations of hits. It has a single-layer bidirectional LSTM encoder with 512 nodes in each direction, a 2-layer LSTM decoder with 256 nodes in each layer, and a Z with 256 dimensions. During training it was given 0 free bits, and had a fixed beta value of 0.8. After 300k steps, the final accuracy is 0.73 and KL divergence is 11 bits.\n",
        "* **drums_2bar_oh_hikl**: This *high* KL model was trained for *better reconstruction and interpolation*. The output is a one-hot encoding of 2^9 combinations of hits. It has a single-layer bidirectional LSTM encoder with 512 nodes in each direction, a 2-layer LSTM decoder with 256 nodes in each layer, and a Z with 256 dimensions. During training it was given 96 free bits and had a fixed beta value of 0.2. It was trained with scheduled sampling with an inverse sigmoid schedule and a rate of 1000. After 300k, steps the final accuracy is 0.97 and KL divergence is 107 bits.\n",
        "* **drums_2bar_nade_reduced**: This model outputs a multi-label \"pianoroll\" with 9 classes. It has a single-layer bidirectional LSTM encoder with 512 nodes in each direction, a 2-layer LSTM-NADE decoder with 512 nodes in each layer and 9-dimensional NADE with 128 hidden units, and a Z with 256 dimensions. During training it was given 96 free bits and has a fixed beta value of 0.2. It was trained with scheduled sampling with an inverse sigmoid schedule and a rate of 1000. After 300k steps, the final accuracy is 0.98 and KL divergence is 110 bits.\n",
        "* **drums_2bar_nade_full**:  The output is a multi-label \"pianoroll\" with 61 classes. A single-layer bidirectional LSTM encoder with 512 nodes in each direction, a 2-layer LSTM-NADE decoder with 512 nodes in each layer and 61-dimensional NADE with 128 hidden units, and a Z with 256 dimensions. During training it was given 0 free bits and has a fixed beta value of 0.2. It was trained with scheduled sampling with an inverse sigmoid schedule and a rate of 1000. After 300k steps, the final accuracy is 0.90 and KL divergence is 116 bits."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0x8YTRDwv8Gk",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Load Pretrained Models\n",
        "\n",
        "drums_models = {}\n",
        "# One-hot encoded.\n",
        "drums_config = configs.CONFIG_MAP['cat-drums_2bar_small']\n",
        "drums_models['drums_2bar_oh_lokl'] = TrainedModel(drums_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/drums_2bar_small.lokl.ckpt')\n",
        "drums_models['drums_2bar_oh_hikl'] = TrainedModel(drums_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/drums_2bar_small.hikl.ckpt')\n",
        "\n",
        "# Multi-label NADE.\n",
        "drums_nade_reduced_config = configs.CONFIG_MAP['nade-drums_2bar_reduced']\n",
        "drums_models['drums_2bar_nade_reduced'] = TrainedModel(drums_nade_reduced_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/drums_2bar_nade.reduced.ckpt')\n",
        "drums_nade_full_config = configs.CONFIG_MAP['nade-drums_2bar_full']\n",
        "drums_models['drums_2bar_nade_full'] = TrainedModel(drums_nade_full_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/drums_2bar_nade.full.ckpt')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lEJptw-V4CEJ",
        "colab_type": "text"
      },
      "source": [
        "## Generate Samples"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zRUlAshMpDnR",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Generate 4 samples from the prior of one of the models listed above.\n",
        "drums_sample_model = \"drums_2bar_oh_lokl\" #@param [\"drums_2bar_oh_lokl\", \"drums_2bar_oh_hikl\", \"drums_2bar_nade_reduced\", \"drums_2bar_nade_full\"]\n",
        "temperature = 0.5 #@param {type:\"slider\", min:0.1, max:1.5, step:0.1}\n",
        "drums_samples = drums_models[drums_sample_model].sample(n=4, length=32, temperature=temperature)\n",
        "for ns in drums_samples:\n",
        "  play(ns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OSwhxkru5mB6",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Optionally download generated MIDI samples.\n",
        "for i, ns in enumerate(drums_samples):\n",
        "  download(ns, '%s_sample_%d.mid' % (drums_sample_model, i))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RTsrpipz4cFc",
        "colab_type": "text"
      },
      "source": [
        "## Generate Interpolations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7cnZfjdGwwZg",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Option 1: Use example MIDI files for interpolation endpoints.\n",
        "input_drums_midi_data = [\n",
        "    tf.io.gfile.GFile(fn, mode='rb').read()\n",
        "    for fn in sorted(tf.io.gfile.glob(BASE_DIR + '/midi/drums_2bar*.mid'))]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yjfyjWPtb8fV",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Option 2: upload your own MIDI files to use for interpolation endpoints instead of those provided.\n",
        "input_drums_midi_data = files.upload().values() or input_drums_midi_data"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zqCJFtHYb-7A",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Extract drums from MIDI files. This will extract all unique 2-bar drum beats using a sliding window with a stride of 1 bar.\n",
        "drums_input_seqs = [mm.midi_to_sequence_proto(m) for m in input_drums_midi_data]\n",
        "extracted_beats = []\n",
        "for ns in drums_input_seqs:\n",
        "  extracted_beats.extend(drums_nade_full_config.data_converter.from_tensors(\n",
        "      drums_nade_full_config.data_converter.to_tensors(ns)[1]))\n",
        "for i, ns in enumerate(extracted_beats):\n",
        "  print(\"Beat\", i)\n",
        "  play(ns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MeAboOS1xDgE",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Interpolate between 2 beats, selected from those in the previous cell.\n",
        "drums_interp_model = \"drums_2bar_oh_hikl\" #@param [\"drums_2bar_oh_lokl\", \"drums_2bar_oh_hikl\", \"drums_2bar_nade_reduced\", \"drums_2bar_nade_full\"]\n",
        "start_beat = 0 #@param {type:\"integer\"}\n",
        "end_beat = 1 #@param {type:\"integer\"}\n",
        "start_beat = extracted_beats[start_beat]\n",
        "end_beat = extracted_beats[end_beat]\n",
        "\n",
        "temperature = 0.5 #@param {type:\"slider\", min:0.1, max:1.5, step:0.1}\n",
        "num_steps = 13 #@param {type:\"integer\"}\n",
        "\n",
        "drums_interp = interpolate(drums_models[drums_interp_model], start_beat, end_beat, num_steps=num_steps, temperature=temperature)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nkKoQwFEcxpi",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Optionally download interpolation MIDI file.\n",
        "download(drums_interp, '%s_interp.mid' % drums_interp_model)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "moLOftFqBS-0",
        "colab_type": "text"
      },
      "source": [
        "# 2-Bar Melody Model\n",
        "\n",
        "The pre-trained model consists of a single-layer bidirectional LSTM encoder with 2048 nodes in each direction, a 3-layer LSTM decoder with 2048 nodes in each layer, and Z with 512 dimensions. The model was given 0 free bits, and had its beta valued annealed at an exponential rate of 0.99999 from 0 to 0.43 over 200k steps. It was trained with scheduled sampling with an inverse sigmoid schedule and a rate of 1000. The final accuracy is 0.95 and KL divergence is 58 bits."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2XCPjwd6BVtm",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Load the pre-trained model.\n",
        "mel_2bar_config = configs.CONFIG_MAP['cat-mel_2bar_big']\n",
        "mel_2bar = TrainedModel(mel_2bar_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/mel_2bar_big.ckpt')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wM6gOe6X3hWB",
        "colab_type": "text"
      },
      "source": [
        "## Generate Samples"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RwXUA74cNkh0",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Generate 4 samples from the prior.\n",
        "temperature = 0.5 #@param {type:\"slider\", min:0.1, max:1.5, step:0.1}\n",
        "mel_2_samples = mel_2bar.sample(n=4, length=32, temperature=temperature)\n",
        "for ns in mel_2_samples:\n",
        "  play(ns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MLg9dQ2D1Xpu",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Optionally download samples.\n",
        "for i, ns in enumerate(mel_2_samples):\n",
        "  download(ns, 'mel_2bar_sample_%d.mid' % i)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8YxEHHI937Oa",
        "colab_type": "text"
      },
      "source": [
        "## Generate Interpolations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H5wCWLMPLfYz",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Option 1: Use example MIDI files for interpolation endpoints.\n",
        "input_mel_midi_data = [\n",
        "    tf.io.gfile.GFile(fn, 'rb').read()\n",
        "    for fn in sorted(tf.io.gfile.glob(BASE_DIR + '/midi/mel_2bar*.mid'))]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hu5SeYFnNEe5",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Option 2: Upload your own MIDI files to use for interpolation endpoints instead of those provided.\n",
        "input_mel_midi_data = files.upload().values() or input_mel_midi_data"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Xy4vizNUH8GJ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Extract melodies from MIDI files. This will extract all unique 2-bar melodies using a sliding window with a stride of 1 bar.\n",
        "mel_input_seqs = [mm.midi_to_sequence_proto(m) for m in input_mel_midi_data]\n",
        "extracted_mels = []\n",
        "for ns in mel_input_seqs:\n",
        "  extracted_mels.extend(\n",
        "      mel_2bar_config.data_converter.from_tensors(\n",
        "          mel_2bar_config.data_converter.to_tensors(ns)[1]))\n",
        "for i, ns in enumerate(extracted_mels):\n",
        "  print(\"Melody\", i)\n",
        "  play(ns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8J4vloU3Pgtz",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Interpolate between 2 melodies, selected from those in the previous cell.\n",
        "start_melody = 0 #@param {type:\"integer\"}\n",
        "end_melody = 1 #@param {type:\"integer\"}\n",
        "start_mel = extracted_mels[start_melody]\n",
        "end_mel = extracted_mels[end_melody]\n",
        "\n",
        "temperature = 0.5 #@param {type:\"slider\", min:0.1, max:1.5, step:0.1}\n",
        "num_steps = 13 #@param {type:\"integer\"}\n",
        "\n",
        "mel_2bar_interp = interpolate(mel_2bar, start_mel, end_mel, num_steps=num_steps, temperature=temperature)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hZVP4JRmTCvB",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Optionally download interpolation MIDI file.\n",
        "download(mel_2bar_interp, 'mel_2bar_interp.mid')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mDKI2rmOk0Dv",
        "colab_type": "text"
      },
      "source": [
        "# 16-bar Melody Models\n",
        "\n",
        "The pre-trained hierarchical model consists of a 2-layer stacked bidirectional LSTM encoder with 2048 nodes in each direction for each layer, a 16-step 2-layer LSTM \"conductor\" decoder with 1024 nodes in each layer, a 2-layer LSTM core decoder with 1024 nodes in each layer, and a Z with 512 dimensions. It was given 256 free bits, and had a fixed beta value of 0.2. After 25k steps, the final accuracy is 0.90 and KL divergence is 277 bits."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9zcfdVjjk3Pp",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Load the pre-trained models.\n",
        "mel_16bar_models = {}\n",
        "hierdec_mel_16bar_config = configs.CONFIG_MAP['hierdec-mel_16bar']\n",
        "mel_16bar_models['hierdec_mel_16bar'] = TrainedModel(hierdec_mel_16bar_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/mel_16bar_hierdec.ckpt')\n",
        "\n",
        "flat_mel_16bar_config = configs.CONFIG_MAP['flat-mel_16bar']\n",
        "mel_16bar_models['baseline_flat_mel_16bar'] = TrainedModel(flat_mel_16bar_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/mel_16bar_flat.ckpt')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l47dxtR82s0t",
        "colab_type": "text"
      },
      "source": [
        "## Generate Samples"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Bptfh7C1njpV",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Generate 4 samples from the selected model prior.\n",
        "mel_sample_model = \"hierdec_mel_16bar\" #@param [\"hierdec_mel_16bar\", \"baseline_flat_mel_16bar\"]\n",
        "temperature = 0.5 #@param {type:\"slider\", min:0.1, max:1.5, step:0.1}\n",
        "mel_16_samples = mel_16bar_models[mel_sample_model].sample(n=4, length=256, temperature=temperature)\n",
        "for ns in mel_16_samples:\n",
        "  play(ns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X4sDzwq623Ei",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Optionally download MIDI samples.\n",
        "for i, ns in enumerate(mel_16_samples):\n",
        "  download(ns, '%s_sample_%d.mid' % (mel_sample_model, i))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ss6V0582zpU",
        "colab_type": "text"
      },
      "source": [
        "## Generate Means"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "38nwpNp_lprY",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Option 1: Use example MIDI files for interpolation endpoints.\n",
        "input_mel_16_midi_data = [\n",
        "    tf.io.gfile.GFile(fn, 'rb').read()\n",
        "    for fn in sorted(tf.io.gfile.glob(BASE_DIR + '/midi/mel_16bar*.mid'))]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z-7VjVcPUpHN",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Option 2: upload your own MIDI files to use for interpolation endpoints instead of those provided.\n",
        "input_mel_16_midi_data = files.upload().values() or input_mel_16_midi_data"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C-WE4Nq2OJxH",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Extract melodies from MIDI files. This will extract all unique 16-bar melodies using a sliding window with a stride of 1 bar.\n",
        "mel_input_seqs = [mm.midi_to_sequence_proto(m) for m in input_mel_16_midi_data]\n",
        "extracted_16_mels = []\n",
        "for ns in mel_input_seqs:\n",
        "  extracted_16_mels.extend(\n",
        "      hierdec_mel_16bar_config.data_converter.from_tensors(\n",
        "          hierdec_mel_16bar_config.data_converter.to_tensors(ns)[1]))\n",
        "for i, ns in enumerate(extracted_16_mels):\n",
        "  print(\"Melody\", i)\n",
        "  play(ns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u_Xp1rpTrayv",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Compute the reconstructions and mean of the two melodies, selected from the previous cell.\n",
        "mel_interp_model = \"hierdec_mel_16bar\" #@param [\"hierdec_mel_16bar\", \"baseline_flat_mel_16bar\"]\n",
        "\n",
        "start_melody = 0 #@param {type:\"integer\"}\n",
        "end_melody = 1 #@param {type:\"integer\"}\n",
        "start_mel = extracted_16_mels[start_melody]\n",
        "end_mel = extracted_16_mels[end_melody]\n",
        "\n",
        "temperature = 0.5 #@param {type:\"slider\", min:0.1, max:1.5, step:0.1}\n",
        "\n",
        "mel_16bar_mean = interpolate(mel_16bar_models[mel_interp_model], start_mel, end_mel, num_steps=3, max_length=256, individual_duration=32, temperature=temperature)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rONYqtlLkyS2",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Optionally download mean MIDI file.\n",
        "download(mel_16bar_mean, '%s_mean.mid' % mel_interp_model)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0d4st_BlUBdl",
        "colab_type": "text"
      },
      "source": [
        "#16-bar \"Trio\" Models (lead, bass, drums)\n",
        "\n",
        "We present two pre-trained models for 16-bar trios: a hierarchical model and a flat (baseline) model.\n",
        "\n",
        "The pre-trained hierarchical model consists of a 2-layer stacked bidirectional LSTM encoder with 2048 nodes in each direction for each layer, a 16-step 2-layer LSTM \"conductor\" decoder with 1024 nodes in each layer, 3 (lead, bass, drums) 2-layer LSTM core decoders with 1024 nodes in each layer, and a Z with 512 dimensions. It was given 1024 free bits, and had a fixed beta value of 0.1.  It was trained with scheduled sampling with an inverse sigmoid schedule and a rate of 1000. After 50k steps, the final accuracy is 0.82 for lead, 0.87 for bass, and 0.90 for drums, and the KL divergence is 1027 bits.\n",
        "\n",
        "The pre-trained flat model consists of a 2-layer stacked bidirectional LSTM encoder with 2048 nodes in each direction for each layer, a 3-layer LSTM decoder with 2048 nodes in each layer, and a Z with 512 dimensions. It was given 1024 free bits, and had a fixed beta value of 0.1.  It was trained with scheduled sampling with an inverse sigmoid schedule and a rate of 1000. After 50k steps, the final accuracy is 0.67 for lead, 0.66 for bass, and 0.79 for drums, and the KL divergence is 1016 bits."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FDW3h0cqUERq",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Load the pre-trained models.\n",
        "trio_models = {}\n",
        "hierdec_trio_16bar_config = configs.CONFIG_MAP['hierdec-trio_16bar']\n",
        "trio_models['hierdec_trio_16bar'] = TrainedModel(hierdec_trio_16bar_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/trio_16bar_hierdec.ckpt')\n",
        "\n",
        "flat_trio_16bar_config = configs.CONFIG_MAP['flat-trio_16bar']\n",
        "trio_models['baseline_flat_trio_16bar'] = TrainedModel(flat_trio_16bar_config, batch_size=4, checkpoint_dir_or_path=BASE_DIR + '/checkpoints/trio_16bar_flat.ckpt')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6PxW0_7Z2fvb",
        "colab_type": "text"
      },
      "source": [
        "## Generate Samples"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XKk8rGihUR6B",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Generate 4 samples from the selected model prior.\n",
        "trio_sample_model = \"hierdec_trio_16bar\" #@param [\"hierdec_trio_16bar\", \"baseline_flat_trio_16bar\"]\n",
        "temperature = 0.5 #@param {type:\"slider\", min:0.1, max:1.5, step:0.1}\n",
        "\n",
        "trio_16_samples = trio_models[trio_sample_model].sample(n=4, length=256, temperature=temperature)\n",
        "for ns in trio_16_samples:\n",
        "  play(ns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fic5W7Z7m7Op",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Optionally download MIDI samples.\n",
        "for i, ns in enumerate(trio_16_samples):\n",
        "  download(ns, '%s_sample_%d.mid' % (trio_sample_model, i))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LGIZPuZc2dIa",
        "colab_type": "text"
      },
      "source": [
        "## Generate Means"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3msZzI89UU_F",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Option 1: Use example MIDI files for interpolation endpoints.\n",
        "input_trio_midi_data = [\n",
        "    tf.io.gfile.GFile(fn, 'rb').read()\n",
        "    for fn in sorted(tf.io.gfile.glob(BASE_DIR + '/midi/trio_16bar*.mid'))]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9ig0w2cSUs9n",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Option 2: Upload your own MIDI files to use for interpolation endpoints instead of those provided.\n",
        "input_trio_midi_data = files.upload().values() or input_trio_midi_data"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mawDY278UZKY",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Extract trios from MIDI files. This will extract all unique 16-bar trios using a sliding window with a stride of 1 bar.\n",
        "trio_input_seqs = [mm.midi_to_sequence_proto(m) for m in input_trio_midi_data]\n",
        "extracted_trios = []\n",
        "for ns in trio_input_seqs:\n",
        "  extracted_trios.extend(\n",
        "      hierdec_trio_16bar_config.data_converter.from_tensors(\n",
        "          hierdec_trio_16bar_config.data_converter.to_tensors(ns)[1]))\n",
        "for i, ns in enumerate(extracted_trios):\n",
        "  print(\"Trio\", i)\n",
        "  play(ns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UQuTQOlZW1LX",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Compute the reconstructions and mean of the two trios, selected from the previous cell.\n",
        "trio_interp_model = \"hierdec_trio_16bar\" #@param [\"hierdec_trio_16bar\", \"baseline_flat_trio_16bar\"]\n",
        "\n",
        "start_trio = 0 #@param {type:\"integer\"}\n",
        "end_trio = 1 #@param {type:\"integer\"}\n",
        "start_trio = extracted_trios[start_trio]\n",
        "end_trio = extracted_trios[end_trio]\n",
        "\n",
        "temperature = 0.5 #@param {type:\"slider\", min:0.1, max:1.5, step:0.1}\n",
        "trio_16bar_mean = interpolate(trio_models[trio_interp_model], start_trio, end_trio, num_steps=3, max_length=256, individual_duration=32, temperature=temperature)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bqSklK8CU_cQ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Optionally download mean MIDI file.\n",
        "download(trio_16bar_mean, '%s_mean.mid' % trio_interp_model)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}