model.py

import torch
from torch import nn
import numpy as np

def build_mlp(layers, final_activation='sigmoid'):
    layers_list = []
    for i in range(len(layers) - 1):
        layers_list.append(nn.Linear(layers[i], layers[i + 1]))
        if i != len(layers) - 2:
            layers_list.append(nn.ReLU())
    if final_activation == 'sigmoid':
        layers_list.append(nn.Sigmoid())
    elif final_activation == 'relu':
        layers_list.append(nn.ReLU())
    
    return nn.Sequential(*layers_list)

class VAE(nn.Module):
    def __init__(self, audio_processor, latent_dim=2, hidden_dims = [2, 4, 8]):
        # Setup
        torch.set_default_dtype(torch.float32)
        super().__init__()

        # Configuration
        self.input_audio_length = 64000

        # Audio Configuration
        self.audio_processor = audio_processor
        
        # Parse args
        self.latent_dim = latent_dim
        self.hidden_dims = hidden_dims

        self.construct()
        
    def construct(self):
        # Create a random mel spectrogram to get the input shape
        test_signal = torch.from_numpy(np.random.rand(64000)).float()
        test_spectrogram = self.audio_processor.signal_to_spectrogram(test_signal)
        (input_channels, input_width, input_height) = test_spectrogram.shape # (128, 128)

        # Encoder
        encoder_modules = []
        input_encoding = input_channels

        for h_dim in self.hidden_dims:
            encoder_modules.append(
                nn.Sequential(
                    nn.Conv2d(input_encoding, out_channels=h_dim, kernel_size=3, stride=2, padding=1),
                    nn.BatchNorm2d(h_dim),
                    nn.LeakyReLU()
                )
            )
            input_encoding = h_dim
        self.spectrogram_encoder = nn.Sequential(*encoder_modules)

        # Find size of encoder output
        spectrogram_encoded_test = self.spectrogram_encoder(torch.randn(1, 1, input_height, input_width))
        self.spectrogram_encoded_shape = spectrogram_encoded_test.shape
        self.spectrogram_encoded_dim = spectrogram_encoded_test.flatten(start_dim=1).shape[1]

        # Set up MLPs for the audio features
        self.family_encoder = build_mlp([11, 16, 32])
        self.instrument_encoder = build_mlp([1006, 512, 128, 64, 32])
        self.source_encoder = build_mlp([3, 8, 16, 32])
        self.note_encoder = build_mlp([2, 8, 16, 32])
        self.qualities_encoder = build_mlp([10, 16, 32])
        self.id_encoder = build_mlp([32, 64, 32])
        # MLP for combined audio features (6*4 is input)
        self.feature_encoder = build_mlp([192, 256, 512])
        
        encoded_dim = self.spectrogram_encoded_dim + 512

        # Latent Bottleneck
        self.fc_mu = nn.Linear(encoded_dim, self.latent_dim)
        self.fc_var = nn.Linear(encoded_dim, self.latent_dim)

        # Decoder

        # Set up MLPs for the audio features
        self.feature_decoder = build_mlp([512, 256, 192])
        self.family_decoder = build_mlp([32, 16, 11])
        self.instrument_decoder = build_mlp([32, 64, 128, 512, 1006])
        self.source_decoder = build_mlp([32, 16, 8, 3])
        self.note_decoder = build_mlp([32, 16, 8, 2])
        self.qualities_decoder = build_mlp([32, 16, 10])
        self.id_decoder = build_mlp([32, 64, 32], final_activation='relu')

        self.hidden_dims.reverse()
        self.decoder_input = nn.Linear(self.latent_dim, encoded_dim)
        decoder_modules = []
        for i in range(len(self.hidden_dims) - 1):
            decoder_modules.append(
                nn.Sequential(
                    nn.ConvTranspose2d(self.hidden_dims[i],
                                       self.hidden_dims[i + 1],
                                       kernel_size=3,
                                       stride=2,
                                       padding=1,
                                       output_padding=1),
                    nn.BatchNorm2d(self.hidden_dims[i + 1]),
                    nn.LeakyReLU())
            )
        self.spectrogram_decoder = nn.Sequential(*decoder_modules)

        self.final_layer = nn.Sequential(
            nn.ConvTranspose2d(self.hidden_dims[-1],
                               self.hidden_dims[-1],
                               kernel_size=3,
                               stride=2,
                               padding=1,
                               output_padding=1),
            nn.BatchNorm2d(self.hidden_dims[-1]),
            nn.LeakyReLU(),
            nn.Conv2d(self.hidden_dims[-1], out_channels=input_channels,
                      kernel_size=3, padding=1),
            nn.Sigmoid())

    def encode(self, input):
        # Encode Audio Features
        family_result = self.family_encoder(input['family'])
        instrument_result = self.instrument_encoder(input['instrument'])
        source_result = self.source_encoder(input['source'])
        note_result = self.note_encoder(input['note'])
        qualities_result = self.qualities_encoder(input['qualities'])
        id_result = self.id_encoder(input['id'])
        feature_input = torch.cat([family_result, instrument_result, source_result, note_result, qualities_result, id_result], dim=1)
        feature_result = self.feature_encoder(feature_input)

        # print("Feature Result Shape: ", feature_result.shape)

        # Encode Spectrogram
        spectrogram_result = self.spectrogram_encoder(input['spectrogram_normalized'])
        spectrogram_result = torch.flatten(spectrogram_result, start_dim=1)

        # print("Spectrogram Result Shape: ", spectrogram_result.shape)

        # Combine Audio Features and Spectrogram
        result = torch.cat([spectrogram_result, feature_result], dim=1)
        
        mu = self.fc_mu(result)
        log_var = self.fc_var(result)
        z = self.reparameterize(mu, log_var)

        # print("Mu Shape: ", mu.shape)
        # print("Log Var Shape: ", log_var.shape)
        # print("Z Shape: ", z.shape)

        return mu, log_var, z

    def decode(self, z):

        # Convert from latent dim to encoded dim
        result = self.decoder_input(z)

        # Separate and Decode Spectrogram
        spectrogram_result = result[:, 0:self.spectrogram_encoded_dim] # Get the spectrogram part of the result
        spectrogram_result = spectrogram_result.view(-1, self.spectrogram_encoded_shape[1], self.spectrogram_encoded_shape[2], self.spectrogram_encoded_shape[3]) # Reshape to the shape of the spectrogram
        spectrogram_result = self.spectrogram_decoder(spectrogram_result) # Decode the spectrogram
        spectrogram_normalized = self.final_layer(spectrogram_result) # Pass through final layer
        spectrogram_denormalized = self.audio_processor.denormalize(spectrogram_normalized)

        feature_result = result[:, self.spectrogram_encoded_dim:] # Get the audio feature part of the result
        feature_result = self.feature_decoder(feature_result) # Decode the audio features
        family_result = self.family_decoder(feature_result[:, 0:32])
        instrument_result = self.instrument_decoder(feature_result[:, 32:64])
        source_result = self.source_decoder(feature_result[:, 64:96])
        note_result = self.note_decoder(feature_result[:, 96:128])
        qualities_result = self.qualities_decoder(feature_result[:, 128:160])
        id_result = self.id_decoder(feature_result[:, 160:192])

        reconstructed = {
            'spectrogram_normalized': spectrogram_normalized,
            'spectrogram': spectrogram_denormalized,
            'family': family_result,
            'instrument': instrument_result,
            'source': source_result,
            'note': note_result,
            'qualities': qualities_result,
            'id': id_result
        }

        return reconstructed

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = eps.mul(std).add_(mu)

        return z
    
    def forward(self, input):
        mu, log_var, z = self.encode(input)
        return self.decode(z), mu, log_var
    
    def generate(self, x):
        return self.forward(x)[0]
    
    def sample(self, num_samples, current_device="cpu"):
        # completely random sample from latent space
        z = torch.randn(num_samples, self.latent_dim)
        z = z.to(current_device)
        samples = self.decode(z)
        return samples
    
    def conditional_sample(self, num_samples, condition, current_device="cpu"):
        # Encoded spectrogram component is completely random
        rand_spectrogram = torch.randn(num_samples, self.spectrogram_encoded_dim)

        # Encoded audio features component is conditioned by input
        family_result = self.family_encoder(condition['family'])
        instrument_result = self.instrument_encoder(condition['instrument'])
        source_result = self.source_encoder(condition['source'])
        note_result = self.note_encoder(condition['note'])
        qualities_result = self.qualities_encoder(condition['qualities'])
        id_result = self.id_encoder(condition['id'])
        feature_input = torch.cat([family_result, instrument_result, source_result, note_result, qualities_result, id_result], dim=1)
        feature_result = self.feature_encoder(feature_input)

        result = torch.cat([rand_spectrogram, feature_result], dim=1)
        mu = self.fc_mu(result)
        log_var = self.fc_var(result)
        z = self.reparameterize(mu, log_var)

        z = self.decoder_input(z)
        z = torch.randn(num_samples, self.latent_dim)
        z = z.to(current_device)
        samples = self.decode(z)
        return samples
    
# if __name__ == "__main__":
#     from audio_processor import AudioProcessor
#     from torch.utils.data import DataLoader
#     from data import get_split
#     test = get_split("test")
#     model = VAE(test.audio_processor)

#     test_loader = DataLoader(test, batch_size=5, shuffle=True)
#     for i, x in enumerate(test_loader):
#         for key,val in x.items():
#             print(key, val.shape)
#         reconstructed, mu, log_var = model(x)
#         for key,val in reconstructed.items():
#             print(key, val.shape)
#         # print(model(x))
#         break