ultimate_accompaniment_transformer.py

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

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/github/asigalov61/Ultimate-Accompaniment-Transformer/blob/main/Ultimate_Accompaniment_Transformer.ipynb

# Ultimate Accompaniment Transformer (ver. 3.0)

***

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 2024

***

# (GPU CHECK)
"""

#@title NVIDIA GPU check
!nvidia-smi

"""# (SETUP ENVIRONMENT)"""

#@title Install dependencies
!git clone --depth 1 https://github.com/asigalov61/Ultimate-Accompaniment-Transformer
!pip install huggingface_hub
!pip install einops
!pip install torch-summary
!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 Ultimate Accompaniment Transformer modules...')

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

print('=' * 70)
print('Loading main Ultimate Accompaniment Transformer modules...')
import torch

# %cd /content/Ultimate-Accompaniment-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 Ultimate Accompaniment 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 Chords Progressions Transformer Pre-Trained Model

#@markdown Choose model

select_model_to_load = "187M-Improved-Very-Fast-Small" # @param ["187M-Original-Very-Fast-Small", "187M-Improved-Very-Fast-Small"]

#@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 Chords Progressions Transformer Pre-Trained Model...')
print('Please wait...')
print('=' * 70)

full_path_to_models_dir = "/content/Ultimate-Accompaniment-Transformer/Models"

if select_model_to_load == '187M-Original-Very-Fast-Small':

  dim = 2048
  depth = 4
  heads = 16

  model_checkpoint_file_name = 'Ultimate_Accompaniment_Transformer_Small_Trained_Model_13414_steps_0.5291_loss_0.8338_acc.pth'
  model_path = full_path_to_models_dir+'/Small/'+model_checkpoint_file_name
  if os.path.isfile(model_path):
    print('Model already exists...')

  else:
    hf_hub_download(repo_id='asigalov61/Ultimate-Accompaniment-Transformer',
                    filename=model_checkpoint_file_name,
                    local_dir='/content/Ultimate-Accompaniment-Transformer/Models/Small',
                    local_dir_use_symlinks=False)
else:

  dim = 2048
  depth = 4
  heads = 16

  model_checkpoint_file_name = 'Ultimate_Accompaniment_Transformer_Small_Improved_Trained_Model_13649_steps_0.3229_loss_0.898_acc.pth'
  model_path = full_path_to_models_dir+'/Small_Improved/'+model_checkpoint_file_name
  if os.path.isfile(model_path):
    print('Model already exists...')

  else:
    hf_hub_download(repo_id='asigalov61/Ultimate-Accompaniment-Transformer',
                    filename=model_checkpoint_file_name,
                    local_dir='/content/Ultimate-Accompaniment-Transformer/Models/Small_Improved',
                    local_dir_use_symlinks=False)

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

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 = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)

SEQ_LEN = 8192 # Models seq len
PAD_IDX = 767 # Models pad index

# instantiate the model

model = TransformerWrapper(
    num_tokens = PAD_IDX+1,
    max_seq_len = SEQ_LEN,
    attn_layers = Decoder(dim = dim, depth = depth, heads = heads, attn_flash = True)
    )

model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)

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

print('Loading model checkpoint...')

model.load_state_dict(torch.load(model_path))
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/Ultimate-Accompaniment-Transformer-Tokens-Embeddings-Plot.png", bbox_inches="tight")

"""# (GENERATE)

# (IMPROV)
"""

#@title Standard Improv Generator

#@markdown Generation settings

melody_MIDI_patch_number = 40 # @param {type:"slider", min:0, max:127, step:1}
accompaniment_MIDI_patch_number = 0 # @param {type:"slider", min:0, max:127, step:1}

number_of_tokens_to_generate = 512 # @param {type:"slider", min:30, max:8190, step:3}
number_of_batches_to_generate = 4 #@param {type:"slider", min:1, max:16, step:1}
temperature = 0.9 # @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('Ultimate Accompaniment Transformer Standard Improv Model Generator')
print('=' * 70)

outy = [0]

print('Selected Improv sequence:')
print(outy)
print('=' * 70)

torch.cuda.empty_cache()

inp = [outy] * 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=True,
                        verbose=True)

out0 = out.tolist()

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

torch.cuda.empty_cache()

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

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(out1) != 0:

      song = out1
      song_f = []

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

      patches = [0] * 16
      patches[3] = melody_MIDI_patch_number
      patches[0] = accompaniment_MIDI_patch_number

      for s in song:

        if 0 <= s < 128:

            time += s * 32

        if 128 <= s < 256:

            dur = (s-128) * 32

        if 256 <= s < 512:

            chan = (s-256) // 128

            if chan == 0:
              channel = 0
              pat = accompaniment_MIDI_patch_number
            else:
              channel = 3
              pat = melody_MIDI_patch_number

            pitch = (s-256) % 128

            vel = max(40, pitch)

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

      data = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,
                                                      output_signature = 'Ultimate Accompaniment Transformer',
                                                      output_file_name = '/content/Ultimate-Accompaniment-Transformer-Composition_'+str(i),
                                                      track_name='Project Los Angeles',
                                                      list_of_MIDI_patches=patches
                                                      )


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

      fname = '/content/Ultimate-Accompaniment-Transformer-Composition_'+str(i)

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

      TMIDIX.plot_ms_SONG(song_f, plot_title=fname)

