RNN_Music_Generation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "RNN_Music-Generation.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "kahm6Z8v_TqC"
      },
      "source": [
        "!sudo apt install -y fluidsynth"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M0lAReB7_Vqb"
      },
      "source": [
        "!pip install --upgrade pyfluidsynth"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G46kKoQZmIa8"
      },
      "source": [
        "!pip install pretty_midi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GsLFq7nsiqcq"
      },
      "source": [
        "import collections\n",
        "import datetime\n",
        "import fluidsynth\n",
        "import glob\n",
        "import numpy as np\n",
        "import pathlib\n",
        "import pandas as pd\n",
        "import pretty_midi\n",
        "import seaborn as sns\n",
        "import tensorflow as tf\n",
        "\n",
        "from IPython import display\n",
        "from matplotlib import pyplot as plt\n",
        "from typing import Dict, List, Optional, Sequence, Tuple"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Efja_OtJNzAM"
      },
      "source": [
        "seed = 42\n",
        "tf.random.set_seed(seed)\n",
        "np.random.seed(seed)\n",
        "\n",
        "_SAMPLING_RATE = 16000"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mwja4SWmibrL"
      },
      "source": [
        "data_dir = pathlib.Path('data/maestro-v2.0.0')\n",
        "if not data_dir.exists():\n",
        "  tf.keras.utils.get_file(\n",
        "      'maestro-v2.0.0-midi.zip',\n",
        "      origin='https://storage.googleapis.com/magentadata/datasets/maestro/v2.0.0/maestro-v2.0.0-midi.zip',\n",
        "      extract=True,\n",
        "      cache_dir='.', cache_subdir='data',\n",
        "  )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "72iFI1bPB9o1"
      },
      "source": [
        "filenames = glob.glob(str(data_dir/'**/*.mid*'))\n",
        "print('Number of files:', len(filenames))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6oSCbHvJNbci"
      },
      "source": [
        "sample_file = filenames[1]\n",
        "print(sample_file)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1YSQ5DjRI2md"
      },
      "source": [
        "pm = pretty_midi.PrettyMIDI(sample_file)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vzoHAaVY_kyY"
      },
      "source": [
        "def display_audio(pm: pretty_midi.PrettyMIDI, seconds=30):\n",
        "  waveform = pm.fluidsynth(fs=_SAMPLING_RATE)\n",
        "  waveform_short = waveform[:seconds*_SAMPLING_RATE]\n",
        "  return display.Audio(waveform_short, rate=_SAMPLING_RATE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GOe-3AAi_sRw"
      },
      "source": [
        "display_audio(pm)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SIGHYQPZQnRo"
      },
      "source": [
        "print('Number of instruments:', len(pm.instruments))\n",
        "instrument = pm.instruments[0]\n",
        "instrument_name = pretty_midi.program_to_instrument_name(instrument.program)\n",
        "print('Instrument name:', instrument_name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nYZm_VehYOTZ"
      },
      "source": [
        "for i, note in enumerate(instrument.notes[:10]):\n",
        "  note_name = pretty_midi.note_number_to_name(note.pitch)\n",
        "  duration = note.end - note.start\n",
        "  print(f'{i}: pitch={note.pitch}, note_name={note_name},'\n",
        "        f' duration={duration:.4f}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wyp_wdcEPWby"
      },
      "source": [
        "def midi_to_notes(midi_file: str) -> pd.DataFrame:\n",
        "  pm = pretty_midi.PrettyMIDI(midi_file)\n",
        "  instrument = pm.instruments[0]\n",
        "  notes = collections.defaultdict(list)\n",
        "\n",
        "  \n",
        "  sorted_notes = sorted(instrument.notes, key=lambda note: note.start)\n",
        "  prev_start = sorted_notes[0].start\n",
        "\n",
        "  for note in sorted_notes:\n",
        "    start = note.start\n",
        "    end = note.end\n",
        "    notes['pitch'].append(note.pitch)\n",
        "    notes['start'].append(start)\n",
        "    notes['end'].append(end)\n",
        "    notes['step'].append(start - prev_start)\n",
        "    notes['duration'].append(end - start)\n",
        "    prev_start = start\n",
        "\n",
        "  return pd.DataFrame({name: np.array(value) for name, value in notes.items()})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X0kPjLBlcnY6"
      },
      "source": [
        "raw_notes = midi_to_notes(sample_file)\n",
        "raw_notes.head()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WE9YXrGZbY2X"
      },
      "source": [
        "get_note_names = np.vectorize(pretty_midi.note_number_to_name)\n",
        "sample_note_names = get_note_names(raw_notes['pitch'])\n",
        "sample_note_names[:10]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "liD2N7x_WOTp"
      },
      "source": [
        "def plot_piano_roll(notes: pd.DataFrame, count: Optional[int] = None):\n",
        "  if count:\n",
        "    title = f'First {count} notes'\n",
        "  else:\n",
        "    title = f'Whole track'\n",
        "    count = len(notes['pitch'])\n",
        "  plt.figure(figsize=(20, 4))\n",
        "  plot_pitch = np.stack([notes['pitch'], notes['pitch']], axis=0)\n",
        "  plot_start_stop = np.stack([notes['start'], notes['end']], axis=0)\n",
        "  plt.plot(\n",
        "      plot_start_stop[:, :count], plot_pitch[:, :count], color=\"b\", marker=\".\")\n",
        "  plt.xlabel('Time [s]')\n",
        "  plt.ylabel('Pitch')\n",
        "  _ = plt.title(title)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vWeUbqmAXjOs"
      },
      "source": [
        "plot_piano_roll(raw_notes, count=100)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G7l76hEDZX8Z"
      },
      "source": [
        "plot_piano_roll(raw_notes)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pq9C9XBBaK7W"
      },
      "source": [
        "def plot_distributions(notes: pd.DataFrame, drop_percentile=2.5):\n",
        "  plt.figure(figsize=[15, 5])\n",
        "  plt.subplot(1, 3, 1)\n",
        "  sns.histplot(notes, x=\"pitch\", bins=20)\n",
        "\n",
        "  plt.subplot(1, 3, 2)\n",
        "  max_step = np.percentile(notes['step'], 100 - drop_percentile)\n",
        "  sns.histplot(notes, x=\"step\", bins=np.linspace(0, max_step, 21))\n",
        "  \n",
        "  plt.subplot(1, 3, 3)\n",
        "  max_duration = np.percentile(notes['duration'], 100 - drop_percentile)\n",
        "  sns.histplot(notes, x=\"duration\", bins=np.linspace(0, max_duration, 21))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Nu2Pw24acFD"
      },
      "source": [
        "plot_distributions(raw_notes)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BD5rsMRARYoV"
      },
      "source": [
        "def notes_to_midi(\n",
        "  notes: pd.DataFrame,\n",
        "  out_file: str, \n",
        "  instrument_name: str,\n",
        "  velocity: int = 100,  \n",
        ") -> pretty_midi.PrettyMIDI:\n",
        "\n",
        "  pm = pretty_midi.PrettyMIDI()\n",
        "  instrument = pretty_midi.Instrument(\n",
        "      program=pretty_midi.instrument_name_to_program(\n",
        "          instrument_name))\n",
        "\n",
        "  prev_start = 0\n",
        "  for i, note in notes.iterrows():\n",
        "    start = float(prev_start + note['step'])\n",
        "    end = float(start + note['duration'])\n",
        "    note = pretty_midi.Note(\n",
        "        velocity=velocity,\n",
        "        pitch=int(note['pitch']),\n",
        "        start=start,\n",
        "        end=end,\n",
        "    )\n",
        "    instrument.notes.append(note)\n",
        "    prev_start = start\n",
        "\n",
        "  pm.instruments.append(instrument)\n",
        "  pm.write(out_file)\n",
        "  return pm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wTazLbuWPIPF"
      },
      "source": [
        "example_file = 'example.midi'\n",
        "example_pm = notes_to_midi(\n",
        "    raw_notes, out_file=example_file, instrument_name=instrument_name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fGRLs-eR_4uK"
      },
      "source": [
        "display_audio(example_pm)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GiaQiTnXSW-T"
      },
      "source": [
        "num_files = 5\n",
        "all_notes = []\n",
        "for f in filenames[:num_files]:\n",
        "  notes = midi_to_notes(f)\n",
        "  all_notes.append(notes)\n",
        "\n",
        "all_notes = pd.concat(all_notes)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F4bMDeRvgWqx"
      },
      "source": [
        "n_notes = len(all_notes)\n",
        "print('Number of notes parsed:', n_notes)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mvNHCHZdXG2P"
      },
      "source": [
        "key_order = ['pitch', 'step', 'duration']\n",
        "train_notes = np.stack([all_notes[key] for key in key_order], axis=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PLC_19tshyFk"
      },
      "source": [
        "notes_ds = tf.data.Dataset.from_tensor_slices(train_notes)\n",
        "notes_ds.element_spec"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZkEC-5s6wJJV"
      },
      "source": [
        "def create_sequences(\n",
        "    dataset: tf.data.Dataset, \n",
        "    seq_length: int,\n",
        "    vocab_size = 128,\n",
        ") -> tf.data.Dataset:\n",
        "  \"\"\"Returns TF Dataset of sequence and label examples.\"\"\"\n",
        "  seq_length = seq_length+1\n",
        "\n",
        " \n",
        "  windows = dataset.window(seq_length, shift=1, stride=1,\n",
        "                              drop_remainder=True)\n",
        "\n",
        "  \n",
        "  flatten = lambda x: x.batch(seq_length, drop_remainder=True)\n",
        "  sequences = windows.flat_map(flatten)\n",
        "  \n",
        "  \n",
        "  def scale_pitch(x):\n",
        "    x = x/[vocab_size,1.0,1.0]\n",
        "    return x\n",
        "\n",
        " \n",
        "  def split_labels(sequences):\n",
        "    inputs = sequences[:-1]\n",
        "    labels_dense = sequences[-1]\n",
        "    labels = {key:labels_dense[i] for i,key in enumerate(key_order)}\n",
        "\n",
        "    return scale_pitch(inputs), labels\n",
        "\n",
        "  return sequences.map(split_labels, num_parallel_calls=tf.data.AUTOTUNE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fGA3VxcFXZ4T"
      },
      "source": [
        "seq_length = 25\n",
        "vocab_size = 128\n",
        "seq_ds = create_sequences(notes_ds, seq_length, vocab_size)\n",
        "seq_ds.element_spec"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ESK9cL7__TF3"
      },
      "source": [
        "for seq, target in seq_ds.take(1):\n",
        "  print('sequence shape:', seq.shape)\n",
        "  print('sequence elements (first 10):', seq[0: 10])\n",
        "  print()\n",
        "  print('target:', target)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fTpFoiM_AV_Y"
      },
      "source": [
        "batch_size = 64\n",
        "buffer_size = n_notes - seq_length  \n",
        "train_ds = (seq_ds\n",
        "            .shuffle(buffer_size)\n",
        "            .batch(batch_size, drop_remainder=True)\n",
        "            .cache()\n",
        "            .prefetch(tf.data.experimental.AUTOTUNE))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LySbjV0GzXQu"
      },
      "source": [
        "train_ds.element_spec"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "erxLOif08e8v"
      },
      "source": [
        "def mse_with_positive_pressure(y_true: tf.Tensor, y_pred: tf.Tensor):\n",
        "  mse = (y_true - y_pred) ** 2\n",
        "  positive_pressure = 10 * tf.maximum(-y_pred, 0.0)\n",
        "  return tf.reduce_mean(mse + positive_pressure)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kNaVWcCzAm5V"
      },
      "source": [
        "input_shape = (seq_length, 3)\n",
        "learning_rate = 0.005\n",
        "\n",
        "inputs = tf.keras.Input(input_shape)\n",
        "x = tf.keras.layers.LSTM(128)(inputs)\n",
        "\n",
        "outputs = {\n",
        "  'pitch': tf.keras.layers.Dense(128, name='pitch')(x),\n",
        "  'step': tf.keras.layers.Dense(1, name='step')(x),\n",
        "  'duration': tf.keras.layers.Dense(1, name='duration')(x),\n",
        "}\n",
        "\n",
        "model = tf.keras.Model(inputs, outputs)\n",
        "\n",
        "loss = {\n",
        "      'pitch': tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "          from_logits=True),\n",
        "      'step': mse_with_positive_pressure,\n",
        "      'duration': mse_with_positive_pressure,\n",
        "}\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n",
        "\n",
        "model.compile(loss=loss, optimizer=optimizer)\n",
        "\n",
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BlATt7Rl0XJl"
      },
      "source": [
        "losses = model.evaluate(train_ds, return_dict=True)\n",
        "losses"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9fQB5SiN3ufX"
      },
      "source": [
        "model.compile(\n",
        "    loss=loss,\n",
        "    loss_weights={\n",
        "        'pitch': 0.05,\n",
        "        'step': 1.0,\n",
        "        'duration':1.0,\n",
        "    },\n",
        "    optimizer=optimizer,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T7CzWmFR38ut"
      },
      "source": [
        "model.evaluate(train_ds, return_dict=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uQA_rwKEgPjp"
      },
      "source": [
        "callbacks = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        filepath='./training_checkpoints/ckpt_{epoch}',\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.EarlyStopping(\n",
        "        monitor='loss',\n",
        "        patience=5,\n",
        "        verbose=1,\n",
        "        restore_best_weights=True),\n",
        "]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aLoYY8-XaPFN"
      },
      "source": [
        "%%time\n",
        "epochs = 50\n",
        "\n",
        "history = model.fit(\n",
        "    train_ds,\n",
        "    epochs=epochs,\n",
        "    callbacks=callbacks,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PYBSjgDWiUfT"
      },
      "source": [
        "plt.plot(history.epoch, history.history['loss'], label='total loss')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1mil8ZyJNe1w"
      },
      "source": [
        "def predict_next_note(\n",
        "    notes: np.ndarray, \n",
        "    keras_model: tf.keras.Model, \n",
        "    temperature: float = 1.0) -> int:\n",
        "  \"\"\"Generates a note IDs using a trained sequence model.\"\"\"\n",
        "\n",
        "  assert temperature > 0\n",
        "\n",
        "  inputs = tf.expand_dims(notes, 0)\n",
        "\n",
        "  predictions = model.predict(inputs)\n",
        "  pitch_logits = predictions['pitch']\n",
        "  step = predictions['step']\n",
        "  duration = predictions['duration']\n",
        " \n",
        "  pitch_logits /= temperature\n",
        "  pitch = tf.random.categorical(pitch_logits, num_samples=1)\n",
        "  pitch = tf.squeeze(pitch, axis=-1)\n",
        "  duration = tf.squeeze(duration, axis=-1)\n",
        "  step = tf.squeeze(step, axis=-1)\n",
        "\n",
        "  step = tf.maximum(0, step)\n",
        "  duration = tf.maximum(0, duration)\n",
        "\n",
        "  return int(pitch), float(step), float(duration)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "87fPl4auPdR3"
      },
      "source": [
        "temperature = 2.0\n",
        "num_predictions = 120\n",
        "\n",
        "sample_notes = np.stack([raw_notes[key] for key in key_order], axis=1)\n",
        "\n",
        "input_notes = (\n",
        "    sample_notes[:seq_length] / np.array([vocab_size, 1, 1]))\n",
        "\n",
        "generated_notes = []\n",
        "prev_start = 0\n",
        "for _ in range(num_predictions):\n",
        "  pitch, step, duration = predict_next_note(input_notes, model, temperature)\n",
        "  start = prev_start + step\n",
        "  end = start + duration\n",
        "  input_note = (pitch, step, duration)\n",
        "  generated_notes.append((*input_note, start, end))\n",
        "  input_notes = np.delete(input_notes, 0, axis=0)\n",
        "  input_notes = np.append(input_notes, np.expand_dims(input_note, 0), axis=0)\n",
        "  prev_start = start\n",
        "\n",
        "generated_notes = pd.DataFrame(\n",
        "    generated_notes, columns=(*key_order, 'start', 'end'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0MK7HmqLuqka"
      },
      "source": [
        "generated_notes.head(10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e9K9KHPaTNnK"
      },
      "source": [
        "out_file = 'output.mid'\n",
        "out_pm = notes_to_midi(\n",
        "    generated_notes, out_file=out_file, instrument_name=instrument_name)\n",
        "display_audio(out_pm)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zv-b6tIBHLe0"
      },
      "source": [
        "from google.colab import files\n",
        "files.download(out_file)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NlNsxcnhvbcK"
      },
      "source": [
        "plot_piano_roll(generated_notes)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j5bco2WVRkAa"
      },
      "source": [
        "plot_distributions(generated_notes)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}