model/musemorphose.py

import torch
from torch import nn
import torch.nn.functional as F
from transformer_encoder import VAETransformerEncoder
from transformer_helpers import (
  weights_init, PositionalEncoding, TokenEmbedding, generate_causal_mask
)

class VAETransformerDecoder(nn.Module):
  def __init__(self, n_layer, n_head, d_model, d_ff, d_seg_emb, dropout=0.1, activation='relu', cond_mode='in-attn'):
    super(VAETransformerDecoder, self).__init__()
    self.n_layer = n_layer
    self.n_head = n_head
    self.d_model = d_model
    self.d_ff = d_ff
    self.d_seg_emb = d_seg_emb
    self.dropout = dropout
    self.activation = activation
    self.cond_mode = cond_mode

    if cond_mode == 'in-attn':
      self.seg_emb_proj = nn.Linear(d_seg_emb, d_model, bias=False)
    elif cond_mode == 'pre-attn':
      self.seg_emb_proj = nn.Linear(d_seg_emb + d_model, d_model, bias=False)

    self.decoder_layers = nn.ModuleList()
    for i in range(n_layer):
      self.decoder_layers.append(
        nn.TransformerEncoderLayer(d_model, n_head, d_ff, dropout, activation)
      )

  def forward(self, x, seg_emb):
    if not hasattr(self, 'cond_mode'):
      self.cond_mode = 'in-attn'
    attn_mask = generate_causal_mask(x.size(0)).to(x.device)
    # print (attn_mask.size())

    if self.cond_mode == 'in-attn':
      seg_emb = self.seg_emb_proj(seg_emb)
    elif self.cond_mode == 'pre-attn':
      x = torch.cat([x, seg_emb], dim=-1)
      x = self.seg_emb_proj(x)

    out = x
    for i in range(self.n_layer):
      if self.cond_mode == 'in-attn':
        out += seg_emb
      out = self.decoder_layers[i](out, src_mask=attn_mask)

    return out

