Sparse_Yoda.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "gradient": {
          "editing": false,
          "id": "ac5a4cf0-d9d2-47b5-9633-b53f8d99a4d2",
          "kernelId": ""
        },
        "id": "SiTIpPjArIyr"
      },
      "source": [
        "# Sparse Yoda (ver. 1.0)\n",
        "\n",
        "***\n",
        "\n",
        "Powered by tegridy-tools: https://github.com/asigalov61/tegridy-tools\n",
        "\n",
        "***\n",
        "\n",
        "Credit for the Sparse Trainsformer implementation used in this colab goes out @lucidrains https://github.com/lucidrains/sinkhorn-transformer\n",
        "\n",
        "***\n",
        "\n",
        "WARNING: This complete implementation is a functioning model of the Artificial Intelligence. Please excercise great humility, care, and respect. https://www.nscai.gov/\n",
        "\n",
        "***\n",
        "\n",
        "#### Project Los Angeles\n",
        "\n",
        "#### Tegridy Code 2022\n",
        "\n",
        "***"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "gradient": {
          "editing": false,
          "id": "fa0a611c-1803-42ae-bdf6-a49b5a4e781b",
          "kernelId": ""
        },
        "id": "gOd93yV0sGd2"
      },
      "source": [
        "# (Setup Environment)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "gradient": {
          "editing": false,
          "id": "39411b40-9e39-416e-8fe4-d40f733e7956",
          "kernelId": ""
        },
        "id": "lw-4aqV3sKQG"
      },
      "outputs": [],
      "source": [
        "#@title nvidia-smi gpu check\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "gradient": {
          "editing": false,
          "id": "a1a45a91-d909-4fd4-b67a-5e16b971d179",
          "kernelId": ""
        },
        "id": "fX12Yquyuihc"
      },
      "outputs": [],
      "source": [
        "#@title Install all dependencies (run only once per session)\n",
        "\n",
        "!git clone https://github.com/asigalov61/Yoda\n",
        "\n",
        "!pip install sinkhorn_transformer\n",
        "\n",
        "!pip install torch\n",
        "!pip install tqdm\n",
        "!pip install pickle5\n",
        "!pip install matplotlib\n",
        "\n",
        "!apt install fluidsynth #Pip does not work for some reason. Only apt works\n",
        "!pip install midi2audio\n",
        "!pip install pretty_midi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "gradient": {
          "editing": false,
          "id": "b8207b76-9514-4c07-95db-95a4742e52c5",
          "kernelId": ""
        },
        "id": "z7n9vnKmug1J"
      },
      "outputs": [],
      "source": [
        "#@title Import all needed modules\n",
        "\n",
        "print('Loading needed modules. Please wait...')\n",
        "import os\n",
        "import copy\n",
        "import random\n",
        "import tqdm\n",
        "import pickle5 as pickle\n",
        "\n",
        "from sinkhorn_transformer import SinkhornTransformerLM\n",
        "from sinkhorn_transformer.autoregressive_wrapper import AutoregressiveWrapper\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.optim as optim\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "\n",
        "print('Loading TMIDIX module...')\n",
        "os.chdir('/content/Yoda')\n",
        "import TMIDIX\n",
        "\n",
        "print('Loading aux modules...')\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from midi2audio import FluidSynth\n",
        "import pretty_midi\n",
        "import librosa.display\n",
        "from IPython.display import display, Javascript, HTML, Audio\n",
        "\n",
        "os.chdir('/content/')\n",
        "print('Done! Enjoy! :)')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ObPxlEutsQBj"
      },
      "source": [
        "# (MODEL)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Download pre-trained Sparse Yoda model and the training data file\n",
        "%cd /content/\n",
        "\n",
        "print('=' * 70)\n",
        "print('Unzipping pre-trained dataset-model...Please wait...')\n",
        "print('=' * 70)\n",
        "\n",
        "!wget --no-check-certificate -O 'Sparse-Yoda-Train-Data.pickle' \"https://onedrive.live.com/download?cid=8A0D502FC99C608F&resid=8A0D502FC99C608F%2118567&authkey=AGy0sliZlnM0rEk\"\n",
        "\n",
        "print('=' * 70)\n",
        "\n",
        "!wget --no-check-certificate -O 'Sparse-Yoda-Trained-Model.pth' \"https://onedrive.live.com/download?cid=8A0D502FC99C608F&resid=8A0D502FC99C608F%2118566&authkey=APcucQuadwRnHtQ\"\n",
        "\n",
        "\n",
        "print('Done! Enjoy! :)')\n",
        "print('=' * 70)\n",
        "%cd /content/"
      ],
      "metadata": {
        "cellView": "form",
        "id": "q2YTwObSiwbi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mdKFoeke9L7H"
      },
      "source": [
        "# (LOAD)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load and prep the original training data which will be used to prime the model\n",
        "full_path_to_original_training_data = \"/content/Sparse-Yoda-Train-Data.pickle\" #@param {type:\"string\"}\n",
        "\n",
        "melody_chords_f = pickle.load(open(full_path_to_original_training_data, 'rb'))\n",
        "\n",
        "randomize_dataset = False\n",
        "\n",
        "print('=' * 70)\n",
        "print('Prepping INTs dataset...')\n",
        "\n",
        "if randomize_dataset:\n",
        "    print('=' * 70)\n",
        "    print('Randomizing the dataset...')\n",
        "    random.shuffle(melody_chords_f)\n",
        "    print('Done!')\n",
        "    \n",
        "print('=' * 70)\n",
        "print('Processing the dataset...')\n",
        "\n",
        "r = 0\n",
        "\n",
        "train_data1 = []\n",
        "\n",
        "itimes = []\n",
        "ipitches = []\n",
        "\n",
        "\n",
        "for chords_list in tqdm.tqdm(melody_chords_f):\n",
        "    \n",
        "    train_data1.extend([0]) # Intro/Zero Token\n",
        "    \n",
        "    for i in chords_list:\n",
        "\n",
        "        if i[0] != 0: # This is the chordification line\n",
        "            train_data1.extend([i[0]]) # start-times\n",
        "            itimes.extend([i[0]])\n",
        "            ipitches.extend([i[1] + (i[2] * 16) + (i[3] * 16 * 128)])\n",
        "        # And this is the main MIDI note line (triple stack)\n",
        "        main_note = [i[1] + (i[2] * 16) + (i[3] * 16 * 128)] # Main note == [duration / pitch / channel]\n",
        "        \n",
        "        if main_note != [0]: # Main note error control...\n",
        "            train_data1.extend(main_note) # Main note == [duration / pitch / channel]\n",
        "\n",
        "print('Done!')        \n",
        "print('=' * 70)\n",
        "        \n",
        "print('Total INTs:', len(train_data1))\n",
        "print('Minimum INT:', min(train_data1))\n",
        "print('Maximum INT:', max(train_data1))\n",
        "print('Unique INTs:', len(set(train_data1)))\n",
        "print('Intro/Zero INTs:', train_data1.count(0))\n",
        "print('=' * 70)"
      ],
      "metadata": {
        "id": "awnEp4NAh3Bl",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gradient": {
          "id": "c83edd89-9a36-430a-9fa7-3a967417c88e",
          "kernelId": ""
        },
        "id": "OaNkGcFo9UP_",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Load/Reload the model\n",
        "\n",
        "from collections import OrderedDict\n",
        "\n",
        "full_path_to_model_checkpoint = \"/content/Sparse-Yoda-Trained-Model.pth\" #@param {type:\"string\"}\n",
        "\n",
        "print('Loading the model...')\n",
        "# constants\n",
        "\n",
        "NUM_BATCHES = int(1e5)\n",
        "BATCH_SIZE = 4\n",
        "GRADIENT_ACCUMULATE_EVERY = 4\n",
        "LEARNING_RATE = 1e-4\n",
        "VALIDATE_EVERY  = 100\n",
        "GENERATE_EVERY  = 500\n",
        "GENERATE_LENGTH = 512\n",
        "SEQ_LEN = 4096\n",
        "\n",
        "model = SinkhornTransformerLM(\n",
        "    num_tokens = max(train_data1)+1,\n",
        "    emb_dim = 128,\n",
        "    dim = 1024,\n",
        "    depth = 16,\n",
        "    max_seq_len = SEQ_LEN,\n",
        "    heads = 16,\n",
        "    bucket_size = 128,\n",
        "    ff_chunks = 2,\n",
        "    causal = True,\n",
        "    reversible = True,\n",
        "    attn_dropout = 0.1,\n",
        "    n_local_attn_heads = 4\n",
        ")\n",
        "\n",
        "model = AutoregressiveWrapper(model)\n",
        "model.cuda()\n",
        "\n",
        "model.load_state_dict(torch.load(full_path_to_model_checkpoint))\n",
        "model.eval()\n",
        "\n",
        "print('Done!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UX1_5y5Fu8AH"
      },
      "source": [
        "# (GENERATE MUSIC)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Custom MIDI option"
      ],
      "metadata": {
        "id": "IAYbC6hC3X0H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Custom MIDI\n",
        "full_path_to_custom_MIDI = \"/content/Yoda/seed4.mid\" #@param {type:\"string\"}\n",
        "\n",
        "score = TMIDIX.midi2ms_score(open(full_path_to_custom_MIDI, 'rb').read())\n",
        "\n",
        "events_matrix = []\n",
        "\n",
        "itrack = 1\n",
        "\n",
        "patches = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
        "\n",
        "patch_map = [[0, 1, 2, 3, 4, 5, 6, 7], # Piano \n",
        "              [24, 25, 26, 27, 28, 29, 30], # Guitar\n",
        "              [32, 33, 34, 35, 36, 37, 38, 39], # Bass\n",
        "              [40, 41], # Violin\n",
        "              [42, 43], # Cello\n",
        "              [46], # Harp\n",
        "              [56, 57, 58, 59, 60], # Trumpet\n",
        "              [71, 72], # Clarinet\n",
        "              [73, 74, 75], # Flute\n",
        "              [-1], # Fake Drums\n",
        "              [52, 53] # Choir\n",
        "            ]\n",
        "\n",
        "while itrack < len(score):\n",
        "    for event in score[itrack]:         \n",
        "        if event[0] == 'note' or event[0] == 'patch_change':\n",
        "            events_matrix.append(event)\n",
        "    itrack += 1\n",
        "\n",
        "events_matrix1 = []\n",
        "for event in events_matrix:\n",
        "        if event[0] == 'patch_change':\n",
        "            patches[event[2]] = event[3]\n",
        "\n",
        "        if event[0] == 'note':\n",
        "            event.extend([patches[event[3]]])\n",
        "            once = False\n",
        "            \n",
        "            for p in patch_map:\n",
        "                if event[6] in p and event[3] != 9: # Except the drums\n",
        "                    event[3] = patch_map.index(p)\n",
        "                    once = True\n",
        "                    \n",
        "            if not once and event[3] != 9: # Except the drums\n",
        "                event[3] = 0 # All other instruments/patches channel\n",
        "                event[5] = max(80, event[5])\n",
        "                \n",
        "            if event[3] < 11: # We won't write chans 11-16 for now...\n",
        "                events_matrix1.append(event)\n",
        "                #stats[event[3]] += 1\n",
        "\n",
        "# recalculating timings\n",
        "\n",
        "for e in events_matrix1:\n",
        "    e[1] = int(e[1] / 16)\n",
        "    e[2] = int(e[2] / 128)\n",
        "\n",
        "# final processing...\n",
        "\n",
        "if len(events_matrix1) > 0:\n",
        "    \n",
        "    events_matrix1.sort(key=lambda x: (x[1], x[4]))\n",
        "\n",
        "    cho = []\n",
        "    pe = events_matrix1[0]\n",
        "    melody_chords = []\n",
        "    for e in events_matrix1:\n",
        "\n",
        "        time = max(0, min(255, e[1]-pe[1]))\n",
        "        dur = max(0, min(15, e[2]))\n",
        "        cha = max(0, min(15, e[3]))\n",
        "        ptc = max(0, min(127, e[4]))\n",
        "        vel = max(0, min(127, e[5]))\n",
        "\n",
        "        melody_chords.append([time, dur, ptc, cha, vel])\n",
        "\n",
        "        pe = e\n",
        "inputs = []\n",
        "for i in melody_chords:\n",
        "\n",
        "  if i[0] != 0: # This is the chordification line\n",
        "      inputs.extend([i[0]]) # start-times\n",
        "      \n",
        "  # And this is the main MIDI note line (triple stack)\n",
        "  main_note = [i[1] + (i[2] * 16) + (i[3] * 16 * 128)] # Main note == [duration / pitch / channel]\n",
        "  \n",
        "  if main_note != [0]: # Main note error control...\n",
        "      inputs.extend(main_note) # Main note == [duration / pitch / channel]\n",
        "\n",
        "print('Done!')"
      ],
      "metadata": {
        "id": "3MnXjDWlab7B",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate"
      ],
      "metadata": {
        "id": "5AJ_QMOd3vd_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Generate\n",
        "priming_type = \"Custom MIDI\" #@param [\"Custom MIDI\", \"Random Dataset Point\"]\n",
        "number_of_prime_tokens = 256 #@param {type:\"slider\", min:32, max:512, step:32}\n",
        "number_of_tokens_to_generate = 256 #@param {type:\"slider\", min:64, max:512, step:64}\n",
        "temperature = 0.8 #@param {type:\"slider\", min:0.1, max:1, step:0.1}\n",
        "\n",
        "\n",
        "if priming_type == 'Random Dataset Point':\n",
        "  # Random dataset point\n",
        "  r = random.randint(0, int(len(train_data1) / 1))\n",
        "  out = train_data1[r:r+number_of_prime_tokens]\n",
        "\n",
        "else:\n",
        "  # Custom MIDI line\n",
        "  out = inputs[:number_of_prime_tokens] \n",
        "\n",
        "out1 = []\n",
        "out1.extend(out)\n",
        "\n",
        "sample = model.generate(torch.LongTensor(out).cuda(), number_of_tokens_to_generate, temperature=temperature)\n",
        "\n",
        "out2 = sample.cpu().numpy().tolist()\n",
        "out1.extend(out2)\n",
        "\n",
        "if len(out1) != 0:\n",
        "    \n",
        "    song = out1\n",
        "    song_f = []\n",
        "    time = 0\n",
        "    dur = 0\n",
        "    vel = 0\n",
        "    pitch = 0\n",
        "    channel = 0\n",
        "    \n",
        "    for s in song:\n",
        "        if s < 256:\n",
        "            time += s * 16\n",
        "            \n",
        "        else:\n",
        "            channel = s // 16 // 128\n",
        "\n",
        "            pitch = (s // 16) % 128\n",
        "            \n",
        "            dur = ((s % 16) * 128) + 128\n",
        "            \n",
        "            # Velocities for each channel:\n",
        "            if channel == 0:  # Piano     \n",
        "                vel = 60\n",
        "            if channel == 1:  # Guitar     \n",
        "                vel = 70            \n",
        "            if channel == 2:  # Bass     \n",
        "                vel = 60            \n",
        "            if channel == 3:  # Violin\n",
        "                vel = 90            \n",
        "            if channel == 4:  # Cello     \n",
        "                vel = 100\n",
        "            if channel == 5:  # Harp     \n",
        "                vel = 80\n",
        "            if channel == 6:  # Trumpet     \n",
        "                vel = 100            \n",
        "            if channel == 7:  # Clarinet     \n",
        "                vel = 100           \n",
        "            if channel == 8:  # Flute\n",
        "                vel = 100                          \n",
        "            if channel == 9:  # Drums\n",
        "                vel = 80            \n",
        "            if channel == 10:  # Choir     \n",
        "                vel = 110                  \n",
        "                               \n",
        "            song_f.append(['note', time, dur, channel, pitch, vel ])\n",
        "\n",
        "    detailed_stats = TMIDIX.Tegridy_SONG_to_MIDI_Converter(song_f,\n",
        "                                                        output_signature = 'Sparse Yoda',  \n",
        "                                                        output_file_name = '/content/Sparse-Yoda-Music-Composition', \n",
        "                                                        track_name='Project Los Angeles',\n",
        "                                                        list_of_MIDI_patches=[0, 24, 32, 40, 42, 46, 56, 71, 73, 0, 53, 0, 0, 0, 0, 0],\n",
        "                                                        number_of_ticks_per_quarter=500)\n",
        "\n",
        "    print('Done!')\n",
        "\n",
        "print('=' * 70)\n",
        "print('Displaying resulting composition...')\n",
        "fname = 'Sparse-Yoda-Music-Composition'\n",
        "\n",
        "pm = pretty_midi.PrettyMIDI(fname + '.mid')\n",
        "\n",
        "# Retrieve piano roll of the MIDI file\n",
        "piano_roll = pm.get_piano_roll()\n",
        "\n",
        "plt.figure(figsize=(14, 5))\n",
        "librosa.display.specshow(piano_roll, x_axis='time', y_axis='cqt_note', fmin=1, hop_length=160, sr=16000, cmap=plt.cm.hot)\n",
        "plt.title(fname)\n",
        "\n",
        "FluidSynth(\"/usr/share/sounds/sf2/FluidR3_GM.sf2\", 16000).midi_to_audio(str(fname + '.mid'), str(fname + '.wav'))\n",
        "Audio(str(fname + '.wav'), rate=16000)"
      ],
      "metadata": {
        "id": "E7uCe36UOBal",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YzCMd94Tu_gz"
      },
      "source": [
        "# Congrats! You did it! :)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Sparse_Yoda.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}