heptabit_music_transformer_composer.py

# -*- coding: utf-8 -*-
"""Heptabit_Music_Transformer_Composer.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/asigalov61/Heptabit-Music-Transformer/blob/main/Heptabit_Music_Transformer_Composer.ipynb

# Heptabit Music Transformer Composer (ver. 1.5)

***

Powered by tegridy-tools: https://github.com/asigalov61/tegridy-tools

***

WARNING: This complete implementation is a functioning model of the Artificial Intelligence. Please excercise great humility, care, and respect. https://www.nscai.gov/

***

#### Project Los Angeles

#### Tegridy Code 2023

***

# (GPU CHECK)
"""

#@title NVIDIA GPU check
!nvidia-smi

"""# (SETUP ENVIRONMENT)"""

#@title Install dependencies
!git clone --depth 1 https://github.com/asigalov61/Heptabit-Music-Transformer
!pip install huggingface_hub
!pip install torch
!pip install einops
!pip install torch-summary
!pip install tqdm
!pip install matplotlib
!apt install fluidsynth #Pip does not work for some reason. Only apt works

# Commented out IPython magic to ensure Python compatibility.
#@title Import modules

print('=' * 70)
print('Loading core Heptabit Music Transformer modules...')

import os
import pickle
import secrets
import statistics
from time import time
import tqdm

print('=' * 70)
print('Loading main Heptabit Music Transformer modules...')
import torch

# %cd /content/Heptabit-Music-Transformer

import TMIDIX

from midi_to_colab_audio import midi_to_colab_audio

from x_transformer_1_23_2 import *

import random

# %cd /content/
print('=' * 70)
print('Loading aux Heptabit Music Transformer modules...')

import matplotlib.pyplot as plt

from torchsummary import summary
from sklearn import metrics

from IPython.display import Audio, display

from huggingface_hub import hf_hub_download

from google.colab import files

print('=' * 70)
print('Done!')
print('Enjoy! :)')
print('=' * 70)

"""# (LOAD MODEL)"""

#@title Load Heptabit Music Transformer Large Model

#@markdown Very fast model, 32 layers, 245k MIDIs training corpus

full_path_to_model_checkpoint = "/content/Heptabit-Music-Transformer/Models/Large/Heptabit_Music_Transformer_Large_Trained_Model_38430_steps_0.4225_loss_0.8744_acc.pth" #@param {type:"string"}

#@markdown Model precision option

model_precision = "bfloat16" # @param ["bfloat16", "float16"]

#@markdown bfloat16 == Half precision/faster speed (if supported, otherwise the model will default to float16)

#@markdown float16 == Full precision/fast speed

plot_tokens_embeddings = False # @param {type:"boolean"}

print('=' * 70)
print('Loading Heptabit Music Transformer Large Pre-Trained Model...')
print('Please wait...')
print('=' * 70)

if os.path.isfile(full_path_to_model_checkpoint):
  print('Model already exists...')

else:
  hf_hub_download(repo_id='asigalov61/Heptabit-Music-Transformer',
                  filename='Heptabit_Music_Transformer_Large_Trained_Model_38430_steps_0.4225_loss_0.8744_acc.pth',
                  local_dir='/content/Heptabit-Music-Transformer/Models/Large',
                  local_dir_use_symlinks=False)

print('=' * 70)
print('Instantiating model...')

torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
device_type = 'cuda'

if model_precision == 'bfloat16' and torch.cuda.is_bf16_supported():
  dtype = 'bfloat16'
else:
  dtype = 'float16'

if model_precision == 'float16':
  dtype = 'float16'

ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)

SEQ_LEN = 8192

# instantiate the model

model = TransformerWrapper(
    num_tokens = 17287,
    max_seq_len = SEQ_LEN,
    attn_layers = Decoder(dim = 1024, depth = 32, heads = 32, attn_flash = True)
)

model = AutoregressiveWrapper(model, ignore_index=17286)

model.cuda()
print('=' * 70)

print('Loading model checkpoint...')

model.load_state_dict(torch.load(full_path_to_model_checkpoint))
print('=' * 70)

model.eval()

print('Done!')
print('=' * 70)

print('Model will use', dtype, 'precision...')
print('=' * 70)

# Model stats
print('Model summary...')
summary(model)

