Move file location

modules/image_restoration/{model_downloader.py → real_esrgan/model_downloader.py}

@@ -1,8 +1,13 @@
 from modules.live_portrait.model_downloader import download_model
 
 MODELS_REALESRGAN_URL = {
-    "RealESRGAN_x2": "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x2.pth",
     "realesr-general-x4v3": "https://huggingface.co/jhj0517/realesr-general-x4v3/resolve/main/realesr-general-x4v3.pth",
+    "RealESRGAN_x2": "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x2.pth",
+}
+
+MODELS_REALESRGAN_SCALABILITY = {
+    "realesr-general-x4v3": [1, 2, 4],
+    "RealESRGAN_x2": [2]
 }
 
 

modules/image_restoration/{real_esrgan_inferencer.py → real_esrgan/real_esrgan_inferencer.py}

@@ -3,10 +3,14 @@ import gradio as gr
 import torch
 from PIL import Image
 import numpy as np
-from typing import Optional
+from typing import Optional, Literal, List, Dict, Tuple, Union
+from realesrgan.utils import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 
 from modules.utils.paths import *
-from .model_downloader import download_resrgan_model, MODELS_REALESRGAN_URL
+from .model_downloader import download_resrgan_model, MODELS_REALESRGAN_URL, MODELS_REALESRGAN_SCALABILITY
+from modules.utils.image_helper import save_image
+from .rrdb_net import RRDBNet
 
 
 class RealESRGANInferencer:
@@ -16,46 +20,88 @@ class RealESRGANInferencer:
         self.model_dir = model_dir
         self.output_dir = output_dir
         self.device = self.get_device()
+        self.arc = None
         self.model = None
-        self.up_sampler = None
         self.face_enhancer = None
 
         self.available_models = list(MODELS_REALESRGAN_URL.keys())
         self.default_model = self.available_models[0]
+        self.model_config = {
+            "model_name": self.default_model,
+            "scale": 1,
+            "half_precision": True
+        }
 
     def load_model(self,
                    model_name: Optional[str] = None,
-                   scale: int = 1,
+                   scale: Literal[1, 2, 4] = 1,
+                   half_precision: bool = True,
                    progress: gr.Progress = gr.Progress()):
+        model_config = {
+            "model_name": model_name,
+            "scale": scale,
+            "half_precision": half_precision
+        }
+        if model_config == self.model_config and self.model is not None:
+            return
+        else:
+            self.model_config = model_config
+
         if model_name is None:
             model_name = self.default_model
-        if not model_name.endswith(".pth"):
-            model_name += ".pth"
+
         model_path = os.path.join(self.model_dir, model_name)
+        if not model_name.endswith(".pth"):
+            model_path += ".pth"
 
         if not os.path.exists(model_path):
             progress(0, f"Downloading RealESRGAN model to : {model_path}")
-            name, ext = os.path.splitext(model_name)
-            download_resrgan_model(model_path, MODELS_REALESRGAN_URL[name])
+            download_resrgan_model(model_path, MODELS_REALESRGAN_URL[model_name])
+
+        name, ext = os.path.splitext(model_name)
+        assert scale in MODELS_REALESRGAN_SCALABILITY[name]
+        if name == 'RealESRGAN_x2':  # x4 RRDBNet model
+            arc = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+            netscale = 4
+        else:  # x4 VGG-style model (S size) : "realesr-general-x4v3"
+            arc = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+            netscale = 4
 
-        if self.model is None:
-            self.model = RealESRGAN(self.device, scale=scale)
-            self.model.load_weights(model_path=model_path, download=False)
+        self.model = RealESRGANer(
+            scale=netscale,
+            model_path=model_path,
+            model=arc,
+            half=half_precision,
+        )
+        self.model.device = torch.device(self.get_device())
 
     def restore_image(self,
                       img_path: str,
+                      model_name: Optional[str] = None,
+                      scale: int = 1,
+                      half_precision: bool = True,
                       overwrite: bool = True):
