cymic/Waifu_Diffusion_Webui
1
1# this file is copied from CodeFormer repository. Please see comment in modules/codeformer_model.py2 3'''4VQGAN code, adapted from the original created by the Unleashing Transformers authors:5https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py6 7'''8import numpy as np9import torch10import torch.nn as nn11import torch.nn.functional as F12import copy13from basicsr.utils import get_root_logger14from basicsr.utils.registry import ARCH_REGISTRY15 16def normalize(in_channels):17 return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)18 19 20@torch.jit.script21def swish(x):22 return x*torch.sigmoid(x)23 24 25# Define VQVAE classes26class VectorQuantizer(nn.Module):27 def __init__(self, codebook_size, emb_dim, beta):28 super(VectorQuantizer, self).__init__()29 self.codebook_size = codebook_size # number of embeddings30 self.emb_dim = emb_dim # dimension of embedding31 self.beta = beta # commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^232 self.embedding = nn.Embedding(self.codebook_size, self.emb_dim)33 self.embedding.weight.data.uniform_(-1.0 / self.codebook_size, 1.0 / self.codebook_size)34 35 def forward(self, z):36 # reshape z -> (batch, height, width, channel) and flatten37 z = z.permute(0, 2, 3, 1).contiguous()38 z_flattened = z.view(-1, self.emb_dim)39 40 # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z41 d = (z_flattened ** 2).sum(dim=1, keepdim=True) + (self.embedding.weight**2).sum(1) - \42 2 * torch.matmul(z_flattened, self.embedding.weight.t())43 44 mean_distance = torch.mean(d)45 # find closest encodings46 # min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)47 min_encoding_scores, min_encoding_indices = torch.topk(d, 1, dim=1, largest=False)48 # [0-1], higher score, higher confidence49 min_encoding_scores = torch.exp(-min_encoding_scores/10)50 51 min_encodings = torch.zeros(min_encoding_indices.shape[0], self.codebook_size).to(z)52 min_encodings.scatter_(1, min_encoding_indices, 1)53 54 # get quantized latent vectors55 z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)56 # compute loss for embedding57 loss = torch.mean((z_q.detach()-z)**2) + self.beta * torch.mean((z_q - z.detach()) ** 2)58 # preserve gradients59 z_q = z + (z_q - z).detach()60 61 # perplexity62 e_mean = torch.mean(min_encodings, dim=0)63 perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))64 # reshape back to match original input shape65 z_q = z_q.permute(0, 3, 1, 2).contiguous()66 67 return z_q, loss, {68 "perplexity": perplexity,69 "min_encodings": min_encodings,70 "min_encoding_indices": min_encoding_indices,71 "min_encoding_scores": min_encoding_scores,72 "mean_distance": mean_distance73 }74 75 def get_codebook_feat(self, indices, shape):76 # input indices: batch*token_num -> (batch*token_num)*177 # shape: batch, height, width, channel78 indices = indices.view(-1,1)79 min_encodings = torch.zeros(indices.shape[0], self.codebook_size).to(indices)80 min_encodings.scatter_(1, indices, 1)81 # get quantized latent vectors82 z_q = torch.matmul(min_encodings.float(), self.embedding.weight)83 84 if shape is not None: # reshape back to match original input shape85 z_q = z_q.view(shape).permute(0, 3, 1, 2).contiguous()86 87 return z_q88 89 90class GumbelQuantizer(nn.Module):91 def __init__(self, codebook_size, emb_dim, num_hiddens, straight_through=False, kl_weight=5e-4, temp_init=1.0):92 super().__init__()93 self.codebook_size = codebook_size # number of embeddings94 self.emb_dim = emb_dim # dimension of embedding95 self.straight_through = straight_through96 self.temperature = temp_init97 self.kl_weight = kl_weight98 self.proj = nn.Conv2d(num_hiddens, codebook_size, 1) # projects last encoder layer to quantized logits99 self.embed = nn.Embedding(codebook_size, emb_dim)100 101 def forward(self, z):102 hard = self.straight_through if self.training else True103 104 logits = self.proj(z)105 106 soft_one_hot = F.gumbel_softmax(logits, tau=self.temperature, dim=1, hard=hard)107 108 z_q = torch.einsum("b n h w, n d -> b d h w", soft_one_hot, self.embed.weight)109 110 # + kl divergence to the prior loss111 qy = F.softmax(logits, dim=1)112 diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.codebook_size + 1e-10), dim=1).mean()113 min_encoding_indices = soft_one_hot.argmax(dim=1)114 115 return z_q, diff, {116 "min_encoding_indices": min_encoding_indices117 }118 119 120class Downsample(nn.Module):121 def __init__(self, in_channels):122 super().__init__()123 self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)124 125 def forward(self, x):126 pad = (0, 1, 0, 1)127 x = torch.nn.functional.pad(x, pad, mode="constant", value=0)128 x = self.conv(x)129 return x130 131 132class Upsample(nn.Module):133 def __init__(self, in_channels):134 super().__init__()135 self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)136 137 def forward(self, x):138 x = F.interpolate(x, scale_factor=2.0, mode="nearest")139 x = self.conv(x)140 141 return x142 143 144class ResBlock(nn.Module):145 def __init__(self, in_channels, out_channels=None):146 super(ResBlock, self).__init__()147 self.in_channels = in_channels148 self.out_channels = in_channels if out_channels is None else out_channels149 self.norm1 = normalize(in_channels)150 self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)151 self.norm2 = normalize(out_channels)152 self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)153 if self.in_channels != self.out_channels:154 self.conv_out = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)155 156 def forward(self, x_in):157 x = x_in158 x = self.norm1(x)159 x = swish(x)160 x = self.conv1(x)161 x = self.norm2(x)162 x = swish(x)163 x = self.conv2(x)164 if self.in_channels != self.out_channels:165 x_in = self.conv_out(x_in)166 167 return x + x_in168 169 170class AttnBlock(nn.Module):171 def __init__(self, in_channels):172 super().__init__()173 self.in_channels = in_channels174 175 self.norm = normalize(in_channels)176 self.q = torch.nn.Conv2d(177 in_channels,178 in_channels,179 kernel_size=1,180 stride=1,181 padding=0182 )183 self.k = torch.nn.Conv2d(184 in_channels,185 in_channels,186 kernel_size=1,187 stride=1,188 padding=0189 )190 self.v = torch.nn.Conv2d(191 in_channels,192 in_channels,193 kernel_size=1,194 stride=1,195 padding=0196 )197 self.proj_out = torch.nn.Conv2d(198 in_channels,199 in_channels,200 kernel_size=1,201 stride=1,202 padding=0203 )204 205 def forward(self, x):206 h_ = x207 h_ = self.norm(h_)208 q = self.q(h_)209 k = self.k(h_)210 v = self.v(h_)211 212 # compute attention213 b, c, h, w = q.shape214 q = q.reshape(b, c, h*w)215 q = q.permute(0, 2, 1) 216 k = k.reshape(b, c, h*w)217 w_ = torch.bmm(q, k) 218 w_ = w_ * (int(c)**(-0.5))219 w_ = F.softmax(w_, dim=2)220 221 # attend to values222 v = v.reshape(b, c, h*w)223 w_ = w_.permute(0, 2, 1) 224 h_ = torch.bmm(v, w_)225 h_ = h_.reshape(b, c, h, w)226 227 h_ = self.proj_out(h_)228 229 return x+h_230 231 232class Encoder(nn.Module):233 def __init__(self, in_channels, nf, emb_dim, ch_mult, num_res_blocks, resolution, attn_resolutions):234 super().__init__()235 self.nf = nf236 self.num_resolutions = len(ch_mult)237 self.num_res_blocks = num_res_blocks238 self.resolution = resolution239 self.attn_resolutions = attn_resolutions240 241 curr_res = self.resolution242 in_ch_mult = (1,)+tuple(ch_mult)243 244 blocks = []245 # initial convultion246 blocks.append(nn.Conv2d(in_channels, nf, kernel_size=3, stride=1, padding=1))247 248 # residual and downsampling blocks, with attention on smaller res (16x16)249 for i in range(self.num_resolutions):250 block_in_ch = nf * in_ch_mult[i]251 block_out_ch = nf * ch_mult[i]252 for _ in range(self.num_res_blocks):253 blocks.append(ResBlock(block_in_ch, block_out_ch))254 block_in_ch = block_out_ch255 if curr_res in attn_resolutions:256 blocks.append(AttnBlock(block_in_ch))257 258 if i != self.num_resolutions - 1:259 blocks.append(Downsample(block_in_ch))260 curr_res = curr_res // 2261 262 # non-local attention block263 blocks.append(ResBlock(block_in_ch, block_in_ch))264 blocks.append(AttnBlock(block_in_ch))265 blocks.append(ResBlock(block_in_ch, block_in_ch))266 267 # normalise and convert to latent size268 blocks.append(normalize(block_in_ch))269 blocks.append(nn.Conv2d(block_in_ch, emb_dim, kernel_size=3, stride=1, padding=1))270 self.blocks = nn.ModuleList(blocks)271 272 def forward(self, x):273 for block in self.blocks:274 x = block(x)275 276 return x277 278 279class Generator(nn.Module):280 def __init__(self, nf, emb_dim, ch_mult, res_blocks, img_size, attn_resolutions):281 super().__init__()282 self.nf = nf 283 self.ch_mult = ch_mult 284 self.num_resolutions = len(self.ch_mult)285 self.num_res_blocks = res_blocks286 self.resolution = img_size 287 self.attn_resolutions = attn_resolutions288 self.in_channels = emb_dim289 self.out_channels = 3290 block_in_ch = self.nf * self.ch_mult[-1]291 curr_res = self.resolution // 2 ** (self.num_resolutions-1)292 293 blocks = []294 # initial conv295 blocks.append(nn.Conv2d(self.in_channels, block_in_ch, kernel_size=3, stride=1, padding=1))296 297 # non-local attention block298 blocks.append(ResBlock(block_in_ch, block_in_ch))299 blocks.append(AttnBlock(block_in_ch))300 blocks.append(ResBlock(block_in_ch, block_in_ch))301 302 for i in reversed(range(self.num_resolutions)):303 block_out_ch = self.nf * self.ch_mult[i]304 305 for _ in range(self.num_res_blocks):306 blocks.append(ResBlock(block_in_ch, block_out_ch))307 block_in_ch = block_out_ch308 309 if curr_res in self.attn_resolutions:310 blocks.append(AttnBlock(block_in_ch))311 312 if i != 0:313 blocks.append(Upsample(block_in_ch))314 curr_res = curr_res * 2315 316 blocks.append(normalize(block_in_ch))317 blocks.append(nn.Conv2d(block_in_ch, self.out_channels, kernel_size=3, stride=1, padding=1))318 319 self.blocks = nn.ModuleList(blocks)320 321 322 def forward(self, x):323 for block in self.blocks:324 x = block(x)325 326 return x327 328 329@ARCH_REGISTRY.register()330class VQAutoEncoder(nn.Module):331 def __init__(self, img_size, nf, ch_mult, quantizer="nearest", res_blocks=2, attn_resolutions=[16], codebook_size=1024, emb_dim=256,332 beta=0.25, gumbel_straight_through=False, gumbel_kl_weight=1e-8, model_path=None):333 super().__init__()334 logger = get_root_logger()335 self.in_channels = 3 336 self.nf = nf 337 self.n_blocks = res_blocks 338 self.codebook_size = codebook_size339 self.embed_dim = emb_dim340 self.ch_mult = ch_mult341 self.resolution = img_size342 self.attn_resolutions = attn_resolutions343 self.quantizer_type = quantizer344 self.encoder = Encoder(345 self.in_channels,346 self.nf,347 self.embed_dim,348 self.ch_mult,349 self.n_blocks,350 self.resolution,351 self.attn_resolutions352 )353 if self.quantizer_type == "nearest":354 self.beta = beta #0.25355 self.quantize = VectorQuantizer(self.codebook_size, self.embed_dim, self.beta)356 elif self.quantizer_type == "gumbel":357 self.gumbel_num_hiddens = emb_dim358 self.straight_through = gumbel_straight_through359 self.kl_weight = gumbel_kl_weight360 self.quantize = GumbelQuantizer(361 self.codebook_size,362 self.embed_dim,363 self.gumbel_num_hiddens,364 self.straight_through,365 self.kl_weight366 )367 self.generator = Generator(368 self.nf, 369 self.embed_dim,370 self.ch_mult, 371 self.n_blocks, 372 self.resolution, 373 self.attn_resolutions374 )375 376 if model_path is not None:377 chkpt = torch.load(model_path, map_location='cpu')378 if 'params_ema' in chkpt:379 self.load_state_dict(torch.load(model_path, map_location='cpu')['params_ema'])380 logger.info(f'vqgan is loaded from: {model_path} [params_ema]')381 elif 'params' in chkpt:382 self.load_state_dict(torch.load(model_path, map_location='cpu')['params'])383 logger.info(f'vqgan is loaded from: {model_path} [params]')384 else:385 raise ValueError(f'Wrong params!')386 387 388 def forward(self, x):389 x = self.encoder(x)390 quant, codebook_loss, quant_stats = self.quantize(x)391 x = self.generator(quant)392 return x, codebook_loss, quant_stats393 394 395 396# patch based discriminator397@ARCH_REGISTRY.register()398class VQGANDiscriminator(nn.Module):399 def __init__(self, nc=3, ndf=64, n_layers=4, model_path=None):400 super().__init__()401 402 layers = [nn.Conv2d(nc, ndf, kernel_size=4, stride=2, padding=1), nn.LeakyReLU(0.2, True)]403 ndf_mult = 1404 ndf_mult_prev = 1405 for n in range(1, n_layers): # gradually increase the number of filters406 ndf_mult_prev = ndf_mult407 ndf_mult = min(2 ** n, 8)408 layers += [409 nn.Conv2d(ndf * ndf_mult_prev, ndf * ndf_mult, kernel_size=4, stride=2, padding=1, bias=False),410 nn.BatchNorm2d(ndf * ndf_mult),411 nn.LeakyReLU(0.2, True)412 ]413 414 ndf_mult_prev = ndf_mult415 ndf_mult = min(2 ** n_layers, 8)416 417 layers += [418 nn.Conv2d(ndf * ndf_mult_prev, ndf * ndf_mult, kernel_size=4, stride=1, padding=1, bias=False),419 nn.BatchNorm2d(ndf * ndf_mult),420 nn.LeakyReLU(0.2, True)421 ]422 423 layers += [424 nn.Conv2d(ndf * ndf_mult, 1, kernel_size=4, stride=1, padding=1)] # output 1 channel prediction map425 self.main = nn.Sequential(*layers)426 427 if model_path is not None:428 chkpt = torch.load(model_path, map_location='cpu')429 if 'params_d' in chkpt:430 self.load_state_dict(torch.load(model_path, map_location='cpu')['params_d'])431 elif 'params' in chkpt:432 self.load_state_dict(torch.load(model_path, map_location='cpu')['params'])433 else:434 raise ValueError(f'Wrong params!')435 436 def forward(self, x):437 return self.main(x)