mini_muse.py

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

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Tyq-n-kH0iKl9Zp1Kbtp6FQKJUw9DOnW

# Mini Muse (ver. 1.0)

***

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

***

Credit for GPT2-RGA code used in this colab goes out @ Sashmark97 https://github.com/Sashmark97/midigen and @ Damon Gwinn https://github.com/gwinndr/MusicTransformer-Pytorch

***

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 2022

***

# (Setup Environment)
"""

#@title nvidia-smi gpu check
!nvidia-smi

#@title Install all dependencies (run only once per session)

!git clone https://github.com/asigalov61/Mini-Muse

!pip install torch

!pip install tqdm
!pip install matplotlib
!pip install torch-summary

!apt install fluidsynth #Pip does not work for some reason. Only apt works
!pip install midi2audio
!pip install pretty_midi

#@title Import all needed modules

print('Loading needed modules. Please wait...')
import os
import random
import copy

from collections import OrderedDict

from tqdm import tqdm

import matplotlib.pyplot as plt

from torchsummary import summary

print('Loading TMIDIX module...')
os.chdir('/content/Mini-Muse')

import TMIDIX
from GPT2RGAX import *

from midi2audio import FluidSynth
import pretty_midi
import librosa.display
from IPython.display import Audio

os.chdir('/content/')

"""# (UNZIP)"""

# Commented out IPython magic to ensure Python compatibility.
#@title Unzip pre-trained Mini Muse Model
# %cd /content/Mini-Muse/Model

print('=' * 70)
print('Unzipping pre-trained model...Please wait...')
print('=' * 70)

!cat /content/Mini-Muse/Model/Mini-Muse-Trained-Model.zip* > Mini-Muse-Trained-Model.zip
print('=' * 70)

!unzip -j Mini-Muse-Trained-Model.zip
print('=' * 70)

print('Done! Enjoy! :)')
print('=' * 70)
# %cd /content/

# Commented out IPython magic to ensure Python compatibility.
#@title Unzip Mini Muse Training Data
# %cd /content/Mini-Muse/Training-Data

print('=' * 70)
print('Unzipping pre-trained model...Please wait...')
print('=' * 70)

!cat /content/Mini-Muse/Training-Data/Mini-Muse-Training-Data.zip* > Mini-Muse-Training-Data.zip
print('=' * 70)

!unzip -j Mini-Muse-Training-Data.zip
print('=' * 70)

print('Done! Enjoy! :)')
print('=' * 70)
# %cd /content/

"""# (LOAD)"""

#@title Load/Reload the model

full_path_to_model_checkpoint = "/content/Mini-Muse/Model/Mini_Muse_Trained_Model_88000_steps_0.6129_loss.pth" #@param {type:"string"}

print('Loading the model...')
config = GPTConfig(512, 
                   1024,
                   dim_feedforward=1024,
                   n_layer=16, 
                   n_head=16, 
                   n_embd=1024,
                   enable_rpr=True,
                   er_len=1024)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = GPT(config)

state_dict = torch.load(full_path_to_model_checkpoint, map_location=device)

new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] #remove 'module'
    new_state_dict[name] = v

model.load_state_dict(new_state_dict)

model.to(device)

model.eval()

print('Done!')

summary(model)

#@title Load/Reload training data
print('Loading training data...')
train_data = TMIDIX.Tegridy_Any_Pickle_File_Reader('/content/Mini-Muse/Training-Data/Mini-Muse-Training-Data')
print('Done!')

"""# (GENERATE)"""

#@title Generate Prime
number_of_prime_tokens = 256 #@param {type:"slider", min:16, max:512, step:16}

#@markdown Select desired combination of instruments (less is better)

Piano = True #@param {type:"boolean"}
Guitar = False #@param {type:"boolean"}
Bass = False #@param {type:"boolean"}
Violin = False #@param {type:"boolean"}
Cello = False #@param {type:"boolean"}
Harp = False #@param {type:"boolean"}
Trumpet = False #@param {type:"boolean"}
Clarinet = False #@param {type:"boolean"}
Flute = False #@param {type:"boolean"}
Drums = False #@param {type:"boolean"}
Choir = False #@param {type:"boolean"}

instruments = []

if Piano:
  instruments += [0]

if Guitar:
  instruments += [1]

if Bass:
  instruments += [2]

if Violin:
  instruments += [3]

if Cello:
  instruments += [4]

if Harp:
  instruments += [5]

if Trumpet:
  instruments += [6]

if Clarinet:
  instruments += [7]

if Flute:
  instruments += [8]

if Drums:
  instruments += [9]

if Choir:
  instruments += [10]

iidx = []

print('Looking for matching primes...')

