import torch import torch.nn as nn from torch.nn.modules.normalization import LayerNorm import random import numpy as np from utilities.constants import * from utilities.device import get_device from .positional_encoding import PositionalEncoding from .rpr import TransformerDecoderRPR, TransformerDecoderLayerRPR from datetime import datetime import json class VideoMusicTransformer(nn.Module): def __init__(self, n_layers=6, num_heads=8, d_model=512, dim_feedforward=1024, dropout=0.1, max_sequence_midi =2048, max_sequence_video=300, max_sequence_chord=300, total_vf_dim = 0, rpr=False): super(VideoMusicTransformer, self).__init__() self.nlayers = n_layers self.nhead = num_heads self.d_model = d_model self.d_ff = dim_feedforward self.dropout = dropout self.max_seq_midi = max_sequence_midi self.max_seq_video = max_sequence_video self.max_seq_chord = max_sequence_chord self.rpr = rpr # Input embedding for video and music features self.embedding = nn.Embedding(CHORD_SIZE, self.d_model) self.embedding_root = nn.Embedding(CHORD_ROOT_SIZE, self.d_model) self.embedding_attr = nn.Embedding(CHORD_ATTR_SIZE, self.d_model) self.total_vf_dim = total_vf_dim self.Linear_vis = nn.Linear(self.total_vf_dim, self.d_model) self.Linear_chord = nn.Linear(self.d_model+1, self.d_model) # Positional encoding self.positional_encoding = PositionalEncoding(self.d_model, self.dropout, self.max_seq_chord) self.positional_encoding_video = PositionalEncoding(self.d_model, self.dropout, self.max_seq_video) # Add condition (minor or major) self.condition_linear = nn.Linear(1, self.d_model) # Base transformer if(not self.rpr): self.transformer = nn.Transformer( d_model=self.d_model, nhead=self.nhead, num_encoder_layers=self.nlayers, num_decoder_layers=self.nlayers, dropout=self.dropout, # activation=self.ff_activ, dim_feedforward=self.d_ff ) # RPR Transformer else: decoder_norm = LayerNorm(self.d_model) decoder_layer = TransformerDecoderLayerRPR(self.d_model, self.nhead, self.d_ff, self.dropout, er_len=self.max_seq_chord) decoder = TransformerDecoderRPR(decoder_layer, self.nlayers, decoder_norm) self.transformer = nn.Transformer( d_model=self.d_model, nhead=self.nhead, num_encoder_layers=self.nlayers, num_decoder_layers=self.nlayers, dropout=self.dropout, # activation=self.ff_activ, dim_feedforward=self.d_ff, custom_decoder=decoder ) self.Wout = nn.Linear(self.d_model, CHORD_SIZE) self.Wout_root = nn.Linear(self.d_model, CHORD_ROOT_SIZE) self.Wout_attr = nn.Linear(self.d_model, CHORD_ATTR_SIZE) self.softmax = nn.Softmax(dim=-1) self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") def forward(self, x, x_root, x_attr, feature_semantic_list, feature_key, feature_scene_offset, feature_motion, feature_emotion, mask=True): if(mask is True): mask = self.transformer.generate_square_subsequent_mask(x.shape[1]).to(self.device) else: mask = None x_root = self.embedding_root(x_root) x_attr = self.embedding_attr(x_attr) x = x_root + x_attr feature_key_padded = torch.full((x.shape[0], x.shape[1], 1), feature_key.item()) feature_key_padded = feature_key_padded.to(self.device) x = torch.cat([x, feature_key_padded], dim=-1) xf = self.Linear_chord(x) ### Video (SemanticList + SceneOffset + Motion + Emotion) (ENCODER) ### vf_concat = feature_semantic_list[0].float() for i in range(1, len(feature_semantic_list)): vf_concat = torch.cat( (vf_concat, feature_semantic_list[i].float()), dim=2) vf_concat = torch.cat([vf_concat, feature_scene_offset.unsqueeze(-1).float()], dim=-1) # -> (max_seq_video, batch_size, d_model+1) vf_concat = torch.cat([vf_concat, feature_motion.unsqueeze(-1).float()], dim=-1) # -> (max_seq_video, batch_size, d_model+1) vf_concat = torch.cat([vf_concat, feature_emotion.float()], dim=-1) # -> (max_seq_video, batch_size, d_model+1) vf = self.Linear_vis(vf_concat) ### POSITIONAL ENCODING ### xf = xf.permute(1,0,2) # -> (max_seq-1, batch_size, d_model) vf = vf.permute(1,0,2) # -> (max_seq_video, batch_size, d_model) xf = self.positional_encoding(xf) vf = self.positional_encoding_video(vf) ### TRANSFORMER ### x_out = self.transformer(src=vf, tgt=xf, tgt_mask=mask) x_out = x_out.permute(1,0,2) if IS_SEPERATED: y_root = self.Wout_root(x_out) y_attr = self.Wout_attr(x_out) del mask return y_root, y_attr else: y = self.Wout(x_out) del mask return y def generate(self, feature_semantic_list = [], feature_key=None, feature_scene_offset=None, feature_motion=None, feature_emotion=None, primer=None, primer_root=None, primer_attr=None, target_seq_length=300, beam=0, beam_chance=1.0, max_conseq_N = 0, max_conseq_chord = 2): assert (not self.training), "Cannot generate while in training mode" print("Generating sequence of max length:", target_seq_length) with open('dataset/vevo_meta/chord_inv.json') as json_file: chordInvDic = json.load(json_file) with open('dataset/vevo_meta/chord_root.json') as json_file: chordRootDic = json.load(json_file) with open('dataset/vevo_meta/chord_attr.json') as json_file: chordAttrDic = json.load(json_file) gen_seq = torch.full((1,target_seq_length), CHORD_PAD, dtype=TORCH_LABEL_TYPE, device=self.device) gen_seq_root = torch.full((1,target_seq_length), CHORD_ROOT_PAD, dtype=TORCH_LABEL_TYPE, device=self.device) gen_seq_attr = torch.full((1,target_seq_length), CHORD_ATTR_PAD, dtype=TORCH_LABEL_TYPE, device=self.device) num_primer = len(primer) gen_seq[..., :num_primer] = primer.type(TORCH_LABEL_TYPE).to(self.device) gen_seq_root[..., :num_primer] = primer_root.type(TORCH_LABEL_TYPE).to(self.device) gen_seq_attr[..., :num_primer] = primer_attr.type(TORCH_LABEL_TYPE).to(self.device) cur_i = num_primer while(cur_i < target_seq_length): y = self.softmax( self.forward( gen_seq[..., :cur_i], gen_seq_root[..., :cur_i], gen_seq_attr[..., :cur_i], feature_semantic_list, feature_key, feature_scene_offset, feature_motion, feature_emotion) )[..., :CHORD_END] token_probs = y[:, cur_i-1, :] if(beam == 0): beam_ran = 2.0 else: beam_ran = random.uniform(0,1) if(beam_ran <= beam_chance): token_probs = token_probs.flatten() top_res, top_i = torch.topk(token_probs, beam) beam_rows = top_i // CHORD_SIZE beam_cols = top_i % CHORD_SIZE gen_seq = gen_seq[beam_rows, :] gen_seq[..., cur_i] = beam_cols else: # token_probs.shape : [1, 157] # 0: N, 1: C, ... , 156: B:maj7 # 157 chordEnd 158 padding if max_conseq_N == 0: token_probs[0][0] = 0.0 isMaxChord = True if cur_i >= max_conseq_chord : preChord = gen_seq[0][cur_i-1].item() for k in range (1, max_conseq_chord): if preChord != gen_seq[0][cur_i-1-k].item(): isMaxChord = False else: isMaxChord = False if isMaxChord: preChord = gen_seq[0][cur_i-1].item() token_probs[0][preChord] = 0.0 distrib = torch.distributions.categorical.Categorical(probs=token_probs) next_token = distrib.sample() gen_seq[:, cur_i] = next_token gen_chord = chordInvDic[ str( next_token.item() ) ] chord_arr = gen_chord.split(":") if len(chord_arr) == 1: chordRootID = chordRootDic[chord_arr[0]] chordAttrID = 1 chordRootID = torch.tensor([chordRootID]).to(self.device) chordAttrID = torch.tensor([chordAttrID]).to(self.device) gen_seq_root[:, cur_i] = chordRootID gen_seq_attr[:, cur_i] = chordAttrID elif len(chord_arr) == 2: chordRootID = chordRootDic[chord_arr[0]] chordAttrID = chordAttrDic[chord_arr[1]] chordRootID = torch.tensor([chordRootID]).to(self.device) chordAttrID = torch.tensor([chordAttrID]).to(self.device) gen_seq_root[:, cur_i] = chordRootID gen_seq_attr[:, cur_i] = chordAttrID # Let the transformer decide to end if it wants to if(next_token == CHORD_END): print("Model called end of sequence at:", cur_i, "/", target_seq_length) break cur_i += 1 if(cur_i % 50 == 0): print(cur_i, "/", target_seq_length) return gen_seq[:, :cur_i]