diff --git a/README.md b/README.md index c3e1580..577a5f4 100644 --- a/README.md +++ b/README.md @@ -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} +} +``` \ No newline at end of file diff --git a/__init__.py b/__init__.py index 834c8f2..747eacb 100644 --- a/__init__.py +++ b/__init__.py @@ -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"] \ No newline at end of file +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/module_convir/ConvIR.py b/module_convir/ConvIR.py new file mode 100644 index 0000000..104c9d6 --- /dev/null +++ b/module_convir/ConvIR.py @@ -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) diff --git a/module_convir/constants_convir.py b/module_convir/constants_convir.py new file mode 100644 index 0000000..f40fa89 --- /dev/null +++ b/module_convir/constants_convir.py @@ -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" diff --git a/module_convir/convir_utils.py b/module_convir/convir_utils.py new file mode 100644 index 0000000..65155d3 --- /dev/null +++ b/module_convir/convir_utils.py @@ -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 diff --git a/module_convir/layers.py b/module_convir/layers.py new file mode 100644 index 0000000..9949dab --- /dev/null +++ b/module_convir/layers.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index f74f332..715019a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/run_convir.py b/run_convir.py new file mode 100644 index 0000000..3a02162 --- /dev/null +++ b/run_convir.py @@ -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,) diff --git a/workflow/convir.png b/workflow/convir.png new file mode 100644 index 0000000..97a998f Binary files /dev/null and b/workflow/convir.png differ diff --git a/workflow/workflow_convir.json b/workflow/workflow_convir.json new file mode 100644 index 0000000..ccb6ee2 --- /dev/null +++ b/workflow/workflow_convir.json @@ -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 +} \ No newline at end of file