for i in tqdm(range(0, len(train_data), 1025)):
  instr = sorted(list(set([y // 10 for y in train_data[i+1:i+1025:4][:number_of_prime_tokens]])))
  if instr == sorted(instruments):
    iidx.append(i)

print('Found', len(iidx), 'matching primes...')

iindex = random.choice(iidx)
print('Selected prime #', iindex // 1025)
print('Prime index:', iindex)

out1 = train_data[iindex:iindex+number_of_prime_tokens+1]

if len(out1) != 0:
    
    song = out1
    song_f = []
    time = 0
    dur = 0
    vel = 0
    pitch = 0
    channel = 0

    son = []

    song1 = []

    for s in song:
      if s > 127:
        son.append(s)

      else:
        if len(son) == 4:
          song1.append(son)
        son = []
        son.append(s)
    
    for s in song1:

        channel = s[0] // 10

        vel = (s[0] % 10) * 16

        time += (s[1]-128) * 16
            
        dur = (s[2] - 256) * 32
        
        pitch = (s[3] - 384)
                                  
        song_f.append(['note', time, dur, channel, pitch, vel ])

    detailed_stats = TMIDIX.Tegridy_SONG_to_MIDI_Converter(song_f,
                                                        output_signature = 'Mini Muse',  
                                                        output_file_name = '/content/Mini-Muse-Music-Composition', 
                                                        track_name='Project Los Angeles',
                                                        list_of_MIDI_patches=[0, 24, 32, 40, 42, 46, 56, 71, 73, 0, 53, 0, 0, 0, 0, 0],
                                                        number_of_ticks_per_quarter=500)

    print('Done!')

print('Displaying resulting composition...')
fname = '/content/Mini-Muse-Music-Composition'

pm = pretty_midi.PrettyMIDI(fname + '.mid')

# Retrieve piano roll of the MIDI file
piano_roll = pm.get_piano_roll()

plt.figure(figsize=(14, 5))
librosa.display.specshow(piano_roll, x_axis='time', y_axis='cqt_note', fmin=1, hop_length=160, sr=16000, cmap=plt.cm.hot)
plt.title(fname)

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

#@title Single Continuation Block Generator

#@markdown NOTE: Play with the settings to get different results

temperature = 0.8 #@param {type:"slider", min:0.1, max:1, step:0.1}
show_stats = True #@param {type:"boolean"}

#===================================================================
print('=' * 70)
print('Mini Muse Music Model Continuation Generator')
print('=' * 70)

print('Generation settings:')
print('=' * 70)
print('Number of prime tokens:', number_of_prime_tokens)
print('Model temperature:', temperature)

print('=' * 70)
print('Generating...')

inputs = train_data[iindex:iindex+number_of_prime_tokens+1]

rand_seq = model.generate(torch.Tensor(inputs), 
                                          target_seq_length=1024,
                                          temperature=temperature,
                                          stop_token=512,
                                          verbose=show_stats)
  
out1 = rand_seq[0].cpu().numpy().tolist()

if len(out1) != 0:
    
    song = out1
    song_f = []
    time = 0
    dur = 0
    vel = 0
    pitch = 0
    channel = 0

    son = []

    song1 = []

    for s in song:
      if s > 127:
        son.append(s)

      else:
        if len(son) == 4:
          song1.append(son)
        son = []
        son.append(s)
    
    for s in song1:

        channel = s[0] // 10

        vel = (s[0] % 10) * 16

        time += (s[1]-128) * 16
            
        dur = (s[2] - 256) * 32
        
        pitch = (s[3] - 384)
                                  
        song_f.append(['note', time, dur, channel, pitch, vel ])

    detailed_stats = TMIDIX.Tegridy_SONG_to_MIDI_Converter(song_f,
                                                        output_signature = 'Mini Muse',  
                                                        output_file_name = '/content/Mini-Muse-Music-Composition', 
                                                        track_name='Project Los Angeles',
                                                        list_of_MIDI_patches=[0, 24, 32, 40, 42, 46, 56, 71, 73, 0, 53, 0, 0, 0, 0, 0],
                                                        number_of_ticks_per_quarter=500)

    print('Done!')

print('Displaying resulting composition...')
fname = '/content/Mini-Muse-Music-Composition'

pm = pretty_midi.PrettyMIDI(fname + '.mid')

# Retrieve piano roll of the MIDI file
piano_roll = pm.get_piano_roll()

plt.figure(figsize=(14, 5))
librosa.display.specshow(piano_roll, x_axis='time', y_axis='cqt_note', fmin=1, hop_length=160, sr=16000, cmap=plt.cm.hot)
plt.title(fname)

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

#@title Auto-continue resulting composition

#@markdown NOTE: This may or may not work well due to long-term structure decay
number_of_continuation_blocks = 1 #@param {type:"slider", min:1, max:5, step:1}

out2 = copy.deepcopy(out1)

for i in range(number_of_continuation_blocks):

  rand_seq = model.generate(torch.Tensor(out2[-64:]), 
                                            target_seq_length=1024,
                                            temperature=temperature,
                                            stop_token=512,
                                            verbose=show_stats)
    
  out = rand_seq[0].cpu().numpy().tolist()

  out2.extend(out[64:])

if len(out2) != 0:
    
    song = out2
    song_f = []
    time = 0
    dur = 0
    vel = 0
    pitch = 0
    channel = 0

    son = []

    song1 = []

    for s in song:
      if s > 127:
        son.append(s)

      else:
        if len(son) == 4:
          song1.append(son)
        son = []
        son.append(s)
    
    for s in song1:

        channel = s[0] // 10

        vel = (s[0] % 10) * 16

        time += (s[1]-128) * 16
            
        dur = (s[2] - 256) * 32
        
        pitch = (s[3] - 384)
                                  
        song_f.append(['note', time, dur, channel, pitch, vel ])

    detailed_stats = TMIDIX.Tegridy_SONG_to_MIDI_Converter(song_f,
                                                        output_signature = 'Mini Muse',  
                                                        output_file_name = '/content/Mini-Muse-Music-Composition', 
                                                        track_name='Project Los Angeles',
                                                        list_of_MIDI_patches=[0, 24, 32, 40, 42, 46, 56, 71, 73, 0, 53, 0, 0, 0, 0, 0],
                                                        number_of_ticks_per_quarter=500)

    print('Done!')

print('Displaying resulting composition...')
fname = '/content/Mini-Muse-Music-Composition'

pm = pretty_midi.PrettyMIDI(fname + '.mid')

# Retrieve piano roll of the MIDI file
piano_roll = pm.get_piano_roll()

plt.figure(figsize=(14, 5))
librosa.display.specshow(piano_roll, x_axis='time', y_axis='cqt_note', fmin=1, hop_length=160, sr=16000, cmap=plt.cm.hot)
plt.title(fname)

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

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