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