Upload inference.py

inference.py

@@ -0,0 +1,398 @@
+import sys
+sys.path.append('code')
+
+
+import numpy as np
+import os
+import pretty_midi as pyd
+import torch
+import sys
+from model import VAE
+from util_tools.format_converter import melody_data2matrix, melody_matrix2data, chord_data2matrix, chord_matrix2data
+from torch.distributions import kl_divergence, Normal
+from nottingham_dataset import Nottingham
+import copy
+
+def chord_grid2data(est_pitch, bpm=60., start=0., max_simu_note=6, pitch_eos=129, num_step=32, min_pitch=0):
+        est_pitch = est_pitch[:, :, 0]#.cpu().detach().numpy() #(32, max_simu_note-1), NO BATCH HERE
+        if est_pitch.shape[1] == max_simu_note:
+            est_pitch = est_pitch[:, 1:]
+        
+        #print(est_pitch.shape)
+        #print(est_pitch)
+        harmonic_rhythm = 1. - (est_pitch[:, 0]==pitch_eos) * 1.
+        #print(harmonic_rhythm)
+
+        pr = np.zeros((32, 128), dtype=int)
+        alpha = 0.25 * 60 / bpm
+        notes = []
+        for t in range(num_step):
+            for n in range(max_simu_note-1):
+                note = est_pitch[t, n]
+                if note == pitch_eos:
+                    break
+                pitch = note + 12*4
+                duration = 1
+                for j in range(t+1, num_step):
+                    if harmonic_rhythm[j] == 1:
+                        break
+                    duration +=1
+                pr[t, pitch] = min(duration, 32 - t)
+                notes.append(
+                    pyd.Note(100, int(pitch), start + t * alpha,
+                                    start + (t + duration) * alpha))
+        chord = pyd.Instrument(program=pyd.instrument_name_to_program('Acoustic Grand Piano'))
+        chord.notes = notes
+        return chord
+
+def melody_matrix2data(melody_matrix, tempo=120, start_time=0.0, get_list=False):
+    HOLD_PITCH = 12
+    REST_PITCH = 13
+    #melodyMatrix = melody_matrix[:, :ROLL_SIZE]
+    chroma = np.concatenate((melody_matrix[:, :12], melody_matrix[:, 15: 17]), axis=-1)
+    register = melody_matrix[:, -10:]
+    #print(chroma.shape)
+    melodySequence = np.argmax(chroma, axis=-1)
+    #print(melodySequence)
+    
+    melody_notes = []
+    minStep = 60 / tempo / 4
+    onset_or_rest = [i for i in range(len(melodySequence)) if not melodySequence[i]==HOLD_PITCH]
+    onset_or_rest.append(len(melodySequence))
+
+    for idx, onset in enumerate(onset_or_rest[:-1]):
+        if melodySequence[onset] == REST_PITCH:
+            continue
+        else:
+            pitch = melodySequence[onset] + 12 * np.argmax(register[onset])
+            #print(pitch)
+            start = onset * minStep
+            end = onset_or_rest[idx+1] * minStep
+            noteRecon = pyd.Note(velocity=100, pitch=pitch, start=start_time+start, end=start_time+end)
+            melody_notes.append(noteRecon)
+    if get_list:
+        return melody_notes
+    else:  
+        melody = pyd.Instrument(program=pyd.instrument_name_to_program('Acoustic Grand Piano'))
+        melody.notes = melody_notes
+        return melody
+
+
+def get_gt(chord, melody):
+    #chord: (num_step, max_simu_note, 1), numpy
+    #melody: (num_step, 28), numpy
+    chord_recon = chord_grid2data(chord, 30, pitch_eos=13)
+    melody_recon = melody_matrix2data(melody, 120)
+    music = pyd.PrettyMIDI(initial_tempo=120)
+    music.instruments.append(melody_recon)
+    music.instruments.append(chord_recon)
+    return music
+
+def shift(original_melody, p_shift):
+    melody = copy.deepcopy(original_melody).cpu().detach().numpy()[0]
+    onsets, pitch = np.nonzero(melody[:, :12])
+    onsets, register = np.nonzero(melody[:, -10:])
+    onset130 = pitch + register*12
+    onset130 += p_shift
+    onset12 = onset130 % 12
+    register = onset130 // 12
+    melody[onsets, :] = 0
+    melody[onsets, onset12] = 1.
+    melody[onsets, register+17] = 1.
+    return torch.from_numpy(melody).float().unsqueeze(0)
+
+
+def reconstruct(chord, melody):
+    #chord: (1, num_step, max_simu_note, 1), torch.LongTensor, cuda()
+    #melody: (1, num_step*4, 28), torch.FloatTensor, cuda()
+    lengths = model.get_len_index_tensor(chord)  # lengths: (B, num_step)
+    chord = model.index_tensor_to_multihot_tensor(chord)
+    chord = model.enc_note_embedding(chord)    #(B, num_step, max_simu_note, note_emb_size)
+    mel_ebd = model.enc_note_embedding(melody)    #(B, num_step*4, note_emb_size)
+    melody_beat_summary = mel_ebd[:, ::4, :] + mel_ebd[:, 1::4, :] + mel_ebd[:, 2::4, :] + mel_ebd[:, 3::4, :]
+    dist, mu, = model.encoder(chord, lengths, melody_beat_summary)
+    z = dist.mean
+    pitch_outs = model.decoder(z, melody_beat_summary, 
+                                    inference=True, x=None, lengths=None, 
+                                    teacher_forcing_ratio1=0., teacher_forcing_ratio2=0.)
+    pitch_outs = pitch_outs.max(-1, keepdim=True)[1]
+    pitch_outs = pitch_outs.cpu().detach().numpy()
+    chord_track = chord_grid2data(pitch_outs[0], bpm=120//4, start=0, pitch_eos=13)
+    
+    melody = melody.cpu().detach().numpy()[0]
+    melody_track = melody_matrix2data(melody, tempo=120)
+
+    music = pyd.PrettyMIDI()
+    music.instruments.append(melody_track)
+    music.instruments.append(chord_track)
+    return music
+
+def melody_control(chord, melody, new_melody):
+    #chord: (B, num_step, max_simu_note, 1), torch.LongTensor, cuda()
+    #melody: (B, num_step*4, 28), torch.FloatTensor, cuda()
+    #new_melody: (B, num_step*4, 28), torch.FloatTensor, cuda()
+
+    lengths = model.get_len_index_tensor(chord)  # lengths: (B, num_step)
+    chord = model.index_tensor_to_multihot_tensor(chord)
+    chord = model.enc_note_embedding(chord)    #(B, num_step, max_simu_note, note_emb_size)
+    mel_ebd = model.enc_note_embedding(melody)    #(B, num_step*4, note_emb_size)
+    melody_beat_summary = mel_ebd[:, ::4, :] + mel_ebd[:, 1::4, :] + mel_ebd[:, 2::4, :] + mel_ebd[:, 3::4, :]
+    new_mel_ebd = model.enc_note_embedding(new_melody)    #(B, num_step*4, note_emb_size)
+    new_melody_beat_summary = new_mel_ebd[:, ::4, :] + new_mel_ebd[:, 1::4, :] + new_mel_ebd[:, 2::4, :] + new_mel_ebd[:, 3::4, :]
+    dist, mu, = model.encoder(chord, lengths, melody_beat_summary)
+    z = dist.mean
+    pitch_outs = model.decoder(z, new_melody_beat_summary, 
+                                    inference=True, x=None, lengths=None, 
+                                    teacher_forcing_ratio1=0., teacher_forcing_ratio2=0.)
+    pitch_outs = pitch_outs.max(-1, keepdim=True)[1]
+    pitch_outs = pitch_outs.cpu().detach().numpy()
+    chord_track = chord_grid2data(pitch_outs[0], bpm=120//4, start=0, pitch_eos=13)
+    
+    new_melody = new_melody.cpu().detach().numpy()[0]
+    melody_track = melody_matrix2data(new_melody, tempo=120)
+
+    music = pyd.PrettyMIDI()
+    music.instruments.append(melody_track)
+    music.instruments.append(chord_track)
+    return music
+
+def melody_prior_control(new_melody):
+    #new_melody: (B, num_step*4, 28), torch.FloatTensor, cuda()
+    new_mel_ebd = model.enc_note_embedding(new_melody)    #(B, num_step*4, note_emb_size)
+    new_melody_beat_summary = new_mel_ebd[:, ::4, :] + new_mel_ebd[:, 1::4, :] + new_mel_ebd[:, 2::4, :] + new_mel_ebd[:, 3::4, :]
+    z = Normal(torch.zeros(128), torch.ones(128)).rsample().unsqueeze(0)
+    pitch_outs = model.decoder(z, new_melody_beat_summary, 
+                                    inference=True, x=None, lengths=None, 
+                                    teacher_forcing_ratio1=0., teacher_forcing_ratio2=0.)
+    pitch_outs = pitch_outs.max(-1, keepdim=True)[1]
+    pitch_outs = pitch_outs.cpu().detach().numpy()
+    chord_track = chord_grid2data(pitch_outs[0], bpm=120//4, start=0, pitch_eos=13)
+    
+    new_melody = new_melody.cpu().detach().numpy()[0]
+    melody_track = melody_matrix2data(new_melody, tempo=120)
+
+    music = pyd.PrettyMIDI()
+    music.instruments.append(melody_track)
+    music.instruments.append(chord_track)
+    return music
+
+
+import utils
+config_fn = './code/model_config.json'
+#train_hyperparams = utils.load_params_dict('train_hyperparams', config_fn)
+model_params = utils.load_params_dict('model_params', config_fn)
+data_repr_params = utils.load_params_dict('data_repr', config_fn)
+#project_params = utils.load_params_dict('project', config_fn)
+#dataset_path = utils.load_params_dict('dataset_paths', config_fn)
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+model = VAE(max_simu_note=data_repr_params['max_simu_note'],
+        max_pitch=data_repr_params['max_pitch'],
+        min_pitch=data_repr_params['min_pitch'],
+        pitch_sos=data_repr_params['pitch_sos'],
+        pitch_eos=data_repr_params['pitch_eos'],
+        pitch_pad=data_repr_params['pitch_pad'],
+        num_step=data_repr_params['num_time_step'],
+
+        note_emb_size=model_params['note_emb_size'],
+        enc_notes_hid_size=model_params['enc_notes_hid_size'],
+        enc_time_hid_size=model_params['enc_time_hid_size'],
+        z_size=model_params['z_size'],
+        dec_emb_hid_size=model_params['dec_emb_hid_size'],
+        dec_time_hid_size=model_params['dec_time_hid_size'],
+        dec_notes_hid_size=model_params['dec_notes_hid_size'],
+        discr_nhead = model_params["discr_nhead"],
+        discr_hid_size = model_params["discr_hid_size"], 
+        discr_dropout = model_params["discr_dropout"], 
+        discr_nlayer = model_params["discr_nlayer"],
+
+        device=device
+        )
+
+
+weight_path = './code/ad-ptvae_param.pt'
+params = torch.load(weight_path,map_location=torch.device(device))
+if 'model_state_dict' in params:
+    params = params['model_state_dict']
+model.load_state_dict(params)
+if torch.cuda.is_available():
+  model.cuda()
+else:
+  model.cpu()
+  
+model.eval()
+print('-'*100)
+print(f'Loaded {weight_path}')
+print('-'*100)
+
+
+dataset = np.load('./code/data.npy', allow_pickle=True).T
+print('-'*100)
+print(f'Loaded ./code/data.npy')
+print('-'*100)
+np.random.seed(0)
+np.random.shuffle(dataset)
+anchor = int(dataset.shape[0] * 0.95)
+val_data = dataset[anchor:, :]
+val_set = Nottingham(dataset=val_data.T, 
+                        length=128, 
+                        step_size=16, 
+                        chord_fomat='pr', shift_high=0, shift_low=0)
+print(len(val_set))
+WRITE_PATH = './code/demo_generate'
+if not os.path.exists(WRITE_PATH):
+    os.makedirs(WRITE_PATH)
+
+
+chord_1, _, melody_1, _ = val_set.__getitem__(338)
+music = get_gt(chord_1, melody_1)
+music.write(os.path.join(WRITE_PATH, 'gt_1.mid'))
+chord_1 = torch.from_numpy(chord_1).long().unsqueeze(0)
+melody_1 = torch.from_numpy(melody_1).float().unsqueeze(0)
+music = reconstruct(chord_1, melody_1)
+music.write(os.path.join(WRITE_PATH, 'recon_1.mid'))
+print(f'Saved to {WRITE_PATH}/recon_1.mid')
+
+chord_2, _, melody_2, _ = val_set.__getitem__(2749)
+music = get_gt(chord_2, melody_2)
+music.write(os.path.join(WRITE_PATH, 'gt_2.mid'))
+chord_2 = torch.from_numpy(chord_2).long().unsqueeze(0)
+melody_2 = torch.from_numpy(melody_2).float().unsqueeze(0)
+music = reconstruct(chord_2, melody_2)
+music.write(os.path.join(WRITE_PATH, 'recon_2.mid'))
+print(f'Saved to {WRITE_PATH}/recon_2.mid')
+
+chord_3, _, melody_3, _ = val_set.__getitem__(3413)
+music = get_gt(chord_3, melody_3)
+music.write(os.path.join(WRITE_PATH, 'gt_3.mid'))
+chord_3 = torch.from_numpy(chord_3).long().unsqueeze(0)
+melody_3 = torch.from_numpy(melody_3).float().unsqueeze(0)
+music = reconstruct(chord_3, melody_3)
+music.write(os.path.join(WRITE_PATH, 'recon_3.mid'))
+print(f'Saved to {WRITE_PATH}/recon_3.mid')
+
+chord_4, _, melody_4, _ = val_set.__getitem__(5126)
+music = get_gt(chord_4, melody_4)
+music.write(os.path.join(WRITE_PATH, 'gt_4.mid'))
+chord_4 = torch.from_numpy(chord_4).long().unsqueeze(0)
+melody_4 = torch.from_numpy(melody_4).float().unsqueeze(0)
+music = reconstruct(chord_4, melody_4)
+music.write(os.path.join(WRITE_PATH, 'recon_4.mid'))
+
+
+midi = pyd.PrettyMIDI(f'{WRITE_PATH}/recon_1.mid')
+melody = melody_data2matrix(midi.instruments[0], midi.get_downbeats())
+melody = val_set.truncate_melody(melody)
+melody_1_modal_change = torch.from_numpy(melody).float().unsqueeze(0)
+
+midi = pyd.PrettyMIDI(f'{WRITE_PATH}/recon_2.mid')
+melody = melody_data2matrix(midi.instruments[0], midi.get_downbeats())
+melody = val_set.truncate_melody(melody)
+melody_2_modal_change = torch.from_numpy(melody).float().unsqueeze(0)
+
+midi = pyd.PrettyMIDI(f'{WRITE_PATH}/recon_3.mid')
+melody = melody_data2matrix(midi.instruments[0], midi.get_downbeats())
+melody = val_set.truncate_melody(melody)
+melody_3_modal_change = torch.from_numpy(melody).float().unsqueeze(0)
+
+midi = pyd.PrettyMIDI(f'{WRITE_PATH}/recon_4.mid')
+melody = melody_data2matrix(midi.instruments[0], midi.get_downbeats())
+melody = val_set.truncate_melody(melody)
+melody_4_modal_change = torch.from_numpy(melody).float().unsqueeze(0)
+
+
+
+
+music = melody_control(chord_1, melody_1, shift(melody_1, 6))
+music.write(os.path.join(WRITE_PATH, 'control_1_transpose.mid'))
+print(f'Saved to {WRITE_PATH}/control_1_transpose.mid')
+
+music = melody_control(chord_2, melody_2, shift(melody_2, 6))
+music.write(os.path.join(WRITE_PATH, 'control_2_transpose.mid'))
+print(f'Saved to {WRITE_PATH}/control_2_transpose.mid')
+
+music = melody_control(chord_3, melody_3, shift(melody_3, 6))
+music.write(os.path.join(WRITE_PATH, 'control_3_transpose.mid'))
+print(f'Saved to {WRITE_PATH}/control_3_transpose.mid')
+
+music = melody_control(chord_4, melody_4, shift(melody_4, 6))
+music.write(os.path.join(WRITE_PATH, 'control_4_transpose1.mid'))
+print(f'Saved to {WRITE_PATH}/control_4_transpose.mid')
+
+
+
+music = melody_control(chord_1, melody_1, melody_1_modal_change)
+music.write(os.path.join(WRITE_PATH, 'control_1_modal_change.mid'))
+print(f'Saved to {WRITE_PATH}/control_1_modal_change.mid')
+
+music = melody_control(chord_2, melody_2, melody_2_modal_change)
+music.write(os.path.join(WRITE_PATH, 'control_2_modal_change.mid'))
+print(f'Saved to {WRITE_PATH}/control_2_modal_change.mid')
+
+music = melody_control(chord_3, melody_3, melody_3_modal_change)
+music.write(os.path.join(WRITE_PATH, 'control_3_modal_change.mid'))
+print(f'Saved to {WRITE_PATH}/control_3_modal_change.mid')
+
+music = melody_control(chord_4, melody_4, melody_4_modal_change)
+music.write(os.path.join(WRITE_PATH, 'control_4_modal_change.mid'))
+print(f'Saved to {WRITE_PATH}/control_4_modal_change.mid')
+
+
+
+
+music = melody_prior_control(melody_1)
+music.write(os.path.join(WRITE_PATH, 'control_1_prior.mid'))
+print(f'Saved to {WRITE_PATH}/control_1_prior.mid')
+
+music = melody_prior_control(melody_2)
+music.write(os.path.join(WRITE_PATH, 'control_2_prior.mid'))
+print(f'Saved to {WRITE_PATH}/control_2_prior.mid')
+
+music = melody_prior_control(melody_3)
+music.write(os.path.join(WRITE_PATH, 'control_3_prior.mid'))
+print(f'Saved to {WRITE_PATH}/control_3_prior.mid')
+
+music = melody_prior_control(melody_4)
+music.write(os.path.join(WRITE_PATH, 'control_4_prior.mid'))
+print(f'Saved to {WRITE_PATH}/control_4_prior.mid')
+
+
+
+
+music = melody_control(chord_1, melody_1, melody_2)
+music.write(os.path.join(WRITE_PATH, 'control_1c+2m.mid'))
+
+music = melody_control(chord_2, melody_2, melody_1)
+music.write(os.path.join(WRITE_PATH, 'control_2c+1m.mid'))
+
+music = melody_control(chord_1, melody_1, melody_3)
+music.write(os.path.join(WRITE_PATH, 'control_1c+3m.mid'))
+
+music = melody_control(chord_3, melody_3, melody_1)
+music.write(os.path.join(WRITE_PATH, 'control_3c+1m.mid'))
+
+music = melody_control(chord_1, melody_1, melody_4)
+music.write(os.path.join(WRITE_PATH, 'control_1c+4m.mid'))
+
+music = melody_control(chord_4, melody_4, melody_1)
+music.write(os.path.join(WRITE_PATH, 'control_4c+1m.mid'))
+
+music = melody_control(chord_2, melody_2, melody_3)
+music.write(os.path.join(WRITE_PATH, 'control_2c+3m.mid'))
+
+music = melody_control(chord_3, melody_3, melody_2)
+music.write(os.path.join(WRITE_PATH, 'control_3c+2m.mid'))
+
+music = melody_control(chord_2, melody_2, melody_4)
+music.write(os.path.join(WRITE_PATH, 'control_2c+4m.mid'))
+
+music = melody_control(chord_4, melody_4, melody_2)
+music.write(os.path.join(WRITE_PATH, 'control_4c+2m.mid'))
+
+music = melody_control(chord_3, melody_3, melody_4)
+music.write(os.path.join(WRITE_PATH, 'control_3c+4m.mid'))
+
+music = melody_control(chord_4, melody_4, melody_3)
+music.write(os.path.join(WRITE_PATH, 'control_4c+3m.mid'))
+