add convir

2024-08-11 18:59:36 +07:00
parent a882d53d4e
commit 7aef6d9a2d
10 changed files with 1424 additions and 7 deletions
@@ -3,12 +3,12 @@
 A collection of ComfyUI nodes.

 Click name to jump to workflow
-1. [**WFEN Face Restore**](#wfen). Paper: [Efficient Face Super-Resolution via Wavelet-based Feature Enhancement Network](https://github.com/PRIS-CV/WFEN)
+1. [**WFEN**](#wfen) Face Restore. Paper: [Efficient Face Super-Resolution via Wavelet-based Feature Enhancement Network](https://github.com/PRIS-CV/WFEN)
 2. [**RealViformer**](#realviformer) - Paper: [Investigating Attention for Real-World Video Super-Resolution](https://github.com/Yuehan717/RealViformer)
 3. [**ProPIH**](#propih). Paper: [Progressive Painterly Image Harmonization from Low-level Styles to High-level Styles](https://github.com/bcmi/ProPIH-Painterly-Image-Harmonization)
 4. [**CoLIE**](#colie). Paper: [Fast Context-Based Low-Light Image Enhancement via Neural Implicit Representations](https://github.com/ctom2/colie)
 5. [**VFIMamba**](#vfimamba). Paper: [Video Frame Interpolation with State Space Models](https://github.com/MCG-NJU/VFIMamba)
-6. [**ConvIR**](#convir) [Revitalizing Convolutional Network for Image Restoration](https://github.com/c-yn/ConvIR)
+6. [**ConvIR**](#convir). Paper: [Revitalizing Convolutional Network for Image Restoration](https://github.com/c-yn/ConvIR)

 ## Workflows

@@ -85,3 +85,97 @@ VFIMamba Example (top: Original, bottom: 5X, 20FPS)

 ### ConvIR

+Download models in the [Pretrained models - gdrive](https://github.com/c-yn/ConvIR?tab=readme-ov-file#pretrained-models) section
+
+[workflow_convir.json](workflow/workflow_convir.json)
+
+![convir](workflow/convir.png)
+
+```
+models\convir
+│ deraining.pkl
+│
+├─defocus
+│   dpdd-base.pkl
+│   dpdd-large.pkl
+│   dpdd-small.pkl
+│
+├─dehaze
+│   densehaze-base.pkl
+│   densehaze-small.pkl
+│   gta5-base.pkl
+│   gta5-small.pkl
+│   haze4k-base.pkl
+│   haze4k-large.pkl
+│   haze4k-small.pkl
+│   ihaze-base.pkl
+│   ihaze-small.pkl
+│   its-base.pkl
+│   its-small.pkl
+│   nhhaze-base.pkl
+│   nhhaze-small.pkl
+│   nhr-base.pkl
+│   nhr-small.pkl
+│   ohaze-base.pkl
+│   ohaze-small.pkl
+│   ots-base.pkl
+│   ots-small.pkl
+│
+├─desnow
+│   csd-base.pkl
+│   csd-small.pkl
+│   snow100k-base.pkl
+│   snow100k-small.pkl
+│   srrs-base.pkl
+│   srrs-small.pkl
+│
+└─modeblur
+    convir_gopro.pkl
+    convir_rsblur.pkl
+```
+
+
+## Credits
+
+```
+@misc{chobola2024fast,
+      title={Fast Context-Based Low-Light Image Enhancement via Neural Implicit Representations}, 
+      author={Tomáš Chobola and Yu Liu and Hanyi Zhang and Julia A. Schnabel and Tingying Peng},
+      year={2024},
+      eprint={2407.12511},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2407.12511}, 
+}
+```
+
+```
+@misc{zhang2024vfimambavideoframeinterpolation,
+      title={VFIMamba: Video Frame Interpolation with State Space Models}, 
+      author={Guozhen Zhang and Chunxu Liu and Yutao Cui and Xiaotong Zhao and Kai Ma and Limin Wang},
+      year={2024},
+      eprint={2407.02315},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2407.02315}, 
+}
+```
+
+```
+@article{cui2024revitalizing,
+  title={Revitalizing Convolutional Network for Image Restoration},
+  author={Cui, Yuning and Ren, Wenqi and Cao, Xiaochun and Knoll, Alois},
+  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
+  year={2024},
+  publisher={IEEE}
+}
+
+@inproceedings{cui2023irnext,
+  title={IRNeXt: Rethinking Convolutional Network Design for Image Restoration},
+  author={Cui, Yuning and Ren, Wenqi and Yang, Sining and Cao, Xiaochun and Knoll, Alois},
+  booktitle={International Conference on Machine Learning},
+  pages={6545--6564},
+  year={2023},
+  organization={PMLR}
+}
+```
@@ -1,5 +1,11 @@
-from .run import (
-    WFEN, RealViFormerSR, ProPIH_Harmonizer, CoLIE_LowLight_Enhance
+from .run import WFEN, RealViFormerSR, ProPIH_Harmonizer, CoLIE_LowLight_Enhance
+
+from .run_convir import (
+    ConvIR_DeHaze,
+    ConvIR_DeRain,
+    ConvIR_DeSnow,
+    ConvIR_MotionDeBlur,
+    ConvIR_DefocusDeblur,
 )

 NODE_CLASS_MAPPINGS = {
@@ -7,6 +13,11 @@ NODE_CLASS_MAPPINGS = {
    "RealViFormerSR": RealViFormerSR,
    "ProPIH_Harmonizer": ProPIH_Harmonizer,
    "CoLIE_LowLight_Enhance": CoLIE_LowLight_Enhance,
+    "ConvIR_DeHaze": ConvIR_DeHaze,
+    "ConvIR_DeRain": ConvIR_DeRain,
+    "ConvIR_DeSnow": ConvIR_DeSnow,
+    "ConvIR_MotionDeBlur": ConvIR_MotionDeBlur,
+    "ConvIR_DefocusDeblur": ConvIR_DefocusDeblur,
 }

 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -14,10 +25,16 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "RealViFormerSR": "RealViFormer Video SR",
    "ProPIH_Harmonizer": "ProPIH Harmonizer",
    "CoLIE_LowLight_Enhance": "CoLIE LowLight Enhance",
+    "ConvIR_DeHaze": "ConvIR DeHaze",
+    "ConvIR_DeRain": "ConvIR DeRain",
+    "ConvIR_DeSnow": "ConvIR DeSnow",
+    "ConvIR_MotionDeBlur": "ConvIR Motion DeBlur",
+    "ConvIR_DefocusDeblur": "ConvIR Defocus Deblur",
 }

 try:
    from .vfi_mamba_run import VFI_MAMBA
+
    NODE_CLASS_MAPPINGS["VFI_MAMBA"] = VFI_MAMBA
    NODE_DISPLAY_NAME_MAPPINGS["VFI_MAMBA"] = "VFI Mamba"
 except ImportError as e:
@@ -25,4 +42,4 @@ except ImportError as e:
    pass


-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
@@ -0,0 +1,196 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .layers import BasicConv, ResBlock
+
+
+class EBlock(nn.Module):
+    def __init__(self, out_channel, num_res, data):
+        super(EBlock, self).__init__()
+
+        layers = [ResBlock(out_channel, out_channel, data) for _ in range(num_res - 1)]
+        layers.append(ResBlock(out_channel, out_channel, data, filter=True))
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class DBlock(nn.Module):
+    def __init__(self, channel, num_res, data):
+        super(DBlock, self).__init__()
+
+        layers = [ResBlock(channel, channel, data) for _ in range(num_res - 1)]
+        layers.append(ResBlock(channel, channel, data, filter=True))
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class SCM(nn.Module):
+    def __init__(self, out_plane):
+        super(SCM, self).__init__()
+        self.main = nn.Sequential(
+            BasicConv(3, out_plane // 4, kernel_size=3, stride=1, relu=True),
+            BasicConv(
+                out_plane // 4, out_plane // 2, kernel_size=1, stride=1, relu=True
+            ),
+            BasicConv(
+                out_plane // 2, out_plane // 2, kernel_size=3, stride=1, relu=True
+            ),
+            BasicConv(out_plane // 2, out_plane, kernel_size=1, stride=1, relu=False),
+            nn.InstanceNorm2d(out_plane, affine=True),
+        )
+
+    def forward(self, x):
+        x = self.main(x)
+        return x
+
+
+class FAM(nn.Module):
+    def __init__(self, channel):
+        super(FAM, self).__init__()
+        self.merge = BasicConv(
+            channel * 2, channel, kernel_size=3, stride=1, relu=False
+        )
+
+    def forward(self, x1, x2):
+        return self.merge(torch.cat([x1, x2], dim=1))
+
+
+class ConvIR(nn.Module):
+    def __init__(self, version, data):
+        super(ConvIR, self).__init__()
+
+        if version == "small":
+            num_res = 4
+        elif version == "base":
+            num_res = 8
+        elif version == "large":
+            num_res = 16
+
+        base_channel = 32
+
+        self.Encoder = nn.ModuleList(
+            [
+                EBlock(base_channel, num_res, data),
+                EBlock(base_channel * 2, num_res, data),
+                EBlock(base_channel * 4, num_res, data),
+            ]
+        )
+
+        self.feat_extract = nn.ModuleList(
+            [
+                BasicConv(3, base_channel, kernel_size=3, relu=True, stride=1),
+                BasicConv(
+                    base_channel, base_channel * 2, kernel_size=3, relu=True, stride=2
+                ),
+                BasicConv(
+                    base_channel * 2,
+                    base_channel * 4,
+                    kernel_size=3,
+                    relu=True,
+                    stride=2,
+                ),
+                BasicConv(
+                    base_channel * 4,
+                    base_channel * 2,
+                    kernel_size=4,
+                    relu=True,
+                    stride=2,
+                    transpose=True,
+                ),
+                BasicConv(
+                    base_channel * 2,
+                    base_channel,
+                    kernel_size=4,
+                    relu=True,
+                    stride=2,
+                    transpose=True,
+                ),
+                BasicConv(base_channel, 3, kernel_size=3, relu=False, stride=1),
+            ]
+        )
+
+        self.Decoder = nn.ModuleList(
+            [
+                DBlock(base_channel * 4, num_res, data),
+                DBlock(base_channel * 2, num_res, data),
+                DBlock(base_channel, num_res, data),
+            ]
+        )
+
+        self.Convs = nn.ModuleList(
+            [
+                BasicConv(
+                    base_channel * 4,
+                    base_channel * 2,
+                    kernel_size=1,
+                    relu=True,
+                    stride=1,
+                ),
+                BasicConv(
+                    base_channel * 2, base_channel, kernel_size=1, relu=True, stride=1
+                ),
+            ]
+        )
+
+        self.ConvsOut = nn.ModuleList(
+            [
+                BasicConv(base_channel * 4, 3, kernel_size=3, relu=False, stride=1),
+                BasicConv(base_channel * 2, 3, kernel_size=3, relu=False, stride=1),
+            ]
+        )
+
+        self.FAM1 = FAM(base_channel * 4)
+        self.SCM1 = SCM(base_channel * 4)
+        self.FAM2 = FAM(base_channel * 2)
+        self.SCM2 = SCM(base_channel * 2)
+
+    def forward(self, x):
+        x_2 = F.interpolate(x, scale_factor=0.5)
+        x_4 = F.interpolate(x_2, scale_factor=0.5)
+        z2 = self.SCM2(x_2)
+        z4 = self.SCM1(x_4)
+
+        outputs = list()
+        # 256
+        x_ = self.feat_extract[0](x)
+        res1 = self.Encoder[0](x_)
+        # 128
+        z = self.feat_extract[1](res1)
+        z = self.FAM2(z, z2)
+        res2 = self.Encoder[1](z)
+        # 64
+        z = self.feat_extract[2](res2)
+        z = self.FAM1(z, z4)
+        z = self.Encoder[2](z)
+
+        z = self.Decoder[0](z)
+        z_ = self.ConvsOut[0](z)
+        # 128
+        z = self.feat_extract[3](z)
+        outputs.append(z_ + x_4)
+
+        z = torch.cat([z, res2], dim=1)
+        z = self.Convs[0](z)
+        z = self.Decoder[1](z)
+        z_ = self.ConvsOut[1](z)
+        # 256
+        z = self.feat_extract[4](z)
+        outputs.append(z_ + x_2)
+
+        z = torch.cat([z, res1], dim=1)
+        z = self.Convs[1](z)
+        z = self.Decoder[2](z)
+        z = self.feat_extract[5](z)
+        outputs.append(z + x)
+
+        return outputs
+
+
+def build_net(version, data):
+    return ConvIR(version, data)
@@ -0,0 +1,66 @@
+MODEL_DIR = "models/convir"
+SIZE_BASE = "base"
+SIZE_SMALL = "small"
+SIZE_LARGE = "large"
+
+########### Image Dehazing ###########
+DEHAZE_FULL = "haze4k"
+DEHAZE_MODELS = [
+    "its",
+    "ots",
+    "densehaze",
+    "haze4k",
+    "ihaze",
+    "nhhaze",
+    "ohaze",
+    "gta5",
+    "nhr",
+]
+
+DEHAZE_DEFAULT = "haze4k"
+
+DEHAZE_SIZES = ["small", "base", "large"]
+
+
+def get_dehaze_model(model_name: str, model_size: str):
+    if model_name != DEHAZE_FULL and model_size == "large":
+        print(f"Invalid size for model {model_name}. Revert to {SIZE_BASE}")
+        model_size = SIZE_BASE
+    return f"dehaze/{model_name}-{model_size}.pkl", model_size
+
+
+########### Image Deraining ###########
+DERAIN_MODEL = "deraining"
+DERAIN_SIZE = "large"
+
+
+def get_derain_model():
+    return f"{DERAIN_MODEL}.pkl"
+
+
+########### Image Desnowing ###########
+DESNOW_MODELS = ["csd", "snow100k", "srrs"]
+DESNOW_DEFAULT = "snow100k"
+DESNOW_SIZES = ["small", "base"]
+
+
+def get_desnow_model(model_name: str, model_size: str):
+    return f"desnow/{model_name}-{model_size}.pkl"
+
+
+########### Defocus Deblurring ###########
+DEFOCUS_MODEL = "dpdd"
+DEFOCUS_SIZES = ["small", "base", "large"]
+
+
+def get_defocus_model(model_size: str):
+    return f"defocus/{DEFOCUS_MODEL}-{model_size}.pkl"
+
+
+########### Motion Deblurring ###########
+DEBLUR_MODELS = ["gopro", "rsblur"]
+DEBLUR_DEFAULT = "gopro"
+
+
+def get_deblur_model(model_name: str):
+    return f"modeblur/convir_{model_name}.pkl"
@@ -0,0 +1,30 @@
+import torch
+import torch.nn.functional as F
+
+
+def preprocess_image(image: torch.Tensor, device, factor=32, min_size=225):
+    # [B, H, W, C]
+    input_img = image.permute(0, 3, 1, 2).to(device)  # [B, C, H, W]
+    _, _, h, w = input_img.shape
+
+    # Resize if either dimension is smaller than min_size
+    if h < min_size or w < min_size:
+        scale = min_size / min(h, w)
+        new_h, new_w = int(h * scale), int(w * scale)
+        input_img = F.interpolate(
+            input_img, size=(new_h, new_w), mode="bilinear", align_corners=False
+        )
+        h, w = new_h, new_w
+
+    # Calculate padding
+    H = ((h + factor - 1) // factor) * factor
+    W = ((w + factor - 1) // factor) * factor
+
+    padh = H - h
+    padw = W - w
+
+    # Pad the image
+    if padh > 0 or padw > 0:
+        input_img = F.pad(input_img, (0, padw, 0, padh), mode="reflect")
+
+    return input_img, h, w
@@ -0,0 +1,271 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicConv(nn.Module):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        stride,
+        bias=True,
+        norm=False,
+        relu=True,
+        transpose=False,
+    ):
+        super(BasicConv, self).__init__()
+        if bias and norm:
+            bias = False
+
+        padding = kernel_size // 2
+        layers = list()
+        if transpose:
+            padding = kernel_size // 2 - 1
+            layers.append(
+                nn.ConvTranspose2d(
+                    in_channel,
+                    out_channel,
+                    kernel_size,
+                    padding=padding,
+                    stride=stride,
+                    bias=bias,
+                )
+            )
+        else:
+            layers.append(
+                nn.Conv2d(
+                    in_channel,
+                    out_channel,
+                    kernel_size,
+                    padding=padding,
+                    stride=stride,
+                    bias=bias,
+                )
+            )
+        if norm:
+            layers.append(nn.BatchNorm2d(out_channel))
+        if relu:
+            layers.append(nn.GELU())
+        self.main = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.main(x)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, data, filter=False):
+        super(ResBlock, self).__init__()
+        self.main = nn.Sequential(
+            BasicConv(in_channel, out_channel, kernel_size=3, stride=1, relu=True),
+            DeepPoolLayer(in_channel, out_channel, data) if filter else nn.Identity(),
+            BasicConv(out_channel, out_channel, kernel_size=3, stride=1, relu=False),
+        )
+
+    def forward(self, x):
+        return self.main(x) + x
+
+
+class DeepPoolLayer(nn.Module):
+    def __init__(self, k, k_out, data):
+        super(DeepPoolLayer, self).__init__()
+        self.pools_sizes = [8, 4, 2]
+
+        if data in [
+            # DeHaze
+            "its",
+            "ots",
+            "densehaze",
+            "haze4k",
+            "ihaze",
+            "nhhaze",
+            "nhr",
+            "ohaze",
+            # DeSnow
+            "csd",
+            "snow100k",
+            "srrs",
+            # DefocusDeBlur
+            "dpdd",
+            # DeBlur
+            "gopro",
+            "rsblur",
+        ]:
+            dilation = [7, 9, 11]
+        elif data in ["gta5"]:
+            dilation = [5, 9, 11]
+        elif data in ["deraining"]:
+            dilation = [3, 7, 9]
+
+        pools, convs, dynas = [], [], []
+        for j, i in enumerate(self.pools_sizes):
+            pools.append(nn.AvgPool2d(kernel_size=i, stride=i))
+            convs.append(nn.Conv2d(k, k, 3, 1, 1, bias=False))
+            dynas.append(MultiShapeKernel(dim=k, kernel_size=3, dilation=dilation[j]))
+        self.pools = nn.ModuleList(pools)
+        self.convs = nn.ModuleList(convs)
+        self.dynas = nn.ModuleList(dynas)
+        self.relu = nn.GELU()
+        self.conv_sum = nn.Conv2d(k, k_out, 3, 1, 1, bias=False)
+
+    def forward(self, x):
+        x_size = x.size()
+        resl = x
+        y_up = None
+        for i in range(len(self.pools_sizes)):
+            if i == 0:
+                y = self.dynas[i](self.convs[i](self.pools[i](x)))
+            else:
+                y = self.dynas[i](self.convs[i](self.pools[i](x) + y_up))
+            resl = torch.add(
+                resl, F.interpolate(y, x_size[2:], mode="bilinear", align_corners=True)
+            )
+            if i != len(self.pools_sizes) - 1:
+                y_up = F.interpolate(
+                    y, scale_factor=2, mode="bilinear", align_corners=True
+                )
+        resl = self.relu(resl)
+        resl = self.conv_sum(resl)
+
+        return resl
+
+
+class dynamic_filter(nn.Module):
+    def __init__(self, inchannels, kernel_size=3, dilation=1, stride=1, group=8):
+        super(dynamic_filter, self).__init__()
+        self.stride = stride
+        self.kernel_size = kernel_size
+        self.group = group
+        self.dilation = dilation
+
+        self.conv = nn.Conv2d(
+            inchannels, group * kernel_size**2, kernel_size=1, stride=1, bias=False
+        )
+        self.bn = nn.BatchNorm2d(group * kernel_size**2)
+        self.act = nn.Tanh()
+
+        nn.init.kaiming_normal_(self.conv.weight, mode="fan_out", nonlinearity="relu")
+        self.lamb_l = nn.Parameter(torch.zeros(inchannels), requires_grad=True)
+        self.lamb_h = nn.Parameter(torch.zeros(inchannels), requires_grad=True)
+        self.pad = nn.ReflectionPad2d(self.dilation * (kernel_size - 1) // 2)
+
+        self.ap = nn.AdaptiveAvgPool2d((1, 1))
+        self.gap = nn.AdaptiveAvgPool2d(1)
+
+        self.inside_all = nn.Parameter(
+            torch.zeros(inchannels, 1, 1), requires_grad=True
+        )
+
+    def forward(self, x):
+        identity_input = x
+        low_filter = self.ap(x)
+        low_filter = self.conv(low_filter)
+        low_filter = self.bn(low_filter)
+
+        n, c, h, w = x.shape
+        x = F.unfold(
+            self.pad(x), kernel_size=self.kernel_size, dilation=self.dilation
+        ).reshape(n, self.group, c // self.group, self.kernel_size**2, h * w)
+
+        n, c1, p, q = low_filter.shape
+        low_filter = low_filter.reshape(
+            n, c1 // self.kernel_size**2, self.kernel_size**2, p * q
+        ).unsqueeze(2)
+
+        low_filter = self.act(low_filter)
+
+        low_part = torch.sum(x * low_filter, dim=3).reshape(n, c, h, w)
+
+        out_low = low_part * (self.inside_all + 1.0) - self.inside_all * self.gap(
+            identity_input
+        )
+
+        out_low = out_low * self.lamb_l[None, :, None, None]
+
+        out_high = (identity_input) * (self.lamb_h[None, :, None, None] + 1.0)
+
+        return out_low + out_high
+
+
+class cubic_attention(nn.Module):
+    def __init__(self, dim, group, dilation, kernel) -> None:
+        super().__init__()
+
+        self.H_spatial_att = spatial_strip_att(
+            dim, dilation=dilation, group=group, kernel=kernel
+        )
+        self.W_spatial_att = spatial_strip_att(
+            dim, dilation=dilation, group=group, kernel=kernel, H=False
+        )
+        self.gamma = nn.Parameter(torch.zeros(dim, 1, 1))
+        self.beta = nn.Parameter(torch.ones(dim, 1, 1))
+
+    def forward(self, x):
+        out = self.H_spatial_att(x)
+        out = self.W_spatial_att(out)
+        return self.gamma * out + x * self.beta
+
+
+class spatial_strip_att(nn.Module):
+    def __init__(self, dim, kernel=3, dilation=1, group=2, H=True) -> None:
+        super().__init__()
+
+        self.k = kernel
+        pad = dilation * (kernel - 1) // 2
+        self.kernel = (1, kernel) if H else (kernel, 1)
+        self.padding = (kernel // 2, 1) if H else (1, kernel // 2)
+        self.dilation = dilation
+        self.group = group
+        self.pad = (
+            nn.ReflectionPad2d((pad, pad, 0, 0))
+            if H
+            else nn.ReflectionPad2d((0, 0, pad, pad))
+        )
+        self.conv = nn.Conv2d(dim, group * kernel, kernel_size=1, stride=1, bias=False)
+        self.ap = nn.AdaptiveAvgPool2d((1, 1))
+        self.filter_act = nn.Tanh()
+        self.inside_all = nn.Parameter(torch.zeros(dim, 1, 1), requires_grad=True)
+        self.lamb_l = nn.Parameter(torch.zeros(dim), requires_grad=True)
+        self.lamb_h = nn.Parameter(torch.zeros(dim), requires_grad=True)
+        gap_kernel = (None, 1) if H else (1, None)
+        self.gap = nn.AdaptiveAvgPool2d(gap_kernel)
+
+    def forward(self, x):
+        identity_input = x.clone()
+        filter = self.ap(x)
+        filter = self.conv(filter)
+        n, c, h, w = x.shape
+        x = F.unfold(
+            self.pad(x), kernel_size=self.kernel, dilation=self.dilation
+        ).reshape(n, self.group, c // self.group, self.k, h * w)
+        n, c1, p, q = filter.shape
+        filter = filter.reshape(n, c1 // self.k, self.k, p * q).unsqueeze(2)
+        filter = self.filter_act(filter)
+        out = torch.sum(x * filter, dim=3).reshape(n, c, h, w)
+
+        out_low = out * (self.inside_all + 1.0) - self.inside_all * self.gap(
+            identity_input
+        )
+        out_low = out_low * self.lamb_l[None, :, None, None]
+        out_high = identity_input * (self.lamb_h[None, :, None, None] + 1.0)
+
+        return out_low + out_high
+
+
+class MultiShapeKernel(nn.Module):
+    def __init__(self, dim, kernel_size=3, dilation=1, group=8):
+        super().__init__()
+
+        self.square_att = dynamic_filter(
+            inchannels=dim, dilation=dilation, group=group, kernel_size=kernel_size
+        )
+        self.strip_att = cubic_attention(
+            dim, group=group, dilation=dilation, kernel=kernel_size
+        )
+
+    def forward(self, x):
+        x1 = self.strip_att(x)
+        x2 = self.square_att(x)
+
+        return x1 + x2
@@ -1,7 +1,7 @@
 [project]
 name = "fm_nodes"
-description = "An expanding collection of ComfyUI nodes focused on new papers. Including: WFEN, RealViFormer, ProPIH, CoLIE, VFIMamba"
-version = "1.0.4"
+description = "An expanding collection of ComfyUI nodes focused on new papers. Including: WFEN, RealViFormer, ProPIH, CoLIE, VFIMamba, ConvIR"
+version = "1.0.5"
 license = {file = "LICENSE"}
 dependencies = ["torch", "einops"]

@@ -0,0 +1,170 @@
+from pathlib import Path
+
+import torch
+from comfy.utils import ProgressBar
+
+from .module_convir.constants_convir import (
+    DEBLUR_DEFAULT,
+    DEBLUR_MODELS,
+    DEFOCUS_MODEL,
+    DEFOCUS_SIZES,
+    DEHAZE_DEFAULT,
+    DEHAZE_MODELS,
+    DEHAZE_SIZES,
+    DERAIN_MODEL,
+    DERAIN_SIZE,
+    DESNOW_DEFAULT,
+    DESNOW_MODELS,
+    DESNOW_SIZES,
+    MODEL_DIR,
+    SIZE_BASE,
+    SIZE_LARGE,
+    get_deblur_model,
+    get_defocus_model,
+    get_dehaze_model,
+    get_derain_model,
+    get_desnow_model,
+)
+from .module_convir.ConvIR import ConvIR
+from .module_convir.convir_utils import preprocess_image
+
+base_dir = Path(__file__).resolve().parent
+
+
+def run_convir(src_img, model_name, model_size, model_file):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    src_img, h, w = preprocess_image(src_img, device)
+
+    model = ConvIR(version=model_size, data=model_name)
+
+    checkpoint_path = f"{base_dir}/{MODEL_DIR}/{model_file}"
+    state_dict = torch.load(checkpoint_path)
+    model.load_state_dict(state_dict["model"])
+    model = model.to(device)
+    model.eval()
+
+    num_frames = src_img.size(0)
+    pbar = ProgressBar(num_frames)
+    result: list[torch.Tensor] = []
+
+    for i in range(num_frames):
+        image = src_img[i].unsqueeze(0)
+        pred = model.forward(image)[2]
+        pred = pred[:, :, :h, :w]
+        pred_clip = torch.clamp(pred, 0, 1)
+        result.append(pred_clip)
+        pbar.update_absolute(i, num_frames)
+
+    result_ts = torch.cat(result, dim=0).permute(0, 2, 3, 1)
+    return result_ts
+
+
+class ConvIR_DeHaze:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_img": ("IMAGE",),
+                "model_name": (DEHAZE_MODELS, {"default": DEHAZE_DEFAULT}),
+                "model_size": (DEHAZE_SIZES, {"default": SIZE_BASE}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_img",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes/ConvIR"
+
+    def todo(self, src_img: torch.Tensor, model_name: str, model_size: str):
+        model_file, model_size = get_dehaze_model(model_name, model_size)
+
+        result_ts = run_convir(src_img, model_name, model_size, model_file)
+        return (result_ts,)
+
+
+class ConvIR_DeRain:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_img": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_img",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes/ConvIR"
+
+    def todo(self, src_img: torch.Tensor):
+        model_file = get_derain_model()
+
+        result_ts = run_convir(src_img, DERAIN_MODEL, DERAIN_SIZE, model_file)
+        return (result_ts,)
+
+
+class ConvIR_DeSnow:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_img": ("IMAGE",),
+                "model_name": (DESNOW_MODELS, {"default": DESNOW_DEFAULT}),
+                "model_size": (DESNOW_SIZES, {"default": SIZE_BASE}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_img",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes/ConvIR"
+
+    def todo(self, src_img: torch.Tensor, model_name: str, model_size: str):
+        model_file = get_desnow_model(model_name, model_size)
+
+        result_ts = run_convir(src_img, model_name, model_size, model_file)
+        return (result_ts,)
+
+
+class ConvIR_MotionDeBlur:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_img": ("IMAGE",),
+                "model_name": (DEBLUR_MODELS, {"default": DEBLUR_DEFAULT}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_img",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes/ConvIR"
+
+    def todo(self, src_img: torch.Tensor, model_name: str):
+        model_file = get_deblur_model(model_name)
+
+        result_ts = run_convir(src_img, model_name, SIZE_LARGE, model_file)
+        return (result_ts,)
+
+
+class ConvIR_DefocusDeblur:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_img": ("IMAGE",),
+                "model_size": (DEFOCUS_SIZES, {"default": SIZE_BASE}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("res_img",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes/ConvIR"
+
+    def todo(self, src_img: torch.Tensor, model_size: str):
+        model_file = get_defocus_model(model_size)
+
+        result_ts = run_convir(src_img, DEFOCUS_MODEL, model_size, model_file)
+        return (result_ts,)
@@ -0,0 +1,573 @@
+{
+  "last_node_id": 16,
+  "last_link_id": 59,
+  "nodes": [
+    {
+      "id": 12,
+      "type": "Reroute",
+      "pos": [
+        330,
+        -43
+      ],
+      "size": [
+        75,
+        26
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "",
+          "type": "*",
+          "link": 30
+        }
+      ],
+      "outputs": [
+        {
+          "name": "",
+          "type": "IMAGE",
+          "links": [
+            31,
+            32,
+            33,
+            40,
+            41
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "showOutputText": false,
+        "horizontal": false
+      }
+    },
+    {
+      "id": 2,
+      "type": "LoadImage",
+      "pos": [
+        -54,
+        155
+      ],
+      "size": {
+        "0": 422.6976623535156,
+        "1": 339.1871643066406
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            30,
+            43
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "cat.jpg",
+        "image"
+      ]
+    },
+    {
+      "id": 13,
+      "type": "Reroute",
+      "pos": [
+        491,
+        447
+      ],
+      "size": [
+        75,
+        26
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "",
+          "type": "*",
+          "link": 43
+        }
+      ],
+      "outputs": [
+        {
+          "name": "",
+          "type": "IMAGE",
+          "links": [
+            44
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "showOutputText": false,
+        "horizontal": false
+      }
+    },
+    {
+      "id": 4,
+      "type": "Image Comparer (rgthree)",
+      "pos": {
+        "0": 1081,
+        "1": 228,
+        "2": 0,
+        "3": 0,
+        "4": 0,
+        "5": 0,
+        "6": 0,
+        "7": 0,
+        "8": 0,
+        "9": 0
+      },
+      "size": {
+        "0": 516.6257934570312,
+        "1": 345.0682067871094
+      },
+      "flags": {},
+      "order": 10,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "image_a",
+          "type": "IMAGE",
+          "link": 26,
+          "dir": 3
+        },
+        {
+          "name": "image_b",
+          "type": "IMAGE",
+          "link": 45,
+          "dir": 3
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "comparer_mode": "Slide"
+      },
+      "widgets_values": [
+        [
+          {
+            "name": "A",
+            "selected": true,
+            "url": "/view?filename=rgthree.compare._temp_ekoah_00015_.png&type=temp&subfolder=&rand=0.05077056793140655"
+          },
+          {
+            "name": "B",
+            "selected": true,
+            "url": "/view?filename=rgthree.compare._temp_ekoah_00016_.png&type=temp&subfolder=&rand=0.8029056834303627"
+          }
+        ]
+      ]
+    },
+    {
+      "id": 10,
+      "type": "Reroute",
+      "pos": [
+        865,
+        114
+      ],
+      "size": [
+        75,
+        26
+      ],
+      "flags": {},
+      "order": 9,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "",
+          "type": "*",
+          "link": 59
+        }
+      ],
+      "outputs": [
+        {
+          "name": "",
+          "type": "IMAGE",
+          "links": [
+            26,
+            54
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "showOutputText": false,
+        "horizontal": false
+      }
+    },
+    {
+      "id": 9,
+      "type": "ConvIR_DeRain",
+      "pos": [
+        591,
+        -16
+      ],
+      "size": {
+        "0": 210,
+        "1": 26
+      },
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "src_img",
+          "type": "IMAGE",
+          "link": 32
+        }
+      ],
+      "outputs": [
+        {
+          "name": "res_img",
+          "type": "IMAGE",
+          "links": [],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ConvIR_DeRain"
+      }
+    },
+    {
+      "id": 11,
+      "type": "ConvIR_DeSnow",
+      "pos": [
+        498,
+        78
+      ],
+      "size": {
+        "0": 315,
+        "1": 82
+      },
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "src_img",
+          "type": "IMAGE",
+          "link": 33
+        }
+      ],
+      "outputs": [
+        {
+          "name": "res_img",
+          "type": "IMAGE",
+          "links": [],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ConvIR_DeSnow"
+      },
+      "widgets_values": [
+        "snow100k",
+        "base"
+      ]
+    },
+    {
+      "id": 14,
+      "type": "ConvIR_MotionDeBlur",
+      "pos": [
+        490,
+        227
+      ],
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "src_img",
+          "type": "IMAGE",
+          "link": 41
+        }
+      ],
+      "outputs": [
+        {
+          "name": "res_img",
+          "type": "IMAGE",
+          "links": [],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ConvIR_MotionDeBlur"
+      },
+      "widgets_values": [
+        "gopro"
+      ]
+    },
+    {
+      "id": 15,
+      "type": "ConvIR_DefocusDeblur",
+      "pos": [
+        487,
+        352
+      ],
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "src_img",
+          "type": "IMAGE",
+          "link": 40
+        }
+      ],
+      "outputs": [
+        {
+          "name": "res_img",
+          "type": "IMAGE",
+          "links": [],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ConvIR_DefocusDeblur"
+      },
+      "widgets_values": [
+        "base"
+      ]
+    },
+    {
+      "id": 1,
+      "type": "ConvIR_DeHaze",
+      "pos": [
+        501,
+        -153
+      ],
+      "size": {
+        "0": 315,
+        "1": 82
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "src_img",
+          "type": "IMAGE",
+          "link": 31,
+          "slot_index": 0
+        }
+      ],
+      "outputs": [
+        {
+          "name": "res_img",
+          "type": "IMAGE",
+          "links": [
+            59
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ConvIR_DeHaze"
+      },
+      "widgets_values": [
+        "haze4k",
+        "base"
+      ]
+    },
+    {
+      "id": 5,
+      "type": "SaveImage",
+      "pos": [
+        1083,
+        -171
+      ],
+      "size": {
+        "0": 517.8594970703125,
+        "1": 340.2516784667969
+      },
+      "flags": {},
+      "order": 11,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 54
+        }
+      ],
+      "properties": {},
+      "widgets_values": [
+        "ConvIR"
+      ]
+    },
+    {
+      "id": 16,
+      "type": "Reroute",
+      "pos": [
+        842,
+        448
+      ],
+      "size": [
+        75,
+        26
+      ],
+      "flags": {},
+      "order": 8,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "",
+          "type": "*",
+          "link": 44
+        }
+      ],
+      "outputs": [
+        {
+          "name": "",
+          "type": "IMAGE",
+          "links": [
+            45
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "showOutputText": false,
+        "horizontal": false
+      }
+    }
+  ],
+  "links": [
+    [
+      26,
+      10,
+      0,
+      4,
+      0,
+      "IMAGE"
+    ],
+    [
+      30,
+      2,
+      0,
+      12,
+      0,
+      "*"
+    ],
+    [
+      31,
+      12,
+      0,
+      1,
+      0,
+      "IMAGE"
+    ],
+    [
+      32,
+      12,
+      0,
+      9,
+      0,
+      "IMAGE"
+    ],
+    [
+      33,
+      12,
+      0,
+      11,
+      0,
+      "IMAGE"
+    ],
+    [
+      40,
+      12,
+      0,
+      15,
+      0,
+      "IMAGE"
+    ],
+    [
+      41,
+      12,
+      0,
+      14,
+      0,
+      "IMAGE"
+    ],
+    [
+      43,
+      2,
+      0,
+      13,
+      0,
+      "*"
+    ],
+    [
+      44,
+      13,
+      0,
+      16,
+      0,
+      "*"
+    ],
+    [
+      45,
+      16,
+      0,
+      4,
+      1,
+      "IMAGE"
+    ],
+    [
+      54,
+      10,
+      0,
+      5,
+      0,
+      "IMAGE"
+    ],
+    [
+      59,
+      1,
+      0,
+      10,
+      0,
+      "*"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        -140,
+        227
+      ]
+    }
+  },
+  "version": 0.4
+}