Tiny_Music_Transformer.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Tiny Music Transformer (ver. 1.0)\n",
        "\n",
        "***\n",
        "\n",
        "Powered by tegridy-tools: https://github.com/asigalov61/tegridy-tools\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 2023\n",
        "\n",
        "***"
      ],
      "metadata": {
        "id": "gpy3qsulqHa5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# (GPU CHECK)"
      ],
      "metadata": {
        "id": "W_So4w8fqPGL"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X3rABEpKCO02",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title NVIDIA GPU check\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# (SETUP ENVIRONMENT)"
      ],
      "metadata": {
        "id": "C0XxnXGFqVyh"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vK40g6V_BTNj",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Install dependencies\n",
        "!git clone --depth 1 https://github.com/asigalov61/Tiny-Music-Transformer\n",
        "!pip install huggingface_hub\n",
        "!pip install torch\n",
        "!pip install einops\n",
        "!pip install torch-summary\n",
        "!pip install tqdm\n",
        "!pip install matplotlib\n",
        "!apt install fluidsynth #Pip does not work for some reason. Only apt works"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DzCOZU_gBiQV",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Import modules\n",
        "\n",
        "print('=' * 70)\n",
        "print('Loading core Tiny Music Transformer modules...')\n",
        "\n",
        "import os\n",
        "import copy\n",
        "import pickle\n",
        "import secrets\n",
        "import statistics\n",
        "from time import time\n",
        "import tqdm\n",
        "\n",
        "print('=' * 70)\n",
        "print('Loading main Tiny Music Transformer modules...')\n",
        "import torch\n",
        "\n",
        "%cd /content/Tiny-Music-Transformer\n",
        "\n",
        "import TMIDIX\n",
        "\n",
        "from midi_to_colab_audio import midi_to_colab_audio\n",
        "\n",
        "from x_transformer_1_23_2 import *\n",
        "\n",
        "import random\n",
        "\n",
        "%cd /content/\n",
        "print('=' * 70)\n",
        "print('Loading aux Tiny Music Transformer modules...')\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from torchsummary import summary\n",
        "from sklearn import metrics\n",
        "\n",
        "from IPython.display import Audio, display\n",
        "\n",
        "from huggingface_hub import hf_hub_download\n",
        "\n",
        "from google.colab import files\n",
        "\n",
        "print('=' * 70)\n",
        "print('Done!')\n",
        "print('Enjoy! :)')\n",
        "print('=' * 70)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eI3aQtHzqSnp"
      },
      "source": [
        "# (LOAD MODEL)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load Tiny Music Transformer Pre-Trained Model\n",
        "\n",
        "#@markdown Choose model\n",
        "\n",
        "select_model_to_load = \"139M-32L-Very-Fast-Tiny\" # @param [\"139M-32L-Very-Fast-Tiny\"]\n",
        "\n",
        "#@markdown Model precision option\n",
        "\n",
        "model_precision = \"bfloat16\" # @param [\"bfloat16\", \"float16\"]\n",
        "\n",
        "#@markdown bfloat16 == Half precision/faster speed (if supported, otherwise the model will default to float16)\n",
        "\n",
        "#@markdown float16 == Full precision/fast speed\n",
        "\n",
        "plot_tokens_embeddings = \"None\" # @param [\"None\", \"Start Times\", \"Durations Velocities\", \"Piano Pitches\", \"Violin Pitches\", \"Aux\"]\n",
        "\n",
        "full_path_to_models_dir = \"/content/Tiny-Music-Transformer/Model/\"\n",
        "\n",
        "print('=' * 70)\n",
        "print('Loading Tiny Music Transformer Training Data...')\n",
        "print('Please wait...')\n",
        "print('=' * 70)\n",
        "\n",
        "training_data_file = 'Tiny_Music_Transformer_Train_Data_Mono_Melodies_Piano_Violin_MIDI_Dataset.pickle'\n",
        "\n",
        "if os.path.isfile(full_path_to_models_dir+training_data_file):\n",
        "  print('Training data already exists...')\n",
        "\n",
        "else:\n",
        "  hf_hub_download(repo_id='asigalov61/Tiny-Music-Transformer',\n",
        "                  filename=training_data_file,\n",
        "                  local_dir='/content/Tiny-Music-Transformer/Model',\n",
        "                  local_dir_use_symlinks=False)\n",
        "\n",
        "print('=' * 70)\n",
        "train_data = TMIDIX.Tegridy_Any_Pickle_File_Reader('/content/Tiny-Music-Transformer/Model/Tiny_Music_Transformer_Train_Data_Mono_Melodies_Piano_Violin_MIDI_Dataset')\n",
        "print('Done!')\n",
        "print('=' * 70)\n",
        "\n",
        "print('Loading Tiny Music Transformer', select_model_to_load,'Pre-Trained Model...')\n",
        "print('Please wait...')\n",
        "print('=' * 70)\n",
        "\n",
        "if select_model_to_load == '139M-32L-Very-Fast-Tiny':\n",
        "\n",
        "  model_checkpoint_file_name = 'Tiny_Music_Transformer_Tiny_Trained_Model_10737_steps_0.4039_loss_0.8729_acc.pth'\n",
        "  model_path = full_path_to_models_dir+model_checkpoint_file_name\n",
        "  if os.path.isfile(model_path):\n",
        "    print('Model already exists...')\n",
        "\n",
        "  else:\n",
        "    hf_hub_download(repo_id='asigalov61/Tiny-Music-Transformer',\n",
        "                    filename=model_checkpoint_file_name,\n",
        "                    local_dir='/content/Tiny-Music-Transformer/Model',\n",
        "                    local_dir_use_symlinks=False)\n",
        "\n",
        "\n",
        "\n",
        "print('=' * 70)\n",
        "print('Instantiating model...')\n",
        "\n",
        "torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul\n",
        "torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn\n",
        "device_type = 'cuda'\n",
        "\n",
        "if model_precision == 'bfloat16' and torch.cuda.is_bf16_supported():\n",
        "  dtype = 'bfloat16'\n",
        "else:\n",
        "  dtype = 'float16'\n",
        "\n",
        "if model_precision == 'float16':\n",
        "  dtype = 'float16'\n",
        "\n",
        "ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]\n",
        "ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)\n",
        "\n",
        "SEQ_LEN = 8192\n",
        "\n",
        "# instantiate the model\n",
        "\n",
        "model = TransformerWrapper(\n",
        "    num_tokens = 644,\n",
        "    max_seq_len = SEQ_LEN,\n",
        "    attn_layers = Decoder(dim = 512, depth = 32, heads = 16, attn_flash = True)\n",
        ")\n",
        "\n",
        "model = AutoregressiveWrapper(model, ignore_index=643)\n",
        "\n",
        "model.cuda()\n",
        "print('=' * 70)\n",
        "\n",
        "print('Loading model checkpoint...')\n",
        "\n",
        "model.load_state_dict(torch.load(model_path))\n",
        "print('=' * 70)\n",
        "\n",
        "model.eval()\n",
        "\n",
        "print('Done!')\n",
        "print('=' * 70)\n",
        "\n",
        "print('Model will use', dtype, 'precision...')\n",
        "print('=' * 70)\n",
        "\n",
        "# Model stats\n",
        "print('Model summary...')\n",
        "summary(model)\n",
        "\n",
        "# Plot Token Embeddings\n",
        "if plot_tokens_embeddings != 'None':\n",
        "  tok_emb = model.net.token_emb.emb.weight.detach().cpu().tolist()\n",
        "\n",
        "if plot_tokens_embeddings == 'Start Times':\n",
        "  tok_range = [0, 128]\n",
        "\n",
        "elif plot_tokens_embeddings == 'Durations Velocities':\n",
        "  tok_range = [128, 256]\n",
        "\n",
        "elif plot_tokens_embeddings == 'Piano Pitches':\n",
        "  tok_range = [256, 384]\n",
        "\n",
        "elif plot_tokens_embeddings == 'Violin Pitches':\n",
        "  tok_range = [384, 512]\n",
        "\n",
        "elif plot_tokens_embeddings == 'Aux':\n",
        "  tok_range = [512, 643]\n",
        "\n",
        "if plot_tokens_embeddings != 'None':\n",
        "\n",
        "  tok_emb1 = []\n",
        "\n",
        "  for t in tok_emb[tok_range[0]:tok_range[1]]:\n",
        "    tok_emb1.append(t)\n",
        "\n",
        "  cos_sim = metrics.pairwise_distances(\n",
        "    tok_emb1, metric='cosine'\n",
        "  )\n",
        "  plt.figure(figsize=(7, 7))\n",
        "  plt.imshow(cos_sim, cmap=\"inferno\", interpolation=\"nearest\")\n",
        "  im_ratio = cos_sim.shape[0] / cos_sim.shape[1]\n",
        "  plt.colorbar(fraction=0.046 * im_ratio, pad=0.04)\n",
        "  plt.xlabel(\"Position\")\n",
        "  plt.ylabel(\"Position\")\n",
        "  plt.tight_layout()\n",
        "  plt.plot()\n",
        "  plt.savefig(\"/content/Tiny-Music-Transformer-Tokens-Embeddings-Plot.png\", bbox_inches=\"tight\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "V4s_G8yUL0cH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# (GENERATE)"
      ],
      "metadata": {
        "id": "7xNyANjZsCOi"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# (PRIMED CONTINUATION)"
      ],
      "metadata": {
        "id": "Gt03VtO6uKkb"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dkvXYwR_qSnx",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Primed Continuation\n",
        "\n",
        "#@markdown Generation settings\n",
        "\n",
        "number_of_prime_tokens = 512 # @param {type:\"slider\", min:3, max:2048, step:3}\n",
        "number_of_tokens_to_generate = 1024 # @param {type:\"slider\", min:30, max:8190, step:3}\n",
        "number_of_batches_to_generate = 8 #@param {type:\"slider\", min:1, max:16, step:1}\n",
        "temperature = 0.9 # @param {type:\"slider\", min:0.1, max:1, step:0.05}\n",
        "\n",
        "#@markdown Other settings\n",
        "include_prime_tokens_in_generated_output = True #@param {type:\"boolean\"}\n",
        "render_MIDI_to_audio = True # @param {type:\"boolean\"}\n",
        "\n",
        "print('=' * 70)\n",
        "print('Tiny Music Transformer Primed Continuation Model Generator')\n",
        "print('=' * 70)\n",
        "\n",
        "random.shuffle(train_data)\n",
        "\n",
        "melody_chords = random.choice(train_data)\n",
        "\n",
        "outy = melody_chords[:number_of_prime_tokens]\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "inp = [outy] * number_of_batches_to_generate\n",
        "\n",
        "inp = torch.LongTensor(inp).cuda()\n",
        "\n",
        "with ctx:\n",
        "  out = model.generate(inp,\n",
        "                        number_of_tokens_to_generate,\n",
        "                        temperature=temperature,\n",
        "                        return_prime=include_prime_tokens_in_generated_output,\n",
        "                        verbose=True)\n",
        "\n",
        "out0 = out.tolist()\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "print('=' * 70)\n",
        "print('Done!')\n",
        "print('=' * 70)\n",
        "\n",
        "#======================================================================\n",
        "print('Rendering results...')\n",
        "\n",
        "for i in range(number_of_batches_to_generate):\n",
        "\n",
        "  print('=' * 70)\n",
        "  print('Batch #', i)\n",
        "  print('=' * 70)\n",
        "\n",
        "  out1 = out0[i]\n",
        "\n",
        "  print('Sample INTs', out1[:12])\n",
        "  print('=' * 70)\n",
        "\n",
        "  if len(out) != 0:\n",
        "\n",
        "      song = out1\n",
        "      song_f = []\n",
        "\n",
        "      time = 0\n",
        "      dur = 0\n",
        "      vel = 90\n",
        "      pitch = 0\n",
        "      channel = 0\n",
        "\n",
        "      patches = [0] * 16\n",
        "      patches[3] = 40\n",
        "\n",
        "      for ss in song:\n",
        "\n",
        "          if 0 <= ss < 128:\n",
        "\n",
        "              time += (ss * 16)\n",
        "\n",
        "          if 128 <= ss < 256:\n",
        "\n",
        "              dur = (ss-128) * 16\n",
        "\n",
        "          if 256 <= ss < 512:\n",
        "\n",
        "              pitch = (ss-256) % 128\n",
        "\n",
        "              channel = (ss-256) // 128\n",
        "\n",
        "              if channel == 1:\n",
        "                channel = 3\n",
        "                vel = 110\n",
        "              else:\n",
        "                channel = 0\n",
        "                vel = 80\n",
        "\n",
        "              song_f.append(['note', time, dur, channel, pitch, vel, patches[channel] ])\n",
        "\n",
        "      patches = [0 if x==-1 else x for x in patches]\n",
        "\n",
        "      detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,\n",
        "                                                                output_signature = 'Tiny Music Transformer',\n",
        "                                                                output_file_name = '/content/Tiny-Music-Transformer-Music-Composition_'+str(i),\n",
        "                                                                track_name='Project Los Angeles',\n",
        "                                                                list_of_MIDI_patches=patches\n",
        "                                                                )\n",
        "      print('=' * 70)\n",
        "      print('Displaying resulting composition...')\n",
        "      print('=' * 70)\n",
        "\n",
        "      fname = '/content/Tiny-Music-Transformer-Music-Composition_'+str(i)\n",
        "\n",
        "      if render_MIDI_to_audio:\n",
        "        midi_audio = midi_to_colab_audio(fname + '.mid')\n",
        "        display(Audio(midi_audio, rate=16000, normalize=False))\n",
        "\n",
        "      TMIDIX.plot_ms_SONG(song_f, plot_title=fname)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Congrats! You did it! :)"
      ],
      "metadata": {
        "id": "eoWDEy6CwDr6"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4",
      "gpuClass": "premium"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.9.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}