doesnt work
This commit is contained in:
FuouM
@@ -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)
|
||||
|
||||
|
||||
|
||||
### 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
|
||||
```
|
||||
+3
-1
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
+7
-1
@@ -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"
|
||||
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"
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 139 KiB |
@@ -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),)
|
||||
|
||||
+342
@@ -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=<AddBackward0>)"
|
||||
]
|
||||
},
|
||||
"execution_count": 17,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"pred"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
Reference in New Issue
Block a user