ldm/models/vqgan_dual.py

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