add wfen node

.gitignore

@@ -1,4 +1,7 @@
+*.pth
+
 # Byte-compiled / optimized / DLL files
+__pycache__
 __pycache__/
 *.py[cod]
 *$py.class

README.md

@@ -1,3 +1,13 @@
 # FM_nodes
 
-A few ComfyUI nodes, starting with WFEN, [Efficient Face Super-Resolution via Wavelet-based Feature Enhancement Network](https://github.com/PRIS-CV/WFEN)
+A collection of ComfyUI nodes. Including:
+1. WFEN Face Restore, [Efficient Face Super-Resolution via Wavelet-based Feature Enhancement Network](https://github.com/PRIS-CV/WFEN)
+
+
+## Workflows
+
+### WFEN
+
+[workflow_wfen_facecrop.json](workflow/workflow_wfen_facecrop.json)
+
+![wfen_facecrop](workflow/wfen_facecrop.png)

__init__.py

@@ -0,0 +1,14 @@
+from .run import (
+    WFEN,
+)
+
+NODE_CLASS_MAPPINGS = {
+    "WFEN": WFEN
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "WFEN": "WFEN Face Restore"
+}
+
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

constants.py

@@ -0,0 +1 @@
+WFEN_MODEL_PATH = "models/wfen/WFEN.pth"

requirements.txt

@@ -0,0 +1,2 @@
+torch
+einops

run.py

@@ -0,0 +1,48 @@
+"""
+@author: Fuou Marinas
+@title: FM Nodes
+@nickname: FM_nodes
+@description: A collection of nodes. WFEN Face Super Resolution.
+"""
+
+from pathlib import Path
+
+import torch
+
+from .constants import WFEN_MODEL_PATH
+from .utils import ensure_size
+from .wfen_module.wfen_model import WFENModel
+
+base_dir = Path(__file__).resolve().parent
+
+
+class WFEN:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_img": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_img",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes"
+
+    def todo(
+        self,
+        src_img: torch.Tensor,
+    ):
+        print(f"{src_img.shape=}")
+        source_tensor = ensure_size(src_img.permute(0, 3, 1, 2))
+        print(f"{source_tensor.shape=}")
+        wfening = WFENModel()
+        wfening.for_load_pretrain_model(f"{base_dir}/{WFEN_MODEL_PATH}")
+
+        out_tensor = wfening.netG.forward(source_tensor)
+        # out_tensor = ensure_size(out_tensor, size=src_img.shape[1:3])
+        print(f"{out_tensor.shape=}")
+        out_tensor = out_tensor.permute(0, 2, 3, 1)
+
+        return (out_tensor,)

utils.py

@@ -0,0 +1,11 @@
+import torch
+import torch.nn.functional as F
+
+
+def ensure_size(tensor: torch.Tensor, size: tuple[int, int] = (128, 128)):
+    if tensor.shape[-2:] == size:
+        return tensor
+    else:
+        return F.interpolate(
+            tensor, size=size, mode="bilinear", align_corners=False
+        )

wfen_module/arch_blocks.py

@@ -0,0 +1,372 @@
+import math
+import numbers
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+
+
+# ---------------------------------------------------------------------------------------------------------------------
+# Layer Norm
+def to_3d(x):
+    return rearrange(x, "b c h w -> b (h w) c")
+
+
+def to_4d(x, h, w):
+    return rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
+
+
+class BiasFree_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(BiasFree_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return x / torch.sqrt(sigma + 1e-5) * self.weight
+
+
+class WithBias_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(WithBias_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        mu = x.mean(-1, keepdim=True)
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return (x - mu) / torch.sqrt(sigma + 1e-5) * self.weight + self.bias
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim, LayerNorm_type):
+        super(LayerNorm, self).__init__()
+        if LayerNorm_type == "BiasFree":
+            self.body = BiasFree_LayerNorm(dim)
+        else:
+            self.body = WithBias_LayerNorm(dim)
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+        return to_4d(self.body(to_3d(x)), h, w)
+
+
+# ---------------------------------------------------------------------------------------------------------------------
+
+
+# FFN
+class FeedForward(nn.Module):
+    def __init__(self, dim, ffn_expansion_factor, bias, input_resolution=None):
+        super(FeedForward, self).__init__()
+
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.ffn_expansion_factor = ffn_expansion_factor
+
+        hidden_features = int(dim * ffn_expansion_factor)
+        self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
+        self.dwconv = nn.Conv2d(
+            hidden_features * 2,
+            hidden_features * 2,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=hidden_features * 2,
+            bias=bias,
+        )
+        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        x = self.project_in(x)
+        x1, x2 = self.dwconv(x).chunk(2, dim=1)
+        x = F.gelu(x1) * x2
+        x = self.project_out(x)
+        return x
+
+
+class NormLayer(nn.Module):
+    """Normalization Layers.
+    ------------
+    # Arguments
+        - channels: input channels, for batch norm and instance norm.
+        - input_size: input shape without batch size, for layer norm.
+    """
+
+    def __init__(self, channels, normalize_shape=None, norm_type="bn"):
+        super(NormLayer, self).__init__()
+        norm_type = norm_type.lower()
+        if norm_type == "bn":
+            self.norm = nn.BatchNorm2d(channels)
+        elif norm_type == "in":
+            self.norm = nn.InstanceNorm2d(channels, affine=True)
+        elif norm_type == "gn":
+            self.norm = nn.GroupNorm(32, channels, affine=True)
+        elif norm_type == "pixel":
+            self.norm = lambda x: F.normalize(x, p=2, dim=1)
+        elif norm_type == "layer":
+            self.norm = nn.LayerNorm(normalize_shape)
+        elif norm_type == "none":
+            self.norm = lambda x: x
+        else:
+            assert 1 == 0, "Norm type {} not support.".format(norm_type)
+
+    def forward(self, x):
+        return self.norm(x)
+
+
+class ReluLayer(nn.Module):
+    """Relu Layer.
+    ------------
+    # Arguments
+        - relu type: type of relu layer, candidates are
+            - ReLU
+            - LeakyReLU: default relu slope 0.2
+            - PRelu
+            - SELU
+            - none: direct pass
+    """
+
+    def __init__(self, channels, relu_type="relu"):
+        super(ReluLayer, self).__init__()
+        relu_type = relu_type.lower()
+        if relu_type == "relu":
+            self.func = nn.ReLU(True)
+        elif relu_type == "leakyrelu":
+            self.func = nn.LeakyReLU(0.2, inplace=True)
+        elif relu_type == "prelu":
+            self.func = nn.PReLU(channels)
+        elif relu_type == "selu":
+            self.func = nn.SELU(True)
+        elif relu_type == "none":
+            self.func = lambda x: x
+        else:
+            assert 1 == 0, "Relu type {} not support.".format(relu_type)
+
+    def forward(self, x):
+        return self.func(x)
+
+
+class ConvLayer(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        scale="none",
+        norm_type="none",
+        relu_type="none",
+        use_pad=True,
+    ):
+        super(ConvLayer, self).__init__()
+        self.use_pad = use_pad
+
+        bias = True if norm_type in ["pixel", "none"] else False
+        stride = 2 if scale == "down" else 1
+
+        self.scale_func = lambda x: x
+        if scale == "up":
+            self.scale_func = lambda x: nn.functional.interpolate(
+                x, scale_factor=2, mode="nearest"
+            )
+
+        self.reflection_pad = nn.ReflectionPad2d(kernel_size // 2)
+        self.conv2d = nn.Conv2d(
+            in_channels, out_channels, kernel_size, stride, bias=bias
+        )
+
+        self.relu = ReluLayer(out_channels, relu_type)
+        self.norm = NormLayer(out_channels, norm_type=norm_type)
+
+    def forward(self, x):
+        out = self.scale_func(x)
+        if self.use_pad:
+            out = self.reflection_pad(out)
+        out = self.conv2d(out)
+        out = self.norm(out)
+        out = self.relu(out)
+        return out
+
+
+class ResidualBlock(nn.Module):
+    """
+    Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
+    ------------------
+    # Args
+        - hg_depth: depth of HourGlassBlock. 0: don't use attention map.
+        - use_pmask: whether use previous mask as HourGlassBlock input.
+    """
+
+    def __init__(
+        self,
+        c_in,
+        c_out,
+        relu_type="prelu",
+        norm_type="bn",
+        scale="none",
+        hg_depth=2,
+        att_name="spar",
+    ):
+        super(ResidualBlock, self).__init__()
+        self.c_in = c_in
+        self.c_out = c_out
+        self.norm_type = norm_type
+        self.relu_type = relu_type
+        self.hg_depth = hg_depth
+
+        kwargs = {"norm_type": norm_type, "relu_type": relu_type}
+
+        if scale == "none" and c_in == c_out:
+            self.shortcut_func = lambda x: x
+        else:
+            self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
+
+        self.preact_func = nn.Sequential(
+            NormLayer(c_in, norm_type=self.norm_type),
+            ReluLayer(c_in, self.relu_type),
+        )
+
+        if scale == "down":
+            scales = ["none", "down"]
+        elif scale == "up":
+            scales = ["up", "none"]
+        elif scale == "none":
+            scales = ["none", "none"]
+
+        self.conv1 = ConvLayer(c_in, c_out, 3, scales[0], **kwargs)
+        self.conv2 = ConvLayer(
+            c_out, c_out, 3, scales[1], norm_type=norm_type, relu_type="none"
+        )
+
+        if att_name.lower() == "spar":
+            c_attn = 1
+        elif att_name.lower() == "spar3d":
+            c_attn = c_out
+        else:
+            raise Exception("Attention type {} not implemented".format(att_name))
+
+        self.att_func = HourGlassBlock(self.hg_depth, c_out, c_attn, **kwargs)
+
+    def forward(self, x):
+        identity = self.shortcut_func(x)
+        out = self.preact_func(x)
+        out = self.conv1(out)
+        out = self.conv2(out)
+        out = identity + self.att_func(out)
+        return out
+
+
+class HourGlassBlock(nn.Module):
+    """Simplified HourGlass block.
+    Reference: https://github.com/1adrianb/face-alignment
+    --------------------------
+    """
+
+    def __init__(
+        self,
+        depth,
+        c_in,
+        c_out,
+        c_mid=64,
+        norm_type="bn",
+        relu_type="prelu",
+    ):
+        super(HourGlassBlock, self).__init__()
+        self.depth = depth
+        self.c_in = c_in
+        self.c_mid = c_mid
+        self.c_out = c_out
+        self.kwargs = {"norm_type": norm_type, "relu_type": relu_type}
+
+        if self.depth:
+            self._generate_network(self.depth)
+            self.out_block = nn.Sequential(
+                ConvLayer(self.c_mid, self.c_out, norm_type="none", relu_type="none"),
+                nn.Sigmoid(),
+            )
+
+    def _generate_network(self, level):
+        if level == self.depth:
+            c1, c2 = self.c_in, self.c_mid
+        else:
+            c1, c2 = self.c_mid, self.c_mid
+
+        self.add_module("b1_" + str(level), ConvLayer(c1, c2, **self.kwargs))
+        self.add_module(
+            "b2_" + str(level), ConvLayer(c1, c2, scale="down", **self.kwargs)
+        )
+        if level > 1:
+            self._generate_network(level - 1)
+        else:
+            self.add_module(
+                "b2_plus_" + str(level),
+                ConvLayer(self.c_mid, self.c_mid, **self.kwargs),
+            )
+
+        self.add_module(
+            "b3_" + str(level),
+            ConvLayer(self.c_mid, self.c_mid, scale="up", **self.kwargs),
+        )
+
+    def _forward(self, level, in_x):
+        up1 = self._modules["b1_" + str(level)](in_x)
+        low1 = self._modules["b2_" + str(level)](in_x)
+        if level > 1:
+            low2 = self._forward(level - 1, low1)
+        else:
+            low2 = self._modules["b2_plus_" + str(level)](low1)
+
+        up2 = self._modules["b3_" + str(level)](low2)
+        if up1.shape[2:] != up2.shape[2:]:
+            up2 = nn.functional.interpolate(up2, up1.shape[2:])
+
+        return up1 + up2
+
+    def forward(self, x, pmask=None):
+        if self.depth == 0:
+            return x
+        input_x = x
+        x = self._forward(self.depth, x)
+        self.att_map = self.out_block(x)
+        x = input_x * self.att_map
+        return x
+
+
+class UpSampler(nn.Sequential):
+    def __init__(self, scale, n_feats, out_feats):
+        m = []
+        if scale == 8:
+            kernel_size = 3
+        elif scale == 16:
+            kernel_size = 5
+        else:
+            kernel_size = 1
+
+        if (scale & (scale - 1)) == 0:  # Is scale = 2^n?
+            for _ in range(int(math.log(scale, 2))):
+                m.append(
+                    nn.Conv2d(
+                        in_channels=n_feats,
+                        out_channels=4 * out_feats,
+                        kernel_size=kernel_size,
+                        stride=1,
+                        padding=kernel_size // 2,
+                    )
+                )
+                m.append(nn.PixelShuffle(upscale_factor=2))
+                m.append(nn.PReLU())
+        super(UpSampler, self).__init__(*m)

wfen_module/models.py

@@ -0,0 +1,468 @@
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+from .arch_blocks import FeedForward, LayerNorm
+
+
+class GSA(nn.Module):
+    def __init__(self, channels, num_heads=8, bias=False):
+        super(GSA, self).__init__()
+        self.channels = channels
+        self.num_heads = num_heads
+
+        self.temperature = nn.Parameter(torch.ones(1, 1, 1))
+        self.act = nn.ReLU()
+
+        self.qkv = nn.Conv2d(channels, channels * 3, kernel_size=1, bias=bias)
+        self.qkv_dwconv = nn.Conv2d(
+            channels * 3,
+            channels * 3,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=channels * 3,
+            bias=bias,
+        )
+        self.project_out = nn.Conv2d(channels, channels, kernel_size=1, bias=bias)
+
+    def forward(self, x, prev_atns=None):
+        b, c, h, w = x.shape
+        if prev_atns is None:
+            qkv = self.qkv_dwconv(self.qkv(x))
+            q, k, v = qkv.chunk(3, dim=1)
+            q = rearrange(q, "b (head c) h w -> b head c (h w)", head=self.num_heads)
+            k = rearrange(k, "b (head c) h w -> b head c (h w)", head=self.num_heads)
+            v = rearrange(v, "b (head c) h w -> b head c (h w)", head=self.num_heads)
+
+            q = torch.nn.functional.normalize(q, dim=-1)
+            k = torch.nn.functional.normalize(k, dim=-1)
+
+            attn = (q @ k.transpose(-2, -1)) * self.temperature
+            attn = self.act(attn)
+            out = attn @ v
+            y = rearrange(
+                out, "b head c (h w) -> b (head c) h w", head=self.num_heads, h=h, w=w
+            )
+            y = rearrange(
+                y, "b (head c) h w -> b (c head) h w", head=self.num_heads, h=h, w=w
+            )
+            y = self.project_out(y)
+            return y, attn
+        else:
+            attn = prev_atns
+            v = rearrange(x, "b (head c) h w -> b head c (h w)", head=self.num_heads)
+            out = attn @ v
+            y = rearrange(
+                out, "b head c (h w) -> b (head c) h w", head=self.num_heads, h=h, w=w
+            )
+            y = rearrange(
+                y, "b (head c) h w -> b (c head) h w", head=self.num_heads, h=h, w=w
+            )
+            y = self.project_out(y)
+            return y
+
+
+class RSA(nn.Module):
+    def __init__(
+        self, channels, num_heads, shifts=1, window_sizes=[4, 8, 12], bias=False
+    ):
+        super(RSA, self).__init__()
+        self.channels = channels
+        self.shifts = shifts
+        self.window_sizes = window_sizes
+
+        self.temperature = nn.Parameter(torch.ones(1, 1, 1))
+        self.act = nn.ReLU()
+
+        self.qkv = nn.Conv2d(channels, channels * 3, kernel_size=1, bias=bias)
+        self.qkv_dwconv = nn.Conv2d(
+            channels * 3,
+            channels * 3,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=channels * 3,
+            bias=bias,
+        )
+        self.project_out = nn.Conv2d(channels, channels, kernel_size=1, bias=bias)
+
+    def forward(self, x, prev_atns=None):
+        b, c, h, w = x.shape
+        if prev_atns is None:
+            wsize = self.window_sizes
+            x_ = x
+            if self.shifts > 0:
+                x_ = torch.roll(x_, shifts=(-wsize // 2, -wsize // 2), dims=(2, 3))
+            qkv = self.qkv_dwconv(self.qkv(x_))
+            q, k, v = qkv.chunk(3, dim=1)
+            q = rearrange(
+                q, "b c (h dh) (w dw) -> b (h w) (dh dw) c", dh=wsize, dw=wsize
+            )
+            k = rearrange(
+                k, "b c (h dh) (w dw) -> b (h w) (dh dw) c", dh=wsize, dw=wsize
+            )
+            v = rearrange(
+                v, "b c (h dh) (w dw) -> b (h w) (dh dw) c", dh=wsize, dw=wsize
+            )
+
+            q = torch.nn.functional.normalize(q, dim=-1)
+            k = torch.nn.functional.normalize(k, dim=-1)
+
+            attn = (
+                q.transpose(-2, -1) @ k
+            ) * self.temperature  # b (h w) (dh dw) (dh dw)
+            attn = self.act(attn)
+            out = v @ attn
+            out = rearrange(
+                out,
+                "b (h w) (dh dw) c-> b (c) (h dh) (w dw)",
+                h=h // wsize,
+                w=w // wsize,
+                dh=wsize,
+                dw=wsize,
+            )
+            if self.shifts > 0:
+                out = torch.roll(out, shifts=(wsize // 2, wsize // 2), dims=(2, 3))
+            y = self.project_out(out)
+            return y, attn
+        else:
+            wsize = self.window_sizes
+            if self.shifts > 0:
+                x = torch.roll(x, shifts=(-wsize // 2, -wsize // 2), dims=(2, 3))
+            atn = prev_atns
+            v = rearrange(
+                x, "b (c) (h dh) (w dw) -> b (h w) (dh dw) c", dh=wsize, dw=wsize
+            )
+            y_ = v @ atn
+            y_ = rearrange(
+                y_,
+                "b (h w) (dh dw) c-> b (c) (h dh) (w dw)",
+                h=h // wsize,
+                w=w // wsize,
+                dh=wsize,
+                dw=wsize,
+            )
+            if self.shifts > 0:
+                y_ = torch.roll(y_, shifts=(wsize // 2, wsize // 2), dims=(2, 3))
+            y = self.project_out(y_)
+            return y
+
+
+class FDT(nn.Module):
+    def __init__(
+        self,
+        inp_channels,
+        window_sizes,
+        shifts,
+        num_heads,
+        shared_depth=1,
+        ffn_expansion_factor=2.66,
+    ):
+        super(FDT, self).__init__()
+        self.shared_depth = shared_depth
+
+        modules_ffd = {}
+        modules_att = {}
+        modules_norm = {}
+        for i in range(shared_depth):
+            modules_ffd["ffd{}".format(i)] = FeedForward(
+                inp_channels, ffn_expansion_factor, bias=False
+            )
+            modules_att["att_{}".format(i)] = RSA(
+                channels=inp_channels,
+                num_heads=num_heads,
+                shifts=shifts,
+                window_sizes=window_sizes,
+            )
+            modules_norm["norm_{}".format(i)] = LayerNorm(inp_channels, "WithBias")
+            modules_norm["norm_{}".format(i + 2)] = LayerNorm(inp_channels, "WithBias")
+        self.modules_ffd = nn.ModuleDict(modules_ffd)
+        self.modules_att = nn.ModuleDict(modules_att)
+        self.modules_norm = nn.ModuleDict(modules_norm)
+
+        modulec_ffd = {}
+        modulec_att = {}
+        modulec_norm = {}
+        for i in range(shared_depth):
+            modulec_ffd["ffd{}".format(i)] = FeedForward(
+                inp_channels, ffn_expansion_factor, bias=False
+            )
+            modulec_att["att_{}".format(i)] = GSA(
+                channels=inp_channels, num_heads=num_heads
+            )
+            modulec_norm["norm_{}".format(i)] = LayerNorm(inp_channels, "WithBias")
+            modulec_norm["norm_{}".format(i + 2)] = LayerNorm(inp_channels, "WithBias")
+        self.modulec_ffd = nn.ModuleDict(modulec_ffd)
+        self.modulec_att = nn.ModuleDict(modulec_att)
+        self.modulec_norm = nn.ModuleDict(modulec_norm)
+
+    def forward(self, x):
+        atn = None
+        B, C, H, W = x.size()
+        for i in range(self.shared_depth):
+            if i == 0:  ## only calculate attention for the 1-st module
+                x_, atn = self.modules_att["att_{}".format(i)](
+                    self.modules_norm["norm_{}".format(i)](x), None
+                )
+                x = (
+                    self.modules_ffd["ffd{}".format(i)](
+                        self.modules_norm["norm_{}".format(i + 2)](x_ + x)
+                    )
+                    + x_
+                )
+            else:
+                x_ = self.modules_att["att_{}".format(i)](
+                    self.modules_norm["norm_{}".format(i)](x), atn
+                )
+                x = (
+                    self.modules_ffd["ffd{}".format(i)](
+                        self.modules_norm["norm_{}".format(i + 2)](x_ + x)
+                    )
+                    + x_
+                )
+
+        for i in range(self.shared_depth):
+            if i == 0:  ## only calculate attention for the 1-st module
+                x_, atn = self.modulec_att["att_{}".format(i)](
+                    self.modulec_norm["norm_{}".format(i)](x), None
+                )
+                x = (
+                    self.modulec_ffd["ffd{}".format(i)](
+                        self.modulec_norm["norm_{}".format(i + 2)](x_ + x)
+                    )
+                    + x_
+                )
+            else:
+                x = self.modulec_att["att_{}".format(i)](
+                    self.modulec_norm["norm_{}".format(i)](x), atn
+                )
+                x = (
+                    self.modulec_ffd["ffd{}".format(i)](
+                        self.modulec_norm["norm_{}".format(i + 2)](x_ + x)
+                    )
+                    + x_
+                )
+
+        return x
+
+
+class HaarWavelet(nn.Module):
+    def __init__(self, in_channels, grad=False):
+        super(HaarWavelet, self).__init__()
+        self.in_channels = in_channels
+
+        self.haar_weights = torch.ones(4, 1, 2, 2)
+        # h
+        self.haar_weights[1, 0, 0, 1] = -1
+        self.haar_weights[1, 0, 1, 1] = -1
+        # v
+        self.haar_weights[2, 0, 1, 0] = -1
+        self.haar_weights[2, 0, 1, 1] = -1
+        # d
+        self.haar_weights[3, 0, 1, 0] = -1
+        self.haar_weights[3, 0, 0, 1] = -1
+
+        self.haar_weights = torch.cat([self.haar_weights] * self.in_channels, 0)
+        self.haar_weights = nn.Parameter(self.haar_weights)
+        self.haar_weights.requires_grad = grad
+
+    def forward(self, x, rev=False):
+        if not rev:
+            out = (
+                F.conv2d(
+                    x, self.haar_weights, bias=None, stride=2, groups=self.in_channels
+                )
+                / 4.0
+            )
+            out = out.reshape(
+                [x.shape[0], self.in_channels, 4, x.shape[2] // 2, x.shape[3] // 2]
+            )
+            out = torch.transpose(out, 1, 2)
+            out = out.reshape(
+                [x.shape[0], self.in_channels * 4, x.shape[2] // 2, x.shape[3] // 2]
+            )
+            return out
+        else:
+            out = x.reshape([x.shape[0], 4, self.in_channels, x.shape[2], x.shape[3]])
+            out = torch.transpose(out, 1, 2)
+            out = out.reshape(
+                [x.shape[0], self.in_channels * 4, x.shape[2], x.shape[3]]
+            )
+            return F.conv_transpose2d(
+                out, self.haar_weights, bias=None, stride=2, groups=self.in_channels
+            )
+
+
+class WFU(nn.Module):
+    def __init__(self, dim_big, dim_small):
+        super(WFU, self).__init__()
+        self.dim = dim_big
+        self.HaarWavelet = HaarWavelet(dim_big, grad=False)
+        self.InverseHaarWavelet = HaarWavelet(dim_big, grad=False)
+        self.RB = nn.Sequential(
+            nn.Conv2d(dim_big, dim_big, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(dim_big, dim_big, kernel_size=3, padding=1),
+        )
+
+        self.channel_tranformation = nn.Sequential(
+            nn.Conv2d(
+                dim_big + dim_small, dim_big + dim_small // 1, kernel_size=1, padding=0
+            ),
+            nn.ReLU(),
+            nn.Conv2d(dim_big + dim_small // 1, dim_big * 3, kernel_size=1, padding=0),
+        )
+
+    def forward(self, x_big, x_small):
+        haar = self.HaarWavelet(x_big, rev=False)
+        a = haar.narrow(1, 0, self.dim)
+        h = haar.narrow(1, self.dim, self.dim)
+        v = haar.narrow(1, self.dim * 2, self.dim)
+        d = haar.narrow(1, self.dim * 3, self.dim)
+
+        hvd = self.RB(h + v + d)
+        a_ = self.channel_tranformation(torch.cat([x_small, a], dim=1))
+        out = self.InverseHaarWavelet(torch.cat([hvd, a_], dim=1), rev=True)
+
+        return out
+
+
+class WFD(nn.Module):
+    def __init__(self, dim_in, dim, need=False):
+        super(WFD, self).__init__()
+        self.need = need
+        if need:
+            self.first_conv = nn.Conv2d(dim_in, dim, kernel_size=1, padding=0)
+            self.HaarWavelet = HaarWavelet(dim, grad=False)
+            self.dim = dim
+        else:
+            self.HaarWavelet = HaarWavelet(dim_in, grad=False)
+            self.dim = dim_in
+
+    def forward(self, x):
+        if self.need:
+            x = self.first_conv(x)
+
+        haar = self.HaarWavelet(x, rev=False)
+        a = haar.narrow(1, 0, self.dim)
+        h = haar.narrow(1, self.dim, self.dim)
+        v = haar.narrow(1, self.dim * 2, self.dim)
+        d = haar.narrow(1, self.dim * 3, self.dim)
+
+        return a, h + v + d
+
+
+class WFEN(nn.Module):
+    def __init__(
+        self,
+        inchannel=3,
+        min_ch=40,
+        res_depth=6,
+    ):
+        super(WFEN, self).__init__()
+        self.first_conv = nn.Conv2d(inchannel, min_ch, kernel_size=3, padding=1)
+
+        self.HaarDownsample1 = WFD(min_ch, min_ch * 2, True)
+        self.HaarDownsample2 = WFD(min_ch * 2, min_ch * 2 * 2, True)
+        self.HaarDownsample3 = WFD(min_ch * 2 * 2, min_ch * 2 * 2, True)
+
+        self.TransformerDown1 = nn.Sequential(
+            FDT(inp_channels=min_ch, window_sizes=8, shifts=0, num_heads=4),
+            FDT(inp_channels=min_ch, window_sizes=8, shifts=0, num_heads=4),
+        )
+        self.TransformerDown2 = nn.Sequential(
+            FDT(inp_channels=min_ch * 2, window_sizes=4, shifts=1, num_heads=4),
+        )
+        self.TransformerDown3 = nn.Sequential(
+            FDT(inp_channels=min_ch * 2 * 2, window_sizes=2, shifts=0, num_heads=8),
+        )
+
+        self.RB1 = nn.Sequential(
+            nn.Conv2d(min_ch * 2, min_ch * 2, kernel_size=1, padding=0, groups=1),
+            nn.ReLU(),
+            nn.Conv2d(min_ch * 2, min_ch * 2, kernel_size=1, padding=0, groups=1),
+        )
+        self.RB2 = nn.Sequential(
+            nn.Conv2d(
+                min_ch * 2 * 2, min_ch * 2 * 2, kernel_size=1, padding=0, groups=1
+            ),
+            nn.ReLU(),
+            nn.Conv2d(
+                min_ch * 2 * 2, min_ch * 2 * 2, kernel_size=1, padding=0, groups=1
+            ),
+        )
+        self.RB3 = nn.Sequential(
+            nn.Conv2d(
+                min_ch * 2 * 2, min_ch * 2 * 2, kernel_size=1, padding=0, groups=1
+            ),
+            nn.ReLU(),
+            nn.Conv2d(
+                min_ch * 2 * 2, min_ch * 2 * 2, kernel_size=1, padding=0, groups=1
+            ),
+        )
+
+        self.Transformer0 = FDT(
+            inp_channels=min_ch * 2 * 2, window_sizes=1, shifts=0, num_heads=8
+        )
+        Transformer = []
+        for i in range(res_depth - 1):
+            Transformer.append(
+                FDT(inp_channels=min_ch * 2 * 2, window_sizes=1, shifts=1, num_heads=8),
+            ) if i % 2 == 0 else Transformer.append(
+                FDT(inp_channels=min_ch * 2 * 2, window_sizes=1, shifts=0, num_heads=8)
+            )
+        self.Transformer = nn.Sequential(*Transformer)
+
+        self.TransformerUp1 = nn.Sequential(
+            FDT(inp_channels=min_ch * 2 * 2, window_sizes=8, shifts=1, num_heads=8),
+        )
+        self.TransformerUp2 = nn.Sequential(
+            FDT(inp_channels=min_ch * 2, window_sizes=4, shifts=0, num_heads=4),
+        )
+        self.TransformerUp3 = nn.Sequential(
+            FDT(inp_channels=min_ch, window_sizes=2, shifts=1, num_heads=4),
+            FDT(inp_channels=min_ch, window_sizes=2, shifts=0, num_heads=4),
+        )
+
+        self.HaarFeatureFusion1 = WFU(min_ch * 2 * 2, min_ch * 2 * 2)
+        self.HaarFeatureFusion2 = WFU(min_ch * 2, min_ch * 2 * 2)
+        self.HaarFeatureFusion3 = WFU(min_ch, min_ch * 2)
+
+        self.out_conv = nn.Conv2d(min_ch, inchannel, kernel_size=3, padding=1)
+
+    def forward(self, input_img):
+        x_first = self.first_conv(input_img)
+
+        ############ encoder ############
+        x1 = self.TransformerDown1(x_first)  # hw:128  c:40
+        x1_a, x1_hvd = self.HaarDownsample1(x1)  # all hw:64 c:80
+
+        x2 = self.TransformerDown2(x1_a)  # hw:64  c:80
+        x2 = x2 + self.RB1(x1_hvd)  # hw:64  c:160
+        x2_a, x2_hvd = self.HaarDownsample2(x2)  # all hw:40 c:160
+
+        x3 = self.TransformerDown3(x2_a)  # hw:40  c:160
+        x3 = x3 + self.RB2(x2_hvd)
+        x3_a, x3_hvd = self.HaarDownsample3(x3)  # all hw:16 c:160
+        ############ encoder ############
+
+        x_trans0 = self.Transformer0(x3_a)  # hw:16 c:160
+        x_trans = self.Transformer(x_trans0)  # hw:16 c:160
+        x_trans = x_trans + self.RB3(x3_hvd)  # hw:16 c:160
+
+        ############ decoder ############
+        x_up1 = self.HaarFeatureFusion1(x3, x_trans)  # hw:40 c:160
+        x_1 = self.TransformerUp1(x_up1)  # hw:40 c:160
+
+        x_up2 = self.HaarFeatureFusion2(x2, x_1)  # hw:64 c:80
+        x_2 = self.TransformerUp2(x_up2)  # hw:64 c:80
+
+        x_up3 = self.HaarFeatureFusion3(x1, x_2)  # hw:128 c:40
+        x_3 = self.TransformerUp3(x_up3)  # hw:32 c:40
+        ############ decoder ############
+
+        out_img = self.out_conv(x_3 + x_first)
+
+        return out_img

wfen_module/wfen_model.py

@@ -0,0 +1,32 @@
+import torch
+import torch.nn.utils as tutils
+from torch import nn
+
+from .models import WFEN
+
+
+def apply_norm(net: nn.Module, weight_norm_type: str):
+    for m in net.modules():
+        if isinstance(m, nn.Conv2d):
+            if weight_norm_type.lower() == "spectral_norm":
+                tutils.spectral_norm(m)
+            elif weight_norm_type.lower() == "weight_norm":
+                tutils.weight_norm(m)
+            else:
+                pass
+
+
+class WFENModel:
+    def __init__(self) -> None:
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.netG = WFEN()
+        self.netG.to(self.device)
+        self.netG = nn.DataParallel(self.netG)
+
+    def for_load_pretrain_model(self, path):
+        print("Loading pretrained model", path)
+        weight = torch.load(path)
+        self.netG.module.load_state_dict(weight)
+
+    def forward(self, img_LR):
+        return self.netG(img_LR)

workflow/workflow_wfen_facecrop.json

@@ -0,0 +1 @@
+{"last_node_id": 13, "last_link_id": 24, "nodes": [{"id": 10, "type": "PreviewImage", "pos": [628, 362], "size": [210, 246], "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 19}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 1, "type": "Image Crop Face", "pos": [599, 211], "size": {"0": 315, "1": 102}, "flags": {}, "order": 1, "mode": 0, "inputs": [{"name": "image", "type": "IMAGE", "link": 1, "slot_index": 0}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [2, 21], "shape": 3, "slot_index": 0}, {"name": "CROP_DATA", "type": "CROP_DATA", "links": [13], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "Image Crop Face"}, "widgets_values": [0.2, "lbpcascade_animeface.xml"]}, {"id": 3, "type": "PreviewImage", "pos": [972, 166], "size": [140, 97.08797341961395], "flags": {}, "order": 2, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 2}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 11, "type": "WFEN", "pos": [893, 370], "size": {"0": 210, "1": 26}, "flags": {}, "order": 3, "mode": 0, "inputs": [{"name": "src_img", "type": "IMAGE", "link": 21}], "outputs": [{"name": "res_img", "type": "IMAGE", "links": [22, 23], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "WFEN"}}, {"id": 12, "type": "PreviewImage", "pos": [1134, 123], "size": [140, 156.58081104083675], "flags": {}, "order": 4, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 22}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 9, "type": "Image Paste Face", "pos": [867, 487], "size": {"0": 315, "1": 122}, "flags": {}, "order": 5, "mode": 0, "inputs": [{"name": "image", "type": "IMAGE", "link": 16}, {"name": "crop_image", "type": "IMAGE", "link": 23}, {"name": "crop_data", "type": "CROP_DATA", "link": 13}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [24], "shape": 3, "slot_index": 0}, {"name": "MASK_IMAGE", "type": "IMAGE", "links": [19], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "Image Paste Face"}, "widgets_values": [0.25, 0]}, {"id": 13, "type": "SaveImage", "pos": [1236, 315], "size": [210, 285.6982048475769], "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 24}], "properties": {}, "widgets_values": ["WFEN"]}, {"id": 2, "type": "LoadImage", "pos": [1456, 267], "size": [210, 333.63426899370893], "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [1, 16], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": null, "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "LoadImage"}, "widgets_values": ["cj.png", "image"]}], "links": [[1, 2, 0, 1, 0, "IMAGE"], [2, 1, 0, 3, 0, "IMAGE"], [13, 1, 1, 9, 2, "CROP_DATA"], [16, 2, 0, 9, 0, "IMAGE"], [19, 9, 1, 10, 0, "IMAGE"], [21, 1, 0, 11, 0, "IMAGE"], [22, 11, 0, 12, 0, "IMAGE"], [23, 11, 0, 9, 1, "IMAGE"], [24, 9, 0, 13, 0, "IMAGE"]], "groups": [], "config": {}, "extra": {"ds": {"scale": 1.464100000000007, "offset": [-547.6249506008551, -46.66978242529448]}}, "version": 0.4}