euterpe_x_ttm_edition.py

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

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/asigalov61/Euterpe-X/blob/main/Euterpe_X_TTM_Edition.ipynb

# Euterpe X TTM Edition (ver. 5.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 2023

***

# (GPU CHECK)
"""

#@title NVIDIA GPU check
!nvidia-smi

"""# (SETUP ENVIRONMENT)"""

#@title Install dependencies
!git clone --depth 1 https://github.com/asigalov61/Euterpe-X
!pip install huggingface_hub
!pip install torch
!pip install einops
!pip install fuzzywuzzy[speedup]
!pip install torch-summary
!pip install tqdm
!pip install matplotlib
!apt install fluidsynth #Pip does not work for some reason. Only apt works
!pip install midi2audio

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

print('=' * 70)
print('Loading core Euterpe X modules...')

import os
import pickle

import secrets
import statistics
from time import time
import tqdm

print('=' * 70)
print('Loading main Euterpe X modules...')
import torch

# %cd /content/Euterpe-X

import TMIDIX

from x_transformer import *

# %cd /content/

from fuzzywuzzy import process

print('=' * 70)
print('Loading aux Euterpe X modules...')

import matplotlib.pyplot as plt

from torchsummary import summary
from sklearn import metrics

from midi2audio import FluidSynth
from IPython.display import Audio, display

import random

from huggingface_hub import hf_hub_download

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

"""# (LOAD MODEL)

## NOTE: Choose one and do not forget to restart the colab runtime if you are switching between models
"""

#@title Load Euterpe X Small Model

#@markdown Fast model, 32 layers, 141k MIDIs training corpus

full_path_to_model_checkpoint = "/content/Euterpe-X/Models/Small/Euterpe_X_Small_Trained_Model_58000_steps_0.6865_loss_0.7964_acc.pth" #@param {type:"string"}

#@markdown Model precision option

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

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

#@markdown float16 == Half precision/double speed

#@markdown float32 == Full precision/normal speed

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

