ldm/models/vqgan_dual_non_dict.py

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