Super_Piano_3_Alt_2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Super Piano 3 Alt 2.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyML/l9/cIXdhZnYU8JuN0LM",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/asigalov61/SuperPiano/blob/master/Super_Piano_3_Alt_2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cj9PFo6cF670"
      },
      "source": [
        "# Super Piano 3\n",
        "## Google Music Transformer\n",
        "\n",
        "### All thanks and credit for this colab go out to Prayag Chatha on whose repo and code it is based: https://github.com/chathasphere/pno-ai"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F-PCBbGB11Dp",
        "cellView": "form"
      },
      "source": [
        "#@title Clone pno-ai repo and install dependencies\n",
        "!git clone 'https://github.com/asigalov61/pno-ai'\n",
        "!pip install pretty_midi\n",
        "!pip install rtmidi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VsR2kUVY6SWs",
        "cellView": "form"
      },
      "source": [
        "#@title Import modules and setup main (hyper)parameters\n",
        "%cd /content/pno-ai/\n",
        "import os, time, datetime\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "torch.set_default_tensor_type(torch.cuda.FloatTensor)\n",
        "from random import shuffle\n",
        "from preprocess import PreprocessingPipeline\n",
        "from train import train\n",
        "from model import MusicTransformer\n",
        "\n",
        "n_epochs = 60\n",
        "batch_size = 128\n",
        "sampling_rate = 125\n",
        "n_velocity_bins = 32\n",
        "seq_length = 256\n",
        "n_tokens = 256 + sampling_rate + n_velocity_bins\n",
        "transformer = MusicTransformer(n_tokens, seq_length, \n",
        "        d_model = 128, n_heads = 8, d_feedforward=512, \n",
        "        depth = 6, positional_encoding=True, relative_pos=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ht3vjmhM2TgU",
        "cellView": "form"
      },
      "source": [
        "#@title Download Google Magenta MAESTRO v.2.0.0 Piano MIDI Dataset (~1300 MIDIs)\n",
        "%cd /content/pno-ai/data\n",
        "!wget 'https://storage.googleapis.com/magentadata/datasets/maestro/v2.0.0/maestro-v2.0.0-midi.zip'\n",
        "!unzip -j maestro-v2.0.0-midi.zip\n",
        "!rm maestro-v2.0.0-midi.zip\n",
        "%cd /content/pno-ai"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JfQ3wCXBEKW1"
      },
      "source": [
        "%cd /content/pno-ai/data\n",
        "!unzip -j '/content/pno-ai/data/Super-Piano-2-Performance-DataSet-CC-BY-NC-SA.zip'\n",
        "!rm '/content/pno-ai/data/Super-Piano-2-Performance-DataSet-CC-BY-NC-SA.zip'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WZpRFWmSC4cD",
        "cellView": "form"
      },
      "source": [
        "#@title Process custom MIDI dataset\n",
        "%cd /content/pno-ai\n",
        "pipeline = PreprocessingPipeline(input_dir=\"data\", stretch_factors=[0.975, 1, 1.025],\n",
        "           split_size=30, sampling_rate=sampling_rate, n_velocity_bins=n_velocity_bins,\n",
        "           transpositions=range(-2,3), training_val_split=0.9, max_encoded_length=seq_length+1,\n",
        "           min_encoded_length=257)\n",
        "\n",
        "pipeline_start = time.time()\n",
        "pipeline.run()\n",
        "runtime = time.time() - pipeline_start\n",
        "print(f\"MIDI pipeline runtime: {runtime / 60 : .1f}m\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4gKYrVKRCpM1",
        "cellView": "form"
      },
      "source": [
        "#@title Train the model\n",
        "%cd '/content/pno-ai'\n",
        "today = datetime.date.today().strftime('%m%d%Y')\n",
        "checkpoint = f\"saved_models/tf_{today}\"\n",
        "\n",
        "training_sequences = pipeline.encoded_sequences['training']\n",
        "validation_sequences = pipeline.encoded_sequences['validation']\n",
        "\n",
        "\n",
        "train(transformer, training_sequences, validation_sequences,\n",
        "               epochs = n_epochs, evaluate_per = 1, custom_schedule=True, custom_loss=True,\n",
        "               batch_size = batch_size, batches_per_print=100,\n",
        "               padding_index=0, checkpoint_path=checkpoint)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S8tuYH63MRXp",
        "cellView": "form"
      },
      "source": [
        "#@title Save the model\n",
        "%cd /content/pno-ai\n",
        "torch.save(transformer, 'trained_model.pth')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t1MTzF3NDKYb",
        "cellView": "form"
      },
      "source": [
        "#@title Load pre-trained checkpoint if needed\n",
        "%cd /content/pno-ai\n",
        "transformer.load_state_dict(torch.load('/content/pno-ai/saved_models/tf_09232020_e4'))\n",
        "print(\"Successfully loaded checkpoint!\")\n",
        "transformer.eval()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7qq3cS3Gmj2G",
        "cellView": "form"
      },
      "source": [
        "#@title Define Helper Functions\n",
        "import numpy as np\n",
        "from pretty_midi import Note, PrettyMIDI, Instrument\n",
        "import torch.nn.functional as F\n",
        "import copy, pathlib\n",
        "\n",
        "def vectorize(sequence):\n",
        "    \"\"\"\n",
        "    Converts a list of pretty_midi Note objects into a numpy array of\n",
        "    dimension (n_notes x 4)\n",
        "    \"\"\"\n",
        "    array = [[note.start, note.end, note.pitch, note.velocity] for\n",
        "            note in sequence]\n",
        "    return np.asarray(array)\n",
        "\n",
        "def devectorize(note_array):\n",
        "    \"\"\"\n",
        "    Converts a vectorized note sequence into a list of pretty_midi Note\n",
        "    objects\n",
        "    \"\"\"\n",
        "    return [Note(start = a[0], end = a[1], pitch=a[2],\n",
        "        velocity=a[3]) for a in note_array.tolist()]\n",
        "\n",
        "\n",
        "def one_hot(sequence, n_states):\n",
        "    \"\"\"\n",
        "    Given a list of integers and the maximal number of unique values found\n",
        "    in the list, return a one-hot encoded tensor of shape (m, n)\n",
        "    where m is sequence length and n is n_states.\n",
        "    \"\"\"\n",
        "    if torch.cuda.is_available():\n",
        "        return torch.eye(n_states)[sequence,:].cuda()\n",
        "    else:\n",
        "        return torch.eye(n_states)[sequence,:]\n",
        "\n",
        "def decode_one_hot(vector):\n",
        "    '''\n",
        "    Given a one-hot encoded vector, return the non-zero index\n",
        "    '''\n",
        "    return vector.nonzero().item()\n",
        "\n",
        "def prepare_batches(sequences, batch_size):\n",
        "    \"\"\"\n",
        "    Splits a list of sequences into batches of a fixed size. Each sequence yields an input sequence\n",
        "    and a target sequence, with the latter one time step ahead. For example, the sequence \"to be or not\n",
        "    to be\" gives an input sequence of \"to be or not to b\" and a target sequence of \"o be or not to be.\"\n",
        "    \"\"\"\n",
        "    n_sequences = len(sequences)\n",
        "    for i in range(0, n_sequences, batch_size):\n",
        "        batch = sequences[i:i+batch_size]\n",
        "\t#needs to be in sorted order for packing batches to work\n",
        "        batch = sorted(batch, key = len, reverse=True)\n",
        "        input_sequences, target_sequences = [], []\n",
        "\n",
        "        for sequence in batch:\n",
        "            input_sequences.append(sequence[:-1])\n",
        "            target_sequences.append(sequence[1:])\n",
        "\n",
        "        yield input_sequences, target_sequences\n",
        "\n",
        "def clones(module, N):\n",
        "    \"Clone N identical layers of a module\"\n",
        "    return torch.nn.ModuleList([copy.deepcopy(module) for i in range(N)])\n",
        "\n",
        "def d(tensor=None):\n",
        "    if tensor is None:\n",
        "        return 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "    return 'cuda' if tensor.is_cuda else 'cpu'\n",
        "\n",
        "def write_midi(note_sequence, output_dir, filename):\n",
        "\n",
        "    #make output directory\n",
        "    pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "    #generate midi\n",
        "    midi = PrettyMIDI()\n",
        "    piano_track = Instrument(program=0, is_drum=False, name=filename)\n",
        "    piano_track.notes = note_sequence\n",
        "    midi.instruments.append(piano_track)\n",
        "    output_name = output_dir + f\"{filename}.midi\"\n",
        "    midi.write(output_name)\n",
        "\n",
        "def sample(model, sample_length, prime_sequence=[], temperature=1):\n",
        "    \"\"\"\n",
        "    Generate a MIDI event sequence of a fixed length by randomly sampling from a model's distribution of sequences. Optionally, \"seed\" the sequence with a prime. A well-trained model will create music that responds to the prime and develops upon it.\n",
        "    \"\"\"\n",
        "    #deactivate training mode\n",
        "    model.eval()\n",
        "    if len(prime_sequence) == 0:\n",
        "        #if no prime is provided, randomly select a starting event\n",
        "        input_sequence = [np.random.randint(model.n_tokens)]\n",
        "    else:\n",
        "        input_sequence = prime_sequence.copy()\n",
        "\n",
        "    #add singleton dimension for the batch\n",
        "    input_tensor = torch.LongTensor(input_sequence).unsqueeze(0).cuda()\n",
        "\n",
        "    for i in range(sample_length):\n",
        "        #select probabilities of *next* token\n",
        "        out = model(input_tensor)[0, -1, :]\n",
        "        #out is a 1d tensor of shape (n_tokens)\n",
        "        probs = F.softmax(out / temperature, dim=0)\n",
        "        #sample prob distribution for next character\n",
        "        c = torch.multinomial(probs,1)\n",
        "        input_tensor = torch.cat([input_tensor[:,1:], c[None]], dim=1)\n",
        "        input_sequence.append(c.item())\n",
        "\n",
        "    return input_sequence"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GqjlxHyTHdhl",
        "cellView": "form"
      },
      "source": [
        "#@title Generate Output\n",
        "import uuid\n",
        "import torch\n",
        "from torch import nn\n",
        "torch.set_default_tensor_type(torch.cuda.FloatTensor)\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(device)\n",
        "from model import MusicTransformer\n",
        "from preprocess import SequenceEncoder\n",
        "import midi_input\n",
        "\n",
        "\n",
        "class GeneratorError(Exception):\n",
        "    pass\n",
        "\n",
        "model_key = ''\n",
        "\n",
        "    \n",
        "n_velocity_events = 32\n",
        "n_time_shift_events = 125\n",
        "\n",
        "decoder = SequenceEncoder(n_time_shift_events, n_velocity_events,\n",
        "           min_events=0)\n",
        "\n",
        "prime_sequence = []\n",
        "\n",
        "\n",
        "\n",
        "temp = 0.4\n",
        "\n",
        "trial_key = str(uuid.uuid4())[:6]\n",
        "n_trials = 1\n",
        "\n",
        "sample_l = 2048\n",
        "\n",
        "keep_g = False\n",
        "stuck_note_d = None\n",
        "note_sequence = []\n",
        "\n",
        "\n",
        "for i in range(n_trials):\n",
        "            print(\"generating sequence\")\n",
        "            output_sequence = sample(transformer, prime_sequence = prime_sequence, sample_length=sample_l, temperature=temp)\n",
        "            note_sequence = decoder.decode_sequence(output_sequence, \n",
        "                verbose=True, stuck_note_duration=stuck_note_d, keep_ghosts=keep_g)\n",
        "\n",
        "            output_dir = \"/content/pno-ai\"\n",
        "            write_midi(note_sequence, output_dir, 'output')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}