"""# (ACCOMPANIMENT GENERATION)"""

#@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", "Ultimate-Accompaniment-Transformer-Melody-Seed-1", "Ultimate-Accompaniment-Transformer-Melody-Seed-2", "Ultimate-Accompaniment-Transformer-Melody-Seed-3", "Ultimate-Accompaniment-Transformer-Melody-Seed-4", "Ultimate-Accompaniment-Transformer-Melody-Seed-5", "Ultimate-Accompaniment-Transformer-Melody-Seed-6", "Ultimate-Accompaniment-Transformer-Melody-Seed-7"]
render_MIDI_to_audio = False # @param {type:"boolean"}

print('=' * 70)
print('Ultimate Accompaniment Transformer Seed MIDI Loader')
print('=' * 70)

f = ''

if select_seed_MIDI != "Upload your own custom MIDI":
  print('Loading seed MIDI...')
  f = '/content/Ultimate-Accompaniment-Transformer/Seeds/'+select_seed_MIDI+'.mid'

else:
  print('Upload your own custom MIDI...')
  print('=' * 70)
  uploaded_MIDI = files.upload()
  if list(uploaded_MIDI.keys()):
    f = list(uploaded_MIDI.keys())[0]

if f != '':

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

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

  #===============================================================================
  # Raw single-track ms score

  raw_score = TMIDIX.midi2single_track_ms_score(f)

  #===============================================================================
  # Enhanced score notes

  escore_notes = TMIDIX.advanced_score_processor(raw_score, return_enhanced_score_notes=True)[0]

  escore_notes = [e for e in escore_notes if e[3] != 9]

  if len(escore_notes) > 0:

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

      #===============================================================================
      # Augmented enhanced score notes

      escore_notes = TMIDIX.augment_enhanced_score_notes(escore_notes, timings_divider=32)

      cscore = TMIDIX.chordify_score([1000, escore_notes])

      melody = TMIDIX.fix_monophonic_score_durations([sorted(e, key=lambda x: x[4], reverse=True)[0] for e in cscore])

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

      melody_chords = []

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

      pe = cscore[0][0]

      mpe = melody[0]

      midx = 1

      for i, c in enumerate(cscore):

          c.sort(key=lambda x: (x[3], x[4]), reverse=True)

          # Next melody note

          if midx < len(melody):

            # Time
            mtime = melody[midx][1]-mpe[1]

            mdur = melody[midx][2]

            mdelta_time = max(0, min(127, mtime))

            # Durations
            mdur = max(0, min(127, mdur))

            # Pitch
            mptc = melody[midx][4]

          else:
            mtime = 127-mpe[1]

            mdur = mpe[2]

            mdelta_time = max(0, min(127, mtime))

            # Durations
            mdur = max(0, min(127, mdur))

            # Pitch
            mptc = mpe[4]


          e = melody[i]

          #=======================================================
          # Timings...

          time = e[1]-pe[1]

          dur = e[2]

          delta_time = max(0, min(127, time))

          # Durations

          dur = max(0, min(127, dur))

          # Pitches

          ptc = max(1, min(127, e[4]))

          if ptc < 60:
            ptc = 60 + (ptc % 12)

          cha = e[3]

          #=======================================================
          # FINAL NOTE SEQ

          if midx < len(melody):
            melody_chords.append([delta_time, dur+128, ptc+384, mdelta_time+512, mptc+640])
            mpe = melody[midx]
            midx += 1
          else:
            melody_chords.append([delta_time, dur+128, ptc+384, mdelta_time+512, mptc+640])

          pe = e

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

  song = melody_chords
  song_f = []



  time = 0
  dur = 128
  vel = 90
  pitch = 0
  pat = 40
  channel = 3

  patches = [0] * 16
  patches[3] = 40
  patches[0] = 0

  for ss in song:



    time += ss[0] * 32

    dur = (ss[1]-128) * 32

    pitch = (ss[2]-256) % 128

    vel = max(40, pitch)

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



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

  print('=' * 70)
  print('Composition stats:')
  print('Composition has', len(cscore), 'chords')
  print('Composition has', len(melody_chords), 'tokens')
  print('Composition MIDI patches:', sorted(set(patches)))
  print('=' * 70)

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

  fname = '/content/Ultimate-Accompaniment-Transformer-Seed-Composition'

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

  TMIDIX.plot_ms_SONG(song_f, plot_title=fname)

else:
  print('=' * 70)

#@title Accompaniment generation

#@markdown NOTE: You can stop the generation at any time to render partial results

#@markdown Output settings

melody_MIDI_patch_number = 40 # @param {type:"slider", min:0, max:127, step:1}
accompaniment_MIDI_patch_number = 0 # @param {type:"slider", min:0, max:127, step:1}

#@markdown Generation settings