class MuseMorphose(nn.Module):
  def __init__(self, enc_n_layer, enc_n_head, enc_d_model, enc_d_ff, 
    dec_n_layer, dec_n_head, dec_d_model, dec_d_ff,
    d_vae_latent, d_embed, n_token,
    enc_dropout=0.1, enc_activation='relu',
    dec_dropout=0.1, dec_activation='relu',
    d_rfreq_emb=32, d_polyph_emb=32,
    n_rfreq_cls=8, n_polyph_cls=8,
    is_training=True, use_attr_cls=True,
    cond_mode='in-attn'
  ):
    super(MuseMorphose, self).__init__()
    self.enc_n_layer = enc_n_layer
    self.enc_n_head = enc_n_head
    self.enc_d_model = enc_d_model
    self.enc_d_ff = enc_d_ff
    self.enc_dropout = enc_dropout
    self.enc_activation = enc_activation

    self.dec_n_layer = dec_n_layer
    self.dec_n_head = dec_n_head
    self.dec_d_model = dec_d_model
    self.dec_d_ff = dec_d_ff
    self.dec_dropout = dec_dropout
    self.dec_activation = dec_activation  

    self.d_vae_latent = d_vae_latent
    self.n_token = n_token
    self.is_training = is_training

    self.cond_mode = cond_mode
    self.token_emb = TokenEmbedding(n_token, d_embed, enc_d_model)
    self.d_embed = d_embed
    self.pe = PositionalEncoding(d_embed)
    self.dec_out_proj = nn.Linear(dec_d_model, n_token)
    self.encoder = VAETransformerEncoder(
      enc_n_layer, enc_n_head, enc_d_model, enc_d_ff, d_vae_latent, enc_dropout, enc_activation
    )

    self.use_attr_cls = use_attr_cls
    if use_attr_cls:
      self.decoder = VAETransformerDecoder(
        dec_n_layer, dec_n_head, dec_d_model, dec_d_ff, d_vae_latent + d_polyph_emb + d_rfreq_emb,
        dropout=dec_dropout, activation=dec_activation,
        cond_mode=cond_mode
      )
    else:
      self.decoder = VAETransformerDecoder(
        dec_n_layer, dec_n_head, dec_d_model, dec_d_ff, d_vae_latent,
        dropout=dec_dropout, activation=dec_activation,
        cond_mode=cond_mode
      )

    if use_attr_cls:
      self.d_rfreq_emb = d_rfreq_emb
      self.d_polyph_emb = d_polyph_emb
      self.rfreq_attr_emb = TokenEmbedding(n_rfreq_cls, d_rfreq_emb, d_rfreq_emb)
      self.polyph_attr_emb = TokenEmbedding(n_polyph_cls, d_polyph_emb, d_polyph_emb)
    else:
      self.rfreq_attr_emb = None
      self.polyph_attr_emb = None

    self.emb_dropout = nn.Dropout(self.enc_dropout)
    self.apply(weights_init)
    

  def reparameterize(self, mu, logvar, use_sampling=True, sampling_var=1.):
    std = torch.exp(0.5 * logvar).to(mu.device)
    if use_sampling:
      eps = torch.randn_like(std).to(mu.device) * sampling_var
    else:
      eps = torch.zeros_like(std).to(mu.device)

    return eps * std + mu

  def get_sampled_latent(self, inp, padding_mask=None, use_sampling=False, sampling_var=0.):
    token_emb = self.token_emb(inp)
    enc_inp = self.emb_dropout(token_emb) + self.pe(inp.size(0))

    _, mu, logvar = self.encoder(enc_inp, padding_mask=padding_mask)
    mu, logvar = mu.reshape(-1, mu.size(-1)), logvar.reshape(-1, mu.size(-1))
    vae_latent = self.reparameterize(mu, logvar, use_sampling=use_sampling, sampling_var=sampling_var)

    return vae_latent

  def generate(self, inp, dec_seg_emb, rfreq_cls=None, polyph_cls=None, keep_last_only=True):
    token_emb = self.token_emb(inp)
    dec_inp = self.emb_dropout(token_emb) + self.pe(inp.size(0))

    if rfreq_cls is not None and polyph_cls is not None:
      dec_rfreq_emb = self.rfreq_attr_emb(rfreq_cls)
      dec_polyph_emb = self.polyph_attr_emb(polyph_cls)
      dec_seg_emb_cat = torch.cat([dec_seg_emb, dec_rfreq_emb, dec_polyph_emb], dim=-1)
    else:
      dec_seg_emb_cat = dec_seg_emb

    out = self.decoder(dec_inp, dec_seg_emb_cat)
    out = self.dec_out_proj(out)

    if keep_last_only:
      out = out[-1, ...]

    return out


  def forward(self, enc_inp, dec_inp, dec_inp_bar_pos, rfreq_cls=None, polyph_cls=None, padding_mask=None):
    # [shape of enc_inp] (seqlen_per_bar, bsize, n_bars_per_sample)
    enc_bt_size, enc_n_bars = enc_inp.size(1), enc_inp.size(2)
    enc_token_emb = self.token_emb(enc_inp)

    # [shape of dec_inp] (seqlen_per_sample, bsize)
    # [shape of rfreq_cls & polyph_cls] same as above 
    # -- (should copy each bar's label to all corresponding indices)
    dec_token_emb = self.token_emb(dec_inp)

    enc_token_emb = enc_token_emb.reshape(
      enc_inp.size(0), -1, enc_token_emb.size(-1)
    )
    enc_inp = self.emb_dropout(enc_token_emb) + self.pe(enc_inp.size(0))
    dec_inp = self.emb_dropout(dec_token_emb) + self.pe(dec_inp.size(0))

    # [shape of padding_mask] (bsize, n_bars_per_sample, seqlen_per_bar)
    # -- should be `True` for padded indices (i.e., those >= seqlen of the bar), `False` otherwise
    if padding_mask is not None:
      padding_mask = padding_mask.reshape(-1, padding_mask.size(-1))

    _, mu, logvar = self.encoder(enc_inp, padding_mask=padding_mask)
    vae_latent = self.reparameterize(mu, logvar)
    vae_latent_reshaped = vae_latent.reshape(enc_bt_size, enc_n_bars, -1)

    dec_seg_emb = torch.zeros(dec_inp.size(0), dec_inp.size(1), self.d_vae_latent).to(vae_latent.device)
    for n in range(dec_inp.size(1)):
      # [shape of dec_inp_bar_pos] (bsize, n_bars_per_sample + 1)
      # -- stores [[start idx of bar #1, sample #1, ..., start idx of bar #K, sample #1, seqlen of sample #1], [same for another sample], ...]
      for b, (st, ed) in enumerate(zip(dec_inp_bar_pos[n, :-1], dec_inp_bar_pos[n, 1:])):
        dec_seg_emb[st:ed, n, :] = vae_latent_reshaped[n, b, :]

    if rfreq_cls is not None and polyph_cls is not None and self.use_attr_cls:
      dec_rfreq_emb = self.rfreq_attr_emb(rfreq_cls)
      dec_polyph_emb = self.polyph_attr_emb(polyph_cls)
      dec_seg_emb_cat = torch.cat([dec_seg_emb, dec_rfreq_emb, dec_polyph_emb], dim=-1)
    else:
      dec_seg_emb_cat = dec_seg_emb

    dec_out = self.decoder(dec_inp, dec_seg_emb_cat)
    dec_logits = self.dec_out_proj(dec_out)

    return mu, logvar, dec_logits

  def compute_loss(self, mu, logvar, beta, fb_lambda, dec_logits, dec_tgt):
    recons_loss = F.cross_entropy(
      dec_logits.view(-1, dec_logits.size(-1)), dec_tgt.contiguous().view(-1), 
      ignore_index=self.n_token - 1, reduction='mean'
    ).float()

    kl_raw = -0.5 * (1 + logvar - mu ** 2 - logvar.exp()).mean(dim=0)
    kl_before_free_bits = kl_raw.mean()
    kl_after_free_bits = kl_raw.clamp(min=fb_lambda)
    kldiv_loss = kl_after_free_bits.mean()

    return {
      'beta': beta,
      'total_loss': recons_loss + beta * kldiv_loss,
      'kldiv_loss': kldiv_loss,
      'kldiv_raw': kl_before_free_bits,
      'recons_loss': recons_loss
    }