# Plot Token Embeddings
if plot_tokens_embeddings:
  tok_emb = model.net.token_emb.emb.weight.detach().cpu().tolist()

  cos_sim = metrics.pairwise_distances(
    tok_emb, metric='cosine'
  )
  plt.figure(figsize=(7, 7))
  plt.imshow(cos_sim, cmap="inferno", interpolation="nearest")
  im_ratio = cos_sim.shape[0] / cos_sim.shape[1]
  plt.colorbar(fraction=0.046 * im_ratio, pad=0.04)
  plt.xlabel("Position")
  plt.ylabel("Position")
  plt.tight_layout()
  plt.plot()
  plt.savefig("/content/Heptabit-Music-Transformer-Large-Tokens-Embeddings-Plot.png", bbox_inches="tight")

"""# (LOAD VELOCITIES MAP)"""

# @title Load velocities map

plot_velocities_map = False # @param {type:"boolean"}

print('=' * 70)
print('Loading Heptabit Music Transformer Velocities Map...')
print('=' * 70)

dynamic_velocities_map = TMIDIX.Tegridy_Any_Pickle_File_Reader('/content/Heptabit-Music-Transformer/Aux-Data/pitches_patches_MIDI_velocities_map')[0]

print('=' * 70)
print('Done!')
print('=' * 70)

if plot_velocities_map:
  cos_sim = metrics.pairwise_distances(
    dynamic_velocities_map, metric='minkowski'
  )
  plt.figure(figsize=(7, 7))
  plt.imshow(cos_sim, cmap="inferno", interpolation="nearest")
  im_ratio = cos_sim.shape[0] / cos_sim.shape[1]
  plt.colorbar(fraction=0.046 * im_ratio, pad=0.04)
  plt.xlabel("Pitches")
  plt.ylabel("Patches")
  plt.tight_layout()
  plt.plot()
  plt.savefig("/content/Heptabit-Music-Transformer-Velocities-Map-Plot.png", bbox_inches="tight")

"""# (LOAD SEED MIDI)"""

#@title Load Seed MIDI

#@markdown Press play button to to upload your own seed MIDI or to load one of the provided sample seed MIDIs from the dropdown list below

select_seed_MIDI = "Upload your own custom MIDI" # @param ["Upload your own custom MIDI", "Heptabit-Music-Transformer-Piano-Seed-1", "Heptabit-Music-Transformer-Piano-Seed-2", "Heptabit-Music-Transformer-Piano-Seed-3", "Heptabit-Music-Transformer-Piano-Seed-4", "Heptabit-Music-Transformer-Piano-Seed-5", "Heptabit-Music-Transformer-MI-Seed-1", "Heptabit-Music-Transformer-MI-Seed-2", "Heptabit-Music-Transformer-MI-Seed-3", "Heptabit-Music-Transformer-MI-Seed-4", "Heptabit-Music-Transformer-MI-Seed-5"]
number_of_prime_tokens = 1023 # @param {type:"slider", min:90, max:8190, step:3}
render_MIDI_to_audio = True # @param {type:"boolean"}

print('=' * 70)
print('Heptabit Music Transformer Seed MIDI Loader')
print('=' * 70)

f = ''

if select_seed_MIDI != "Upload your own custom MIDI":
  print('Loading seed MIDI...')
  f = '/content/Heptabit-Music-Transformer/Seeds/'+select_seed_MIDI+'.mid'
  score = TMIDIX.midi2single_track_ms_score(open(f, 'rb').read(), recalculate_channels=False)

else:
  print('Upload your own custom MIDI...')
  print('=' * 70)
  uploaded_MIDI = files.upload()
  if list(uploaded_MIDI.keys()):
    score = TMIDIX.midi2single_track_ms_score(list(uploaded_MIDI.values())[0], recalculate_channels=False)
    f = list(uploaded_MIDI.keys())[0]

