| import torch
|
| from audioldm.latent_diffusion.ema import *
|
| from audioldm.variational_autoencoder.modules import Encoder, Decoder
|
| from audioldm.variational_autoencoder.distributions import DiagonalGaussianDistribution
|
|
|
| from audioldm.hifigan.utilities import get_vocoder, vocoder_infer
|
|
|
|
|
| class AutoencoderKL(nn.Module):
|
| def __init__(
|
| self,
|
| ddconfig=None,
|
| lossconfig=None,
|
| image_key="fbank",
|
| embed_dim=None,
|
| time_shuffle=1,
|
| subband=1,
|
| ckpt_path=None,
|
| reload_from_ckpt=None,
|
| ignore_keys=[],
|
| colorize_nlabels=None,
|
| monitor=None,
|
| base_learning_rate=1e-5,
|
| ):
|
| super().__init__()
|
|
|
| self.encoder = Encoder(**ddconfig)
|
| self.decoder = Decoder(**ddconfig)
|
|
|
| self.subband = int(subband)
|
|
|
| if self.subband > 1:
|
| print("Use subband decomposition %s" % self.subband)
|
|
|
| self.quant_conv = torch.nn.Conv2d(2 * ddconfig["z_channels"], 2 * embed_dim, 1)
|
| self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
|
|
|
| self.vocoder = get_vocoder(None, "cpu")
|
| self.embed_dim = embed_dim
|
|
|
| if monitor is not None:
|
| self.monitor = monitor
|
|
|
| self.time_shuffle = time_shuffle
|
| self.reload_from_ckpt = reload_from_ckpt
|
| self.reloaded = False
|
| self.mean, self.std = None, None
|
|
|
| def encode(self, x):
|
|
|
| x = self.freq_split_subband(x)
|
| h = self.encoder(x)
|
| moments = self.quant_conv(h)
|
| posterior = DiagonalGaussianDistribution(moments)
|
| return posterior
|
|
|
| def decode(self, z):
|
| z = self.post_quant_conv(z)
|
| dec = self.decoder(z)
|
| dec = self.freq_merge_subband(dec)
|
| return dec
|
|
|
| def decode_to_waveform(self, dec):
|
| dec = dec.squeeze(1).permute(0, 2, 1)
|
| wav_reconstruction = vocoder_infer(dec, self.vocoder)
|
| return wav_reconstruction
|
|
|
| def forward(self, input, sample_posterior=True):
|
| posterior = self.encode(input)
|
| if sample_posterior:
|
| z = posterior.sample()
|
| else:
|
| z = posterior.mode()
|
|
|
| if self.flag_first_run:
|
| print("Latent size: ", z.size())
|
| self.flag_first_run = False
|
|
|
| dec = self.decode(z)
|
|
|
| return dec, posterior
|
|
|
| def freq_split_subband(self, fbank):
|
| if self.subband == 1 or self.image_key != "stft":
|
| return fbank
|
|
|
| bs, ch, tstep, fbins = fbank.size()
|
|
|
| assert fbank.size(-1) % self.subband == 0
|
| assert ch == 1
|
|
|
| return (
|
| fbank.squeeze(1)
|
| .reshape(bs, tstep, self.subband, fbins // self.subband)
|
| .permute(0, 2, 1, 3)
|
| )
|
|
|
| def freq_merge_subband(self, subband_fbank):
|
| if self.subband == 1 or self.image_key != "stft":
|
| return subband_fbank
|
| assert subband_fbank.size(1) == self.subband
|
| bs, sub_ch, tstep, fbins = subband_fbank.size()
|
| return subband_fbank.permute(0, 2, 1, 3).reshape(bs, tstep, -1).unsqueeze(1)
|
|
|
