import torch import torch.nn.functional as F import pytorch_lightning as pl from main 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 self.encoder = {} self.decoder = {} self.quantize = {} self.quant_conv = {} self.post_quant_conv = {} self.loss = {} for i in range(2): self.encoder[i+1] = Encoder(**ddconfig) self.decoder[i+1] = Decoder(**ddconfig) self.quantize[i+1] = VectorQuantizer(n_embed, embed_dim, beta=0.25, remap=remap, sane_index_shape=sane_index_shape) self.quant_conv[i+1] = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1) self.post_quant_conv[i+1] = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) self.loss[i+1] = 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, x, model_key): h = self.encoder[model_key](x) h = self.quant_conv[model_key](h) quant, emb_loss, info = self.quantize[model_key](h) return quant, emb_loss, info def decode(self, quant, model_key): quant = self.post_quant_conv[model_key](quant) dec = self.decoder[model_key](quant) return dec 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, input, model_key): quant, diff, _ = self.encode(input, model_key) dec = self.decode(quant, model_key) return dec, diff 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.float() def training_step(self, batch, batch_idx, optimizer_idx): breakpoint() x1 = self.get_input(batch, self.image_key1) x2 = self.get_input(batch, self.image_key2) xrec1, qloss1 = self.forward(x1, model_key=1) xrec2, qloss2 = self.forward(x2, model_key=2) if optimizer_idx == 0: # autoencoder 1 aeloss1, log_dict_ae1 = self.loss[1](qloss1, x1, xrec1, optimizer_idx, self.global_step, last_layer=self.get_last_layer(), 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.loss[2](qloss2, x2, xrec2, optimizer_idx, self.global_step, last_layer=self.get_last_layer(), 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.loss[1](qloss1, x1, xrec1, optimizer_idx, self.global_step, last_layer=self.get_last_layer(), 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.loss[2](qloss2, x2, xrec2, optimizer_idx, self.global_step, last_layer=self.get_last_layer(), 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): breakpoint() x1 = self.get_input(batch, self.image_key1) x2 = self.get_input(batch, self.image_key2) xrec1, qloss1 = self.forward(x1, model_key=1) xrec2, qloss2 = self.forward(x2, model_key=2) aeloss1, log_dict_ae1 = self.loss[1](qloss1, x1, xrec1, 0, self.global_step, last_layer=self.get_last_layer(model_key=1), split="val") aeloss2, log_dict_ae2 = self.loss[2](qloss2, x2, xrec2, 0, self.global_step, last_layer=self.get_last_layer(model_key=2), split="val") discloss1, log_dict_disc1 = self.loss[1](qloss1, x1, xrec1, 1, self.global_step, last_layer=self.get_last_layer(model_key=1), split="val") discloss2, log_dict_disc2 = self.loss[2](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.encoder[1].parameters())+ list(self.decoder[1].parameters())+ list(self.quantize[1].parameters())+ list(self.quant_conv[1].parameters())+ list(self.post_quant_conv[1].parameters())+ list(self.encoder[2].parameters())+ list(self.decoder[2].parameters())+ list(self.quantize[2].parameters())+ list(self.quant_conv[2].parameters())+ list(self.post_quant_conv[2].parameters()), lr=lr, betas=(0.5, 0.9)) opt_disc = torch.optim.Adam(list(self.loss[1].discriminator.parameters())+ list(self.loss[2].discriminator.parameters()), lr=lr, betas=(0.5, 0.9)) return [opt_ae, opt_disc], [] def get_last_layer(self, model_key): return self.decoder[model_key].conv_out.weight def log_images(self, batch, **kwargs): log = dict() ## log 1 x = self.get_input(batch, self.image_key1) x = x.to(self.device) xrec, _ = self(x) if x.shape[1] > 3: # colorize with random projection assert xrec.shape[1] > 3 x = self.to_rgb(x) xrec = self.to_rgb(xrec) log["inputs1"] = x log["reconstructions1"] = xrec ## log 2 x = self.get_input(batch, self.image_key2) x = x.to(self.device) xrec, _ = self(x) if x.shape[1] > 3: # colorize with random projection assert xrec.shape[1] > 3 x = self.to_rgb(x) xrec = self.to_rgb(xrec) log["inputs2"] = x log["reconstructions2"] = xrec 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