Delete .github/workflows/no-response.yml

.github/workflows/no-response.yml

@@ -1,33 +0,0 @@
-name: No Response
-
-# TODO: it seems not to work
-# Modified from: https://raw.githubusercontent.com/github/docs/main/.github/workflows/no-response.yaml
-
-# **What it does**: Closes issues that don't have enough information to be actionable.
-# **Why we have it**: To remove the need for maintainers to remember to check back on issues periodically
-#                     to see if contributors have responded.
-# **Who does it impact**: Everyone that works on docs or docs-internal.
-
-on:
-  issue_comment:
-    types: [created]
-
-  schedule:
-    # Schedule for five minutes after the hour every hour
-    - cron: '5 * * * *'
-
-jobs:
-  noResponse:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: lee-dohm/no-response@v0.5.0
-        with:
-          token: ${{ github.token }}
-          closeComment: >
-            This issue has been automatically closed because there has been no response
-            to our request for more information from the original author. With only the
-            information that is currently in the issue, we don't have enough information
-            to take action. Please reach out if you have or find the answers we need so
-            that we can investigate further.
-            If you still have questions, please improve your description and re-open it.
-            Thanks :-)

README.md

@@ -35,7 +35,7 @@ It leverages rich and diverse priors encapsulated in a pretrained face GAN (*e.g
 
 :triangular_flag_on_post: **Updates**
 
-- :white_check_mark: Add CodeFormer ([CC BY-NC-SA 4.0 License](https://creativecommons.org/licenses/by-nc-sa/4.0/)) and RestoreFormer.
+- :white_check_mark: Add [RestoreFormer](https://github.com/wzhouxiff/RestoreFormer) inference codes.
 - :white_check_mark: Add [V1.4 model](https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth), which produces slightly more details and better identity than V1.3.
 - :white_check_mark: Add **[V1.3 model](https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth)**, which produces **more natural** restoration results, and better results on *very low-quality* / *high-quality* inputs. See more in [Model zoo](#european_castle-model-zoo), [Comparisons.md](Comparisons.md)
 - :white_check_mark: Integrated to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See [Gradio Web Demo](https://huggingface.co/spaces/akhaliq/GFPGAN).

VERSION

@@ -1 +1 @@
-1.3.6
+1.3.8

cog_predict.py

@@ -42,6 +42,10 @@ class Predictor(BasePredictor):
         if not os.path.exists('gfpgan/weights/GFPGANv1.4.pth'):
             os.system(
                 'wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P ./gfpgan/weights')
+        if not os.path.exists('gfpgan/weights/RestoreFormer.pth'):
+            os.system(
+                'wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P ./gfpgan/weights'
+            )
 
         # background enhancer with RealESRGAN
         model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
@@ -60,15 +64,16 @@ class Predictor(BasePredictor):
         self.current_version = 'v1.4'
 
     def predict(
-        self,
-        img: Path = Input(description='Input'),
-        version: str = Input(
-            description='GFPGAN version. v1.3: better quality. v1.4: more details and better identity.',
-            choices=['v1.2', 'v1.3', 'v1.4'],
-            default='v1.4'),
-        scale: float = Input(description='Rescaling factor', default=2)
+            self,
+            img: Path = Input(description='Input'),
+            version: str = Input(
+                description='GFPGAN version. v1.3: better quality. v1.4: more details and better identity.',
+                choices=['v1.2', 'v1.3', 'v1.4', 'RestoreFormer'],
+                default='v1.4'),
+            scale: float = Input(description='Rescaling factor', default=2),
     ) -> Path:
-        print(img, version, scale)
+        weight = 0.5
+        print(img, version, scale, weight)
         try:
             extension = os.path.splitext(os.path.basename(str(img)))[1]
             img = cv2.imread(str(img), cv2.IMREAD_UNCHANGED)
@@ -109,14 +114,19 @@ class Predictor(BasePredictor):
                         channel_multiplier=2,
                         bg_upsampler=self.upsampler)
                     self.current_version = 'v1.4'
+                elif version == 'RestoreFormer':
+                    self.face_enhancer = GFPGANer(
+                        model_path='gfpgan/weights/RestoreFormer.pth',
+                        upscale=2,
+                        arch='RestoreFormer',
+                        channel_multiplier=2,
+                        bg_upsampler=self.upsampler)
 
             try:
                 _, _, output = self.face_enhancer.enhance(
-                    img, has_aligned=False, only_center_face=False, paste_back=True)
+                    img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
             except RuntimeError as error:
                 print('Error', error)
-            else:
-                extension = 'png'
 
             try:
                 if scale != 2:

gfpgan/archs/codeformer_arch.py

@@ -1,630 +0,0 @@
-"""
-Modified from https://github.com/sczhou/CodeFormer
-VQGAN code, adapted from the original created by the Unleashing Transformers authors:
-https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py
-"""
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from basicsr.utils import get_root_logger
-from basicsr.utils.registry import ARCH_REGISTRY
-from torch import Tensor
-from typing import Optional
-
-
-class VectorQuantizer(nn.Module):
-
-    def __init__(self, codebook_size, emb_dim, beta):
-        super(VectorQuantizer, self).__init__()
-        self.codebook_size = codebook_size  # number of embeddings
-        self.emb_dim = emb_dim  # dimension of embedding
-        self.beta = beta  # commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
-        self.embedding = nn.Embedding(self.codebook_size, self.emb_dim)
-        self.embedding.weight.data.uniform_(-1.0 / self.codebook_size, 1.0 / self.codebook_size)
-
-    def forward(self, z):
-        # reshape z -> (batch, height, width, channel) and flatten
-        z = z.permute(0, 2, 3, 1).contiguous()
-        z_flattened = z.view(-1, self.emb_dim)
-
-        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
-        d = (z_flattened ** 2).sum(dim=1, keepdim=True) + (self.embedding.weight**2).sum(1) - \
-            2 * torch.matmul(z_flattened, self.embedding.weight.t())
-
-        mean_distance = torch.mean(d)
-        # find closest encodings
-        # min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
-        min_encoding_scores, min_encoding_indices = torch.topk(d, 1, dim=1, largest=False)
-        # [0-1], higher score, higher confidence
-        min_encoding_scores = torch.exp(-min_encoding_scores / 10)
-
-        min_encodings = torch.zeros(min_encoding_indices.shape[0], self.codebook_size).to(z)
-        min_encodings.scatter_(1, min_encoding_indices, 1)
-
-        # get quantized latent vectors
-        z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)
-        # compute loss for embedding
-        loss = torch.mean((z_q.detach() - z)**2) + self.beta * torch.mean((z_q - z.detach())**2)
-        # preserve gradients
-        z_q = z + (z_q - z).detach()
-
-        # perplexity
-        e_mean = torch.mean(min_encodings, dim=0)
-        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))
-        # reshape back to match original input shape
-        z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
-        return z_q, loss, {
-            'perplexity': perplexity,
-            'min_encodings': min_encodings,
-            'min_encoding_indices': min_encoding_indices,
-            'min_encoding_scores': min_encoding_scores,
-            'mean_distance': mean_distance
-        }
-
-    def get_codebook_feat(self, indices, shape):
-        # input indices: batch*token_num -> (batch*token_num)*1
-        # shape: batch, height, width, channel
-        indices = indices.view(-1, 1)
-        min_encodings = torch.zeros(indices.shape[0], self.codebook_size).to(indices)
-        min_encodings.scatter_(1, indices, 1)
-        # get quantized latent vectors
-        z_q = torch.matmul(min_encodings.float(), self.embedding.weight)
-
-        if shape is not None:  # reshape back to match original input shape
-            z_q = z_q.view(shape).permute(0, 3, 1, 2).contiguous()
-
-        return z_q
-
-
-class GumbelQuantizer(nn.Module):
-
-    def __init__(self, codebook_size, emb_dim, num_hiddens, straight_through=False, kl_weight=5e-4, temp_init=1.0):
-        super().__init__()
-        self.codebook_size = codebook_size  # number of embeddings
-        self.emb_dim = emb_dim  # dimension of embedding
-        self.straight_through = straight_through
-        self.temperature = temp_init
-        self.kl_weight = kl_weight
-        self.proj = nn.Conv2d(num_hiddens, codebook_size, 1)  # projects last encoder layer to quantized logits
-        self.embed = nn.Embedding(codebook_size, emb_dim)
-
-    def forward(self, z):
-        hard = self.straight_through if self.training else True
-
-        logits = self.proj(z)
-
-        soft_one_hot = F.gumbel_softmax(logits, tau=self.temperature, dim=1, hard=hard)
-
-        z_q = torch.einsum('b n h w, n d -> b d h w', soft_one_hot, self.embed.weight)
-
-        # + kl divergence to the prior loss
-        qy = F.softmax(logits, dim=1)
-        diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.codebook_size + 1e-10), dim=1).mean()
-        min_encoding_indices = soft_one_hot.argmax(dim=1)
-
-        return z_q, diff, {'min_encoding_indices': min_encoding_indices}
-
-
-class Downsample(nn.Module):
-
-    def __init__(self, in_channels):
-        super().__init__()
-        self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
-
-    def forward(self, x):
-        pad = (0, 1, 0, 1)
-        x = torch.nn.functional.pad(x, pad, mode='constant', value=0)
-        x = self.conv(x)
-        return x
-
-
-class Upsample(nn.Module):
-
-    def __init__(self, in_channels):
-        super().__init__()
-        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
-
-    def forward(self, x):
-        x = F.interpolate(x, scale_factor=2.0, mode='nearest')
-        x = self.conv(x)
-
-        return x
-
-
-class AttnBlock(nn.Module):
-
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = normalize(in_channels)
-        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b, c, h, w = q.shape
-        q = q.reshape(b, c, h * w)
-        q = q.permute(0, 2, 1)
-        k = k.reshape(b, c, h * w)
-        w_ = torch.bmm(q, k)
-        w_ = w_ * (int(c)**(-0.5))
-        w_ = F.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b, c, h * w)
-        w_ = w_.permute(0, 2, 1)
-        h_ = torch.bmm(v, w_)
-        h_ = h_.reshape(b, c, h, w)
-
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-
-class Encoder(nn.Module):
-
-    def __init__(self, in_channels, nf, out_channels, ch_mult, num_res_blocks, resolution, attn_resolutions):
-        super().__init__()
-        self.nf = nf
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.attn_resolutions = attn_resolutions
-
-        curr_res = self.resolution
-        in_ch_mult = (1, ) + tuple(ch_mult)
-
-        blocks = []
-        # initial convultion
-        blocks.append(nn.Conv2d(in_channels, nf, kernel_size=3, stride=1, padding=1))
-
-        # residual and downsampling blocks, with attention on smaller res (16x16)
-        for i in range(self.num_resolutions):
-            block_in_ch = nf * in_ch_mult[i]
-            block_out_ch = nf * ch_mult[i]
-            for _ in range(self.num_res_blocks):
-                blocks.append(ResBlock(block_in_ch, block_out_ch))
-                block_in_ch = block_out_ch
-                if curr_res in attn_resolutions:
-                    blocks.append(AttnBlock(block_in_ch))
-
-            if i != self.num_resolutions - 1:
-                blocks.append(Downsample(block_in_ch))
-                curr_res = curr_res // 2
-
-        # non-local attention block
-        blocks.append(ResBlock(block_in_ch, block_in_ch))
-        blocks.append(AttnBlock(block_in_ch))
-        blocks.append(ResBlock(block_in_ch, block_in_ch))
-
-        # normalise and convert to latent size
-        blocks.append(normalize(block_in_ch))
-        blocks.append(nn.Conv2d(block_in_ch, out_channels, kernel_size=3, stride=1, padding=1))
-        self.blocks = nn.ModuleList(blocks)
-
-    def forward(self, x):
-        for block in self.blocks:
-            x = block(x)
-
-        return x
-
-
-class Generator(nn.Module):
-
-    def __init__(self, nf, ch_mult, res_blocks, img_size, attn_resolutions, emb_dim):
-        super().__init__()
-        self.nf = nf
-        self.ch_mult = ch_mult
-        self.num_resolutions = len(self.ch_mult)
-        self.num_res_blocks = res_blocks
-        self.resolution = img_size
-        self.attn_resolutions = attn_resolutions
-        self.in_channels = emb_dim
-        self.out_channels = 3
-        block_in_ch = self.nf * self.ch_mult[-1]
-        curr_res = self.resolution // 2**(self.num_resolutions - 1)
-
-        blocks = []
-        # initial conv
-        blocks.append(nn.Conv2d(self.in_channels, block_in_ch, kernel_size=3, stride=1, padding=1))
-
-        # non-local attention block
-        blocks.append(ResBlock(block_in_ch, block_in_ch))
-        blocks.append(AttnBlock(block_in_ch))
-        blocks.append(ResBlock(block_in_ch, block_in_ch))
-
-        for i in reversed(range(self.num_resolutions)):
-            block_out_ch = self.nf * self.ch_mult[i]
-
-            for _ in range(self.num_res_blocks):
-                blocks.append(ResBlock(block_in_ch, block_out_ch))
-                block_in_ch = block_out_ch
-
-                if curr_res in self.attn_resolutions:
-                    blocks.append(AttnBlock(block_in_ch))
-
-            if i != 0:
-                blocks.append(Upsample(block_in_ch))
-                curr_res = curr_res * 2
-
-        blocks.append(normalize(block_in_ch))
-        blocks.append(nn.Conv2d(block_in_ch, self.out_channels, kernel_size=3, stride=1, padding=1))
-
-        self.blocks = nn.ModuleList(blocks)
-
-    def forward(self, x):
-        for block in self.blocks:
-            x = block(x)
-
-        return x
-
-
-class VQAutoEncoder(nn.Module):
-
-    def __init__(self,
-                 img_size,
-                 nf,
-                 ch_mult,
-                 quantizer='nearest',
-                 res_blocks=2,
-                 attn_resolutions=[16],
-                 codebook_size=1024,
-                 emb_dim=256,
-                 beta=0.25,
-                 gumbel_straight_through=False,
-                 gumbel_kl_weight=1e-8,
-                 model_path=None):
-        super().__init__()
-        logger = get_root_logger()
-        self.in_channels = 3
-        self.nf = nf
-        self.n_blocks = res_blocks
-        self.codebook_size = codebook_size
-        self.embed_dim = emb_dim
-        self.ch_mult = ch_mult
-        self.resolution = img_size
-        self.attn_resolutions = attn_resolutions
-        self.quantizer_type = quantizer
-        self.encoder = Encoder(self.in_channels, self.nf, self.embed_dim, self.ch_mult, self.n_blocks, self.resolution,
-                               self.attn_resolutions)
-        if self.quantizer_type == 'nearest':
-            self.beta = beta  # 0.25
-            self.quantize = VectorQuantizer(self.codebook_size, self.embed_dim, self.beta)
-        elif self.quantizer_type == 'gumbel':
-            self.gumbel_num_hiddens = emb_dim
-            self.straight_through = gumbel_straight_through
-            self.kl_weight = gumbel_kl_weight
-            self.quantize = GumbelQuantizer(self.codebook_size, self.embed_dim, self.gumbel_num_hiddens,
-                                            self.straight_through, self.kl_weight)
-        self.generator = Generator(nf, ch_mult, res_blocks, img_size, attn_resolutions, emb_dim)
-
-        if model_path is not None:
-            chkpt = torch.load(model_path, map_location='cpu')
-            if 'params_ema' in chkpt:
-                self.load_state_dict(torch.load(model_path, map_location='cpu')['params_ema'])
-                logger.info(f'vqgan is loaded from: {model_path} [params_ema]')
-            elif 'params' in chkpt:
-                self.load_state_dict(torch.load(model_path, map_location='cpu')['params'])
-                logger.info(f'vqgan is loaded from: {model_path} [params]')
-            else:
-                raise ValueError('Wrong params!')
-
-    def forward(self, x):
-        x = self.encoder(x)
-        quant, codebook_loss, quant_stats = self.quantize(x)
-        x = self.generator(quant)
-        return x, codebook_loss, quant_stats
-
-
-def calc_mean_std(feat, eps=1e-5):
-    """Calculate mean and std for adaptive_instance_normalization.
-
-    Args:
-        feat (Tensor): 4D tensor.
-        eps (float): A small value added to the variance to avoid
-            divide-by-zero. Default: 1e-5.
-    """
-    size = feat.size()
-    assert len(size) == 4, 'The input feature should be 4D tensor.'
-    b, c = size[:2]
-    feat_var = feat.view(b, c, -1).var(dim=2) + eps
-    feat_std = feat_var.sqrt().view(b, c, 1, 1)
-    feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
-    return feat_mean, feat_std
-
-
-def adaptive_instance_normalization(content_feat, style_feat):
-    """Adaptive instance normalization.
-
-    Adjust the reference features to have the similar color and illuminations
-    as those in the degradate features.
-
-    Args:
-        content_feat (Tensor): The reference feature.
-        style_feat (Tensor): The degradate features.
-    """
-    size = content_feat.size()
-    style_mean, style_std = calc_mean_std(style_feat)
-    content_mean, content_std = calc_mean_std(content_feat)
-    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
-    return normalized_feat * style_std.expand(size) + style_mean.expand(size)
-
-
-class PositionEmbeddingSine(nn.Module):
-    """
-    This is a more standard version of the position embedding, very similar to the one
-    used by the Attention is all you need paper, generalized to work on images.
-    """
-
-    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
-        super().__init__()
-        self.num_pos_feats = num_pos_feats
-        self.temperature = temperature
-        self.normalize = normalize
-        if scale is not None and normalize is False:
-            raise ValueError('normalize should be True if scale is passed')
-        if scale is None:
-            scale = 2 * math.pi
-        self.scale = scale
-
-    def forward(self, x, mask=None):
-        if mask is None:
-            mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
-        not_mask = ~mask
-        y_embed = not_mask.cumsum(1, dtype=torch.float32)
-        x_embed = not_mask.cumsum(2, dtype=torch.float32)
-        if self.normalize:
-            eps = 1e-6
-            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
-            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
-
-        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
-        dim_t = self.temperature**(2 * (dim_t // 2) / self.num_pos_feats)
-
-        pos_x = x_embed[:, :, :, None] / dim_t
-        pos_y = y_embed[:, :, :, None] / dim_t
-        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
-        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
-        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
-        return pos
-
-
-def _get_activation_fn(activation):
-    """Return an activation function given a string"""
-    if activation == 'relu':
-        return F.relu
-    if activation == 'gelu':
-        return F.gelu
-    if activation == 'glu':
-        return F.glu
-    raise RuntimeError(F'activation should be relu/gelu, not {activation}.')
-
-
-class TransformerSALayer(nn.Module):
-
-    def __init__(self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation='gelu'):
-        super().__init__()
-        self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)
-        # Implementation of Feedforward model - MLP
-        self.linear1 = nn.Linear(embed_dim, dim_mlp)
-        self.dropout = nn.Dropout(dropout)
-        self.linear2 = nn.Linear(dim_mlp, embed_dim)
-
-        self.norm1 = nn.LayerNorm(embed_dim)
-        self.norm2 = nn.LayerNorm(embed_dim)
-        self.dropout1 = nn.Dropout(dropout)
-        self.dropout2 = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward(self,
-                tgt,
-                tgt_mask: Optional[Tensor] = None,
-                tgt_key_padding_mask: Optional[Tensor] = None,
-                query_pos: Optional[Tensor] = None):
-
-        # self attention
-        tgt2 = self.norm1(tgt)
-        q = k = self.with_pos_embed(tgt2, query_pos)
-        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
-        tgt = tgt + self.dropout1(tgt2)
-
-        # ffn
-        tgt2 = self.norm2(tgt)
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
-        tgt = tgt + self.dropout2(tgt2)
-        return tgt
-
-
-def normalize(in_channels):
-    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-
-
-@torch.jit.script
-def swish(x):
-    return x * torch.sigmoid(x)
-
-
-class ResBlock(nn.Module):
-
-    def __init__(self, in_channels, out_channels=None):
-        super(ResBlock, self).__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels if out_channels is None else out_channels
-        self.norm1 = normalize(in_channels)
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        self.norm2 = normalize(out_channels)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        if self.in_channels != self.out_channels:
-            self.conv_out = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x_in):
-        x = x_in
-        x = self.norm1(x)
-        x = swish(x)
-        x = self.conv1(x)
-        x = self.norm2(x)
-        x = swish(x)
-        x = self.conv2(x)
-        if self.in_channels != self.out_channels:
-            x_in = self.conv_out(x_in)
-
-        return x + x_in
-
-
-class Fuse_sft_block(nn.Module):
-
-    def __init__(self, in_ch, out_ch):
-        super().__init__()
-        self.encode_enc = ResBlock(2 * in_ch, out_ch)
-
-        self.scale = nn.Sequential(
-            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1), nn.LeakyReLU(0.2, True),
-            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))
-
-        self.shift = nn.Sequential(
-            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1), nn.LeakyReLU(0.2, True),
-            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))
-
-    def forward(self, enc_feat, dec_feat, w=1):
-        enc_feat = self.encode_enc(torch.cat([enc_feat, dec_feat], dim=1))
-        scale = self.scale(enc_feat)
-        shift = self.shift(enc_feat)
-        residual = w * (dec_feat * scale + shift)
-        out = dec_feat + residual
-        return out
-
-
-@ARCH_REGISTRY.register()
-class CodeFormer(VQAutoEncoder):
-
-    def __init__(self,
-                 dim_embd=512,
-                 n_head=8,
-                 n_layers=9,
-                 codebook_size=1024,
-                 latent_size=256,
-                 connect_list=['32', '64', '128', '256'],
-                 fix_modules=['quantize', 'generator']):
-        super(CodeFormer, self).__init__(512, 64, [1, 2, 2, 4, 4, 8], 'nearest', 2, [16], codebook_size)
-
-        if fix_modules is not None:
-            for module in fix_modules:
-                for param in getattr(self, module).parameters():
-                    param.requires_grad = False
-
-        self.connect_list = connect_list
-        self.n_layers = n_layers
-        self.dim_embd = dim_embd
-        self.dim_mlp = dim_embd * 2
-
-        self.position_emb = nn.Parameter(torch.zeros(latent_size, self.dim_embd))
-        self.feat_emb = nn.Linear(256, self.dim_embd)
-
-        # transformer
-        self.ft_layers = nn.Sequential(*[
-            TransformerSALayer(embed_dim=dim_embd, nhead=n_head, dim_mlp=self.dim_mlp, dropout=0.0)
-            for _ in range(self.n_layers)
-        ])
-
-        # logits_predict head
-        self.idx_pred_layer = nn.Sequential(nn.LayerNorm(dim_embd), nn.Linear(dim_embd, codebook_size, bias=False))
-
-        self.channels = {'16': 512, '32': 256, '64': 256, '128': 128, '256': 128, '512': 64}
-
-        # after second residual block for > 16, before attn layer for ==16
-        self.fuse_encoder_block = {'512': 2, '256': 5, '128': 8, '64': 11, '32': 14, '16': 18}
-        # after first residual block for > 16, before attn layer for ==16
-        self.fuse_generator_block = {'16': 6, '32': 9, '64': 12, '128': 15, '256': 18, '512': 21}
-
-        # fuse_convs_dict
-        self.fuse_convs_dict = nn.ModuleDict()
-        for f_size in self.connect_list:
-            in_ch = self.channels[f_size]
-            self.fuse_convs_dict[f_size] = Fuse_sft_block(in_ch, in_ch)
-
-    def _init_weights(self, module):
-        if isinstance(module, (nn.Linear, nn.Embedding)):
-            module.weight.data.normal_(mean=0.0, std=0.02)
-            if isinstance(module, nn.Linear) and module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-    def forward(self, x, weight=0.5, **kwargs):
-        detach_16 = True
-        code_only = False
-        adain = True
-        # ################### Encoder #####################
-        enc_feat_dict = {}
-        out_list = [self.fuse_encoder_block[f_size] for f_size in self.connect_list]
-        for i, block in enumerate(self.encoder.blocks):
-            x = block(x)
-            if i in out_list:
-                enc_feat_dict[str(x.shape[-1])] = x.clone()
-
-        lq_feat = x
-        # ################# Transformer ###################
-        # quant_feat, codebook_loss, quant_stats = self.quantize(lq_feat)
-        pos_emb = self.position_emb.unsqueeze(1).repeat(1, x.shape[0], 1)
-        # BCHW -> BC(HW) -> (HW)BC
-        feat_emb = self.feat_emb(lq_feat.flatten(2).permute(2, 0, 1))
-        query_emb = feat_emb
-        # Transformer encoder
-        for layer in self.ft_layers:
-            query_emb = layer(query_emb, query_pos=pos_emb)
-
-        # output logits
-        logits = self.idx_pred_layer(query_emb)  # (hw)bn
-        logits = logits.permute(1, 0, 2)  # (hw)bn -> b(hw)n
-
-        if code_only:  # for training stage II
-            # logits doesn't need softmax before cross_entropy loss
-            return logits, lq_feat
-
-        # ################# Quantization ###################
-        # if self.training:
-        #     quant_feat = torch.einsum('btn,nc->btc', [soft_one_hot, self.quantize.embedding.weight])
-        #     # b(hw)c -> bc(hw) -> bchw
-        #     quant_feat = quant_feat.permute(0,2,1).view(lq_feat.shape)
-        # ------------
-        soft_one_hot = F.softmax(logits, dim=2)
-        _, top_idx = torch.topk(soft_one_hot, 1, dim=2)
-        quant_feat = self.quantize.get_codebook_feat(top_idx, shape=[x.shape[0], 16, 16, 256])
-        # preserve gradients
-        # quant_feat = lq_feat + (quant_feat - lq_feat).detach()
-
-        if detach_16:
-            quant_feat = quant_feat.detach()  # for training stage III
-        if adain:
-            quant_feat = adaptive_instance_normalization(quant_feat, lq_feat)
-
-        # ################## Generator ####################
-        x = quant_feat
-        fuse_list = [self.fuse_generator_block[f_size] for f_size in self.connect_list]
-
-        for i, block in enumerate(self.generator.blocks):
-            x = block(x)
-            if i in fuse_list:  # fuse after i-th block
-                f_size = str(x.shape[-1])
-                if weight > 0:
-                    x = self.fuse_convs_dict[f_size](enc_feat_dict[f_size].detach(), x, weight)
-        out = x
-        # logits doesn't need softmax before cross_entropy loss
-        # return out, logits, lq_feat
-        return out, logits

