ai-art

  1import torch
  2import torch.nn as nn
  3
  4from core.taming.modules.diffusion import Encoder, Decoder
  5from core.taming.modules.vqvae import VectorQuantizer
  6from core.taming.modules.losses import VQLPIPSWithDiscriminator, DummyLoss
  7
  8from core.utils.loader import safe_load
  9
 10
 11class VQModel(nn.Module):
 12    def __init__(self,
 13                 ddconfig,
 14                 n_embed,
 15                 embed_dim,
 16                 lossconfig=None,
 17                 ckpt_path=None,
 18                 model_dir=None,
 19                 ignore_keys=[],
 20                 image_key="image",
 21                 colorize_nlabels=None,
 22                 monitor=None,
 23                 remap=None,
 24                 sane_index_shape=False,  # tell vector quantizer to return indices as bhw
 25                 ):
 26        super().__init__()
 27        self.image_key = image_key
 28
 29        self.encoder = Encoder(**ddconfig)
 30        self.decoder = Decoder(**ddconfig)
 31
 32        self.loss = DummyLoss()
 33        if lossconfig is not None:
 34            self.loss = VQLPIPSWithDiscriminator(model_dir=model_dir, **lossconfig["params"])
 35
 36        self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
 37                                        remap=remap, sane_index_shape=sane_index_shape)
 38        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
 39        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
 40
 41        if ckpt_path is not None:
 42            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
 43
 44        self.image_key = image_key
 45
 46        if colorize_nlabels is not None:
 47            assert type(colorize_nlabels) == int
 48            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
 49        if monitor is not None:
 50            self.monitor = monitor
 51
 52    def init_from_ckpt(self, path, ignore_keys=list()):
 53        try:
 54            sd = torch.load(path, map_location="cpu")["state_dict"]
 55        except Exception:
 56            sd = safe_load(path, map_location="cpu")["state_dict"]
 57
 58        keys = list(sd.keys())
 59        for k in keys:
 60            for ik in ignore_keys:
 61                if k.startswith(ik):
 62                    print("Deleting key {} from state_dict.".format(k))
 63                    del sd[k]
 64
 65        if "first_stage_model.encoder.conv_in.weight" in sd:
 66            stripped_state_dict = {}
 67            for key in sd:
 68                if key.startswith("first_stage_model."):
 69                    stripped_state_dict[key[18:]] = sd[key]
 70            sd = stripped_state_dict
 71
 72        self.load_state_dict(sd, strict=False)
 73        print(f"Restored from {path}")
 74
 75    def encode(self, x):
 76        h = self.encoder(x)
 77        h = self.quant_conv(h)
 78        quant, emb_loss, info = self.quantize(h)
 79        return quant, emb_loss, info
 80
 81    def decode(self, quant):
 82        quant = self.post_quant_conv(quant)
 83        dec = self.decoder(quant)
 84        return dec
 85
 86    def decode_code(self, code_b):
 87        quant_b = self.quantize.embed_code(code_b)
 88        dec = self.decode(quant_b)
 89        return dec
 90
 91    def forward(self, input):
 92        quant, diff, _ = self.encode(input)
 93        dec = self.decode(quant)
 94        return dec, diff
 95
 96    def get_input(self, batch, device):
 97        x = batch
 98        if len(x.shape) == 3:
 99            x = x[..., None]
100        x = x.to(device, memory_format=torch.contiguous_format)
101        return x.float()
102
103    def training_step(self, batch, batch_idx, optimizer_idx, device='cpu'):
104        x = self.get_input(batch, device)
105        xrec, qloss = self(x)
106
107        if optimizer_idx == 0:
108            # autoencode
109            aeloss = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, last_layer=self.get_last_layer(), split="train")
110            return aeloss
111
112        if optimizer_idx == 1:
113            # discriminator
114            discloss = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, last_layer=self.get_last_layer(), split="train")
115            return discloss
116
117    def configure_optimizers(self):
118        lr = self.learning_rate
119        opt_ae = torch.optim.Adam(list(self.encoder.parameters()) +
120                                  list(self.decoder.parameters()) +
121                                  list(self.quantize.parameters()) +
122                                  list(self.quant_conv.parameters()) +
123                                  list(self.post_quant_conv.parameters()),
124                                  lr=lr, betas=(0.5, 0.9))
125        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
126                                    lr=lr, betas=(0.5, 0.9))
127        return [opt_ae, opt_disc], []
128
129    def get_last_layer(self):
130        return self.decoder.conv_out.weight