number_of_prime_melody_notes = 1 # @param {type:"slider", min:0, max:64, step:1}
use_harmonized_melody = False # @param {type:"boolean"}
harmonized_melody_octave = 4 # @param {type:"slider", min:3, max:5, step:1}
force_accompaniment_generation_for_each_melody_note = False # @param {type:"boolean"}
tokens_sampling_rate = 1 # @param {type:"slider", min:1, max:16, step:1}
number_of_memory_tokens = 4096 # @param {type:"slider", min:32, max:8188, step:16}
temperature = 0.9 # @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('Ultimate Accompaniment Transformer Model Generator')
print('=' * 70)

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

def generate_acc(input_seq,
                  next_note_time,
                  force_acc = False,
                  num_samples = 2,
                  num_batches = 8,
                  num_memory_tokens = 4096,
                  temperature=0.9
                  ):

    input_seq = input_seq[-num_memory_tokens:]

    if force_acc:
      x = torch.LongTensor([input_seq+[0]] * num_batches).cuda()
    else:
      x = torch.LongTensor([input_seq] * num_batches).cuda()

    cur_time = 0
    ctime = 0

    o = 0

    while cur_time < next_note_time and o < 384:

      samples = []

      for _ in range(num_samples):

        with ctx:
          out = model.generate(x,
                              1,
                              temperature=temperature,
                              return_prime=True,
                              verbose=False)

          with torch.no_grad():
            test_loss, test_acc = model(out)

        samples.append([out[:,-1].tolist(), test_acc.tolist()])

      accs = [y[1] for y in samples]
      max_acc = max(accs)
      o = statistics.mode(samples[accs.index(max_acc)][0])

      if 0 <= o < 128:
        cur_time += o

      if cur_time < next_note_time and o < 384:

        ctime = cur_time

        out = torch.LongTensor([[o]] * num_batches).cuda()
        x = torch.cat((x, out), dim=1)

    return list(statistics.mode([tuple(t) for t in x[:, len(input_seq):].tolist()])), ctime

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

torch.cuda.empty_cache()

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

melody_tones_chords = TMIDIX.harmonize_enhanced_melody_score_notes(melody)

harm_melody_chords =  []

for i, m in enumerate(melody_tones_chords):
  cho = []
  mm = sorted(m, reverse=True)
  for mmm in mm:
    cho.extend([0, max(0, min(127,melody[i][2]))+128, (harmonized_melody_octave*12)+mmm+256])

  harm_melody_chords.append(cho)

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

output1 = []
output2 = []

for i in range(number_of_prime_melody_notes):

  output1.extend(melody_chords[i])

  if use_harmonized_melody:
    output1.extend(harm_melody_chords[i])
    output2.append(harm_melody_chords[i])
  else:
    output2.append([])

hidx = number_of_prime_melody_notes
cur_time = 0

for m in tqdm.tqdm(melody_chords[number_of_prime_melody_notes:]):

  try:
    mel = copy.deepcopy(m)
    mel[0] = mel[0] - cur_time
    next_note_time = mel[3]-512

    if use_harmonized_melody:
      input = mel + harm_melody_chords[hidx]
    else:
      input = mel

    output1.extend(input)

    out, cur_time = generate_acc(output1,
                                next_note_time,
                                force_acc=force_accompaniment_generation_for_each_melody_note,
                                num_samples=tokens_sampling_rate,
                                num_batches=tokens_sampling_rate,
                                num_memory_tokens=number_of_memory_tokens,
                                temperature=temperature
                                )
    output1.extend(out)

    if use_harmonized_melody:
      output2.append(harm_melody_chords[hidx] + out)
    else:
      output2.append(out)

    hidx += 1

  except KeyboardInterrupt:
    print('Stopping...')
    print('=' * 70)
    print('Stopping generation...')
    break

  except Exception as e:
    print('=' * 70)
    print('Error:', e)
    break

torch.cuda.empty_cache()

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

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

#===============================================================================
print('Rendering results...')

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

out1 = output2

if len(out1) != 0:

    song = out1
    song_f = []

    time = 0
    ntime = 0
    ndur = 0
    vel = 90
    npitch = 0
    channel = 0

    patches = [0] * 16
    patches[0] = accompaniment_MIDI_patch_number
    patches[3] = melody_MIDI_patch_number

    for i, ss in enumerate(song):

            ntime += melody_chords[i][0] * 32
            ndur = (melody_chords[i][1]-128) * 32
            nchannel = 1
            npitch = (melody_chords[i][2]-256) % 128
            vel = max(40, npitch)+20

            song_f.append(['note', ntime, ndur, 3, npitch, vel, melody_MIDI_patch_number ])

            time = ntime

            for s in ss:

              if 0 <= s < 128:

                  time += s * 32

              if 128 <= s < 256:

                  dur = (s-128) * 32

              if 256 <= s < 384:

                  pitch = (s-256)

                  vel = max(40, pitch)

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

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

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

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

fname = '/content/Ultimate-Accompaniment-Transformer-Composition'

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

TMIDIX.plot_ms_SONG(song_f, plot_title=fname)

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