|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
from contextlib import contextmanager |
|
from einops import rearrange |
|
|
|
from core.models.common.get_model import get_model, register |
|
|
|
from .autokl_modules.diffusion_modules import Encoder, Decoder |
|
from .autokl_modules.distributions import DiagonalGaussianDistribution |
|
|
|
|
|
@register('autoencoderkl') |
|
class AutoencoderKL(nn.Module): |
|
def __init__(self, |
|
ddconfig, |
|
lossconfig, |
|
embed_dim, |
|
ckpt_path=None, |
|
ignore_keys=[], |
|
image_key="image", |
|
colorize_nlabels=None, |
|
monitor=None,): |
|
super().__init__() |
|
self.image_key = image_key |
|
self.encoder = Encoder(**ddconfig) |
|
self.decoder = Decoder(**ddconfig) |
|
assert ddconfig["double_z"] |
|
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.embed_dim = embed_dim |
|
if colorize_nlabels is not None: |
|
assert type(colorize_nlabels)==int |
|
self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) |
|
if monitor is not None: |
|
self.monitor = monitor |
|
|
|
def encode(self, x): |
|
if x.ndim == 5: |
|
is_video = True |
|
num_frames = x.shape[2] |
|
x = rearrange(x, 'b c f h w -> (b f) c h w') |
|
else: |
|
is_video = False |
|
if x.shape[1] == 1: |
|
x = torch.cat([x, x, x], dim=1) |
|
|
|
h = self.encoder(x) |
|
moments = self.quant_conv(h) |
|
if is_video: |
|
moments = rearrange(moments, '(b f) c h w -> b c f h w', f=num_frames) |
|
posterior = DiagonalGaussianDistribution(moments) |
|
return posterior |
|
|
|
def decode(self, z): |
|
if z.ndim == 5: |
|
is_video = True |
|
num_frames = z.shape[2] |
|
z = rearrange(z, 'b c f h w -> (b f) c h w') |
|
else: |
|
is_video = False |
|
num_frames = 1 |
|
|
|
z = self.post_quant_conv(z) |
|
dec = self.decoder(z, num_frames=num_frames) |
|
|
|
if is_video: |
|
dec = rearrange(dec, '(b f) c h w -> b c f h w', f=num_frames) |
|
return dec |
|
|
|
def forward(self, input, sample_posterior=True): |
|
posterior = self.encode(input) |
|
if sample_posterior: |
|
z = posterior.sample().to(input.dtype) |
|
else: |
|
z = posterior.mode().to(input.dtype) |
|
dec = self.decode(z) |
|
return dec, posterior |
|
|
|
def get_input(self, batch, k): |
|
x = batch[k] |
|
if len(x.shape) == 3: |
|
x = x[..., None] |
|
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format) |
|
return x |
|
|
|
def training_step(self, batch, batch_idx, optimizer_idx): |
|
inputs = self.get_input(batch, self.image_key) |
|
reconstructions, posterior = self(inputs) |
|
|
|
if optimizer_idx == 0: |
|
|
|
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, |
|
last_layer=self.get_last_layer(), split="train") |
|
self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) |
|
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False) |
|
return aeloss |
|
|
|
if optimizer_idx == 1: |
|
|
|
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, |
|
last_layer=self.get_last_layer(), split="train") |
|
|
|
self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) |
|
self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False) |
|
return discloss |
|
|
|
def validation_step(self, batch, batch_idx): |
|
inputs = self.get_input(batch, self.image_key) |
|
reconstructions, posterior = self(inputs) |
|
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step, |
|
last_layer=self.get_last_layer(), split="val") |
|
|
|
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step, |
|
last_layer=self.get_last_layer(), split="val") |
|
|
|
self.log("val/rec_loss", log_dict_ae["val/rec_loss"]) |
|
self.log_dict(log_dict_ae) |
|
self.log_dict(log_dict_disc) |
|
return self.log_dict |
|
|
|
def configure_optimizers(self): |
|
lr = self.learning_rate |
|
opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ |
|
list(self.decoder.parameters())+ |
|
list(self.quant_conv.parameters())+ |
|
list(self.post_quant_conv.parameters()), |
|
lr=lr, betas=(0.5, 0.9)) |
|
opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), |
|
lr=lr, betas=(0.5, 0.9)) |
|
return [opt_ae, opt_disc], [] |
|
|
|
def get_last_layer(self): |
|
return self.decoder.conv_out.weight |
|
|
|
@torch.no_grad() |
|
def log_images(self, batch, only_inputs=False, **kwargs): |
|
log = dict() |
|
x = self.get_input(batch, self.image_key) |
|
x = x.to(self.device) |
|
if not only_inputs: |
|
xrec, posterior = self(x) |
|
if x.shape[1] > 3: |
|
|
|
assert xrec.shape[1] > 3 |
|
x = self.to_rgb(x) |
|
xrec = self.to_rgb(xrec) |
|
log["samples"] = self.decode(torch.randn_like(posterior.sample())) |
|
log["reconstructions"] = xrec |
|
log["inputs"] = x |
|
return log |
|
|
|
def to_rgb(self, x): |
|
assert self.image_key == "segmentation" |
|
if not hasattr(self, "colorize"): |
|
self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) |
|
x = F.conv2d(x, weight=self.colorize) |
|
x = 2.*(x-x.min())/(x.max()-x.min()) - 1. |
|
return x |
|
|
|
|
|
class IdentityFirstStage(nn.Module): |
|
def __init__(self, *args, vq_interface=False, **kwargs): |
|
self.vq_interface = vq_interface |
|
super().__init__() |
|
|
|
def encode(self, x, *args, **kwargs): |
|
return x |
|
|
|
def decode(self, x, *args, **kwargs): |
|
return x |
|
|
|
def quantize(self, x, *args, **kwargs): |
|
if self.vq_interface: |
|
return x, None, [None, None, None] |
|
return x |
|
|
|
def forward(self, x, *args, **kwargs): |
|
return x |
|
|