if f != '':

  print('=' * 70)
  print('File:', f)
  print('=' * 70)

  #=======================================================
  # START PROCESSING

  # Convering MIDI to ms score with MIDI.py module
  score = TMIDIX.midi2single_track_ms_score(open(f, 'rb').read(), recalculate_channels=False)

  # INSTRUMENTS CONVERSION CYCLE
  events_matrix = []
  itrack = 1
  patches = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

  while itrack < len(score):
      for event in score[itrack]:
          if event[0] == 'note' or event[0] == 'patch_change':
              events_matrix.append(event)
      itrack += 1

  events_matrix.sort(key=lambda x: x[1])

  events_matrix1 = []

  for event in events_matrix:
          if event[0] == 'patch_change':
                patches[event[2]] = event[3]

          if event[0] == 'note':
                event.extend([patches[event[3]]])

                events_matrix1.append(event)

  if len(events_matrix1) > 0:
    if min([e[1] for e in events_matrix1]) >= 0 and min([e[2] for e in events_matrix1]) >= 0:

      #=======================================================
      # PRE-PROCESSING

      # checking number of instruments in a composition
      instruments_list_without_drums = list(set([y[3] for y in events_matrix1 if y[3] != 9]))
      instruments_list = list(set([y[3] for y in events_matrix1]))

      if len(events_matrix1) > 0 and len(instruments_list_without_drums) > 0:

        # recalculating timings
        for e in events_matrix1:
            e[1] = int(e[1] / 32) # Max 4 seconds for start-times
            e[2] = int(e[2] / 32) # Max 4 seconds for durations

        # Sorting by pitch, then by start-time
        events_matrix1.sort(key=lambda x: x[6])
        events_matrix1.sort(key=lambda x: x[4])
        events_matrix1.sort(key=lambda x: x[1])

        #=======================================================
        # FINAL PRE-PROCESSING

        melody_chords = []

        pe = events_matrix1[0]

        for e in events_matrix1:

            # Cliping all values...
            time = max(0, min(127, e[1]-pe[1]))
            dur = max(0, min(127, e[2]))
            cha = max(0, min(15, e[3]))
            ptc = max(1, min(127, e[4]))

            pat = max(0, min(127, e[6]))

            # Writing final note
            melody_chords.append([time, dur, cha, ptc, pat])

            pe = e

        #=======================================================
        # FINAL PROCESSING
        #=======================================================

        melody_chords2 = []

        # Break between compositions / Intro seq

        if 9 in instruments_list:
          drums_present = 17155 # Yes
        else:
          drums_present = 17154 # No

        if melody_chords[0][2] != 9:
            pat = melody_chords[0][4]
        else:
            pat = 128

        melody_chords2.extend([17285, drums_present, 17156+pat]) # Intro seq

        #=======================================================

        # TOTAL DICTIONARY SIZE 17285+1=17286

        #=======================================================
        # MAIN PROCESSING CYCLE
        #=======================================================

        chords_counter = 1

        comp_chords_len = len([y for y in melody_chords if y[0] != 0])

        for m in melody_chords:

            # if ((comp_chords_len - chords_counter) == 50) and (m[0] != 0):
            #    melody_chords2.extend([17153, 17153, 17153]) # Outro token seq

            if (chords_counter % 50 == 0) and (m[0] != 0):
                nct = 16897+min(255, ((chords_counter // 50)-1)) # chords counter token
                melody_chords2.extend([nct, nct, nct])
                chords_counter += 1
            else:
                if m[0] != 0:
                    chords_counter += 1

            # Drums patch
            if m[2] == 9: # Drums patch will be == 128
                pat = 128
            else:
                pat = m[4]

            cha_pat = (129 * pat) + m[3]

            melody_chords2.extend([m[0], m[1]+128, cha_pat+256])

        # melody_chords2.extend([17285, 17285, 17285]) # EOS

  melody_chords_f = melody_chords2[:number_of_prime_tokens]

  #=======================================================

  song = melody_chords_f

  song_f = []

  time = 0
  dur = 0
  channel = 0
  pitch = 0
  vel = 90

  patches = [-1] * 16

  channels = [0] * 16
  channels[9] = 1

  for ss in song:

      if 0 <= ss < 128:

          time += (ss * 32)

      if 128 <= ss < 256:

          dur = ((ss-128) * 32)

      if 256 <= ss < 16897:
          patch = (ss-256) // 129

          if patch < 128:

              if patch not in patches:
                if 0 in channels:
                    cha = channels.index(0)
                    channels[cha] = 1
                else:
                    cha = 15

                patches[cha] = patch
                channel = patches.index(patch)
              else:
                channel = patches.index(patch)

          if patch == 128:
              channel = 9

          pitch = (ss-256) % 129

          vel = round(dynamic_velocities_map[pitch][patch])

          song_f.append(['note', time, dur, channel, pitch, vel ])

  patches = [0 if x==-1 else x for x in patches]

  detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,
                                                            output_signature = 'Heptabit Music Transformer',
                                                            output_file_name = '/content/Heptabit-Music-Transformer-Seed-Composition',
                                                            track_name='Project Los Angeles',
                                                            list_of_MIDI_patches=patches
                                                            )

  #=======================================================

  print('=' * 70)
  print('Composition stats:')
  print('Composition has', len(melody_chords), 'notes')
  print('Composition has', len(melody_chords2), 'tokens')
  print('Composition MIDI patches:', sorted(list(set([((y-256) // 129) for y in melody_chords_f if 256 <= y < 16897]))))
  print('=' * 70)

  print('Displaying resulting composition...')
  print('=' * 70)

  fname = '/content/Heptabit-Music-Transformer-Seed-Composition'

  x = []
  y =[]
  c = []

  colors = ['red', 'yellow', 'green', 'cyan',
            'blue', 'pink', 'orange', 'purple',
            'gray', 'white', 'gold', 'silver',
            'lightgreen', 'indigo', 'maroon', 'turquoise']

  block_lines = [(song_f[-1][1] / 1000)]
  block_tokens = [min(len(melody_chords_f), number_of_prime_tokens)]

  for s in song_f:
    x.append(s[1] / 1000)
    y.append(s[4])
    c.append(colors[s[3]])

  if render_MIDI_to_audio:
    midi_audio = midi_to_colab_audio(fname + '.mid')
    display(Audio(midi_audio, rate=16000, normalize=False))

  plt.figure(figsize=(14,5))
  ax=plt.axes(title=fname)
  ax.set_facecolor('black')

  plt.scatter(x,y, c=c)
  plt.xlabel("Time")
  plt.ylabel("Pitch")
  plt.show()

else:
  print('=' * 70)

"""# (COMPOSITION LOOP)

## Run the cells below in a loop to generate endless continuation
"""

#@title Standard Continuation Generator

#@markdown Generation settings

try_to_generate_outro = False #@param {type:"boolean"}
number_of_tokens_to_generate = 408 # @param {type:"slider", min:33, max:8190, step:3}
number_of_batches_to_generate = 4 #@param {type:"slider", min:1, max:16, step:1}
preview_length_in_tokens = 120 # @param {type:"slider", min:33, max:240, step:3}
number_of_memory_tokens = 8190 # @param {type:"slider", min:300, max:8190, step:3}
temperature = 0.85 # @param {type:"slider", min:0.1, max:1, step:0.05}

#@markdown Other settings

render_MIDI_to_audio = True # @param {type:"boolean"}

print('=' * 70)
print('Heptabit Music Transformer Standard Continuation Model Generator')
print('=' * 70)

preview = melody_chords_f[-preview_length_in_tokens:]

mel_cho = melody_chords_f[-number_of_memory_tokens:]

if try_to_generate_outro:
  mel_cho.extend([17153, 17153, 17153])

inp = [mel_cho] * number_of_batches_to_generate

inp = torch.LongTensor(inp).cuda()

with ctx:
  out = model.generate(inp,
                        number_of_tokens_to_generate,
                        temperature=temperature,
                        return_prime=False,
                        eos_token=None,
                        verbose=True)

out0 = out.tolist()

print('=' * 70)
print('Done!')
print('=' * 70)
#======================================================================
print('Rendering results...')

for i in range(number_of_batches_to_generate):

  print('=' * 70)
  print('Batch #', i)
  print('=' * 70)

  out1 = out0[i]

  print('Sample INTs', out1[:12])
  print('=' * 70)

  if len(out) != 0:

      song = preview + out1
      song_f = []

      time = 0
      dur = 0
      channel = 0
      pitch = 0
      vel = 90

      patches = [-1] * 16

      channels = [0] * 16
      channels[9] = 1

      for ss in song:

          if 0 <= ss < 128:

              time += (ss * 32)

          if 128 <= ss < 256:

              dur = ((ss-128) * 32)

          if 256 <= ss < 16897:
              patch = (ss-256) // 129

              if patch < 128:

                  if patch not in patches:
                    if 0 in channels:
                        cha = channels.index(0)
                        channels[cha] = 1
                    else:
                        cha = 15

                    patches[cha] = patch
                    channel = patches.index(patch)
                  else:
                    channel = patches.index(patch)

              if patch == 128:
                  channel = 9

              pitch = (ss-256) % 129

              vel = round(dynamic_velocities_map[pitch][patch])

              song_f.append(['note', time, dur, channel, pitch, vel ])

      patches = [0 if x==-1 else x for x in patches]

      detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,
                                                          output_signature = 'Heptabit Music Transformer',
                                                          output_file_name = '/content/Heptabit-Music-Transformer-Music-Composition_'+str(i),
                                                          track_name='Project Los Angeles',
                                                          list_of_MIDI_patches=patches
                                                          )
      print('=' * 70)
      print('Displaying resulting composition...')
      print('=' * 70)

      fname = '/content/Heptabit-Music-Transformer-Music-Composition_'+str(i)

      x = []
      y =[]
      c = []

      colors = ['red', 'yellow', 'green', 'cyan',
                'blue', 'pink', 'orange', 'purple',
                'gray', 'white', 'gold', 'silver',
                'lightgreen', 'indigo', 'maroon', 'turquoise']

      for s in song_f:
        x.append(s[1] / 1000)
        y.append(s[4])
        c.append(colors[s[3]])

      if render_MIDI_to_audio:
        midi_audio = midi_to_colab_audio(fname + '.mid')
        display(Audio(midi_audio, rate=16000, normalize=False))

      plt.figure(figsize=(14,5))
      ax=plt.axes(title=fname)
      ax.set_facecolor('black')

      plt.scatter(x,y, c=c)

      pbl = song_f[(int(preview_length_in_tokens / 3))][1] / 1000

      ax.axvline(x=pbl, c='w')

      plt.xlabel("Time")
      plt.ylabel("Pitch")
      plt.show()

#@title Choose one generated block to add to the composition
block_action = "add_last_generated_block" #@param ["add_last_generated_block", "remove_last_added_block"]
add_block_with_batch_number = 0 #@param {type:"slider", min:0, max:15, step:1}
render_MIDI_to_audio = True # @param {type:"boolean"}

print('=' * 70)

if block_action == 'add_last_generated_block':
  melody_chords_f.extend(out0[min(len(out0)-1, add_block_with_batch_number)])
  print('Block added!')
else:
  if len(block_tokens) > 1:
    melody_chords_f = melody_chords_f[:(len(melody_chords_f)-block_tokens[-1])]
    print('Block removed!')
  else:
    print('Nothing to remove!!!')

print('=' * 70)
print('Composition now has', (len(melody_chords_f) // 3), 'notes')
print('Composition now has', len(melody_chords_f), 'tokens')
print('Composition now has MIDI patches:', sorted(list(set([((y-256) // 129) for y in melody_chords_f if 256 <= y < 16897]))))

print('=' * 70)
print('Sample INTs', out1[:12])
print('=' * 70)

if len(melody_chords_f) != 0:

    song = melody_chords_f
    song_f = []

    time = 0
    dur = 0
    channel = 0
    pitch = 0
    vel = 90

    patches = [-1] * 16

    channels = [0] * 16
    channels[9] = 1

    for ss in song:

        if 0 <= ss < 128:

            time += (ss * 32)

        if 128 <= ss < 256:

            dur = ((ss-128) * 32)

        if 256 <= ss < 16897:
            patch = (ss-256) // 129

            if patch < 128:

                if patch not in patches:
                  if 0 in channels:
                      cha = channels.index(0)
                      channels[cha] = 1
                  else:
                      cha = 15

                  patches[cha] = patch
                  channel = patches.index(patch)
                else:
                  channel = patches.index(patch)

            if patch == 128:
                channel = 9

            pitch = (ss-256) % 129

            vel = round(dynamic_velocities_map[pitch][patch])

            song_f.append(['note', time, dur, channel, pitch, vel ])

    patches = [0 if x==-1 else x for x in patches]

    detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,
                                                              output_signature = 'Heptabit Music Transformer',
                                                              output_file_name = '/content/Heptabit-Music-Transformer-Music-Composition',
                                                              track_name='Project Los Angeles',
                                                              list_of_MIDI_patches=patches
                                                              )
    print('=' * 70)
    print('Displaying resulting composition...')
    print('=' * 70)

    fname = '/content/Heptabit-Music-Transformer-Music-Composition'

    x = []
    y =[]
    c = []

    colors = ['red', 'yellow', 'green', 'cyan',
              'blue', 'pink', 'orange', 'purple',
              'gray', 'white', 'gold', 'silver',
              'lightgreen', 'indigo', 'maroon', 'turquoise']

    if block_action == 'add_last_generated_block':
      block_lines.append((song_f[-1][1] / 1000))
      block_tokens.append(len(out0[min(len(out0)-1, add_block_with_batch_number)]))

    else:
      if len(block_tokens) > 1:
        block_lines.pop()
        block_tokens.pop()

    for s in song_f:
      x.append(s[1] / 1000)
      y.append(s[4])
      c.append(colors[s[3]])

    if render_MIDI_to_audio:
      midi_audio = midi_to_colab_audio(fname + '.mid')
      display(Audio(midi_audio, rate=16000, normalize=False))

    plt.figure(figsize=(14,5))
    ax=plt.axes(title=fname)
    ax.set_facecolor('black')

    plt.scatter(x,y, c=c)

    for bl in block_lines:
      ax.axvline(x=bl, c='w')

    plt.xlabel("Time")
    plt.ylabel("Pitch")
    plt.show()

"""# Congrats! You did it! :)"""