gfpgan/archs/gfpganv1_cleanonnx_arch.py

@@ -1,325 +0,0 @@
-import math
-import random
-import torch
-from basicsr.utils.registry import ARCH_REGISTRY
-from torch import nn
-from torch.nn import functional as F
-
-from .stylegan2_cleanonnx_arch import StyleGAN2GeneratorCleanONNX
-
-
-class StyleGAN2GeneratorCSFT(StyleGAN2GeneratorCleanONNX):
-    """StyleGAN2 Generator with SFT modulation (Spatial Feature Transform).
-
-    It is the clean version without custom compiled CUDA extensions used in StyleGAN2.
-
-    Args:
-        out_size (int): The spatial size of outputs.
-        num_style_feat (int): Channel number of style features. Default: 512.
-        num_mlp (int): Layer number of MLP style layers. Default: 8.
-        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
-        narrow (float): The narrow ratio for channels. Default: 1.
-        sft_half (bool): Whether to apply SFT on half of the input channels. Default: False.
-    """
-
-    def __init__(self, out_size, num_style_feat=512, num_mlp=8, channel_multiplier=2, narrow=1, sft_half=False):
-        super(StyleGAN2GeneratorCSFT, self).__init__(
-            out_size,
-            num_style_feat=num_style_feat,
-            num_mlp=num_mlp,
-            channel_multiplier=channel_multiplier,
-            narrow=narrow)
-        self.sft_half = sft_half
-
-    def forward(self,
-                styles,
-                conditions,
-                input_is_latent=False,
-                noise=None,
-                randomize_noise=True,
-                truncation=1,
-                truncation_latent=None,
-                inject_index=None,
-                return_latents=False):
-        """Forward function for StyleGAN2GeneratorCSFT.
-
-        Args:
-            styles (list[Tensor]): Sample codes of styles.
-            conditions (list[Tensor]): SFT conditions to generators.
-            input_is_latent (bool): Whether input is latent style. Default: False.
-            noise (Tensor | None): Input noise or None. Default: None.
-            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
-            truncation (float): The truncation ratio. Default: 1.
-            truncation_latent (Tensor | None): The truncation latent tensor. Default: None.
-            inject_index (int | None): The injection index for mixing noise. Default: None.
-            return_latents (bool): Whether to return style latents. Default: False.
-        """
-        # style codes -> latents with Style MLP layer
-        if not input_is_latent:
-            styles = [self.style_mlp(s) for s in styles]
-        # noises
-        if noise is None:
-            if randomize_noise:
-                noise = [None] * self.num_layers  # for each style conv layer
-            else:  # use the stored noise
-                noise = [getattr(self.noises, f'noise{i}') for i in range(self.num_layers)]
-        # style truncation
-        if truncation < 1:
-            style_truncation = []
-            for style in styles:
-                style_truncation.append(truncation_latent + truncation * (style - truncation_latent))
-            styles = style_truncation
-        # get style latents with injection
-        if len(styles) == 1:
-            inject_index = self.num_latent
-
-            if styles[0].ndim < 3:
-                # repeat latent code for all the layers
-                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
-            else:  # used for encoder with different latent code for each layer
-                latent = styles[0]
-        elif len(styles) == 2:  # mixing noises
-            if inject_index is None:
-                inject_index = random.randint(1, self.num_latent - 1)
-            latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
-            latent2 = styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1)
-            latent = torch.cat([latent1, latent2], 1)
-
-        # main generation
-        out = self.constant_input(latent.shape[0])
-        out = self.style_conv1(out, latent[:, 0], noise=noise[0])
-        skip = self.to_rgb1(out, latent[:, 1])
-
-        i = 1
-        for conv1, conv2, noise1, noise2, to_rgb in zip(self.style_convs[::2], self.style_convs[1::2], noise[1::2],
-                                                        noise[2::2], self.to_rgbs):
-            out = conv1(out, latent[:, i], noise=noise1)
-
-            # the conditions may have fewer levels
-            if i < len(conditions):
-                # SFT part to combine the conditions
-                if self.sft_half:  # only apply SFT to half of the channels
-                    out_same, out_sft = torch.split(out, int(out.size(1) // 2), dim=1)
-                    # print(out_sft.size(), conditions[i - 1].size(), conditions[i].size())
-                    out_sft = out_sft * conditions[i - 1] + conditions[i]
-                    out = torch.cat([out_same, out_sft], dim=1)
-                else:  # apply SFT to all the channels
-                    out = out * conditions[i - 1] + conditions[i]
-
-            out = conv2(out, latent[:, i + 1], noise=noise2)
-            skip = to_rgb(out, latent[:, i + 2], skip)  # feature back to the rgb space
-            i += 2
-
-        image = skip
-
-        if return_latents:
-            return image, latent
-        else:
-            return image, None
-
-
-class ResBlock(nn.Module):
-    """Residual block with bilinear upsampling/downsampling.
-
-    Args:
-        in_channels (int): Channel number of the input.
-        out_channels (int): Channel number of the output.
-        mode (str): Upsampling/downsampling mode. Options: down | up. Default: down.
-    """
-
-    def __init__(self, in_channels, out_channels, mode='down'):
-        super(ResBlock, self).__init__()
-
-        self.conv1 = nn.Conv2d(in_channels, in_channels, 3, 1, 1)
-        self.conv2 = nn.Conv2d(in_channels, out_channels, 3, 1, 1)
-        self.skip = nn.Conv2d(in_channels, out_channels, 1, bias=False)
-        if mode == 'down':
-            self.scale_factor = 0.5
-        elif mode == 'up':
-            self.scale_factor = 2
-
-    def forward(self, x):
-        out = F.leaky_relu_(self.conv1(x), negative_slope=0.2)
-        # upsample/downsample
-        out = F.interpolate(out, scale_factor=self.scale_factor, mode='bilinear', align_corners=False)
-        out = F.leaky_relu_(self.conv2(out), negative_slope=0.2)
-        # skip
-        x = F.interpolate(x, scale_factor=self.scale_factor, mode='bilinear', align_corners=False)
-        skip = self.skip(x)
-        out = out + skip
-        return out
-
-
-@ARCH_REGISTRY.register()
-class GFPGANv1CleanONNX(nn.Module):
-    """The GFPGAN architecture: Unet + StyleGAN2 decoder with SFT.
-
-    It is the clean version without custom compiled CUDA extensions used in StyleGAN2.
-
-    Ref: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior.
-
-    Args:
-        out_size (int): The spatial size of outputs.
-        num_style_feat (int): Channel number of style features. Default: 512.
-        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
-        decoder_load_path (str): The path to the pre-trained decoder model (usually, the StyleGAN2). Default: None.
-        fix_decoder (bool): Whether to fix the decoder. Default: True.
-
-        num_mlp (int): Layer number of MLP style layers. Default: 8.
-        input_is_latent (bool): Whether input is latent style. Default: False.
-        different_w (bool): Whether to use different latent w for different layers. Default: False.
-        narrow (float): The narrow ratio for channels. Default: 1.
-        sft_half (bool): Whether to apply SFT on half of the input channels. Default: False.
-    """
-
-    def __init__(
-            self,
-            out_size,
-            num_style_feat=512,
-            channel_multiplier=1,
-            decoder_load_path=None,
-            fix_decoder=True,
-            # for stylegan decoder
-            num_mlp=8,
-            input_is_latent=False,
-            different_w=False,
-            narrow=1,
-            sft_half=False):
-
-        super(GFPGANv1CleanONNX, self).__init__()
-        self.input_is_latent = input_is_latent
-        self.different_w = different_w
-        self.num_style_feat = num_style_feat
-
-        unet_narrow = narrow * 0.5  # by default, use a half of input channels
-        channels = {
-            '4': int(512 * unet_narrow),
-            '8': int(512 * unet_narrow),
-            '16': int(512 * unet_narrow),
-            '32': int(512 * unet_narrow),
-            '64': int(256 * channel_multiplier * unet_narrow),
-            '128': int(128 * channel_multiplier * unet_narrow),
-            '256': int(64 * channel_multiplier * unet_narrow),
-            '512': int(32 * channel_multiplier * unet_narrow),
-            '1024': int(16 * channel_multiplier * unet_narrow)
-        }
-
-        self.log_size = int(math.log(out_size, 2))
-        first_out_size = 2**(int(math.log(out_size, 2)))
-
-        self.conv_body_first = nn.Conv2d(3, channels[f'{first_out_size}'], 1)
-
-        # downsample
-        in_channels = channels[f'{first_out_size}']
-        self.conv_body_down = nn.ModuleList()
-        for i in range(self.log_size, 2, -1):
-            out_channels = channels[f'{2**(i - 1)}']
-            self.conv_body_down.append(ResBlock(in_channels, out_channels, mode='down'))
-            in_channels = out_channels
-
-        self.final_conv = nn.Conv2d(in_channels, channels['4'], 3, 1, 1)
-
-        # upsample
-        in_channels = channels['4']
-        self.conv_body_up = nn.ModuleList()
-        for i in range(3, self.log_size + 1):
-            out_channels = channels[f'{2**i}']
-            self.conv_body_up.append(ResBlock(in_channels, out_channels, mode='up'))
-            in_channels = out_channels
-
-        # to RGB
-        self.toRGB = nn.ModuleList()
-        for i in range(3, self.log_size + 1):
-            self.toRGB.append(nn.Conv2d(channels[f'{2**i}'], 3, 1))
-
-        if different_w:
-            linear_out_channel = (int(math.log(out_size, 2)) * 2 - 2) * num_style_feat
-        else:
-            linear_out_channel = num_style_feat
-
-        self.final_linear = nn.Linear(channels['4'] * 4 * 4, linear_out_channel)
-
-        # the decoder: stylegan2 generator with SFT modulations
-        self.stylegan_decoder = StyleGAN2GeneratorCSFT(
-            out_size=out_size,
-            num_style_feat=num_style_feat,
-            num_mlp=num_mlp,
-            channel_multiplier=channel_multiplier,
-            narrow=narrow,
-            sft_half=sft_half)
-
-        # load pre-trained stylegan2 model if necessary
-        if decoder_load_path:
-            self.stylegan_decoder.load_state_dict(
-                torch.load(decoder_load_path, map_location=lambda storage, loc: storage)['params_ema'])
-        # fix decoder without updating params
-        if fix_decoder:
-            for _, param in self.stylegan_decoder.named_parameters():
-                param.requires_grad = False
-
-        # for SFT modulations (scale and shift)
-        self.condition_scale = nn.ModuleList()
-        self.condition_shift = nn.ModuleList()
-        for i in range(3, self.log_size + 1):
-            out_channels = channels[f'{2**i}']
-            if sft_half:
-                sft_out_channels = out_channels
-            else:
-                sft_out_channels = out_channels * 2
-            self.condition_scale.append(
-                nn.Sequential(
-                    nn.Conv2d(out_channels, out_channels, 3, 1, 1), nn.LeakyReLU(0.2, True),
-                    nn.Conv2d(out_channels, sft_out_channels, 3, 1, 1)))
-            self.condition_shift.append(
-                nn.Sequential(
-                    nn.Conv2d(out_channels, out_channels, 3, 1, 1), nn.LeakyReLU(0.2, True),
-                    nn.Conv2d(out_channels, sft_out_channels, 3, 1, 1)))
-
-    def forward(self, x, return_latents=False, return_rgb=True, randomize_noise=True):
-        """Forward function for GFPGANv1Clean.
-
-        Args:
-            x (Tensor): Input images.
-            return_latents (bool): Whether to return style latents. Default: False.
-            return_rgb (bool): Whether return intermediate rgb images. Default: True.
-            randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True.
-        """
-        conditions = []
-        unet_skips = []
-        out_rgbs = []
-
-        # encoder
-        feat = F.leaky_relu_(self.conv_body_first(x), negative_slope=0.2)
-        for i in range(self.log_size - 2):
-            feat = self.conv_body_down[i](feat)
-            unet_skips.insert(0, feat)
-        feat = F.leaky_relu_(self.final_conv(feat), negative_slope=0.2)
-
-        # style code
-        style_code = self.final_linear(feat.view(feat.size(0), -1))
-        if self.different_w:
-            style_code = style_code.view(style_code.size(0), -1, self.num_style_feat)
-
-        # decode
-        for i in range(self.log_size - 2):
-            # add unet skip
-            feat = feat + unet_skips[i]
-            # ResUpLayer
-            feat = self.conv_body_up[i](feat)
-            # generate scale and shift for SFT layers
-            scale = self.condition_scale[i](feat)
-            conditions.append(scale.clone())
-            shift = self.condition_shift[i](feat)
-            conditions.append(shift.clone())
-            # generate rgb images
-            if return_rgb:
-                out_rgbs.append(self.toRGB[i](feat))
-
-        # decoder
-        image, _ = self.stylegan_decoder([style_code],
-                                         conditions,
-                                         return_latents=return_latents,
-                                         input_is_latent=self.input_is_latent,
-                                         randomize_noise=randomize_noise)
-
-        return image

gfpgan/utils.py

@@ -75,10 +75,6 @@ class GFPGANer():
         elif arch == 'RestoreFormer':
             from gfpgan.archs.restoreformer_arch import RestoreFormer
             self.gfpgan = RestoreFormer()
-        elif arch == 'CodeFormer':
-            from gfpgan.archs.codeformer_arch import CodeFormer
-            self.gfpgan = CodeFormer(
-                dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256'])
         # initialize face helper
         self.face_helper = FaceRestoreHelper(
             upscale,

inference_gfpgan.py

@@ -41,7 +41,7 @@ def main():
         type=str,
         default='auto',
         help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs. Default: auto')
-    parser.add_argument('-w', '--weight', type=float, default=0.5, help='Adjustable weights for CodeFormer.')
+    parser.add_argument('-w', '--weight', type=float, default=0.5, help='Adjustable weights.')
     args = parser.parse_args()
 
     args = parser.parse_args()
@@ -104,11 +104,6 @@ def main():
         channel_multiplier = 2
         model_name = 'RestoreFormer'
         url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth'
-    elif args.version == 'CodeFormer':
-        arch = 'CodeFormer'
-        channel_multiplier = 2
-        model_name = 'CodeFormer'
-        url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth'
     else:
         raise ValueError(f'Wrong model version {args.version}.')