| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from mapping_network import AdaptiveInstanceNorm, MappingNetowrk |
| from helpers.imle_helpers import get_1x1 |
| from collections import defaultdict |
| from rtm_core import RTMMappingNetwork |
| import numpy as np |
| import itertools |
|
|
| def parse_layer_string(s): |
| layers = [] |
| for ss in s.split(','): |
| if 'x' in ss: |
| res, num = ss.split('x') |
| count = int(num) |
| layers += [(int(res), None) for _ in range(count)] |
| elif 'm' in ss: |
| res, mixin = [int(a) for a in ss.split('m')] |
| layers.append((res, mixin)) |
| elif 'd' in ss: |
| res, down_rate = [int(a) for a in ss.split('d')] |
| layers.append((res, down_rate)) |
| else: |
| res = int(ss) |
| layers.append((res, None)) |
| return layers |
|
|
| def get_width_settings(width, s): |
| mapping = defaultdict(lambda: width) |
| if s: |
| s = s.split(',') |
| for ss in s: |
| k, v = ss.split(':') |
| mapping[int(k)] = int(v) |
| return mapping |
|
|
| class SEBlock(nn.Module): |
| def __init__(self, channels, reduction=16): |
| super().__init__() |
| self.avg_pool = nn.AdaptiveAvgPool2d(1) |
| self.fc = nn.Sequential( |
| nn.Linear(channels, channels // reduction, bias=False), |
| nn.ReLU(inplace=True), |
| nn.Linear(channels // reduction, channels, bias=False), |
| nn.Sigmoid() |
| ) |
|
|
| def forward(self, x): |
| b, c, _, _ = x.size() |
| y = self.avg_pool(x).view(b, c) |
| y = self.fc(y).view(b, c, 1, 1) |
| return x * y.expand_as(x) |
|
|
|
|
| class ConvNeXtBlock(nn.Module): |
| def __init__(self, dim, H, expansion=4, kernel_size=7, use_se=True, reduction=16, dropout=0.0): |
| super().__init__() |
| self.dw_conv = nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size//2, groups=dim) |
|
|
| if(H.convnext_norm == 'layernorm'): |
| self.norm = nn.LayerNorm(dim, eps=H.convnext_norm_eps) |
| elif(H.convnext_norm == 'rmsnorm'): |
| self.norm = nn.RMSNorm(dim, eps=H.convnext_norm_eps) |
| |
| self.pw_conv1 = nn.Linear(dim, expansion * dim) |
| self.gelu = nn.GELU() |
| self.pw_conv2 = nn.Linear(expansion * dim, dim) |
|
|
| |
| self.use_se = use_se |
| if use_se: |
| self.se = SEBlock(dim, reduction=reduction) |
| else: |
| |
| self.se = nn.Identity() |
|
|
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, m): |
| if isinstance(m, (nn.Conv2d, nn.Linear)): |
| |
| if m.bias is not None: |
| nn.init.constant_(m.bias, 0) |
|
|
| |
| def forward(self, x): |
| |
| x = self.dw_conv(x) |
| |
| x = x.permute(0, 2, 3, 1) |
| x = self.norm(x) |
| x = self.pw_conv1(x) |
| x = self.gelu(x) |
| x = self.pw_conv2(x) |
| x = x.permute(0, 3, 1, 2) |
|
|
| x = self.se(x) |
|
|
| return x |
|
|
| class DecBlock(nn.Module): |
| def __init__(self, H, res, mixin, n_blocks): |
| super().__init__() |
| self.base = res |
| self.mixin = mixin |
| self.H = H |
| self.widths = get_width_settings(H.width, H.custom_width_str) |
| width = self.widths[res] |
|
|
| if mixin is not None and self.widths[mixin] != width: |
| self.proj = get_1x1(self.widths[mixin], width) |
| else: |
| self.proj = nn.Identity() |
|
|
| self.adaIN = AdaptiveInstanceNorm(width, H.latent_dim) |
| self.resnet = ConvNeXtBlock(width, H, kernel_size=7, |
| expansion=H.convnext_expansion, |
| use_se=H.use_se, |
| reduction=H.se_reduction, |
| dropout=H.dropout_p) |
|
|
| self.residual_ratio = nn.Parameter(torch.tensor(H.residual_ratio)) |
| self.residual_type = H.residual_type |
| self.sigmoid = nn.Sigmoid() |
|
|
|
|
| def forward(self, x, w): |
| if self.mixin is not None: |
| x = F.interpolate(x, scale_factor=self.base / self.mixin, mode='bicubic') |
| x = self.proj(x) |
| |
| residual = x |
| x = self.adaIN(x, w) |
| x = self.resnet(x) |
|
|
| if self.residual_type == 'normal': |
| return x * self.sigmoid(self.residual_ratio) + residual |
| |
| elif self.residual_type == 'convex': |
| return x * self.sigmoid(self.residual_ratio) + residual * (1 - self.sigmoid(self.residual_ratio)) |
|
|
| def stopgrad_keep_graph(x): |
| return x.detach() + 0.0 * x |
|
|
| class Decoder(nn.Module): |
| def __init__(self, H): |
| super().__init__() |
| self.H = H |
| self.use_rtm = getattr(H, "use_rtm", False) |
|
|
| if self.use_rtm: |
| self.mapping_network = RTMMappingNetwork( |
| code_dim=H.latent_dim, |
| num_tokens=getattr(H, "num_tokens", 1), |
| H_cycles=getattr(H, "H_cycles", 1), |
| L_cycles=getattr(H, "L_cycles", 1), |
| H_layers=getattr(H, "H_layers", 2), |
| L_layers=getattr(H, "L_layers", 2), |
| hidden_size=getattr(H, "rtm_hidden_size", 256), |
| expansion=getattr(H, "rtm_expansion", 4.0), |
| refinement_steps=getattr(H, "refinement_steps", 1), |
| with_grad=getattr(H, "rtm_with_grad", False), |
| cycle_noise_std=getattr(H, "rtm_cycle_noise_std", 0.0), |
| ) |
| else: |
| self.mapping_network = MappingNetowrk( |
| code_dim=H.latent_dim, n_mlp=H.n_mpl, |
| mapping_lr_multiplier=getattr(H, "mapping_lr_multiplier", 1.0)) |
| resos = set() |
| dec_blocks = [] |
| self.widths = get_width_settings(H.width, H.custom_width_str) |
| blocks = parse_layer_string(H.dec_blocks) |
| for idx, (res, mixin) in enumerate(blocks): |
| dec_blocks.append(DecBlock(H, res, mixin, n_blocks=len(blocks))) |
| resos.add(res) |
| self.resolutions = sorted(resos) |
| self.dec_blocks = nn.ModuleList(dec_blocks) |
| first_res = self.resolutions[0] |
| last_res = self.resolutions[-1] |
| self.constant = nn.Parameter(torch.randn(1, self.widths[first_res], first_res, first_res)) |
| resnets = {} |
|
|
| for res in self.resolutions: |
| key = str(res) |
|
|
| if res < 8: |
| resnets[key] = nn.Identity() |
| else: |
| resnets[key] = get_1x1(self.widths[res], H.image_channels) |
|
|
|
|
| self.resnets = nn.ModuleDict(resnets) |
| self.gains = nn.Parameter(torch.ones(1, H.image_channels, 1, 1)) |
| self.biases = nn.Parameter(torch.zeros(1, H.image_channels, 1, 1)) |
|
|
|
|
| def forward(self, latent_code, spatial_noise=None, input_is_w=False, train=False): |
| if not input_is_w: |
| w = self.mapping_network(latent_code) |
| if isinstance(w, tuple) or isinstance(w, list): |
| w = w[0] |
| else: |
| w = latent_code |
| targets = [] |
| x = self.constant.repeat(latent_code.shape[0], 1, 1, 1) |
|
|
| for idx, block in enumerate(self.dec_blocks): |
| if(block.mixin is not None): |
| intermediate = self.resnets[str(block.mixin)](x) |
| targets.append(intermediate) |
| if(block.mixin >= 8 and self.H.use_stopgrad_for_intermediate): |
| x = x.detach() |
| x = block(x, w) |
| x = self.resnets[str(self.resolutions[-1])](x) |
| x = self.gains * x + self.biases |
| targets.append(x) |
| if(train): |
| return targets |
| else: |
| return targets[-1] |
|
|
| class IMLE(nn.Module): |
| def __init__(self, H): |
| super().__init__() |
| self.decoder = Decoder(H) |
|
|
| def forward(self, latents, spatial_noise=None, input_is_w=False, train=False): |
| return self.decoder.forward(latents, spatial_noise, input_is_w, train) |
|
|
