Shellbrady/LivePortrait5
0
1# coding: utf-82 3"""4This file defines various neural network modules and utility functions, including convolutional and residual blocks,5normalizations, and functions for spatial transformation and tensor manipulation.6"""7 8from torch import nn9import torch.nn.functional as F10import torch11import torch.nn.utils.spectral_norm as spectral_norm12import math13import warnings14 15 16def kp2gaussian(kp, spatial_size, kp_variance):17 """18 Transform a keypoint into gaussian like representation19 """20 mean = kp21 22 coordinate_grid = make_coordinate_grid(spatial_size, mean)23 number_of_leading_dimensions = len(mean.shape) - 124 shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape25 coordinate_grid = coordinate_grid.view(*shape)26 repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 1)27 coordinate_grid = coordinate_grid.repeat(*repeats)28 29 # Preprocess kp shape30 shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3)31 mean = mean.view(*shape)32 33 mean_sub = (coordinate_grid - mean)34 35 out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance)36 37 return out38 39 40def make_coordinate_grid(spatial_size, ref, **kwargs):41 d, h, w = spatial_size42 x = torch.arange(w).type(ref.dtype).to(ref.device)43 y = torch.arange(h).type(ref.dtype).to(ref.device)44 z = torch.arange(d).type(ref.dtype).to(ref.device)45 46 # NOTE: must be right-down-in47 x = (2 * (x / (w - 1)) - 1) # the x axis faces to the right48 y = (2 * (y / (h - 1)) - 1) # the y axis faces to the bottom49 z = (2 * (z / (d - 1)) - 1) # the z axis faces to the inner50 51 yy = y.view(1, -1, 1).repeat(d, 1, w)52 xx = x.view(1, 1, -1).repeat(d, h, 1)53 zz = z.view(-1, 1, 1).repeat(1, h, w)54 55 meshed = torch.cat([xx.unsqueeze_(3), yy.unsqueeze_(3), zz.unsqueeze_(3)], 3)56 57 return meshed58 59 60class ConvT2d(nn.Module):61 """62 Upsampling block for use in decoder.63 """64 65 def __init__(self, in_features, out_features, kernel_size=3, stride=2, padding=1, output_padding=1):66 super(ConvT2d, self).__init__()67 68 self.convT = nn.ConvTranspose2d(in_features, out_features, kernel_size=kernel_size, stride=stride,69 padding=padding, output_padding=output_padding)70 self.norm = nn.InstanceNorm2d(out_features)71 72 def forward(self, x):73 out = self.convT(x)74 out = self.norm(out)75 out = F.leaky_relu(out)76 return out77 78 79class ResBlock3d(nn.Module):80 """81 Res block, preserve spatial resolution.82 """83 84 def __init__(self, in_features, kernel_size, padding):85 super(ResBlock3d, self).__init__()86 self.conv1 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding)87 self.conv2 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding)88 self.norm1 = nn.BatchNorm3d(in_features, affine=True)89 self.norm2 = nn.BatchNorm3d(in_features, affine=True)90 91 def forward(self, x):92 out = self.norm1(x)93 out = F.relu(out)94 out = self.conv1(out)95 out = self.norm2(out)96 out = F.relu(out)97 out = self.conv2(out)98 out += x99 return out100 101 102class UpBlock3d(nn.Module):103 """104 Upsampling block for use in decoder.105 """106 107 def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):108 super(UpBlock3d, self).__init__()109 110 self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,111 padding=padding, groups=groups)112 self.norm = nn.BatchNorm3d(out_features, affine=True)113 114 def forward(self, x):115 out = F.interpolate(x, scale_factor=(1, 2, 2))116 out = self.conv(out)117 out = self.norm(out)118 out = F.relu(out)119 return out120 121 122class DownBlock2d(nn.Module):123 """124 Downsampling block for use in encoder.125 """126 127 def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):128 super(DownBlock2d, self).__init__()129 self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups)130 self.norm = nn.BatchNorm2d(out_features, affine=True)131 self.pool = nn.AvgPool2d(kernel_size=(2, 2))132 133 def forward(self, x):134 out = self.conv(x)135 out = self.norm(out)136 out = F.relu(out)137 out = self.pool(out)138 return out139 140 141class DownBlock3d(nn.Module):142 """143 Downsampling block for use in encoder.144 """145 146 def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):147 super(DownBlock3d, self).__init__()148 '''149 self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,150 padding=padding, groups=groups, stride=(1, 2, 2))151 '''152 self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,153 padding=padding, groups=groups)154 self.norm = nn.BatchNorm3d(out_features, affine=True)155 self.pool = nn.AvgPool3d(kernel_size=(1, 2, 2))156 157 def forward(self, x):158 out = self.conv(x)159 out = self.norm(out)160 out = F.relu(out)161 out = self.pool(out)162 return out163 164 165class SameBlock2d(nn.Module):166 """167 Simple block, preserve spatial resolution.168 """169 170 def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False):171 super(SameBlock2d, self).__init__()172 self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups)173 self.norm = nn.BatchNorm2d(out_features, affine=True)174 if lrelu:175 self.ac = nn.LeakyReLU()176 else:177 self.ac = nn.ReLU()178 179 def forward(self, x):180 out = self.conv(x)181 out = self.norm(out)182 out = self.ac(out)183 return out184 185 186class Encoder(nn.Module):187 """188 Hourglass Encoder189 """190 191 def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):192 super(Encoder, self).__init__()193 194 down_blocks = []195 for i in range(num_blocks):196 down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), min(max_features, block_expansion * (2 ** (i + 1))), kernel_size=3, padding=1))197 self.down_blocks = nn.ModuleList(down_blocks)198 199 def forward(self, x):200 outs = [x]201 for down_block in self.down_blocks:202 outs.append(down_block(outs[-1]))203 return outs204 205 206class Decoder(nn.Module):207 """208 Hourglass Decoder209 """210 211 def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):212 super(Decoder, self).__init__()213 214 up_blocks = []215 216 for i in range(num_blocks)[::-1]:217 in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1)))218 out_filters = min(max_features, block_expansion * (2 ** i))219 up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1))220 221 self.up_blocks = nn.ModuleList(up_blocks)222 self.out_filters = block_expansion + in_features223 224 self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1)225 self.norm = nn.BatchNorm3d(self.out_filters, affine=True)226 227 def forward(self, x):228 out = x.pop()229 for up_block in self.up_blocks:230 out = up_block(out)231 skip = x.pop()232 out = torch.cat([out, skip], dim=1)233 out = self.conv(out)234 out = self.norm(out)235 out = F.relu(out)236 return out237 238 239class Hourglass(nn.Module):240 """241 Hourglass architecture.242 """243 244 def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):245 super(Hourglass, self).__init__()246 self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features)247 self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features)248 self.out_filters = self.decoder.out_filters249 250 def forward(self, x):251 return self.decoder(self.encoder(x))252 253 254class SPADE(nn.Module):255 def __init__(self, norm_nc, label_nc):256 super().__init__()257 258 self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)259 nhidden = 128260 261 self.mlp_shared = nn.Sequential(262 nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1),263 nn.ReLU())264 self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)265 self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)266 267 def forward(self, x, segmap):268 normalized = self.param_free_norm(x)269 segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')270 actv = self.mlp_shared(segmap)271 gamma = self.mlp_gamma(actv)272 beta = self.mlp_beta(actv)273 out = normalized * (1 + gamma) + beta274 return out275 276 277class SPADEResnetBlock(nn.Module):278 def __init__(self, fin, fout, norm_G, label_nc, use_se=False, dilation=1):279 super().__init__()280 # Attributes281 self.learned_shortcut = (fin != fout)282 fmiddle = min(fin, fout)283 self.use_se = use_se284 # create conv layers285 self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=dilation, dilation=dilation)286 self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=dilation, dilation=dilation)287 if self.learned_shortcut:288 self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)289 # apply spectral norm if specified290 if 'spectral' in norm_G:291 self.conv_0 = spectral_norm(self.conv_0)292 self.conv_1 = spectral_norm(self.conv_1)293 if self.learned_shortcut:294 self.conv_s = spectral_norm(self.conv_s)295 # define normalization layers296 self.norm_0 = SPADE(fin, label_nc)297 self.norm_1 = SPADE(fmiddle, label_nc)298 if self.learned_shortcut:299 self.norm_s = SPADE(fin, label_nc)300 301 def forward(self, x, seg1):302 x_s = self.shortcut(x, seg1)303 dx = self.conv_0(self.actvn(self.norm_0(x, seg1)))304 dx = self.conv_1(self.actvn(self.norm_1(dx, seg1)))305 out = x_s + dx306 return out307 308 def shortcut(self, x, seg1):309 if self.learned_shortcut:310 x_s = self.conv_s(self.norm_s(x, seg1))311 else:312 x_s = x313 return x_s314 315 def actvn(self, x):316 return F.leaky_relu(x, 2e-1)317 318 319def filter_state_dict(state_dict, remove_name='fc'):320 new_state_dict = {}321 for key in state_dict:322 if remove_name in key:323 continue324 new_state_dict[key] = state_dict[key]325 return new_state_dict326 327 328class GRN(nn.Module):329 """ GRN (Global Response Normalization) layer330 """331 332 def __init__(self, dim):333 super().__init__()334 self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))335 self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))336 337 def forward(self, x):338 Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)339 Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)340 return self.gamma * (x * Nx) + self.beta + x341 342 343class LayerNorm(nn.Module):344 r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.345 The ordering of the dimensions in the inputs. channels_last corresponds to inputs with346 shape (batch_size, height, width, channels) while channels_first corresponds to inputs347 with shape (batch_size, channels, height, width).348 """349 350 def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):351 super().__init__()352 self.weight = nn.Parameter(torch.ones(normalized_shape))353 self.bias = nn.Parameter(torch.zeros(normalized_shape))354 self.eps = eps355 self.data_format = data_format356 if self.data_format not in ["channels_last", "channels_first"]:357 raise NotImplementedError358 self.normalized_shape = (normalized_shape, )359 360 def forward(self, x):361 if self.data_format == "channels_last":362 return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)363 elif self.data_format == "channels_first":364 u = x.mean(1, keepdim=True)365 s = (x - u).pow(2).mean(1, keepdim=True)366 x = (x - u) / torch.sqrt(s + self.eps)367 x = self.weight[:, None, None] * x + self.bias[:, None, None]368 return x369 370 371def _no_grad_trunc_normal_(tensor, mean, std, a, b):372 # Cut & paste from PyTorch official master until it's in a few official releases - RW373 # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf374 def norm_cdf(x):375 # Computes standard normal cumulative distribution function376 return (1. + math.erf(x / math.sqrt(2.))) / 2.377 378 if (mean < a - 2 * std) or (mean > b + 2 * std):379 warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "380 "The distribution of values may be incorrect.",381 stacklevel=2)382 383 with torch.no_grad():384 # Values are generated by using a truncated uniform distribution and385 # then using the inverse CDF for the normal distribution.386 # Get upper and lower cdf values387 l = norm_cdf((a - mean) / std)388 u = norm_cdf((b - mean) / std)389 390 # Uniformly fill tensor with values from [l, u], then translate to391 # [2l-1, 2u-1].392 tensor.uniform_(2 * l - 1, 2 * u - 1)393 394 # Use inverse cdf transform for normal distribution to get truncated395 # standard normal396 tensor.erfinv_()397 398 # Transform to proper mean, std399 tensor.mul_(std * math.sqrt(2.))400 tensor.add_(mean)401 402 # Clamp to ensure it's in the proper range403 tensor.clamp_(min=a, max=b)404 return tensor405 406 407def drop_path(x, drop_prob=0., training=False, scale_by_keep=True):408 """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).409 410 This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,411 the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...412 See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for413 changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use414 'survival rate' as the argument.415 416 """417 if drop_prob == 0. or not training:418 return x419 keep_prob = 1 - drop_prob420 shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets421 random_tensor = x.new_empty(shape).bernoulli_(keep_prob)422 if keep_prob > 0.0 and scale_by_keep:423 random_tensor.div_(keep_prob)424 return x * random_tensor425 426 427class DropPath(nn.Module):428 """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).429 """430 431 def __init__(self, drop_prob=None, scale_by_keep=True):432 super(DropPath, self).__init__()433 self.drop_prob = drop_prob434 self.scale_by_keep = scale_by_keep435 436 def forward(self, x):437 return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)438 439 440def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):441 return _no_grad_trunc_normal_(tensor, mean, std, a, b)442