Spaces:
Running
Running
| 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 |