import torch
import torch.nn.functional as F
import pytorch_lightning as pl
from ldm.util import instantiate_from_config
from ldm.modules.diffusionmodules.model import Encoder, Decoder
from ldm.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
class VQModelDual(pl.LightningModule):
def __init__(self,
ddconfig,
lossconfig,
n_embed,
embed_dim,
ckpt_path=None,
ignore_keys=[],
image1_key="image1",
image2_key="image2",
colorize_nlabels=None,
monitor=None,
remap=None,
sane_index_shape=False, # tell vector quantizer to return indices as bhw
):
super().__init__()
self.image1_key = image1_key
self.image2_key = image2_key
## model 1
self.encoder1 = Encoder(**ddconfig)
self.decoder1 = Decoder(**ddconfig)
self.quantize1 = VectorQuantizer(n_embed, embed_dim, beta=0.25,
remap=remap, sane_index_shape=sane_index_shape)
self.quant_conv1= torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
self.post_quant_conv1 = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
self.loss1 = instantiate_from_config(lossconfig)
## model 2
self.encoder2 = Encoder(**ddconfig)
self.decoder2 = Decoder(**ddconfig)
self.quantize2 = VectorQuantizer(n_embed, embed_dim, beta=0.25,
remap=remap, sane_index_shape=sane_index_shape)
self.quant_conv2 = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
self.post_quant_conv2 = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
self.loss2 = instantiate_from_config(lossconfig)
if ckpt_path is not None:
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
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 init_from_ckpt(self, path, ignore_keys=list()):
sd = torch.load(path, map_location="cpu")["state_dict"]
keys = list(sd.keys())
for k in keys:
for ik in ignore_keys:
if k.startswith(ik):
print("Deleting key {} from state_dict.".format(k))
del sd[k]
self.load_state_dict(sd, strict=False)
print(f"Restored from {path}")
def encode(self, x1, x2):
h1 = self.encoder1(x1)
h1 = self.quant_conv1(h1)
quant1, emb_loss1, info1 = self.quantize1(h1)
h2 = self.encoder2(x2)
h2 = self.quant_conv2(h2)
quant2, emb_loss2, info2 = self.quantize2(h2)
return quant1, emb_loss1, info1, quant2, emb_loss2, info2
def decode(self, quant1, quant2):
quant1 = self.post_quant_conv1(quant1)
dec1 = self.decoder1(quant1)
quant2 = self.post_quant_conv2(quant2)
dec2 = self.decoder2(quant2)
return dec1, dec2
# def decode_code(self, code_b, model_key):
# quant_b = self.quantize[model_key].embed_code(code_b)
# dec = self.decode(quant_b,model_key)
# return dec
def forward(self, input1, input2):
# quant, diff, _ = self.encode(input, model_key)
quant1, diff1, _, quant2, diff2, _ = self.encode(input1, input2)
dec1, dec2 = self.decode(quant1, quant2)
# dec = self.decode(quant, model_key)
return dec1, dec2, diff1, diff2
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)
x = x.to(memory_format=torch.contiguous_format)
return x.float()
def training_step(self, batch, batch_idx, optimizer_idx):
x1 = self.get_input(batch, self.image1_key)
x2 = self.get_input(batch, self.image2_key)
xrec1, xrec2, qloss1, qloss2 = self.forward(x1, x2)
if optimizer_idx == 0:
# autoencoder 1
aeloss1, log_dict_ae1 = self.loss1(qloss1, x1, xrec1, optimizer_idx, self.global_step,
last_layer=self.get_last_layer(model_key=1), split="train")
self.log("train/aeloss1", aeloss1, prog_bar=True, logger=True, on_step=True, on_epoch=True)
self.log_dict(log_dict_ae1, prog_bar=False, logger=True, on_step=True, on_epoch=True)
# autoencoder 2
aeloss2, log_dict_ae2 = self.loss2(qloss2, x2, xrec2, optimizer_idx, self.global_step,
last_layer=self.get_last_layer(model_key=2), split="train")
self.log("train/aeloss2", aeloss2, prog_bar=True, logger=True, on_step=True, on_epoch=True)
self.log_dict(log_dict_ae2, prog_bar=False, logger=True, on_step=True, on_epoch=True)
return aeloss1 + aeloss2
if optimizer_idx == 1:
# discriminator 1
discloss1, log_dict_disc1 = self.loss1(qloss1, x1, xrec1, optimizer_idx, self.global_step,
last_layer=self.get_last_layer(model_key=1), split="train")
self.log("train/discloss1", discloss1, prog_bar=True, logger=True, on_step=True, on_epoch=True)
self.log_dict(log_dict_disc1, prog_bar=False, logger=True, on_step=True, on_epoch=True)
# discriminator 2
discloss2, log_dict_disc2 = self.loss2(qloss2, x2, xrec2, optimizer_idx, self.global_step,
last_layer=self.get_last_layer(model_key=2), split="train")
self.log("train/discloss", discloss2, prog_bar=True, logger=True, on_step=True, on_epoch=True)
self.log_dict(log_dict_disc2, prog_bar=False, logger=True, on_step=True, on_epoch=True)
return discloss1 + discloss2
def validation_step(self, batch, batch_idx):
x1 = self.get_input(batch, self.image1_key)
x2 = self.get_input(batch, self.image2_key)
xrec1, xrec2, qloss1, qloss2 = self.forward(x1, x2)
aeloss1, log_dict_ae1 = self.loss1(qloss1, x1, xrec1, 0, self.global_step,
last_layer=self.get_last_layer(model_key=1), split="val")
aeloss2, log_dict_ae2 = self.loss2(qloss2, x2, xrec2, 0, self.global_step,
last_layer=self.get_last_layer(model_key=2), split="val")
discloss1, log_dict_disc1 = self.loss1(qloss1, x1, xrec1, 1, self.global_step,
last_layer=self.get_last_layer(model_key=1), split="val")
discloss2, log_dict_disc2 = self.loss2(qloss2, x2, xrec2, 1, self.global_step,
last_layer=self.get_last_layer(model_key=2), split="val")
rec_loss1 = log_dict_ae1["val/rec_loss"]
rec_loss2 = log_dict_ae2["val/rec_loss"]
self.log("val/rec_loss1", rec_loss1,
prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
self.log("val/rec_loss2", rec_loss2,
prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
self.log("val/aeloss1", aeloss1,
prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
self.log("val/aeloss2", aeloss2,
prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
self.log_dict(log_dict_ae1)
self.log_dict(log_dict_disc1)
self.log_dict(log_dict_ae2)
self.log_dict(log_dict_disc2)
return self.log_dict
def configure_optimizers(self):
lr = self.learning_rate
opt_ae = torch.optim.Adam(list(self.encoder1.parameters())+
list(self.decoder1.parameters())+
list(self.quantize1.parameters())+
list(self.quant_conv1.parameters())+
list(self.post_quant_conv1.parameters())+
list(self.encoder2.parameters())+
list(self.decoder2.parameters())+
list(self.quantize2.parameters())+
list(self.quant_conv2.parameters())+
list(self.post_quant_conv2.parameters()),
lr=lr, betas=(0.5, 0.9))
opt_disc = torch.optim.Adam(list(self.loss1.discriminator.parameters())+
list(self.loss2.discriminator.parameters()),
lr=lr, betas=(0.5, 0.9))
return [opt_ae, opt_disc], []
def get_last_layer(self, model_key):
if model_key==1:
return self.decoder2.conv_out.weight
elif model_key==2:
return self.decoder2.conv_out.weight
def log_images(self, batch, **kwargs):
log = dict()
x1 = self.get_input(batch, self.image1_key)
x2 = self.get_input(batch, self.image2_key)
x1 = x1.to(self.device)
x2 = x2.to(self.device)
xrec1, xrec2, _, _ = self.forward(x1, x2)
## log 1
if x1.shape[1] > 3:
# colorize with random projection
assert xrec1.shape[1] > 3
x1 = self.to_rgb(x1)
xrec1 = self.to_rgb(xrec1)
log["inputs1"] = x1
log["reconstructions1"] = xrec1
## log 2
if x2.shape[1] > 3:
# colorize with random projection
assert xrec2.shape[1] > 3
x2 = self.to_rgb(x2)
xrec2 = self.to_rgb(xrec2)
log["inputs2"] = x2
log["reconstructions2"] = xrec2
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 VQModelDualInterface(VQModelDual):
def __init__(self, embed_dim, *args, **kwargs):
super().__init__(embed_dim=embed_dim, *args, **kwargs)
self.embed_dim = embed_dim
def encode(self, x1, x2):
h1 = self.encoder1(x1)
h1 = self.quant_conv1(h1)
h2 = self.encoder2(x2)
h2 = self.quant_conv2(h2)
return h1, h2
def decode(self, h1, h2, force_not_quantize=False):
# also go through quantization layer
if not force_not_quantize:
quant1, emb_loss1, info1 = self.quantize1(h1)
quant2, emb_loss2, info2 = self.quantize2(h2)
else:
quant1 = h1
quant2 = h2
quant1 = self.post_quant_conv1(quant1)
dec1 = self.decoder1(quant1)
quant2 = self.post_quant_conv2(quant2)
dec2 = self.decoder2(quant2)
return dec1, dec2
def decode1(self, h1, force_not_quantize=False):
# also go through quantization layer
if not force_not_quantize:
quant1, emb_loss1, info1 = self.quantize1(h1)
else:
quant1 = h1
quant1 = self.post_quant_conv1(quant1)
dec1 = self.decoder1(quant1)
return dec1
def decode2(self, h2, force_not_quantize=False):
# also go through quantization layer
if not force_not_quantize:
quant2, emb_loss2, info2 = self.quantize2(h2)
else:
quant2 = h2
quant2 = self.post_quant_conv2(quant2)
dec2 = self.decoder2(quant2)
return dec2