Spaces:
Running on Zero
Running on Zero
| from functools import partial | |
| from itertools import repeat | |
| import collections.abc | |
| import torch | |
| import torch.nn as nn | |
| def _ntuple(n): | |
| def parse(x): | |
| if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): | |
| return tuple(x) | |
| return tuple(repeat(x, n)) | |
| return parse | |
| to_1tuple = _ntuple(1) | |
| to_2tuple = _ntuple(2) | |
| to_3tuple = _ntuple(3) | |
| to_4tuple = _ntuple(4) | |
| to_ntuple = _ntuple | |
| def get_conv(in_ch, out_ch, kernel_size, stride, padding, zero_bias=True, zero_weights=False): | |
| conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride, padding) | |
| if zero_bias: | |
| conv.bias.data.mul_(0.0) | |
| if zero_weights: | |
| conv.weight.data.mul_(0.0) | |
| return conv | |
| def conv_k1s1(in_ch, out_ch, zero_bias=True, zero_weights=False): | |
| return get_conv(in_ch, out_ch, 1, 1, 0, zero_bias, zero_weights) | |
| def conv_k3s1(in_ch, out_ch, zero_bias=True, zero_weights=False): | |
| return get_conv(in_ch, out_ch, 3, 1, 1, zero_bias, zero_weights) | |
| def patch_downsample(in_ch, out_ch, rate=2): | |
| return get_conv(in_ch, out_ch, kernel_size=rate, stride=rate, padding=0) | |
| def patch_upsample(in_ch, out_ch, rate=2): | |
| conv = nn.Sequential( | |
| get_conv(in_ch, out_ch * (rate ** 2), kernel_size=1, stride=1, padding=0), | |
| nn.PixelShuffle(rate) | |
| ) | |
| return conv | |
| class SetKey(nn.Module): | |
| """ A dummy layer that is used to mark the position of a layer in the network. | |
| """ | |
| def __init__(self, key): | |
| super().__init__() | |
| self.key = key | |
| def forward(self, x): | |
| return x | |
| class CompressionStopFlag(nn.Module): | |
| """ A dummy layer that is used to mark the stop position of encoding bits. | |
| """ | |
| def __init__(self): | |
| super().__init__() | |
| def forward(self, x): | |
| return x | |
| class FeatureExtractorWithEmbedding(nn.Module): | |
| def __init__(self, blocks): | |
| super().__init__() | |
| self.enc_blocks = nn.ModuleList(blocks) | |
| def forward(self, x, emb=None): | |
| features = {} | |
| for i, block in enumerate(self.enc_blocks): | |
| if isinstance(block, SetKey): | |
| features[block.key] = x | |
| elif getattr(block, 'requires_embedding', False): | |
| x = block(x, emb) | |
| else: | |
| x = block(x) | |
| return x, features | |
| class Mlp(nn.Module): | |
| def __init__( | |
| self, | |
| in_features, | |
| hidden_features=None, | |
| out_features=None, | |
| act_layer=nn.GELU, | |
| norm_layer=None, | |
| bias=True, | |
| use_conv=False, | |
| ): | |
| super().__init__() | |
| out_features = out_features or in_features | |
| hidden_features = hidden_features or in_features | |
| bias = to_2tuple(bias) | |
| linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear | |
| self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0]) | |
| self.act = act_layer() | |
| self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity() | |
| self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1]) | |
| def forward(self, x): | |
| x = self.fc1(x) | |
| x = self.act(x) | |
| x = self.norm(x) | |
| x = self.fc2(x) | |
| return x | |
| class ConvNeXtBlockAdaLN(nn.Module): | |
| def __init__(self, dim, embed_dim=256, out_dim=None, kernel_size=7, mlp_ratio: float = 2, | |
| residual=True, ls_init_value=1e-6): | |
| super().__init__() | |
| # depthwise conv | |
| pad = (kernel_size - 1) // 2 | |
| self.conv_dw = nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=pad, groups=dim) | |
| # layer norm | |
| self.norm = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) | |
| self.norm.affine = False # for FLOPs computing | |
| # AdaLN | |
| self.embedding_layer = nn.Sequential( | |
| nn.Sequential(), | |
| nn.Linear(embed_dim, 2*dim), | |
| nn.Unflatten(1, unflattened_size=(1, 1, 2*dim)) | |
| ) | |
| # MLP | |
| hidden = int(mlp_ratio * dim) | |
| out_dim = out_dim or dim | |
| self.mlp = Mlp(dim, hidden_features=hidden, out_features=out_dim, act_layer=nn.GELU) | |
| # layer scaling | |
| if ls_init_value >= 0: | |
| self.gamma = nn.Parameter(torch.full(size=(1, out_dim, 1, 1), fill_value=1e-6)) | |
| else: | |
| self.gamma = None | |
| self.residual = residual | |
| self.requires_embedding = True | |
| for m in self.modules(): | |
| if isinstance(m, (nn.Conv2d, nn.Linear)): | |
| nn.init.trunc_normal_(m.weight, std=0.02) | |
| if m.bias is not None: | |
| nn.init.zeros_(m.bias) | |
| def forward(self, x, emb): | |
| shortcut = x | |
| # depthwise conv | |
| x = self.conv_dw(x) | |
| # layer norm | |
| x = x.permute(0, 2, 3, 1).contiguous() | |
| x = self.norm(x) | |
| # AdaLN | |
| embedding = self.embedding_layer(emb) | |
| shift, scale = torch.chunk(embedding, chunks=2, dim=-1) | |
| x = x * (1 + scale) + shift | |
| # MLP | |
| x = self.mlp(x) | |
| # scaling | |
| if self.gamma is not None: | |
| x = x.mul(self.gamma.reshape(-1)) | |
| x = x.permute(0, 3, 1, 2).contiguous() | |
| if self.residual: | |
| x = x + shortcut | |
| return x | |