memef4rmer/edit_anything

1import torch2import pytorch_lightning as pl3import torch.nn.functional as F4from contextlib import contextmanager5 6from ldm.modules.diffusionmodules.model import Encoder, Decoder7from ldm.modules.distributions.distributions import DiagonalGaussianDistribution8 9from ldm.util import instantiate_from_config10from ldm.modules.ema import LitEma11 12 13class AutoencoderKL(pl.LightningModule):14    def __init__(self,15                 ddconfig,16                 lossconfig,17                 embed_dim,18                 ckpt_path=None,19                 ignore_keys=[],20                 image_key="image",21                 colorize_nlabels=None,22                 monitor=None,23                 ema_decay=None,24                 learn_logvar=False25                 ):26        super().__init__()27        self.learn_logvar = learn_logvar28        self.image_key = image_key29        self.encoder = Encoder(**ddconfig)30        self.decoder = Decoder(**ddconfig)31        self.loss = instantiate_from_config(lossconfig)32        assert ddconfig["double_z"]33        self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)34        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)35        self.embed_dim = embed_dim36        if colorize_nlabels is not None:37            assert type(colorize_nlabels)==int38            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))39        if monitor is not None:40            self.monitor = monitor41 42        self.use_ema = ema_decay is not None43        if self.use_ema:44            self.ema_decay = ema_decay45            assert 0. < ema_decay < 1.46            self.model_ema = LitEma(self, decay=ema_decay)47            print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")48 49        if ckpt_path is not None:50            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)51 52    def init_from_ckpt(self, path, ignore_keys=list()):53        sd = torch.load(path, map_location="cpu")["state_dict"]54        keys = list(sd.keys())55        for k in keys:56            for ik in ignore_keys:57                if k.startswith(ik):58                    print("Deleting key {} from state_dict.".format(k))59                    del sd[k]60        self.load_state_dict(sd, strict=False)61        print(f"Restored from {path}")62 63    @contextmanager64    def ema_scope(self, context=None):65        if self.use_ema:66            self.model_ema.store(self.parameters())67            self.model_ema.copy_to(self)68            if context is not None:69                print(f"{context}: Switched to EMA weights")70        try:71            yield None72        finally:73            if self.use_ema:74                self.model_ema.restore(self.parameters())75                if context is not None:76                    print(f"{context}: Restored training weights")77 78    def on_train_batch_end(self, *args, **kwargs):79        if self.use_ema:80            self.model_ema(self)81 82    def encode(self, x):83        h = self.encoder(x)84        moments = self.quant_conv(h)85        posterior = DiagonalGaussianDistribution(moments)86        return posterior87 88    def decode(self, z):89        z = self.post_quant_conv(z)90        dec = self.decoder(z)91        return dec92 93    def forward(self, input, sample_posterior=True):94        posterior = self.encode(input)95        if sample_posterior:96            z = posterior.sample()97        else:98            z = posterior.mode()99        dec = self.decode(z)100        return dec, posterior101 102    def get_input(self, batch, k):103        x = batch[k]104        if len(x.shape) == 3:105            x = x[..., None]106        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()107        return x108 109    def training_step(self, batch, batch_idx, optimizer_idx):110        inputs = self.get_input(batch, self.image_key)111        reconstructions, posterior = self(inputs)112 113        if optimizer_idx == 0:114            # train encoder+decoder+logvar115            aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,116                                            last_layer=self.get_last_layer(), split="train")117            self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)118            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)119            return aeloss120 121        if optimizer_idx == 1:122            # train the discriminator123            discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,124                                                last_layer=self.get_last_layer(), split="train")125 126            self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)127            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)128            return discloss129 130    def validation_step(self, batch, batch_idx):131        log_dict = self._validation_step(batch, batch_idx)132        with self.ema_scope():133            log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")134        return log_dict135 136    def _validation_step(self, batch, batch_idx, postfix=""):137        inputs = self.get_input(batch, self.image_key)138        reconstructions, posterior = self(inputs)139        aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,140                                        last_layer=self.get_last_layer(), split="val"+postfix)141 142        discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,143                                            last_layer=self.get_last_layer(), split="val"+postfix)144 145        self.log(f"val{postfix}/rec_loss", log_dict_ae[f"val{postfix}/rec_loss"])146        self.log_dict(log_dict_ae)147        self.log_dict(log_dict_disc)148        return self.log_dict149 150    def configure_optimizers(self):151        lr = self.learning_rate152        ae_params_list = list(self.encoder.parameters()) + list(self.decoder.parameters()) + list(153            self.quant_conv.parameters()) + list(self.post_quant_conv.parameters())154        if self.learn_logvar:155            print(f"{self.__class__.__name__}: Learning logvar")156            ae_params_list.append(self.loss.logvar)157        opt_ae = torch.optim.Adam(ae_params_list,158                                  lr=lr, betas=(0.5, 0.9))159        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),160                                    lr=lr, betas=(0.5, 0.9))161        return [opt_ae, opt_disc], []162 163    def get_last_layer(self):164        return self.decoder.conv_out.weight165 166    @torch.no_grad()167    def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):168        log = dict()169        x = self.get_input(batch, self.image_key)170        x = x.to(self.device)171        if not only_inputs:172            xrec, posterior = self(x)173            if x.shape[1] > 3:174                # colorize with random projection175                assert xrec.shape[1] > 3176                x = self.to_rgb(x)177                xrec = self.to_rgb(xrec)178            log["samples"] = self.decode(torch.randn_like(posterior.sample()))179            log["reconstructions"] = xrec180            if log_ema or self.use_ema:181                with self.ema_scope():182                    xrec_ema, posterior_ema = self(x)183                    if x.shape[1] > 3:184                        # colorize with random projection185                        assert xrec_ema.shape[1] > 3186                        xrec_ema = self.to_rgb(xrec_ema)187                    log["samples_ema"] = self.decode(torch.randn_like(posterior_ema.sample()))188                    log["reconstructions_ema"] = xrec_ema189        log["inputs"] = x190        return log191 192    def to_rgb(self, x):193        assert self.image_key == "segmentation"194        if not hasattr(self, "colorize"):195            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))196        x = F.conv2d(x, weight=self.colorize)197        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.198        return x199 200 201class IdentityFirstStage(torch.nn.Module):202    def __init__(self, *args, vq_interface=False, **kwargs):203        self.vq_interface = vq_interface204        super().__init__()205 206    def encode(self, x, *args, **kwargs):207        return x208 209    def decode(self, x, *args, **kwargs):210        return x211 212    def quantize(self, x, *args, **kwargs):213        if self.vq_interface:214            return x, None, [None, None, None]215        return x216 217    def forward(self, x, *args, **kwargs):218        return x219 220