File size: 2,262 Bytes
feae2c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import os
import torch
import logging
import pytorch_lightning as pl
from ldm.modules.bigvgan.bigvgan import VocoderBigVGAN
from ldm.modules.latent_encoder.model import Encoder, Decoder
from ldm.modules.diffusionmodules.distributions import (
DiagonalGaussianDistribution,
)
import soundfile as sf
class AutoencoderKL(pl.LightningModule):
def __init__(
self,
ddconfig=None,
embed_dim=None,
reload_from_ckpt=None,
ignore_keys=[],
vocoder_path=None,
):
super().__init__()
self.encoder = Encoder(**ddconfig)
self.decoder = Decoder(**ddconfig)
self.vocoder_path = vocoder_path
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 = VocoderBigVGAN(self.vocoder_path, device="cuda")
self.reload_from_ckpt = reload_from_ckpt
if self.reload_from_ckpt is not None:
checkpoint = torch.load(self.reload_from_ckpt)
load_todo_keys = {}
pretrained_state_dict = checkpoint
current_state_dict = self.state_dict()
for key in current_state_dict:
if (
key in pretrained_state_dict.keys()
and pretrained_state_dict[key].size()
== current_state_dict[key].size()
):
load_todo_keys[key] = pretrained_state_dict[key]
else:
logging.info("Key %s mismatch during loading, seems fine" % key)
self.load_state_dict(load_todo_keys, strict=False)
self.reloaded = True
else:
logging.info("Train from scratch")
def encode(self, 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)
fbk = self.decoder(z)
return fbk
def forward(self, input, sample_posterior=True):
posterior = self.encode(input)
z = posterior.sample()
dec = self.decode(z)
return dec, posterior
|