diff --git a/README.md b/README.md index 0ef2865..f235d9a 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,8 @@ Click name to jump to workflow 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. [**AnySR**](#anysr). Paper: [Realizing Image Super-Resolution as Any-Scale, Any-Resource](https://github.com/CrispyFeSo4/AnySR) + ## Workflows ### WFEN @@ -52,3 +53,12 @@ No model needed to be downloaded. Lower `loss_mean` seems to result in brighter [workflow_colie_lowlight.json](workflow/workflow_colie_lowlight.json) ![colie_lowlight](workflow/colie_lowlight.png) + +### AnySR + +Download `anysr_edsr_500.pth(updated)` and `anysr_rdn_500.pth(updated)` from [Checkpoints](https://github.com/CrispyFeSo4/AnySR?tab=readme-ov-file#checkpoints) + +``` +models/anysr/anysr_edsr_500.pth +models/anysr/anysr_rdn_500.pth +``` \ No newline at end of file diff --git a/__init__.py b/__init__.py index 2bdd7e6..96793cd 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,5 @@ from .run import ( - WFEN, RealViFormerSR, ProPIH_Harmonizer, CoLIE_LowLight_Enhance + WFEN, RealViFormerSR, ProPIH_Harmonizer, CoLIE_LowLight_Enhance, AnySR ) NODE_CLASS_MAPPINGS = { @@ -7,6 +7,7 @@ NODE_CLASS_MAPPINGS = { "RealViFormerSR": RealViFormerSR, "ProPIH_Harmonizer": ProPIH_Harmonizer, "CoLIE_LowLight_Enhance": CoLIE_LowLight_Enhance, + "AnySR": AnySR, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -14,6 +15,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "RealViFormerSR": "RealViFormer Video SR", "ProPIH_Harmonizer": "ProPIH Harmonizer", "CoLIE_LowLight_Enhance": "CoLIE LowLight Enhance", + "AnySR": "AnySR", } diff --git a/anysr_module/anysr_model.py b/anysr_module/anysr_model.py new file mode 100644 index 0000000..f6310e8 --- /dev/null +++ b/anysr_module/anysr_model.py @@ -0,0 +1,87 @@ +import torch +from .edsr_anysr import make_edsr +from .rdn_anysr import make_rdn + + +def run_anysr( + image: torch.Tensor, model_arch: str, model_path: str, scale: float, entire_net=True +): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + scale_max = 4 + + model_load = torch.load(model_path)['model'] + model_args = model_load['args']['encoder_spec']['args'] + pretrain_dict = model_load['sd'] + + print(f"{model_args=}") + + image_tensor = image.permute(0, 3, 1, 2).squeeze(dim=0).to(device) + print(f"{image_tensor.shape=}") + if model_arch == "edsr": + model = make_edsr(scale=scale, **model_args) + elif model_arch == "rdn": + model = make_rdn(scale=scale, **model_args) + + model_dict = model.state_dict() + model_dict = replace(pretrain_dict, model_dict) + model.load_state_dict(model_dict, strict=False) + + + # model.load_state_dict(torch.load(model_path)['model']['sd']) + + model = model.to(device) + + h = int(image_tensor.shape[-2] * int(scale)) + w = int(image_tensor.shape[-1] * int(scale)) + scale = h / image_tensor.shape[-2] + + # coord = make_coord((h, w), flatten=False).to(device) + # cell = torch.ones(1, 2).to(device) + # cell[:, 0] *= 2 / h + # cell[:, 1] *= 2 / w + + # cell_factor = max(scale / scale_max, 1) + + # model.encoder.scale = scale + # model.encoder.scale2 = scale + + pred = model( + image_tensor.cuda().unsqueeze(0), + ).squeeze(0) + print(f"{pred.shape=}") + pred = (pred * 0.5 + 0.5).clamp(0, 1).reshape(3, h, w).cpu() + print(f"{pred.shape=}") + pred = pred.permute(1, 2, 0).unsqueeze(dim=0) + + return pred + +def replace(pretrain_dict, model_dict): + new_dict = {} + for k,v in pretrain_dict.items(): + if 'body.0.' in k[15:]: + newk = k[:15] + k[15:].replace('body.0','conv1') + new_dict[newk] = v + elif 'body.2' in k[15:]: + newk = k[:15] + k[15:].replace('body.2','conv2') + new_dict[newk] = v + pretrain_dict.update(new_dict) + common_dict = {k:v for k,v in pretrain_dict.items() if k in model_dict} + model_dict.update(common_dict) + return model_dict + +def make_coord(shape, ranges=None, flatten=True): + """Make coordinates at grid centers.""" + coord_seqs = [] + for i, n in enumerate(shape): + if ranges is None: + v0, v1 = -1, 1 + else: + v0, v1 = ranges[i] + r = (v1 - v0) / (2 * n) + seq = v0 + r + (2 * r) * torch.arange(n).float() + coord_seqs.append(seq) + + ret = torch.stack(torch.meshgrid(*coord_seqs, indexing="ij"), dim=-1) + if flatten: + ret = ret.view(-1, ret.shape[-1]) + return ret diff --git a/anysr_module/arch_util.py b/anysr_module/arch_util.py new file mode 100644 index 0000000..2818d83 --- /dev/null +++ b/anysr_module/arch_util.py @@ -0,0 +1,126 @@ +import torch +import torch.nn as nn +import torch.nn.init as init +import torch.nn.functional as F +# import math + +from torch import Tensor +from typing import Optional, List + + +def make_divisible(v, divisor=8, min_value=8): + if min_value is None: + min_value = divisor + new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_v < 0.9 * v: + new_v += divisor + return int(new_v) + +def initialize_weights(net_l, scale=1): + if not isinstance(net_l, list): + net_l = [net_l] + for net in net_l: + for m in net.modules(): + if isinstance(m, nn.Conv2d): + init.kaiming_normal_(m.weight, a=0, mode='fan_in') + m.weight.data *= scale # for residual block + if m.bias is not None: + m.bias.data.zero_() + elif isinstance(m, nn.Linear): + init.kaiming_normal_(m.weight, a=0, mode='fan_in') + m.weight.data *= scale + if m.bias is not None: + m.bias.data.zero_() + elif isinstance(m, nn.BatchNorm2d): + init.constant_(m.weight, 1) + init.constant_(m.bias.data, 0.0) + +class MeanShift(nn.Conv2d): + def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1): + super(MeanShift, self).__init__(3, 3, kernel_size=1) + std = torch.Tensor(rgb_std) + self.weight.data = torch.eye(3).view(3, 3, 1, 1) + self.weight.data.div_(std.view(3, 1, 1, 1)) + self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean) + self.bias.data.div_(std) + # self.requires_grad = False + for p in self.parameters(): + p.requires_grad = False + +class USConv2d(nn.Conv2d): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, us=[False, False]): + super(USConv2d, self).__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias) + self.width_mult = 1.0 + self.us = us + + def set_width_mult(self, width_mult): + self.width_mult = width_mult + + def forward(self, inputs): + in_channels = inputs.shape[1] // self.groups if self.us[0] else self.in_channels // self.groups + out_channels = self.out_channels + if self.us[1]: + if self.width_mult is not None: + out_channels = int(self.out_channels * self.width_mult) + else: + self.width_mult = 1.0 + out_channels = self.out_channels + + weight = self.weight[:out_channels, :in_channels, :, :] + if self.bias is not None: + bias = self.bias[:out_channels] + else: + bias = self.bias + + y = F.conv2d(inputs, weight, bias, self.stride, self.padding, self.dilation, self.groups) + return y + + +class USConvTranspose2d(nn.ConvTranspose2d): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, us=[False, False]): + super(USConvTranspose2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, output_padding = output_padding) + self.width_mult = None + self.us = us + + def forward(self, input: Tensor, output_size: Optional[List[int]] = None) -> Tensor: + # in_channels = make_divisible(self.in_channels * self.width_mult) if self.us[0] else self.in_channels + in_channels = int(self.in_channels * self.width_mult) if self.us[0] else self.in_channels + out_channels = input.shape[1] if self.us[1] else self.out_channels + + + weight = self.weight[:in_channels, :out_channels, :, :] + + assert isinstance(self.padding, tuple) + output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, self.dilation) # type: ignore[arg-type] + + return F.conv_transpose2d(input, weight, self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + + +class USBatchNorm2d(nn.BatchNorm2d): + def __init__(self, num_features, width_list = None): + super(USBatchNorm2d, self).__init__(num_features, affine=True, track_running_stats=False) + self.width_id = None + + self.bn = nn.ModuleList([ + nn.BatchNorm2d(self.num_features, affine=False) for _ in range(len(width_list)) + ]) + # raise NotImplementedError + + def forward(self, inputs): + num_features = inputs.size(1) + y = F.batch_norm( + inputs, + self.bn[self.width_id].running_mean[:num_features], + self.bn[self.width_id].running_var[:num_features], + self.weight[:num_features], + self.bias[:num_features], + self.training, + self.momentum, + self.eps) + return y + + + + + diff --git a/anysr_module/edsr_anysr.py b/anysr_module/edsr_anysr.py new file mode 100644 index 0000000..f575262 --- /dev/null +++ b/anysr_module/edsr_anysr.py @@ -0,0 +1,296 @@ +# modified from: https://github.com/thstkdgus35/EDSR-PyTorch + +import argparse + +# from models import register +import math +import random +from argparse import Namespace + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# torch.autograd.set_detect_anomaly(True) +# import torch.nn as nn +# import torch.nn.functional as F +# from torch.autograd import Variable + +from . import arch_util +from .arch_util import USConv2d + + +def default_conv(in_channels, out_channels, kernel_size, bias=True): + return nn.Conv2d( + in_channels, out_channels, kernel_size, + padding=(kernel_size//2), bias=bias) + +class MeanShift(nn.Conv2d): + def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1): + super(MeanShift, self).__init__(3, 3, kernel_size=1) + std = torch.Tensor(rgb_std) + self.weight.data = torch.eye(3).view(3, 3, 1, 1) + self.weight.data.div_(std.view(3, 1, 1, 1)) + self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean) + self.bias.data.div_(std) + self.requires_grad = False + +class ResBlock(nn.Module): + def __init__( + self, conv, n_feats, kernel_size, + bias=True, bn=False, act=nn.ReLU(True), res_scale=1): + + super(ResBlock, self).__init__() + m = [] + for i in range(2): + m.append(conv(n_feats, n_feats, kernel_size, bias=bias)) + if bn: + m.append(nn.BatchNorm2d(n_feats)) + if i == 0: + m.append(act) + + self.body = nn.Sequential(*m) + self.res_scale = res_scale + + def forward(self, x): + res = self.body(x).mul(self.res_scale) + res += x + return res + +class ResidualBlock_noBN_anysr(nn.Module): + '''Residual block w/o BN + ---Conv-ReLU-Conv-+- + |________________| + ''' + def __init__(self, nf=64): + super(ResidualBlock_noBN_anysr, self).__init__() + self.conv1 = USConv2d(nf, nf, 3, 1, 1, bias=True, us=[False,True]) + self.conv2 = USConv2d(nf, nf, 3, 1, 1, bias=True, us=[True,False]) + self.fc1 = USConv2d(nf+8, nf*2, 1, 1, 0, bias=True, us=[True,False]) + self.fc2 = USConv2d(nf*2, nf, 1, 1, 0, bias=True, us=[False,True]) + self.gap = nn.AdaptiveAvgPool2d((1, 1)) + self.nf = nf + + # initialization + arch_util.initialize_weights([self.conv1, self.conv2, self.fc1, self.fc2], 0.1) + + def set_width(self, width): + self.conv1.width_mult = width + self.conv2.width_mult = width + self.fc1.width_mult = width + self.fc2.width_mult = width + + def set_scale(self, scale, scale2): + self.scale = scale + self.scale2 = scale2 + + def scale_aware_attn(self,feature,r): + B,C,H,W = feature.size() + global_avg_pooled = self.gap(feature) + feature = global_avg_pooled.view(B, -1) + r0 = r.to('cuda') + max_scale = max(r[0],r[1]) + if max_scale*10%2 == 1: + flag = -1 + else: + flag = 1 + r0 = r0 * flag + + r0 = r0.expand(B, -1) + interleaved_feature = [] + for i in range(0,C,16):# 0-C,step=16 + if i+16 <= C: + feature_trunk = feature[:,i:i+16] + interleaved_feature.append(feature_trunk) + interleaved_feature.append(r0) + else: + remain_channel = C - i + feature_trunk = feature[:,-remain_channel:] + interleaved_feature.append(feature_trunk) + + scale_fea = torch.cat(interleaved_feature, dim=1) + + res = self.fc1(scale_fea.view(B,-1,1,1)) + res = F.relu(res) + weight = self.fc2(res) + result = torch.sigmoid(weight) + + return result + + + def forward(self, x): + identity = x + b = identity.shape[0] + out = self.conv1(x) + + r_info = torch.tensor([self.scale,self.scale2]) + + attn_weight = self.scale_aware_attn(out,r_info).view(b,-1,1,1) + + out = out * attn_weight.expand_as(out) + tmp = F.relu(out) + out = self.conv2(tmp) + return identity + out + + +class Upsampler(nn.Sequential): + def __init__(self, conv, scale, n_feats, bn=False, act=False, bias=True): + + m = [] + if (scale & (scale - 1)) == 0: + for _ in range(int(math.log(scale, 2))): + m.append(conv(n_feats, 4 * n_feats, 3, bias)) + m.append(nn.PixelShuffle(2)) + if bn: + m.append(nn.BatchNorm2d(n_feats)) + if act == 'relu': + m.append(nn.ReLU(True)) + elif act == 'prelu': + m.append(nn.PReLU(n_feats)) + + elif scale == 3: + m.append(conv(n_feats, 9 * n_feats, 3, bias)) + m.append(nn.PixelShuffle(3)) + if bn: + m.append(nn.BatchNorm2d(n_feats)) + if act == 'relu': + m.append(nn.ReLU(True)) + elif act == 'prelu': + m.append(nn.PReLU(n_feats)) + else: + raise NotImplementedError + + super(Upsampler, self).__init__(*m) + +url = { + 'r16f64x2': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_baseline_x2-1bc95232.pt', + 'r16f64x3': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_baseline_x3-abf2a44e.pt', + 'r16f64x4': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_baseline_x4-6b446fab.pt', + 'r32f256x2': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_x2-0edfb8a3.pt', + 'r32f256x3': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_x3-ea3ef2c6.pt', + 'r32f256x4': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_x4-4f62e9ef.pt' +} + +class EDSR_anysr(nn.Module): + def __init__(self, args, conv=default_conv): + super(EDSR_anysr, self).__init__() + + n_resblock = args.n_resblocks + n_resblocks = args.n_resblocks + n_feats = args.n_feats + kernel_size = 3 + self.scale = args.scale[0] + self.scale2 = self.scale + + act = nn.ReLU(True) + self.scale_list = [round(x * 0.1 + 1, 1) for x in range(1, 31)] + + self.args = args + self.scale_idx = 0 + rgb_mean = (0.4488, 0.4371, 0.4040) + rgb_std = (1.0, 1.0, 1.0) + url_name = 'r{}f{}x{}'.format(n_resblocks, n_feats, self.scale) + if url_name in url: + self.url = url[url_name] + else: + self.url = None + self.sub_mean = MeanShift(args.rgb_range,rgb_mean,rgb_std) + self.add_mean = MeanShift(args.rgb_range,rgb_mean,rgb_std) + + m_head = [conv(args.n_colors, n_feats, kernel_size)] + + m_body = [] + for _ in range(n_resblock): + m_body.append(ResidualBlock_noBN_anysr(n_feats)) + + m_body.append(conv(n_feats, n_feats, kernel_size)) + + self.add_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std, 1) + + self.head = nn.Sequential(*m_head) + self.body = nn.Sequential(*m_body) + + if args.no_upsampling: + self.out_dim = n_feats + else: + self.out_dim = args.n_colors + + + def forward(self, x): + + x = self.sub_mean(x) + max_scale = max(self.scale,self.scale2) + if max_scale in self.scale_list[:7]: + width = 0.5 + elif max_scale in self.scale_list[7:15]: + width = 0.7 + elif max_scale in self.scale_list[15:22]: + width = 0.9 + else: + width = 1.0 + # if args.entire_net == 1: + # width = 1 + + x = self.head(x) + res = x + + for i in range(self.args.n_resblocks): + self.body[i].set_width(width) + self.body[i].set_scale(self.scale,self.scale2) + res = self.body[i](res) + + res = self.body[-1](res) + res += x + return res + + + def load_state_dict(self, state_dict, strict=True): + own_state = self.state_dict() + for name, param in state_dict.items(): + if name in own_state: + if isinstance(param, nn.Parameter): + param = param.data + try: + own_state[name].copy_(param) + except Exception: + if name.find('tail') == -1: + raise RuntimeError('While copying the parameter named {}, ' + 'whose dimensions in the model are {} and ' + 'whose dimensions in the checkpoint are {}.' + .format(name, own_state[name].size(), param.size())) + elif strict: + if name.find('tail') == -1: + raise KeyError('unexpected key "{}" in state_dict' + .format(name)) + + +# @register('edsr-anysr-baseline') +def make_edsr_baseline(n_resblocks=16, n_feats=64, res_scale=1, + scale=2, no_upsampling=False, rgb_range=1): + args = Namespace() + args.n_resblocks = n_resblocks + args.n_feats = n_feats + args.res_scale = res_scale + + args.scale = [scale] + args.no_upsampling = no_upsampling + + args.rgb_range = rgb_range + args.n_colors = 3 + return EDSR_anysr(args) + + +# @register('edsr-anysr') +def make_edsr(n_resblocks=32, n_feats=256, res_scale=0.1, + scale=2, no_upsampling=False, rgb_range=1): + args = Namespace() + args.n_resblocks = n_resblocks + args.n_feats = n_feats + args.res_scale = res_scale + + args.scale = [scale] + args.no_upsampling = no_upsampling + + args.rgb_range = rgb_range + args.n_colors = 3 + return EDSR_anysr(args) diff --git a/anysr_module/rdn_anysr.py b/anysr_module/rdn_anysr.py new file mode 100644 index 0000000..5f32fd2 --- /dev/null +++ b/anysr_module/rdn_anysr.py @@ -0,0 +1,460 @@ +# Residual Dense Network for Image Super-Resolution +# https://arxiv.org/abs/1802.08797 +# modified from: https://github.com/thstkdgus35/EDSR-PyTorch + +import argparse +# import math +import random + +# from models import register +# import time +from argparse import Namespace + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# torch.autograd.set_detect_anomaly(True) +# from torch.autograd import Variable + +# from . import arch_util +# from .arch_util import USConv2d + + +def default_conv(in_channels, out_channels, kernel_size, bias=True): + return nn.Conv2d( + in_channels, out_channels, kernel_size, padding=(kernel_size // 2), bias=bias + ) + + +class MeanShift(nn.Conv2d): + def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1): + super(MeanShift, self).__init__(3, 3, kernel_size=1) + std = torch.Tensor(rgb_std) + self.weight.data = torch.eye(3).view(3, 3, 1, 1) + self.weight.data.div_(std.view(3, 1, 1, 1)) + self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean) + self.bias.data.div_(std) + self.requires_grad = False + + +class USConv2d(nn.Conv2d): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + groups=1, + bias=True, + us=[False, False], + ): + super(USConv2d, self).__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + ) + self.width_mult = 1.0 + self.us = us + + def set_width_mult(self, width_mult): + self.width_mult = width_mult + + def forward(self, inputs): + in_channels = ( + inputs.shape[1] // self.groups + if self.us[0] + else self.in_channels // self.groups + ) + out_channels = self.out_channels + if self.us[1]: + if self.width_mult is not None: + out_channels = int(self.out_channels * self.width_mult) + else: + self.width_mult = 1.0 + out_channels = self.out_channels + + weight = self.weight[:out_channels, :in_channels, :, :] + if self.bias is not None: + bias = self.bias[:out_channels] + else: + bias = self.bias + + y = F.conv2d( + inputs, weight, bias, self.stride, self.padding, self.dilation, self.groups + ) + return y + + +class USConv2d_fc2(nn.Conv2d): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + groups=1, + bias=True, + us=[False, False], + ): + super(USConv2d_fc2, self).__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + ) + self.width_mult = 1.0 + self.us = us + + def set_width_mult(self, width_mult): + self.width_mult = width_mult + + def forward(self, inputs): + in_channels = ( + inputs.shape[1] // self.groups + if self.us[0] + else self.in_channels // self.groups + ) + out_channels = self.out_channels + if self.us[1]: + if self.width_mult is not None: + out_channels = int(self.out_channels // 2 * self.width_mult) * 2 + else: + self.width_mult = 1.0 + out_channels = self.out_channels + + weight = self.weight[:out_channels, :in_channels, :, :] + if self.bias is not None: + bias = self.bias[:out_channels] + else: + bias = self.bias + + y = F.conv2d( + inputs, weight, bias, self.stride, self.padding, self.dilation, self.groups + ) + return y + + +class RDB_USConv2d_former_first(nn.Module): + def __init__(self, inChannels, growRate, kSize=3): + super(RDB_USConv2d_former_first, self).__init__() + Cin = inChannels + G = growRate + self.conv = nn.Sequential( + *[ + USConv2d( + Cin, + G, + kSize, + padding=(kSize - 1) // 2, + stride=1, + bias=True, + us=[False, True], + ), + nn.ReLU(), + ] + ) + + def forward(self, x): + out = self.conv(x) + return torch.cat((x, out), 1) + + +class RDB_USConv2d_former(nn.Module): + def __init__(self, inChannels, growRate, kSize=3): + super(RDB_USConv2d_former, self).__init__() + Cin = inChannels + G = growRate + self.conv = nn.Sequential( + *[ + USConv2d( + Cin, + G, + kSize, + padding=(kSize - 1) // 2, + stride=1, + bias=True, + us=[True, True], + ), + nn.ReLU(), + ] + ) + + def forward(self, x): + out = self.conv(x) + return torch.cat((x, out), 1) + + +class RDB_USConv2d_latter(nn.Module): + def __init__(self, inChannels, growRate, kSize=3): + super(RDB_USConv2d_latter, self).__init__() + Cin = inChannels + G = growRate + self.conv = nn.Sequential( + *[ + USConv2d( + Cin, + G, + kSize, + padding=(kSize - 1) // 2, + stride=1, + bias=True, + us=[True, False], + ), + nn.ReLU(), + ] + ) + + def forward(self, x): + out = self.conv(x) + return torch.cat((x, out), 1) + + +class RDB_Conv(nn.Module): + def __init__(self, inChannels, growRate, kSize=3): + super(RDB_Conv, self).__init__() + Cin = inChannels + G = growRate + self.conv = nn.Sequential( + *[nn.Conv2d(Cin, G, kSize, padding=(kSize - 1) // 2, stride=1), nn.ReLU()] + ) + + def forward(self, x): + out = self.conv(x) + return torch.cat((x, out), 1) + + +class RDB_any(nn.Module): + def __init__(self, growRate0, growRate, nConvLayers, kSize=3): + super(RDB_any, self).__init__() + G0 = growRate0 + G = growRate + C = nConvLayers + self.C = nConvLayers + self.G0 = G0 + self.G = G + + convs = [] + for i in range(C // 2): # 0,1,2,3 + convs.append(RDB_Conv(G0 + i * G, G)) + + convs.append(RDB_USConv2d_former_first(G0 + (C // 2) * G, G)) # 4 + for c in range(C - 1 - C // 2): # c=0, block=5 + if (c) % 2 == 1: # former + convs.append(RDB_USConv2d_former(G0 + (c + 5) * G, G)) + else: + convs.append(RDB_USConv2d_latter(G0 + (c + 5) * G, G)) + + self.convs = nn.Sequential(*convs) + + # Local Feature Fusion + self.LFF = USConv2d( + G0 + C * G, G0, 1, padding=0, stride=1, bias=True, us=[True, False] + ) + self.fc1 = USConv2d(2 * (G + 8), 256, 1, 1, 0, bias=True, us=[True, False]) + self.fc2 = USConv2d_fc2(256, 2 * G, 1, 1, 0, bias=True, us=[False, True]) + self.gap = nn.AdaptiveAvgPool2d((1, 1)) + + def set_scale(self, scale, scale2): + self.scale = scale + self.scale2 = scale2 + + def set_width(self, width): + for c in range(self.C): + self.convs[c].conv[0].width_mult = width + self.LFF.width_mult = width + self.fc1.width_mult = width + self.fc2.width_mult = width + + def scale_aware_attn(self, feature, r): + B, C, H, W = feature.size() + global_avg_pooled = self.gap(feature) + feature = global_avg_pooled.view(B, -1) + r0 = r.to("cuda") + + max_scale = max(r[0], r[1]) + if max_scale * 10 % 2 == 1: + flag = -1 + else: + flag = 1 + r0 = r0 * flag + + r0 = r0.expand(B, -1) + + interleaved_feature = [] + for i in range(0, C, 16): # 0-C,step=16 + if i + 16 <= C: + feature_trunk = feature[:, i : i + 16] + interleaved_feature.append(feature_trunk) + interleaved_feature.append(r0) + + else: + remain_channel = C - i + feature_trunk = feature[:, -remain_channel:] + interleaved_feature.append(feature_trunk) + + scale_fea = torch.cat(interleaved_feature, dim=1) + + res = self.fc1(scale_fea.view(B, -1, 1, 1)) + res = F.relu(res) + weight = self.fc2(res) + + result = torch.sigmoid(weight) + return result + + def forward(self, x): + B = x.shape[0] + res = self.convs(x) + r_info = torch.tensor([self.scale, self.scale2]) + + former_fea = res[:, : self.G0 + 6 * self.G, :, :] + attn_fea = res[:, self.G0 + 6 * self.G :, :, :] + + attn_weight = self.scale_aware_attn(attn_fea, r_info).view(B, -1, 1, 1) + + out = attn_fea * attn_weight.expand_as(attn_fea) + res = torch.cat((former_fea, out), dim=1).expand_as(res) + + tmp = self.LFF(res) + res = tmp + x + return res + + +class RDN_anysr(nn.Module): + def __init__(self, args): + super(RDN_anysr, self).__init__() + self.scale = args.scale[0] + self.scale2 = self.scale + G0 = args.G0 + kSize = args.RDNkSize + self.scale_list = [round(x * 0.1 + 1, 1) for x in range(1, 31)] + self.args = args + self.scale_idx = 0 + + self.D, C, G = { + "A": (20, 6, 32), + "B": (16, 8, 64), + }[args.RDNconfig] + + rgb_mean = (0.4488, 0.4371, 0.4040) + rgb_std = (1.0, 1.0, 1.0) + self.sub_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std) + self.add_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std, 1) + + self.SFENet1 = nn.Conv2d( + args.n_colors, G0, kSize, padding=(kSize - 1) // 2, stride=1 + ) + self.SFENet2 = nn.Conv2d(G0, G0, kSize, padding=(kSize - 1) // 2, stride=1) + + self.RDBs = nn.ModuleList() + for i in range(self.D): + self.RDBs.append(RDB_any(growRate0=G0, growRate=G, nConvLayers=C)) + self.GFF = nn.Sequential( + *[ + nn.Conv2d(self.D * G0, G0, 1, padding=0, stride=1), + nn.Conv2d(G0, G0, kSize, padding=(kSize - 1) // 2, stride=1), + ] + ) + + if args.no_upsampling: + self.out_dim = G0 + else: + r = self.scale + self.out_dim = args.n_colors + if r == 2 or r == 3: + self.UPNet = nn.Sequential( + *[ + nn.Conv2d( + G0, G * r * r, kSize, padding=(kSize - 1) // 2, stride=1 + ), + nn.PixelShuffle(r), + nn.Conv2d( + G, args.n_colors, kSize, padding=(kSize - 1) // 2, stride=1 + ), + ] + ) + elif r == 4: + self.UPNet = nn.Sequential( + *[ + nn.Conv2d(G0, G * 4, kSize, padding=(kSize - 1) // 2, stride=1), + nn.PixelShuffle(2), + nn.Conv2d(G, G * 4, kSize, padding=(kSize - 1) // 2, stride=1), + nn.PixelShuffle(2), + nn.Conv2d( + G, args.n_colors, kSize, padding=(kSize - 1) // 2, stride=1 + ), + ] + ) + else: + raise ValueError("scale must be 2 or 3 or 4.") + + def forward(self, x): + max_scale = max(self.scale, self.scale2) + + parser = argparse.ArgumentParser() + parser.add_argument("--test_only", type=int) + parser.add_argument("--entire_net", type=int) + args, _ = parser.parse_known_args() + + if max_scale in self.scale_list[:7]: + width = 0.7 + if args.test_only == 0: # train + if random.random() < 0.6: + width = 1.0 + elif max_scale in self.scale_list[7:15]: + width = 0.8 + elif max_scale in self.scale_list[15:22]: + width = 0.9 + else: + width = 1.0 + if args.entire_net == 1: + width = 1 + + x = self.sub_mean(x) + f__1 = self.SFENet1(x) + x = self.SFENet2(f__1) + + RDBs_out = [] + for i in range(self.D): + self.RDBs[i].set_width(width) + self.RDBs[i].set_scale(self.scale, self.scale2) + x = self.RDBs[i](x) + RDBs_out.append(x) + + x = self.GFF(torch.cat(RDBs_out, 1)) + x += f__1 + + if self.args.no_upsampling: + return x + else: + return self.UPNet(x) + + +# @register("rdn-anysr") +def make_rdn( + G0=64, RDNkSize=3, RDNconfig="B", scale=2, no_upsampling=False, rgb_range=1 +): + args = Namespace() + args.G0 = G0 + args.RDNkSize = RDNkSize + args.RDNconfig = RDNconfig + + args.scale = [scale] + args.no_upsampling = no_upsampling + args.rgb_range = rgb_range + + args.n_colors = 3 + return RDN_anysr(args) diff --git a/constants.py b/constants.py index 8f31860..a283db0 100644 --- a/constants.py +++ b/constants.py @@ -1,4 +1,10 @@ WFEN_MODEL_PATH = "models/wfen/WFEN.pth" REALVIFORMER_MODEL_PATH = "models/realviformer/weights.pth" PROPIH_VGG_MODEL_PATH = "models/propih/vgg_normalised.pth" -PROPIH_G_MODEL_PATH = "models/propih/latest_net_G.pth" \ No newline at end of file +PROPIH_G_MODEL_PATH = "models/propih/latest_net_G.pth" +ANYSR_MODELS = ["edsr", "rdn"] +ANYSR_MODEL_PATHS = { + "edsr": "models/anysr/anysr_edsr_500.pth", + "rdn": "models/anysr/anysr_rdn_500.pth" +} +ANYSR_DEFAULT = "edsr" \ No newline at end of file diff --git a/examples/006.jpg b/examples/006.jpg new file mode 100644 index 0000000..59393cd Binary files /dev/null and b/examples/006.jpg differ diff --git a/run.py b/run.py index cb4f6e2..b29f695 100644 --- a/run.py +++ b/run.py @@ -12,8 +12,12 @@ import torch.nn.functional as F import tqdm from comfy.utils import ProgressBar +from .anysr_module.anysr_model import run_anysr from .colie_module.colie_model import CoLIE_Config, run_colie from .constants import ( + ANYSR_DEFAULT, + ANYSR_MODEL_PATHS, + ANYSR_MODELS, PROPIH_G_MODEL_PATH, PROPIH_VGG_MODEL_PATH, REALVIFORMER_MODEL_PATH, @@ -283,3 +287,45 @@ class CoLIE_LowLight_Enhance: pbar.update_absolute(i, num_frames) return (torch.cat(result, dim=0),) + + +class AnySR: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "src_img": ("IMAGE",), + "scale": ( + "FLOAT", + {"default": 2.00, "min": 0.01}, + ), + "model": (ANYSR_MODELS, {"default": ANYSR_DEFAULT}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("res_img",) + FUNCTION = "todo" + CATEGORY = "FM_nodes" + + def todo( + self, + src_img: torch.Tensor, + scale: float, + model: str, + ): + result: list[torch.Tensor] = [] + num_frames = src_img.size(0) + pbar = ProgressBar(num_frames) + for i in range(num_frames): + image = src_img[i].unsqueeze(0) + res_tensor = run_anysr( + image, + scale=scale, + model_arch=model, + model_path=f"{base_dir}/{ANYSR_MODEL_PATHS[model]}", + ) + result.append(res_tensor) + pbar.update_absolute(i, num_frames) + + return (torch.cat(result, dim=0),) diff --git a/test.ipynb b/test.ipynb new file mode 100644 index 0000000..6ae320c --- /dev/null +++ b/test.ipynb @@ -0,0 +1,342 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "j:\\AI\\ComfyUI_windows_portable\\python_embeded\\Lib\\site-packages\\transformers\\utils\\hub.py:124: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.\n", + " warnings.warn(\n" + ] + } + ], + "source": [ + "import torch\n", + "from torchvision import transforms\n", + "from PIL import Image\n", + "\n", + "# Function to load and preprocess the image\n", + "def load_image_as_tensor(image_path, batch_size=1):\n", + " # Open the image using PIL\n", + " image = Image.open(image_path)\n", + " \n", + " # Define the transformation pipeline\n", + " transform = transforms.Compose([\n", + " transforms.Resize((256, 256)),\n", + " transforms.ToTensor(), # Convert PIL Image to tensor and scale to [0, 1]\n", + " # transforms.Lambda(lambda x: x.unsqueeze(0).repeat(batch_size, 1, 1, 1)) # Add batch dimension and repeat\n", + " ])\n", + " \n", + " # Apply the transformation\n", + " tensor = transform(image)\n", + " \n", + " return tensor" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "image = load_image_as_tensor(\"examples/006.jpg\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([3, 256, 256])" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "image.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", + "scale_max = 2" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "model_path = \"models/anysr/anysr_edsr_500.pth\"" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "model_load = torch.load(model_path)['model']\n", + "model_args = model_load['args']['encoder_spec']['args']\n", + "pretrain_dict = model_load['sd']" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "image_tensor = image.squeeze(dim=0).to(device)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([3, 256, 256])" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "image_tensor.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "from anysr_module.edsr_anysr import make_edsr\n", + "\n", + "\n", + "model = make_edsr(scale=2, **model_args)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "from anysr_module.anysr_model import replace\n", + "\n", + "\n", + "model_dict = model.state_dict()\n", + "model_dict = replace(pretrain_dict, model_dict)\n", + "model.load_state_dict(model_dict, strict=False)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "model = model.to(device)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "model.scale = 2\n", + "model.scale2 = 2" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([3, 256, 256])" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "image_tensor.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "pred = model(\n", + " image_tensor.unsqueeze(dim=0).cuda(),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([1, 256, 256, 256])" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pred.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([1, 256, 256, 256])" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pred.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "ename": "RuntimeError", + "evalue": "shape '[1, 256, 256, 3, 9, 9]' is invalid for input of size 16777216", + "output_type": "error", + "traceback": [ + "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[1;31mRuntimeError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[1;32mIn[14], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m reshaped_tensor \u001b[38;5;241m=\u001b[39m \u001b[43mpred\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreshape\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m256\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m256\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m9\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m9\u001b[39;49m\u001b[43m)\u001b[49m\n", + "\u001b[1;31mRuntimeError\u001b[0m: shape '[1, 256, 256, 3, 9, 9]' is invalid for input of size 16777216" + ] + } + ], + "source": [ + "reshaped_tensor = pred.reshape(1, 256, 256, 3, 9, 9)" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([[[[-0.2023, -0.2465, -0.1883, ..., 0.0145, 0.0291, 0.0131],\n", + " [-0.2058, -0.3171, -0.2691, ..., 0.0181, 0.0362, 0.0219],\n", + " [-0.1518, -0.3128, -0.2633, ..., 0.0039, 0.0277, 0.0250],\n", + " ...,\n", + " [-0.0476, 0.0612, 0.0954, ..., -0.1948, -0.2188, -0.1279],\n", + " [-0.0426, 0.0376, 0.0445, ..., -0.1631, -0.1892, -0.0890],\n", + " [ 0.0051, 0.0672, 0.0702, ..., -0.0010, -0.0048, 0.0314]],\n", + "\n", + " [[-0.1818, -0.2641, -0.2487, ..., -0.1725, -0.1746, -0.1650],\n", + " [-0.0194, -0.1230, -0.1304, ..., -0.2200, -0.2187, -0.1621],\n", + " [-0.0632, -0.1623, -0.1519, ..., -0.2183, -0.2122, -0.1603],\n", + " ...,\n", + " [-0.1712, -0.1563, -0.1632, ..., -0.1246, -0.1534, -0.1306],\n", + " [-0.2195, -0.2493, -0.2216, ..., -0.1596, -0.1834, -0.1382],\n", + " [-0.1708, -0.2002, -0.1785, ..., -0.1240, -0.1355, -0.1496]],\n", + "\n", + " [[ 0.0384, -0.1395, -0.1434, ..., 0.0056, -0.0025, -0.0254],\n", + " [ 0.0819, -0.0770, -0.1296, ..., -0.0080, -0.0228, -0.0414],\n", + " [ 0.0429, -0.0804, -0.1248, ..., -0.0178, -0.0256, -0.0443],\n", + " ...,\n", + " [ 0.1172, 0.0857, 0.0598, ..., -0.0916, -0.1025, -0.0719],\n", + " [ 0.0778, 0.0410, 0.0348, ..., -0.0644, -0.0767, -0.0587],\n", + " [ 0.0415, 0.0601, 0.0726, ..., 0.0271, 0.0418, 0.0895]],\n", + "\n", + " ...,\n", + "\n", + " [[ 0.0076, -0.1249, -0.1511, ..., -0.0041, -0.0041, -0.0131],\n", + " [-0.1826, -0.2713, -0.2433, ..., 0.0897, 0.0770, 0.0489],\n", + " [-0.1696, -0.2405, -0.2122, ..., 0.0730, 0.0656, 0.0382],\n", + " ...,\n", + " [ 0.0398, 0.0766, 0.1059, ..., -0.0038, -0.0599, -0.0317],\n", + " [ 0.0230, 0.0949, 0.1537, ..., -0.0311, -0.0921, -0.0622],\n", + " [ 0.0083, 0.0128, 0.0458, ..., -0.1552, -0.1739, -0.1277]],\n", + "\n", + " [[ 0.1756, 0.2339, 0.2451, ..., 0.1755, 0.1732, 0.1503],\n", + " [ 0.0185, -0.0176, 0.0165, ..., 0.2317, 0.2341, 0.1946],\n", + " [-0.0086, -0.0330, 0.0083, ..., 0.2167, 0.2240, 0.1849],\n", + " ...,\n", + " [ 0.1407, 0.1760, 0.2159, ..., 0.0839, 0.0529, 0.0533],\n", + " [ 0.1521, 0.2473, 0.3145, ..., 0.0638, 0.0380, 0.0335],\n", + " [ 0.1696, 0.2355, 0.2771, ..., 0.0139, -0.0056, 0.0147]],\n", + "\n", + " [[-0.1895, -0.2206, -0.2065, ..., -0.0522, -0.0691, -0.0413],\n", + " [-0.1096, -0.1073, -0.1033, ..., -0.1391, -0.1615, -0.1086],\n", + " [-0.1121, -0.1327, -0.1424, ..., -0.1380, -0.1547, -0.1007],\n", + " ...,\n", + " [ 0.0065, -0.0212, -0.0236, ..., -0.1827, -0.1751, -0.0828],\n", + " [-0.0189, -0.0623, -0.0892, ..., -0.1393, -0.1310, -0.0676],\n", + " [-0.0020, -0.0198, -0.0286, ..., -0.0580, -0.0464, -0.0142]]]],\n", + " device='cuda:0', grad_fn=)" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pred" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}