| import torch |
| import numpy as np |
| import torch.nn.functional as F |
| import torch.nn as nn |
|
|
| from torch import Tensor, nn, no_grad |
| from .autoencoders import OobleckDecoder, OobleckEncoder |
|
|
| from .transformer import ContinuousTransformer |
| LRELU_SLOPE = 0.1 |
| padding_mode = "zeros" |
| sample_eps = 1e-6 |
|
|
| def vae_sample(mean, scale): |
| stdev = nn.functional.softplus(scale) |
| var = stdev * stdev + sample_eps |
| logvar = torch.log(var) |
| latents = torch.randn_like(mean) * stdev + mean |
|
|
| kl = (mean * mean + var - logvar - 1).sum(1).mean() |
| |
| return latents, kl |
|
|
|
|
| class EAR_VAE(nn.Module): |
|
|
| def __init__(self, model_config: dict = None): |
| super().__init__() |
|
|
| if model_config is None: |
| model_config = { |
| "encoder": { |
| "config": { |
| "in_channels": 2, |
| "channels": 128, |
| "c_mults": [1, 2, 4, 8, 16], |
| "strides": [2, 4, 4, 4, 8], |
| "latent_dim": 128, |
| "use_snake": True |
| } |
| }, |
| "decoder": { |
| "config": { |
| "out_channels": 2, |
| "channels": 128, |
| "c_mults": [1, 2, 4, 8, 16], |
| "strides": [2, 4, 4, 4, 8], |
| "latent_dim": 64, |
| "use_nearest_upsample": False, |
| "use_snake": True, |
| "final_tanh": False, |
| }, |
| }, |
| "latent_dim": 64, |
| "downsampling_ratio": 1024, |
| "io_channels": 2, |
| } |
| else: |
| model_config = model_config |
|
|
| if model_config.get("transformer") is not None: |
| self.transformers = ContinuousTransformer( |
| dim=model_config["decoder"]["config"]["latent_dim"], |
| depth=model_config["transformer"]["depth"], |
| **model_config["transformer"].get("config", {}), |
| ) |
| else: |
| self.transformers = None |
|
|
| self.encoder = OobleckEncoder(**model_config["encoder"]["config"]) |
| self.decoder = OobleckDecoder(**model_config["decoder"]["config"]) |
|
|
| def forward(self, audio) -> Tensor: |
| """ |
| audio: Input audio tensor [B,C,T] |
| """ |
| status = self.encoder(audio) |
| mean, scale = status.chunk(2, dim=1) |
| z, kl = vae_sample(mean, scale) |
| |
| if self.transformers is not None: |
| z = z.permute(0, 2, 1) |
| z = self.transformers(z) |
| z = z.permute(0, 2, 1) |
|
|
| x = self.decoder(z) |
| return x, kl |
|
|
| def encode(self, audio, use_sample=True): |
| x = self.encoder(audio) |
| mean, scale = x.chunk(2, dim=1) |
| if use_sample: |
| z, _ = vae_sample(mean, scale) |
| else: |
| z = mean |
| return z |
|
|
| def decode(self, z): |
| |
| if self.transformers is not None: |
| z = z.permute(0, 2, 1) |
| z = self.transformers(z) |
| z = z.permute(0, 2, 1) |
| |
| x = self.decoder(z) |
| return x |
|
|
| @no_grad() |
| def inference(self, audio): |
| z = self.encode(audio) |
| recon_audio = self.decode(z) |
| return recon_audio |
|
|
