freefirehay/codeformer

1import math2import numpy as np3import torch4from torch import nn, Tensor5import torch.nn.functional as F6from typing import Optional, List7 8from basicsr.archs.vqgan_arch import *9from basicsr.utils import get_root_logger10from basicsr.utils.registry import ARCH_REGISTRY11 12def calc_mean_std(feat, eps=1e-5):13    """Calculate mean and std for adaptive_instance_normalization.14 15    Args:16        feat (Tensor): 4D tensor.17        eps (float): A small value added to the variance to avoid18            divide-by-zero. Default: 1e-5.19    """20    size = feat.size()21    assert len(size) == 4, 'The input feature should be 4D tensor.'22    b, c = size[:2]23    feat_var = feat.view(b, c, -1).var(dim=2) + eps24    feat_std = feat_var.sqrt().view(b, c, 1, 1)25    feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)26    return feat_mean, feat_std27 28 29def adaptive_instance_normalization(content_feat, style_feat):30    """Adaptive instance normalization.31 32    Adjust the reference features to have the similar color and illuminations33    as those in the degradate features.34 35    Args:36        content_feat (Tensor): The reference feature.37        style_feat (Tensor): The degradate features.38    """39    size = content_feat.size()40    style_mean, style_std = calc_mean_std(style_feat)41    content_mean, content_std = calc_mean_std(content_feat)42    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)43    return normalized_feat * style_std.expand(size) + style_mean.expand(size)44 45 46class PositionEmbeddingSine(nn.Module):47    """48    This is a more standard version of the position embedding, very similar to the one49    used by the Attention is all you need paper, generalized to work on images.50    """51 52    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):53        super().__init__()54        self.num_pos_feats = num_pos_feats55        self.temperature = temperature56        self.normalize = normalize57        if scale is not None and normalize is False:58            raise ValueError("normalize should be True if scale is passed")59        if scale is None:60            scale = 2 * math.pi61        self.scale = scale62 63    def forward(self, x, mask=None):64        if mask is None:65            mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)66        not_mask = ~mask67        y_embed = not_mask.cumsum(1, dtype=torch.float32)68        x_embed = not_mask.cumsum(2, dtype=torch.float32)69        if self.normalize:70            eps = 1e-671            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale72            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale73 74        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)75        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)76 77        pos_x = x_embed[:, :, :, None] / dim_t78        pos_y = y_embed[:, :, :, None] / dim_t79        pos_x = torch.stack(80            (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=481        ).flatten(3)82        pos_y = torch.stack(83            (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=484        ).flatten(3)85        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)86        return pos87 88def _get_activation_fn(activation):89    """Return an activation function given a string"""90    if activation == "relu":91        return F.relu92    if activation == "gelu":93        return F.gelu94    if activation == "glu":95        return F.glu96    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")97 98 99class TransformerSALayer(nn.Module):100    def __init__(self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="gelu"):101        super().__init__()102        self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)103        # Implementation of Feedforward model - MLP104        self.linear1 = nn.Linear(embed_dim, dim_mlp)105        self.dropout = nn.Dropout(dropout)106        self.linear2 = nn.Linear(dim_mlp, embed_dim)107 108        self.norm1 = nn.LayerNorm(embed_dim)109        self.norm2 = nn.LayerNorm(embed_dim)110        self.dropout1 = nn.Dropout(dropout)111        self.dropout2 = nn.Dropout(dropout)112 113        self.activation = _get_activation_fn(activation)114 115    def with_pos_embed(self, tensor, pos: Optional[Tensor]):116        return tensor if pos is None else tensor + pos117 118    def forward(self, tgt,119                tgt_mask: Optional[Tensor] = None,120                tgt_key_padding_mask: Optional[Tensor] = None,121                query_pos: Optional[Tensor] = None):122        123        # self attention124        tgt2 = self.norm1(tgt)125        q = k = self.with_pos_embed(tgt2, query_pos)126        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,127                              key_padding_mask=tgt_key_padding_mask)[0]128        tgt = tgt + self.dropout1(tgt2)129 130        # ffn131        tgt2 = self.norm2(tgt)132        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))133        tgt = tgt + self.dropout2(tgt2)134        return tgt135 136class Fuse_sft_block(nn.Module):137    def __init__(self, in_ch, out_ch):138        super().__init__()139        self.encode_enc = ResBlock(2*in_ch, out_ch)140 141        self.scale = nn.Sequential(142                    nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),143                    nn.LeakyReLU(0.2, True),144                    nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))145 146        self.shift = nn.Sequential(147                    nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),148                    nn.LeakyReLU(0.2, True),149                    nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))150 151    def forward(self, enc_feat, dec_feat, w=1):152        enc_feat = self.encode_enc(torch.cat([enc_feat, dec_feat], dim=1))153        scale = self.scale(enc_feat)154        shift = self.shift(enc_feat)155        residual = w * (dec_feat * scale + shift)156        out = dec_feat + residual157        return out158 159 160@ARCH_REGISTRY.register()161class CodeFormer(VQAutoEncoder):162    def __init__(self, dim_embd=512, n_head=8, n_layers=9, 163                codebook_size=1024, latent_size=256,164                connect_list=['32', '64', '128', '256'],165                fix_modules=['quantize','generator']):166        super(CodeFormer, self).__init__(512, 64, [1, 2, 2, 4, 4, 8], 'nearest',2, [16], codebook_size)167 168        if fix_modules is not None:169            for module in fix_modules:170                for param in getattr(self, module).parameters():171                    param.requires_grad = False172 173        self.connect_list = connect_list174        self.n_layers = n_layers175        self.dim_embd = dim_embd176        self.dim_mlp = dim_embd*2177 178        self.position_emb = nn.Parameter(torch.zeros(latent_size, self.dim_embd))179        self.feat_emb = nn.Linear(256, self.dim_embd)180 181        # transformer182        self.ft_layers = nn.Sequential(*[TransformerSALayer(embed_dim=dim_embd, nhead=n_head, dim_mlp=self.dim_mlp, dropout=0.0) 183                                    for _ in range(self.n_layers)])184 185        # logits_predict head186        self.idx_pred_layer = nn.Sequential(187            nn.LayerNorm(dim_embd),188            nn.Linear(dim_embd, codebook_size, bias=False))189        190        self.channels = {191            '16': 512,192            '32': 256,193            '64': 256,194            '128': 128,195            '256': 128,196            '512': 64,197        }198 199        # after second residual block for > 16, before attn layer for ==16200        self.fuse_encoder_block = {'512':2, '256':5, '128':8, '64':11, '32':14, '16':18}201        # after first residual block for > 16, before attn layer for ==16202        self.fuse_generator_block = {'16':6, '32': 9, '64':12, '128':15, '256':18, '512':21}203 204        # fuse_convs_dict205        self.fuse_convs_dict = nn.ModuleDict()206        for f_size in self.connect_list:207            in_ch = self.channels[f_size]208            self.fuse_convs_dict[f_size] = Fuse_sft_block(in_ch, in_ch)209 210    def _init_weights(self, module):211        if isinstance(module, (nn.Linear, nn.Embedding)):212            module.weight.data.normal_(mean=0.0, std=0.02)213            if isinstance(module, nn.Linear) and module.bias is not None:214                module.bias.data.zero_()215        elif isinstance(module, nn.LayerNorm):216            module.bias.data.zero_()217            module.weight.data.fill_(1.0)218 219    def forward(self, x, w=0, detach_16=True, code_only=False, adain=False):220        # ################### Encoder #####################221        enc_feat_dict = {}222        out_list = [self.fuse_encoder_block[f_size] for f_size in self.connect_list]223        for i, block in enumerate(self.encoder.blocks):224            x = block(x) 225            if i in out_list:226                enc_feat_dict[str(x.shape[-1])] = x.clone()227 228        lq_feat = x229        # ################# Transformer ###################230        # quant_feat, codebook_loss, quant_stats = self.quantize(lq_feat)231        pos_emb = self.position_emb.unsqueeze(1).repeat(1,x.shape[0],1)232        # BCHW -> BC(HW) -> (HW)BC233        feat_emb = self.feat_emb(lq_feat.flatten(2).permute(2,0,1))234        query_emb = feat_emb235        # Transformer encoder236        for layer in self.ft_layers:237            query_emb = layer(query_emb, query_pos=pos_emb)238 239        # output logits240        logits = self.idx_pred_layer(query_emb) # (hw)bn241        logits = logits.permute(1,0,2) # (hw)bn -> b(hw)n242 243        if code_only: # for training stage II244          # logits doesn't need softmax before cross_entropy loss245            return logits, lq_feat246 247        # ################# Quantization ###################248        # if self.training:249        #     quant_feat = torch.einsum('btn,nc->btc', [soft_one_hot, self.quantize.embedding.weight])250        #     # b(hw)c -> bc(hw) -> bchw251        #     quant_feat = quant_feat.permute(0,2,1).view(lq_feat.shape)252        # ------------253        soft_one_hot = F.softmax(logits, dim=2)254        _, top_idx = torch.topk(soft_one_hot, 1, dim=2)255        quant_feat = self.quantize.get_codebook_feat(top_idx, shape=[x.shape[0],16,16,256])256        # preserve gradients257        # quant_feat = lq_feat + (quant_feat - lq_feat).detach()258 259        if detach_16:260            quant_feat = quant_feat.detach() # for training stage III261        if adain:262            quant_feat = adaptive_instance_normalization(quant_feat, lq_feat)263 264        # ################## Generator ####################265        x = quant_feat266        fuse_list = [self.fuse_generator_block[f_size] for f_size in self.connect_list]267 268        for i, block in enumerate(self.generator.blocks):269            x = block(x) 270            if i in fuse_list: # fuse after i-th block271                f_size = str(x.shape[-1])272                if w>0:273                    x = self.fuse_convs_dict[f_size](enc_feat_dict[f_size].detach(), x, w)274        out = x275        # logits doesn't need softmax before cross_entropy loss276        return out, logits, lq_feat