-        if self.model is None:
-            self.load_model()
+        model_config = {
+            "model_name": self.model_config["model_name"],
+            "scale": scale,
+            "half_precision": half_precision
+        }
+        if self.model is None or self.model_config != model_config:
+            self.load_model(
+                model_name=self.default_model if model_name is None else model_name,
+                scale=scale,
+                half_precision=half_precision
+            )
 
         try:
-            img = Image.open(img_path).convert('RGB')
-            sr_img = self.model.predict(img)
+            output, img_mode = self.model.enhance(img_path, outscale=scale)
+
             if overwrite:
                 output_path = img_path
             else:
                 output_path = get_auto_incremental_file_path(self.output_dir, extension="png")
-            sr_img.save(output_path)
+
+            output_path = save_image(output, output_path=output_path)
             return output_path
         except Exception as e:
             raise

modules/image_restoration/real_esrgan/rrdb_net.py

@@ -0,0 +1,182 @@
+from torch import nn as nn
+import torch
+from torch.nn import init as init
+from torch.nn import functional as F
+from torch.nn.modules.batchnorm import _BatchNorm
+
+
+class ResidualDenseBlock(nn.Module):
+    """Residual Dense Block.
+
+    Used in RRDB block in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat=64, num_grow_ch=32):
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
+        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        # Empirically, we use 0.2 to scale the residual for better performance
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    """Residual in Residual Dense Block.
+
+    Used in RRDB-Net in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat, num_grow_ch=32):
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
+
+    def forward(self, x):
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        # Empirically, we use 0.2 to scale the residual for better performance
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    """Networks consisting of Residual in Residual Dense Block, which is used
+    in ESRGAN.
+
+    ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
+
+    We extend ESRGAN for scale x2 and scale x1.
+    Note: This is one option for scale 1, scale 2 in RRDBNet.
+    We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size
+    and enlarge the channel size before feeding inputs into the main ESRGAN architecture.
+
+    Args:
+        num_in_ch (int): Channel number of inputs.
+        num_out_ch (int): Channel number of outputs.
+        num_feat (int): Channel number of intermediate features.
+            Default: 64
+        num_block (int): Block number in the trunk network. Defaults: 23
+        num_grow_ch (int): Channels for each growth. Default: 32.
+    """
+
+    def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32):
+        super(RRDBNet, self).__init__()
+        self.scale = scale
+        if scale == 2:
+            num_in_ch = num_in_ch * 4
+        elif scale == 1:
+            num_in_ch = num_in_ch * 16
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        # upsample
+        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        if self.scale == 2:
+            feat = pixel_unshuffle(x, scale=2)
+        elif self.scale == 1:
+            feat = pixel_unshuffle(x, scale=4)
+        else:
+            feat = x
+        feat = self.conv_first(feat)
+        body_feat = self.conv_body(self.body(feat))
+        feat = feat + body_feat
+        # upsample
+        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        return out
+
+
+def make_layer(basic_block, num_basic_block, **kwarg):
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (nn.module): nn.module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+
+def pixel_unshuffle(x, scale):
+    """ Pixel unshuffle.
+
+    Args:
+        x (Tensor): Input feature with shape (b, c, hh, hw).
+        scale (int): Downsample ratio.
+
+    Returns:
+        Tensor: the pixel unshuffled feature.
+    """
+    b, c, hh, hw = x.size()
+    out_channel = c * (scale**2)
+    assert hh % scale == 0 and hw % scale == 0
+    h = hh // scale
+    w = hw // scale
+    x_view = x.view(b, c, h, scale, w, scale)
+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)
+
+@torch.no_grad()
+def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
+    """Initialize network weights.
+
+    Args:
+        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
+        scale (float): Scale initialized weights, especially for residual
+            blocks. Default: 1.
+        bias_fill (float): The value to fill bias. Default: 0
+        kwargs (dict): Other arguments for initialization function.
+    """
+    if not isinstance(module_list, list):
+        module_list = [module_list]
+    for module in module_list:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, _BatchNorm):
+                init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)