print('=' * 70)
print('Loading Euterpe X Small 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/Euterpe-X',
                  filename='Euterpe_X_Small_Trained_Model_58000_steps_0.6865_loss_0.7964_acc.pth',
                  local_dir='/content/Euterpe-X/Models/Small/',
                  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'

if model_precision == 'float32':
  dtype = 'float32'

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

SEQ_LEN = 2048

# instantiate the model

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

model = AutoregressiveWrapper(model)

model = torch.nn.DataParallel(model)

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.module.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/Euterpe-X-Small-Tokens-Embeddings-Plot.png", bbox_inches="tight")

#@title Load Euterpe X Large Model

#@markdown Slow model, 60 layers, 314k MIDIs training corpus

full_path_to_model_checkpoint = "/content/Euterpe-X/Models/Large/Euterpe_X_Large_Trained_Model_100000_steps_0.477_loss_0.8533_acc.pth" #@param {type:"string"}

#@markdown Model precision option

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

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

#@markdown float16 == Half precision/double speed

#@markdown float32 == Full precision/normal speed

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

print('=' * 70)
print('Loading Euterpe X 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/Euterpe-X',
                  filename='Euterpe_X_Large_Trained_Model_100000_steps_0.477_loss_0.8533_acc.pth',
                  local_dir='/content/Euterpe-X/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'

if model_precision == 'float32':
  dtype = 'float32'

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

SEQ_LEN = 2048

# instantiate the model

model = TransformerWrapper(
    num_tokens = 3344,
    max_seq_len = SEQ_LEN,
    attn_layers = Decoder(dim = 1024, depth = 60, heads = 8, attn_flash=True)
)

model = AutoregressiveWrapper(model)

model = torch.nn.DataParallel(model)

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.module.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/Euterpe-X-Large-Tokens-Embeddings-Plot.png", bbox_inches="tight")

"""# (LOAD AUX DATA)"""

#@title Load Euterpe X Aux Data

print('=' * 70)
print('Loading Euterpe X Aux Data. Please wait...')
print('=' * 70)

if os.path.isfile('/content/Euterpe-X/Aux-Data/Euterpe_X_Text_To_Music_Aux_Data.pickle'):
  print('Aux Data already exists...')

else:
  hf_hub_download(repo_id='asigalov61/Euterpe-X',
                  filename='Euterpe_X_Text_To_Music_Aux_Data.pickle',
                  local_dir='/content/Euterpe-X/Aux-Data/',
                  local_dir_use_symlinks=False)
print('=' * 70)
AUX_DATA = TMIDIX.Tegridy_Any_Pickle_File_Reader('/content/Euterpe-X/Aux-Data/Euterpe_X_Text_To_Music_Aux_Data')

print('Done!')

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

# @title Load velocities map

#@markdown Select velocity type to apply to the generated output

velocity_type_to_use = "Velocity Map" # @param ["Constant Velocity", "Velocity Map"]

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

print('=' * 70)
print('Loading Euterpe X Velocities Map...')
print('=' * 70)

instruments_patches = [0, 24, 32, 40, 42, 46, 56, 65, 73, 0, 53, 19, 0, 0, 0, 0]

dynamic_velocities_map = TMIDIX.Tegridy_Any_Pickle_File_Reader('/content/Euterpe-X/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='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("Pitches")
  plt.ylabel("Patches")
  plt.tight_layout()
  plt.plot()
  plt.savefig("/content/Euterpe-X-Velocities-Map-Plot.png", bbox_inches="tight")

"""# (GENERATE)"""

#@title Text-To-Music Generation

#@markdown Text-To-Music Settings

#@markdown NOTE: You can enter any desired title or artist, or both

enter_desired_song_title = "Family Guy" #@param {type:"string"}
enter_desired_artist = "TV Themes" #@param {type:"string"}

#@markdown Generation Settings

number_of_tokens_to_generate = 300 # @param {type:"slider", min:30, max:2046, 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.1}

#@markdown Other settings

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

print('=' * 70)
print('Euterpe X TTM Model Generator')
print('=' * 70)

print('Searching titles...Please wait...')
random.shuffle(AUX_DATA)

titles_index = []

for A in AUX_DATA:
  titles_index.append(A[0])

search_string = ''

if enter_desired_song_title != '' and enter_desired_artist != '':
  search_string = enter_desired_song_title + ' --- ' + enter_desired_artist

else:
  search_string = enter_desired_song_title + enter_desired_artist

search_match = process.extract(query=search_string, choices=titles_index, limit=1)
search_index = titles_index.index(search_match[0][0])

print('Done!')
print('=' * 70)
print('Selected title:', AUX_DATA[search_index][0])
print('=' * 70)

if allow_model_to_stop_generation_if_needed:
  min_stop_token = 3343
else:
  min_stop_token = None

# Velocities
velocities_map = [80, 80, 70, 100, 90, 80, 100, 100, 100, 90, 110, 100]
vel_map = AUX_DATA[search_index][1]

for i in range(12):
  if vel_map[i] != 0:
    velocities_map[i] = vel_map[i]

# Loading data...
outy = AUX_DATA[search_index][2][3:]

block_marker = sum([(y * 8) for y in outy if y < 256]) / 1000

inp = [outy] * number_of_batches_to_generate

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

with ctx:
  out = model.module.generate(inp,
                        number_of_tokens_to_generate,
                        temperature=temperature,
                        return_prime=True,
                        eos_token=min_stop_token,
                        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 = out1
      song_f = []

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

      for ss in song:

        if ss > 0 and ss < 256:

            time += ss * 8

        if ss >= 256 and ss < 256+(12*128):

            dur = ((ss-256) % 128) * 30

        if ss >= 256+(12*128) and ss < 256+(12*128)+(12*128):
            channel = (ss-(256+(12*128))) // 128
            pitch = (ss-(256+(12*128))) % 128

            if velocity_type_to_use == 'Constant Velocity':
              vel = velocities_map[channel]
            else:
              vel = round(dynamic_velocities_map[pitch][instruments_patches[channel]])


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

      detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,
                                                          output_signature = 'Euterpe X',
                                                          output_file_name = '/content/Euterpe-X-Music-Composition_'+str(i),
                                                          track_name='Project Los Angeles',
                                                          list_of_MIDI_patches=[0, 24, 32, 40, 42, 46, 56, 65, 73, 0, 53, 19, 0, 0, 0, 0]
                                                          )
      print('=' * 70)
      print('Displaying resulting composition...')
      print('=' * 70)

      fname = '/content/Euterpe-X-Music-Composition_'+str(i)

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

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

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

      if render_MIDI_to_audio:
        FluidSynth("/usr/share/sounds/sf2/FluidR3_GM.sf2", 16000).midi_to_audio(str(fname + '.mid'), str(fname + '.wav'))
        display(Audio(str(fname + '.wav'), rate=16000))

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

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

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

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

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