parson/audioEditing

1import torch2from audioldm.latent_diffusion.ema import *3from audioldm.variational_autoencoder.modules import Encoder, Decoder4from audioldm.variational_autoencoder.distributions import DiagonalGaussianDistribution5 6from audioldm.hifigan.utilities import get_vocoder, vocoder_infer7 8 9class AutoencoderKL(nn.Module):10    def __init__(11        self,12        ddconfig=None,13        lossconfig=None,14        image_key="fbank",15        embed_dim=None,16        time_shuffle=1,17        subband=1,18        ckpt_path=None,19        reload_from_ckpt=None,20        ignore_keys=[],21        colorize_nlabels=None,22        monitor=None,23        base_learning_rate=1e-5,24    ):25        super().__init__()26 27        self.encoder = Encoder(**ddconfig)28        self.decoder = Decoder(**ddconfig)29 30        self.subband = int(subband)31 32        if self.subband > 1:33            print("Use subband decomposition %s" % self.subband)34 35        self.quant_conv = torch.nn.Conv2d(2 * ddconfig["z_channels"], 2 * embed_dim, 1)36        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)37 38        self.vocoder = get_vocoder(None, "cpu")39        self.embed_dim = embed_dim40 41        if monitor is not None:42            self.monitor = monitor43 44        self.time_shuffle = time_shuffle45        self.reload_from_ckpt = reload_from_ckpt46        self.reloaded = False47        self.mean, self.std = None, None48 49    def encode(self, x):50        # x = self.time_shuffle_operation(x)51        x = self.freq_split_subband(x)52        h = self.encoder(x)53        moments = self.quant_conv(h)54        posterior = DiagonalGaussianDistribution(moments)55        return posterior56 57    def decode(self, z):58        z = self.post_quant_conv(z)59        dec = self.decoder(z)60        dec = self.freq_merge_subband(dec)61        return dec62 63    def decode_to_waveform(self, dec):64        dec = dec.squeeze(1).permute(0, 2, 1)65        wav_reconstruction = vocoder_infer(dec, self.vocoder)66        return wav_reconstruction67 68    def forward(self, input, sample_posterior=True):69        posterior = self.encode(input)70        if sample_posterior:71            z = posterior.sample()72        else:73            z = posterior.mode()74 75        if self.flag_first_run:76            print("Latent size: ", z.size())77            self.flag_first_run = False78 79        dec = self.decode(z)80 81        return dec, posterior82 83    def freq_split_subband(self, fbank):84        if self.subband == 1 or self.image_key != "stft":85            return fbank86 87        bs, ch, tstep, fbins = fbank.size()88 89        assert fbank.size(-1) % self.subband == 090        assert ch == 191 92        return (93            fbank.squeeze(1)94            .reshape(bs, tstep, self.subband, fbins // self.subband)95            .permute(0, 2, 1, 3)96        )97 98    def freq_merge_subband(self, subband_fbank):99        if self.subband == 1 or self.image_key != "stft":100            return subband_fbank101        assert subband_fbank.size(1) == self.subband  # Channel dimension102        bs, sub_ch, tstep, fbins = subband_fbank.size()103        return subband_fbank.permute(0, 2, 1, 3).reshape(bs, tstep, -1).unsqueeze(1)104