Fuou Marinas
@@ -10,6 +10,8 @@ Click name to jump to workflow
|
||||
5. [**UniST**](#unist). Paper: [Two Birds, One Stone: A Unified Framework for Joint Learning of Image and Video Style Transfers](https://github.com/NevSNev/UniST)
|
||||
6. [**AesPA-Net**](#aespa-net). Paper: [Aesthetic Pattern-Aware Style Transfer Networks](https://github.com/Kibeom-Hong/AesPA-Net)
|
||||
7. [**TSSAT**](#tssat). Paper: [Two-Stage Statistics-Aware Transformation for Artistic Style Transfer](https://github.com/HalbertCH/TSSAT)
|
||||
8. [**AesFA**](#aesfa). Paper: [An Aesthetic Feature-Aware Arbitrary Neural Style Transfer](https://github.com/Sooyyoungg/AesFA)
|
||||
|
||||
|
||||
[**Extra nodes**](#extra-nodes)
|
||||
|
||||
@@ -26,8 +28,8 @@ All nodes support batched input (i.e video) but is generally not recommended. On
|
||||
|  |  |
|
||||
| **UniST (2023) - Medium - Only squares** | **AesPA-Net (2023) - Medium - Only squares** |
|
||||
|  |  |
|
||||
| **TSSAT (2023) - Slow** | **X** |
|
||||
|  | X |
|
||||
| **TSSAT (2023) - Slow** | **AesFA (2024) - Fast - Only squares** |
|
||||
|  |  |
|
||||
|
||||
### Neural Neighbor
|
||||
|
||||
@@ -167,6 +169,22 @@ models/vgg_normalised.pth
|
||||
|
||||
|
||||
|
||||
|
||||
### AesFA
|
||||
|
||||
Download `main.pth` in the [Getting Started](https://github.com/Sooyyoungg/AesFA?tab=readme-ov-file#getting-started) and place it in `models/aesfa/main.pth`.
|
||||
|
||||
[workflow_aesfa_w_blend.json](workflows/workflow_aesfa_w_blend.json)
|
||||
|
||||
|
||||
|
||||
Authors' Notes:
|
||||
|
||||
> [...] style blending, i.e., using the low-frequency and high-frequency style information from different images.
|
||||
The style-transferred outputs finely keep the color information from the low-frequency image and change the texture information based on the high frequency image.
|
||||
|
||||
> The vertical line-shape artifacts alongside the images are often observed. We reason that these appear because the content features are being convolved directly with the predicted aesthetic feature-aware kernels and biases in our model. In addition, the upsampling operation could be the ones that create artifacts.
|
||||
|
||||
## Extra nodes
|
||||
|
||||
### Coral Color Transfer
|
||||
|
||||
+16
-1
@@ -1,4 +1,15 @@
|
||||
from .run import NeuralNeighbor, CAST, EFDM, MicroAST, UniST, UniST_Video, AesPA, TSSAT
|
||||
from .run import (
|
||||
AesFA,
|
||||
AesFAStyleBlend,
|
||||
CAST,
|
||||
EFDM,
|
||||
TSSAT,
|
||||
AesPA,
|
||||
MicroAST,
|
||||
NeuralNeighbor,
|
||||
UniST,
|
||||
UniST_Video,
|
||||
)
|
||||
from .run_extra import CoralColorTransfer
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
@@ -11,6 +22,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
"UniST_Video": UniST_Video,
|
||||
"AesPA": AesPA,
|
||||
"TSSAT": TSSAT,
|
||||
"AESFA": AesFA,
|
||||
"AesFAStyleBlend": AesFAStyleBlend,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -23,6 +36,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UniST_Video": "UniST Video",
|
||||
"AesPA": "AesPA-Net",
|
||||
"TSSAT": "TSSAT",
|
||||
"AESFA": "AESFA",
|
||||
"AesFAStyleBlend": "AesFA Styles Blending",
|
||||
}
|
||||
|
||||
|
||||
|
||||
+3
-1
@@ -22,4 +22,6 @@ VGG_NORM_CONV51_PATH = "models/vgg_normalised_conv5_1.pth"
|
||||
AESPA_DEC_PATH = "models/aespa/dec_model.pth"
|
||||
AESPA_TRANSFORMER_PATH = "models/aespa/transformer_model.pth"
|
||||
|
||||
TSSAT_DEC_PATH = "models/tssat/decoder_iter_160000.pth"
|
||||
TSSAT_DEC_PATH = "models/tssat/decoder_iter_160000.pth"
|
||||
|
||||
AESFA_PATH = "models/aesfa/main.pth"
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 782 KiB |
@@ -0,0 +1,87 @@
|
||||
import torch
|
||||
import torchvision.transforms as transforms
|
||||
|
||||
from .model import AesFA_test
|
||||
|
||||
|
||||
class MODEL_AESFA:
|
||||
def __init__(self, checkpoint_path: str, device) -> None:
|
||||
self.aesfa_model = AesFA_test(
|
||||
alpha_in=0.5, # input ratio of low-frequency channel
|
||||
alpha_out=0.5, # output ratio of low-frequency channel
|
||||
style_kernel=3, # size of style kernel
|
||||
input_nc=3, # of input image channel
|
||||
nf=64, # of feature map channel after Encoder first layer
|
||||
output_nc=3, # of output image channel
|
||||
freq_ratio=[1, 1], # [high, low] ratio at the last layer
|
||||
)
|
||||
|
||||
dict_model = torch.load(checkpoint_path)
|
||||
|
||||
self.aesfa_model.netE.load_state_dict(dict_model["netE"])
|
||||
self.aesfa_model.netS.load_state_dict(dict_model["netS"])
|
||||
self.aesfa_model.netG.load_state_dict(dict_model["netG"])
|
||||
|
||||
self.aesfa_model.to(device)
|
||||
|
||||
|
||||
def inference_aesfa(
|
||||
src_img: torch.Tensor,
|
||||
style_img: torch.Tensor,
|
||||
device,
|
||||
size: int,
|
||||
do_crop: bool,
|
||||
model: MODEL_AESFA,
|
||||
):
|
||||
content_tf = test_transform(
|
||||
size, do_crop, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
|
||||
)
|
||||
content = content_tf(src_img).to(device)
|
||||
style = content_tf(style_img).to(device)
|
||||
|
||||
stylized_image = model.aesfa_model.forward(content, style)
|
||||
output = un_normalize_batch(stylized_image)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def inference_aesfa_style_blend(
|
||||
src_img: torch.Tensor,
|
||||
style_hi: torch.Tensor,
|
||||
style_lo: torch.Tensor,
|
||||
device,
|
||||
size: int,
|
||||
do_crop: bool,
|
||||
model: MODEL_AESFA,
|
||||
):
|
||||
content_tf = test_transform(
|
||||
size, do_crop, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
|
||||
)
|
||||
content = content_tf(src_img).to(device)
|
||||
style_high = content_tf(style_hi).to(device)
|
||||
style_low = content_tf(style_lo).to(device)
|
||||
|
||||
stylized_image = model.aesfa_model.style_blending(content, style_high, style_low)
|
||||
output = un_normalize_batch(stylized_image)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def test_transform(size, crop, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
|
||||
transform_list = []
|
||||
|
||||
if crop and size != 0:
|
||||
transform_list.append(transforms.Resize(size))
|
||||
transform_list.append(transforms.CenterCrop(size))
|
||||
elif size != 0:
|
||||
transform_list.append(transforms.Resize((size, size)))
|
||||
|
||||
transform_list.append(transforms.Normalize(mean, std))
|
||||
|
||||
return transforms.Compose(transform_list)
|
||||
|
||||
|
||||
def un_normalize_batch(tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
|
||||
for t, m, s in zip(tensor, mean, std):
|
||||
t.mul_(s).add_(m)
|
||||
return tensor
|
||||
@@ -0,0 +1,343 @@
|
||||
import os
|
||||
# import glob
|
||||
# from path import Path
|
||||
import math
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from torch.optim import lr_scheduler
|
||||
# from blocks import *
|
||||
|
||||
def model_save(ckpt_dir, model, optim_E, optim_S, optim_G, epoch, itr=None):
|
||||
if not os.path.exists(ckpt_dir):
|
||||
os.makedirs(ckpt_dir)
|
||||
|
||||
torch.save({'netE': model.netE.state_dict(),
|
||||
'netS': model.netS.state_dict(),
|
||||
'netG': model.netG.state_dict(),
|
||||
'optim_E': optim_E.state_dict(),
|
||||
'optim_S': optim_S.state_dict(),
|
||||
'optim_G': optim_G.state_dict()},
|
||||
'%s/model_iter_%d_epoch_%d.pth' % (ckpt_dir, itr+1, epoch+1))
|
||||
|
||||
def model_load(checkpoint, ckpt_path, model, optim_E, optim_S, optim_G):
|
||||
# if not os.path.exists(ckpt_dir):
|
||||
# epoch = -1
|
||||
# return model, optim_E, optim_S, optim_G, epoch
|
||||
|
||||
# ckpt_path = Path(ckpt_dir)
|
||||
if checkpoint:
|
||||
model_ckpt = ckpt_path + '/' + checkpoint
|
||||
else:
|
||||
ckpt_lst = ckpt_path.glob('model_iter_*')
|
||||
ckpt_lst.sort(key=lambda x: int(x.split('iter_')[1].split('_epoch')[0]))
|
||||
model_ckpt = ckpt_lst[-1]
|
||||
itr = int(model_ckpt.split('iter_')[1].split('_epoch_')[0])
|
||||
epoch = int(model_ckpt.split('iter_')[1].split('_epoch_')[1].split('.')[0])
|
||||
print(model_ckpt)
|
||||
|
||||
dict_model = torch.load(model_ckpt)
|
||||
|
||||
model.netE.load_state_dict(dict_model['netE'])
|
||||
model.netS.load_state_dict(dict_model['netS'])
|
||||
model.netG.load_state_dict(dict_model['netG'])
|
||||
optim_E.load_state_dict(dict_model['optim_E'])
|
||||
optim_S.load_state_dict(dict_model['optim_S'])
|
||||
optim_G.load_state_dict(dict_model['optim_G'])
|
||||
|
||||
return model, optim_E, optim_S, optim_G, epoch, itr
|
||||
|
||||
def test_model_load(checkpoint, model):
|
||||
dict_model = torch.load(checkpoint)
|
||||
model.netE.load_state_dict(dict_model['netE'])
|
||||
model.netS.load_state_dict(dict_model['netS'])
|
||||
model.netG.load_state_dict(dict_model['netG'])
|
||||
return model
|
||||
|
||||
def get_scheduler(optimizer, config):
|
||||
if config.lr_policy == 'lambda':
|
||||
def lambda_rule(epoch):
|
||||
lr_l = 1.0 - max(0, epoch + config.n_epoch - config.n_iter) / float(config.n_iter_decay + 1)
|
||||
return lr_l
|
||||
|
||||
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
||||
elif config.lr_policy == 'step':
|
||||
scheduler = lr_scheduler.StepLR(optimizer, step_size=config.lr_decay_iters, gamma=0.1)
|
||||
elif config.lr_policy == 'plateau':
|
||||
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
|
||||
elif config.lr_policy == 'cosine':
|
||||
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.n_iter, eta_min=0)
|
||||
else:
|
||||
return NotImplementedError('learning rate policy [%s] is not implemented', config.lr_policy)
|
||||
return scheduler
|
||||
|
||||
def update_learning_rate(scheduler, optimizer):
|
||||
scheduler.step()
|
||||
lr = optimizer.param_groups[0]['lr']
|
||||
print('learning rate = %.7f' % lr)
|
||||
|
||||
class Oct_Conv_aftup(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, pad_type, alpha_in, alpha_out):
|
||||
super(Oct_Conv_aftup, self).__init__()
|
||||
lf_in = int(in_channels*alpha_in)
|
||||
lf_out = int(out_channels*alpha_out)
|
||||
hf_in = in_channels - lf_in
|
||||
hf_out = out_channels - lf_out
|
||||
|
||||
self.conv_h = nn.Conv2d(in_channels=hf_in, out_channels=hf_out, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type)
|
||||
self.conv_l = nn.Conv2d(in_channels=lf_in, out_channels=lf_out, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type)
|
||||
|
||||
def forward(self, x):
|
||||
hf, lf = x
|
||||
hf = self.conv_h(hf)
|
||||
lf = self.conv_l(lf)
|
||||
return hf, lf
|
||||
|
||||
class Oct_conv_reLU(nn.ReLU):
|
||||
def forward(self, x):
|
||||
hf, lf = x
|
||||
hf = super(Oct_conv_reLU, self).forward(hf)
|
||||
lf = super(Oct_conv_reLU, self).forward(lf)
|
||||
return hf, lf
|
||||
|
||||
class Oct_conv_lreLU(nn.LeakyReLU):
|
||||
def forward(self, x):
|
||||
hf, lf = x
|
||||
hf = super(Oct_conv_lreLU, self).forward(hf)
|
||||
lf = super(Oct_conv_lreLU, self).forward(lf)
|
||||
return hf, lf
|
||||
|
||||
class Oct_conv_up(nn.Upsample):
|
||||
def forward(self, x):
|
||||
hf, lf = x
|
||||
hf = super(Oct_conv_up, self).forward(hf)
|
||||
lf = super(Oct_conv_up, self).forward(lf)
|
||||
return hf, lf
|
||||
|
||||
|
||||
############## Encoder ##############
|
||||
class OctConv(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride=1,
|
||||
padding=0, groups=1, pad_type='reflect', alpha_in=0.5, alpha_out=0.5, type='normal', freq_ratio = [1, 1]):
|
||||
super(OctConv, self).__init__()
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.type = type
|
||||
self.alpha_in = alpha_in
|
||||
self.alpha_out = alpha_out
|
||||
self.freq_ratio = freq_ratio
|
||||
|
||||
hf_ch_in = int(in_channels * (1 - self.alpha_in))
|
||||
hf_ch_out = int(out_channels * (1 -self. alpha_out))
|
||||
lf_ch_in = in_channels - hf_ch_in
|
||||
lf_ch_out = out_channels - hf_ch_out
|
||||
|
||||
self.avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)
|
||||
self.upsample = nn.Upsample(scale_factor=2)
|
||||
|
||||
self.is_dw = groups == in_channels
|
||||
|
||||
if type == 'first':
|
||||
self.convh = nn.Conv2d(in_channels, hf_ch_out, kernel_size=kernel_size,
|
||||
stride=stride, padding=padding, padding_mode=pad_type, bias = False)
|
||||
self.convl = nn.Conv2d(in_channels, lf_ch_out,
|
||||
kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type, bias=False)
|
||||
elif type == 'last':
|
||||
self.convh = nn.Conv2d(hf_ch_in, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type, bias=False)
|
||||
self.convl = nn.Conv2d(lf_ch_in, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=pad_type, bias=False)
|
||||
else:
|
||||
self.L2L = nn.Conv2d(
|
||||
lf_ch_in, lf_ch_out,
|
||||
kernel_size=kernel_size, stride=stride, padding=padding, groups=math.ceil(alpha_in * groups), padding_mode=pad_type, bias=False
|
||||
)
|
||||
if self.is_dw:
|
||||
self.L2H = None
|
||||
self.H2L = None
|
||||
else:
|
||||
self.L2H = nn.Conv2d(
|
||||
lf_ch_in, hf_ch_out,
|
||||
kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, padding_mode=pad_type, bias=False
|
||||
)
|
||||
self.H2L = nn.Conv2d(
|
||||
hf_ch_in, lf_ch_out,
|
||||
kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, padding_mode=pad_type, bias=False
|
||||
)
|
||||
self.H2H = nn.Conv2d(
|
||||
hf_ch_in, hf_ch_out,
|
||||
kernel_size=kernel_size, stride=stride, padding=padding, groups=math.ceil(groups - alpha_in * groups), padding_mode=pad_type, bias=False
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
if self.type == 'first':
|
||||
hf = self.convh(x)
|
||||
lf = self.avg_pool(x)
|
||||
lf = self.convl(lf)
|
||||
return hf, lf
|
||||
elif self.type == 'last':
|
||||
hf, lf = x
|
||||
out_h = self.convh(hf)
|
||||
out_l = self.convl(self.upsample(lf))
|
||||
output = out_h * self.freq_ratio[0] + out_l * self.freq_ratio[1]
|
||||
return output, out_h, out_l
|
||||
else:
|
||||
hf, lf = x
|
||||
if self.is_dw:
|
||||
hf, lf = self.H2H(hf), self.L2L(lf)
|
||||
else:
|
||||
hf, lf = self.H2H(hf) + self.L2H(self.upsample(lf)), self.L2L(lf) + self.H2L(self.avg_pool(hf))
|
||||
return hf, lf
|
||||
|
||||
|
||||
############## Decoder ##############
|
||||
class AdaOctConv(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, group_div, style_channels, kernel_size,
|
||||
stride, padding, oct_groups, alpha_in, alpha_out, type='normal'):
|
||||
super(AdaOctConv, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.alpha_in = alpha_in
|
||||
self.alpha_out = alpha_out
|
||||
self.type = type
|
||||
|
||||
h_in = int(in_channels * (1 - self.alpha_in))
|
||||
l_in = in_channels - h_in
|
||||
|
||||
n_groups_h = h_in // group_div
|
||||
n_groups_l = l_in // group_div
|
||||
|
||||
style_channels_h = int(style_channels * (1 - self.alpha_in))
|
||||
style_channels_l = int(style_channels - style_channels_h)
|
||||
|
||||
kernel_size_h = kernel_size[0]
|
||||
kernel_size_l = kernel_size[1]
|
||||
kernel_size_A = kernel_size[2]
|
||||
|
||||
self.kernelPredictor_h = KernelPredictor(in_channels=h_in,
|
||||
out_channels=h_in,
|
||||
n_groups=n_groups_h,
|
||||
style_channels=style_channels_h,
|
||||
kernel_size=kernel_size_h)
|
||||
self.kernelPredictor_l = KernelPredictor(in_channels=l_in,
|
||||
out_channels=l_in,
|
||||
n_groups=n_groups_l,
|
||||
style_channels=style_channels_l,
|
||||
kernel_size=kernel_size_l)
|
||||
|
||||
self.AdaConv_h = AdaConv2d(in_channels=h_in, out_channels=h_in, n_groups=n_groups_h)
|
||||
self.AdaConv_l = AdaConv2d(in_channels=l_in, out_channels=l_in, n_groups=n_groups_l)
|
||||
|
||||
self.OctConv = OctConv(in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size_A, stride=stride, padding=padding, groups=oct_groups,
|
||||
alpha_in=alpha_in, alpha_out=alpha_out, type=type)
|
||||
|
||||
self.relu = Oct_conv_lreLU()
|
||||
|
||||
def forward(self, content, style, cond='train'):
|
||||
c_hf, c_lf = content
|
||||
s_hf, s_lf = style
|
||||
h_w_spatial, h_w_pointwise, h_bias = self.kernelPredictor_h(s_hf)
|
||||
l_w_spatial, l_w_pointwise, l_bias = self.kernelPredictor_l(s_lf)
|
||||
|
||||
if cond == 'train':
|
||||
output_h = self.AdaConv_h(c_hf, h_w_spatial, h_w_pointwise, h_bias)
|
||||
output_l = self.AdaConv_l(c_lf, l_w_spatial, l_w_pointwise, l_bias)
|
||||
output = output_h, output_l
|
||||
|
||||
output = self.relu(output)
|
||||
|
||||
output = self.OctConv(output)
|
||||
if self.type != 'last':
|
||||
output = self.relu(output)
|
||||
return output
|
||||
|
||||
if cond == 'test':
|
||||
output_h = self.AdaConv_h(c_hf, h_w_spatial, h_w_pointwise, h_bias)
|
||||
output_l = self.AdaConv_l(c_lf, l_w_spatial, l_w_pointwise, l_bias)
|
||||
output = output_h, output_l
|
||||
output = self.relu(output)
|
||||
output = self.OctConv(output)
|
||||
if self.type != 'last':
|
||||
output = self.relu(output)
|
||||
return output
|
||||
|
||||
class KernelPredictor(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, n_groups, style_channels, kernel_size):
|
||||
super(KernelPredictor, self).__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.n_groups = n_groups
|
||||
self.w_channels = style_channels
|
||||
self.kernel_size = kernel_size
|
||||
|
||||
padding = (kernel_size - 1) / 2
|
||||
self.spatial = nn.Conv2d(style_channels,
|
||||
in_channels * out_channels // n_groups,
|
||||
kernel_size=kernel_size,
|
||||
padding=(math.ceil(padding), math.ceil(padding)),
|
||||
padding_mode='reflect')
|
||||
self.pointwise = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d((1, 1)),
|
||||
nn.Conv2d(style_channels,
|
||||
out_channels * out_channels // n_groups,
|
||||
kernel_size=1)
|
||||
)
|
||||
self.bias = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d((1, 1)),
|
||||
nn.Conv2d(style_channels,
|
||||
out_channels,
|
||||
kernel_size=1)
|
||||
)
|
||||
|
||||
def forward(self, w):
|
||||
w_spatial = self.spatial(w)
|
||||
w_spatial = w_spatial.reshape(len(w),
|
||||
self.out_channels,
|
||||
self.in_channels // self.n_groups,
|
||||
self.kernel_size, self.kernel_size)
|
||||
|
||||
w_pointwise = self.pointwise(w)
|
||||
w_pointwise = w_pointwise.reshape(len(w),
|
||||
self.out_channels,
|
||||
self.out_channels // self.n_groups,
|
||||
1, 1)
|
||||
bias = self.bias(w)
|
||||
bias = bias.reshape(len(w), self.out_channels)
|
||||
return w_spatial, w_pointwise, bias
|
||||
|
||||
class AdaConv2d(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size=3, n_groups=None):
|
||||
super(AdaConv2d, self).__init__()
|
||||
self.n_groups = in_channels if n_groups is None else n_groups
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
padding = (kernel_size - 1) / 2
|
||||
self.conv = nn.Conv2d(in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=(kernel_size, kernel_size),
|
||||
padding=(math.ceil(padding), math.floor(padding)),
|
||||
padding_mode='reflect')
|
||||
|
||||
def forward(self, x, w_spatial, w_pointwise, bias):
|
||||
assert len(x) == len(w_spatial) == len(w_pointwise) == len(bias)
|
||||
x = F.instance_norm(x)
|
||||
|
||||
ys = []
|
||||
for i in range(len(x)):
|
||||
y = self.forward_single(x[i:i+1], w_spatial[i], w_pointwise[i], bias[i])
|
||||
ys.append(y)
|
||||
ys = torch.cat(ys, dim=0)
|
||||
|
||||
ys = self.conv(ys)
|
||||
return ys
|
||||
|
||||
def forward_single(self, x, w_spatial, w_pointwise, bias):
|
||||
assert w_spatial.size(-1) == w_spatial.size(-2)
|
||||
padding = (w_spatial.size(-1) - 1) / 2
|
||||
pad = (math.ceil(padding), math.floor(padding), math.ceil(padding), math.floor(padding))
|
||||
|
||||
x = F.pad(x, pad=pad, mode='reflect')
|
||||
x = F.conv2d(x, w_spatial, groups=self.n_groups)
|
||||
x = F.conv2d(x, w_pointwise, groups=self.n_groups, bias=bias)
|
||||
return x
|
||||
@@ -0,0 +1,66 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from . import networks
|
||||
|
||||
|
||||
class AesFA_test(nn.Module):
|
||||
def __init__(
|
||||
self, alpha_in, alpha_out, style_kernel, input_nc, nf, output_nc, freq_ratio
|
||||
):
|
||||
super(AesFA_test, self).__init__()
|
||||
|
||||
self.netE = networks.define_network(
|
||||
"Encoder",
|
||||
alpha_in,
|
||||
alpha_out,
|
||||
style_kernel,
|
||||
input_nc,
|
||||
nf,
|
||||
output_nc,
|
||||
freq_ratio,
|
||||
)
|
||||
self.netS = networks.define_network(
|
||||
"Encoder",
|
||||
alpha_in,
|
||||
alpha_out,
|
||||
style_kernel,
|
||||
input_nc,
|
||||
nf,
|
||||
output_nc,
|
||||
freq_ratio,
|
||||
)
|
||||
self.netG = networks.define_network(
|
||||
"Generator",
|
||||
alpha_in,
|
||||
alpha_out,
|
||||
style_kernel,
|
||||
input_nc,
|
||||
nf,
|
||||
output_nc,
|
||||
freq_ratio,
|
||||
)
|
||||
|
||||
def forward(self, real_A, real_B):
|
||||
with torch.no_grad():
|
||||
content_A = self.netE.forward_test(real_A, "content")
|
||||
style_B = self.netS.forward_test(real_B, "style")
|
||||
# if freq:
|
||||
# trs_AtoB, trs_AtoB_high, trs_AtoB_low = self.netG(content_A, style_B)
|
||||
# end = time.time()
|
||||
# during = end - start
|
||||
# return trs_AtoB, trs_AtoB_high, trs_AtoB_low, during
|
||||
# else:
|
||||
trs_AtoB = self.netG.forward_test(content_A, style_B)
|
||||
return trs_AtoB
|
||||
|
||||
def style_blending(self, real_A, real_B_1, real_B_2):
|
||||
with torch.no_grad():
|
||||
content_A = self.netE.forward_test(real_A, "content")
|
||||
style_B1_h = self.netS.forward_test(real_B_1, "style")[0]
|
||||
style_B2_l = self.netS.forward_test(real_B_2, "style")[1]
|
||||
style_B = style_B1_h, style_B2_l
|
||||
|
||||
trs_AtoB = self.netG.forward_test(content_A, style_B)
|
||||
|
||||
return trs_AtoB
|
||||
@@ -0,0 +1,390 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
# import torch.nn.functional as F
|
||||
from .blocks import AdaOctConv, Oct_Conv_aftup, Oct_conv_lreLU, Oct_conv_up, OctConv
|
||||
|
||||
# from blocks import *
|
||||
|
||||
|
||||
def define_network(
|
||||
net_type, alpha_in, alpha_out, style_kernel, input_nc, nf, output_nc, freq_ratio
|
||||
):
|
||||
net = None
|
||||
sk = style_kernel
|
||||
|
||||
if net_type == "Encoder":
|
||||
net = Encoder(
|
||||
in_dim=input_nc,
|
||||
nf=nf,
|
||||
style_kernel=[sk, sk],
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
)
|
||||
elif net_type == "Generator":
|
||||
net = Decoder(
|
||||
nf=nf,
|
||||
out_dim=output_nc,
|
||||
style_channel=256,
|
||||
style_kernel=[sk, sk, 3],
|
||||
alpha_in=alpha_in,
|
||||
freq_ratio=freq_ratio,
|
||||
alpha_out=alpha_out,
|
||||
)
|
||||
return net
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, in_dim, nf=64, style_kernel=[3, 3], alpha_in=0.5, alpha_out=0.5):
|
||||
super(Encoder, self).__init__()
|
||||
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels=in_dim, out_channels=nf, kernel_size=7, stride=1, padding=3
|
||||
)
|
||||
|
||||
self.OctConv1_1 = OctConv(
|
||||
in_channels=nf,
|
||||
out_channels=nf,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
groups=64,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="first",
|
||||
)
|
||||
self.OctConv1_2 = OctConv(
|
||||
in_channels=nf,
|
||||
out_channels=2 * nf,
|
||||
kernel_size=1,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
self.OctConv1_3 = OctConv(
|
||||
in_channels=2 * nf,
|
||||
out_channels=2 * nf,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
|
||||
self.OctConv2_1 = OctConv(
|
||||
in_channels=2 * nf,
|
||||
out_channels=2 * nf,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
groups=128,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
self.OctConv2_2 = OctConv(
|
||||
in_channels=2 * nf,
|
||||
out_channels=4 * nf,
|
||||
kernel_size=1,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
self.OctConv2_3 = OctConv(
|
||||
in_channels=4 * nf,
|
||||
out_channels=4 * nf,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
|
||||
self.pool_h = nn.AdaptiveAvgPool2d((style_kernel[0], style_kernel[0]))
|
||||
self.pool_l = nn.AdaptiveAvgPool2d((style_kernel[1], style_kernel[1]))
|
||||
|
||||
self.relu = Oct_conv_lreLU()
|
||||
|
||||
def forward(self, x):
|
||||
enc_feat = []
|
||||
out = self.conv(x)
|
||||
|
||||
out = self.OctConv1_1(out)
|
||||
out = self.relu(out)
|
||||
out = self.OctConv1_2(out)
|
||||
out = self.relu(out)
|
||||
out = self.OctConv1_3(out)
|
||||
out = self.relu(out)
|
||||
enc_feat.append(out)
|
||||
|
||||
out = self.OctConv2_1(out)
|
||||
out = self.relu(out)
|
||||
out = self.OctConv2_2(out)
|
||||
out = self.relu(out)
|
||||
out = self.OctConv2_3(out)
|
||||
out = self.relu(out)
|
||||
enc_feat.append(out)
|
||||
|
||||
out_high, out_low = out
|
||||
out_sty_h = self.pool_h(out_high)
|
||||
out_sty_l = self.pool_l(out_low)
|
||||
out_sty = out_sty_h, out_sty_l
|
||||
|
||||
return out, out_sty, enc_feat
|
||||
|
||||
def forward_test(self, x, cond):
|
||||
out = self.conv(x)
|
||||
|
||||
out = self.OctConv1_1(out)
|
||||
out = self.relu(out)
|
||||
out = self.OctConv1_2(out)
|
||||
out = self.relu(out)
|
||||
out = self.OctConv1_3(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.OctConv2_1(out)
|
||||
out = self.relu(out)
|
||||
out = self.OctConv2_2(out)
|
||||
out = self.relu(out)
|
||||
out = self.OctConv2_3(out)
|
||||
out = self.relu(out)
|
||||
|
||||
if cond == "style":
|
||||
out_high, out_low = out
|
||||
out_sty_h = self.pool_h(out_high)
|
||||
out_sty_l = self.pool_l(out_low)
|
||||
return out_sty_h, out_sty_l
|
||||
else:
|
||||
return out
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
nf=64,
|
||||
out_dim=3,
|
||||
style_channel=512,
|
||||
style_kernel=[3, 3, 3],
|
||||
alpha_in=0.5,
|
||||
alpha_out=0.5,
|
||||
freq_ratio=[1, 1],
|
||||
pad_type="reflect",
|
||||
):
|
||||
super(Decoder, self).__init__()
|
||||
|
||||
group_div = [1, 2, 4, 8]
|
||||
self.up_oct = Oct_conv_up(scale_factor=2)
|
||||
|
||||
self.AdaOctConv1_1 = AdaOctConv(
|
||||
in_channels=4 * nf,
|
||||
out_channels=4 * nf,
|
||||
group_div=group_div[0],
|
||||
style_channels=style_channel,
|
||||
kernel_size=style_kernel,
|
||||
stride=1,
|
||||
padding=1,
|
||||
oct_groups=4 * nf,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
self.OctConv1_2 = OctConv(
|
||||
in_channels=4 * nf,
|
||||
out_channels=2 * nf,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
self.oct_conv_aftup_1 = Oct_Conv_aftup(
|
||||
in_channels=2 * nf,
|
||||
out_channels=2 * nf,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
pad_type=pad_type,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
)
|
||||
|
||||
self.AdaOctConv2_1 = AdaOctConv(
|
||||
in_channels=2 * nf,
|
||||
out_channels=2 * nf,
|
||||
group_div=group_div[1],
|
||||
style_channels=style_channel,
|
||||
kernel_size=style_kernel,
|
||||
stride=1,
|
||||
padding=1,
|
||||
oct_groups=2 * nf,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
self.OctConv2_2 = OctConv(
|
||||
in_channels=2 * nf,
|
||||
out_channels=nf,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
self.oct_conv_aftup_2 = Oct_Conv_aftup(
|
||||
nf, nf, 3, 1, 1, pad_type, alpha_in, alpha_out
|
||||
)
|
||||
|
||||
self.AdaOctConv3_1 = AdaOctConv(
|
||||
in_channels=nf,
|
||||
out_channels=nf,
|
||||
group_div=group_div[2],
|
||||
style_channels=style_channel,
|
||||
kernel_size=style_kernel,
|
||||
stride=1,
|
||||
padding=1,
|
||||
oct_groups=nf,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="normal",
|
||||
)
|
||||
self.OctConv3_2 = OctConv(
|
||||
in_channels=nf,
|
||||
out_channels=nf // 2,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
alpha_in=alpha_in,
|
||||
alpha_out=alpha_out,
|
||||
type="last",
|
||||
freq_ratio=freq_ratio,
|
||||
)
|
||||
|
||||
self.conv4 = nn.Conv2d(in_channels=nf // 2, out_channels=out_dim, kernel_size=1)
|
||||
|
||||
def forward(self, content, style):
|
||||
out = self.AdaOctConv1_1(content, style)
|
||||
out = self.OctConv1_2(out)
|
||||
out = self.up_oct(out)
|
||||
out = self.oct_conv_aftup_1(out)
|
||||
|
||||
out = self.AdaOctConv2_1(out, style)
|
||||
out = self.OctConv2_2(out)
|
||||
out = self.up_oct(out)
|
||||
out = self.oct_conv_aftup_2(out)
|
||||
|
||||
out = self.AdaOctConv3_1(out, style)
|
||||
out = self.OctConv3_2(out)
|
||||
out, out_high, out_low = out
|
||||
|
||||
out = self.conv4(out)
|
||||
out_high = self.conv4(out_high)
|
||||
out_low = self.conv4(out_low)
|
||||
|
||||
return out, out_high, out_low
|
||||
|
||||
def forward_test(self, content, style):
|
||||
out = self.AdaOctConv1_1(content, style, "test")
|
||||
out = self.OctConv1_2(out)
|
||||
out = self.up_oct(out)
|
||||
out = self.oct_conv_aftup_1(out)
|
||||
|
||||
out = self.AdaOctConv2_1(out, style, "test")
|
||||
out = self.OctConv2_2(out)
|
||||
out = self.up_oct(out)
|
||||
out = self.oct_conv_aftup_2(out)
|
||||
|
||||
out = self.AdaOctConv3_1(out, style, "test")
|
||||
out = self.OctConv3_2(out)
|
||||
|
||||
out = self.conv4(out[0])
|
||||
return out
|
||||
|
||||
|
||||
############## Contrastive Loss function ##############
|
||||
def calc_mean_std(feat, eps=1e-5):
|
||||
# eps is a small value added to the variance to avoid divide-by-zero.
|
||||
size = feat.size()
|
||||
assert len(size) == 4
|
||||
N, C = size[:2]
|
||||
feat_var = feat.view(N, C, -1).var(dim=2) + eps
|
||||
feat_std = feat_var.sqrt().view(N, C, 1, 1)
|
||||
feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
|
||||
return feat_mean, feat_std
|
||||
|
||||
|
||||
def calc_content_loss(input, target):
|
||||
assert input.size() == target.size()
|
||||
mse_loss = nn.MSELoss()
|
||||
return mse_loss(input, target)
|
||||
|
||||
|
||||
def calc_style_loss(input, target):
|
||||
assert input.size() == target.size()
|
||||
mse_loss = nn.MSELoss()
|
||||
input_mean, input_std = calc_mean_std(input)
|
||||
target_mean, target_std = calc_mean_std(target)
|
||||
|
||||
loss = mse_loss(input_mean, target_mean) + mse_loss(input_std, target_std)
|
||||
return loss
|
||||
|
||||
|
||||
class EFDM_loss(nn.Module):
|
||||
def __init__(self):
|
||||
super(EFDM_loss, self).__init__()
|
||||
self.mse_loss = nn.MSELoss()
|
||||
|
||||
def efdm_single(self, style, trans):
|
||||
B, C, W, H = style.size(0), style.size(1), style.size(2), style.size(3)
|
||||
|
||||
value_style, index_style = torch.sort(style.view(B, C, -1))
|
||||
value_trans, index_trans = torch.sort(trans.view(B, C, -1))
|
||||
inverse_index = index_trans.argsort(-1)
|
||||
|
||||
return self.mse_loss(
|
||||
trans.view(B, C, -1), value_style.gather(-1, inverse_index)
|
||||
)
|
||||
|
||||
def forward(self, style_E, style_S, translate_E, translate_S, neg_idx):
|
||||
loss = 0.0
|
||||
batch = style_E[0][0].shape[0]
|
||||
for b in range(batch):
|
||||
poss_loss = 0.0
|
||||
neg_loss = 0.0
|
||||
|
||||
# Positive loss
|
||||
for i in range(len(style_E)):
|
||||
poss_loss += self.efdm_single(
|
||||
style_E[i][0][b].unsqueeze(0), translate_E[i][0][b].unsqueeze(0)
|
||||
) + self.efdm_single(
|
||||
style_E[i][1][b].unsqueeze(0), translate_E[i][1][b].unsqueeze(0)
|
||||
)
|
||||
for i in range(len(style_S)):
|
||||
poss_loss += self.efdm_single(
|
||||
style_S[i][0][b].unsqueeze(0), translate_S[i][0][b].unsqueeze(0)
|
||||
) + self.efdm_single(
|
||||
style_S[i][1][b].unsqueeze(0), translate_S[i][1][b].unsqueeze(0)
|
||||
)
|
||||
|
||||
# Negative loss
|
||||
for nb in neg_idx[b]:
|
||||
for i in range(len(style_E)):
|
||||
neg_loss += self.efdm_single(
|
||||
style_E[i][0][nb].unsqueeze(0),
|
||||
translate_E[i][0][b].unsqueeze(0),
|
||||
) + self.efdm_single(
|
||||
style_E[i][1][nb].unsqueeze(0),
|
||||
translate_E[i][1][b].unsqueeze(0),
|
||||
)
|
||||
for i in range(len(style_S)):
|
||||
neg_loss += self.efdm_single(
|
||||
style_S[i][0][nb].unsqueeze(0),
|
||||
translate_S[i][0][b].unsqueeze(0),
|
||||
) + self.efdm_single(
|
||||
style_S[i][1][nb].unsqueeze(0),
|
||||
translate_S[i][1][b].unsqueeze(0),
|
||||
)
|
||||
|
||||
loss += poss_loss / neg_loss
|
||||
|
||||
return loss
|
||||
@@ -0,0 +1,58 @@
|
||||
import torch.nn as nn
|
||||
|
||||
vgg = nn.Sequential(
|
||||
nn.Conv2d(3, 3, (1, 1)),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(3, 64, (3, 3)),
|
||||
nn.ReLU(), # relu1-1
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(64, 64, (3, 3)),
|
||||
nn.ReLU(), # relu1-2
|
||||
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(64, 128, (3, 3)),
|
||||
nn.ReLU(), # relu2-1
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(128, 128, (3, 3)),
|
||||
nn.ReLU(), # relu2-2
|
||||
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(128, 256, (3, 3)),
|
||||
nn.ReLU(), # relu3-1
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(), # relu3-2
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(), # relu3-3
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(), # relu3-4
|
||||
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 512, (3, 3)),
|
||||
nn.ReLU(), # relu4-1, this is the last layer used
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu4-2
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu4-3
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu4-4
|
||||
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu5-1
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu5-2
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu5-3
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu5-4
|
||||
)
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui-styletransferplus"
|
||||
description = "Nodes: Neural Neighbor, CAST, EFDM, MicroAST, UniST, AesPA, TSSAT ..."
|
||||
version = "1.0.2"
|
||||
description = "Nodes: Neural Neighbor, CAST, EFDM, MicroAST, UniST, AesPA, TSSAT, AesFA, ..."
|
||||
version = "1.0.3"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["torch", "torchvision", "scikit-image", "einops"]
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ import tqdm
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
from .constants import (
|
||||
AESFA_PATH,
|
||||
AESPA_DEC_PATH,
|
||||
AESPA_TRANSFORMER_PATH,
|
||||
CAST_DEFAULT,
|
||||
@@ -33,6 +34,13 @@ from .constants import (
|
||||
VGG_NORM_CONV51_PATH,
|
||||
)
|
||||
|
||||
# AesFA
|
||||
from .module_aesfa.aesfa_model import (
|
||||
MODEL_AESFA,
|
||||
inference_aesfa,
|
||||
inference_aesfa_style_blend,
|
||||
)
|
||||
|
||||
# AesPA
|
||||
from .module_aespa.aespa_model import MODEL_AESPA, inference_aespa
|
||||
|
||||
@@ -562,7 +570,7 @@ class TSSAT:
|
||||
),
|
||||
"size": (
|
||||
"INT",
|
||||
{"default": 512, "min": 1, "max": 1024, "step": 1},
|
||||
{"default": 512, "min": 1, "step": 1},
|
||||
),
|
||||
"max_iter": (
|
||||
"INT",
|
||||
@@ -615,6 +623,128 @@ class TSSAT:
|
||||
return (torch.cat(result, dim=0),)
|
||||
|
||||
|
||||
class AesFA:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"src_img": ("IMAGE",),
|
||||
"style_img": ("IMAGE",),
|
||||
"do_crop": (
|
||||
"BOOLEAN",
|
||||
{"default": False},
|
||||
),
|
||||
"size": (
|
||||
"INT",
|
||||
{"default": 512, "min": 1, "step": 1},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("out_img",)
|
||||
FUNCTION = "todo"
|
||||
CATEGORY = "StyleTransferPlus"
|
||||
|
||||
def todo(
|
||||
self,
|
||||
src_img: torch.Tensor,
|
||||
style_img: torch.Tensor,
|
||||
do_crop: bool,
|
||||
size: int,
|
||||
):
|
||||
print(f"{src_img.shape=}")
|
||||
print(f"{style_img.shape=}")
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = MODEL_AESFA(
|
||||
checkpoint_path=f"{base_dir}/{AESFA_PATH}",
|
||||
device=device,
|
||||
)
|
||||
|
||||
params = {
|
||||
"style_img": style_img.permute(0, 3, 1, 2),
|
||||
"device": device,
|
||||
"size": size,
|
||||
"do_crop": do_crop,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
num_frames = src_img.size(0)
|
||||
|
||||
result: list[torch.Tensor] = []
|
||||
with torch.no_grad():
|
||||
for i in range(num_frames):
|
||||
params["src_img"] = src_img[i].unsqueeze(0).permute(0, 3, 1, 2)
|
||||
res_tensor = inference_aesfa(**params)
|
||||
result.append(res_tensor.permute(0, 2, 3, 1))
|
||||
|
||||
return (torch.cat(result, dim=0),)
|
||||
|
||||
|
||||
class AesFAStyleBlend:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"src_img": ("IMAGE",),
|
||||
"style_hi": ("IMAGE",),
|
||||
"style_lo": ("IMAGE",),
|
||||
"do_crop": (
|
||||
"BOOLEAN",
|
||||
{"default": False},
|
||||
),
|
||||
"size": (
|
||||
"INT",
|
||||
{"default": 512, "min": 1, "step": 1},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("out_img",)
|
||||
FUNCTION = "todo"
|
||||
CATEGORY = "StyleTransferPlus"
|
||||
|
||||
def todo(
|
||||
self,
|
||||
src_img: torch.Tensor,
|
||||
style_hi: torch.Tensor,
|
||||
style_lo: torch.Tensor,
|
||||
do_crop: bool,
|
||||
size: int,
|
||||
):
|
||||
print(f"{src_img.shape=}")
|
||||
print(f"{style_hi.shape=}")
|
||||
print(f"{style_lo.shape=}")
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = MODEL_AESFA(
|
||||
checkpoint_path=f"{base_dir}/{AESFA_PATH}",
|
||||
device=device,
|
||||
)
|
||||
|
||||
params = {
|
||||
"style_hi": style_hi.permute(0, 3, 1, 2),
|
||||
"style_lo": style_lo.permute(0, 3, 1, 2),
|
||||
"device": device,
|
||||
"size": size,
|
||||
"do_crop": do_crop,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
num_frames = src_img.size(0)
|
||||
|
||||
result: list[torch.Tensor] = []
|
||||
with torch.no_grad():
|
||||
for i in range(num_frames):
|
||||
params["src_img"] = src_img[i].unsqueeze(0).permute(0, 3, 1, 2)
|
||||
res_tensor = inference_aesfa_style_blend(**params)
|
||||
result.append(res_tensor.permute(0, 2, 3, 1))
|
||||
|
||||
return (torch.cat(result, dim=0),)
|
||||
|
||||
|
||||
def serialize_floats(lst: list[float]):
|
||||
return ",".join(map(str, lst))
|
||||
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 810 KiB |
@@ -0,0 +1,394 @@
|
||||
{
|
||||
"last_node_id": 13,
|
||||
"last_link_id": 34,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 9,
|
||||
"type": "AESFA",
|
||||
"pos": [
|
||||
678,
|
||||
313
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 102
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "src_img",
|
||||
"type": "IMAGE",
|
||||
"link": 33
|
||||
},
|
||||
{
|
||||
"name": "style_img",
|
||||
"type": "IMAGE",
|
||||
"link": 30
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "out_img",
|
||||
"type": "IMAGE",
|
||||
"links": [],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "AESFA"
|
||||
},
|
||||
"widgets_values": [
|
||||
true,
|
||||
1024
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 11,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
682,
|
||||
672
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 314
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
25
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"popova.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 10,
|
||||
"type": "AesFAStyleBlend",
|
||||
"pos": [
|
||||
682,
|
||||
501
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 122
|
||||
},
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "src_img",
|
||||
"type": "IMAGE",
|
||||
"link": 34
|
||||
},
|
||||
{
|
||||
"name": "style_hi",
|
||||
"type": "IMAGE",
|
||||
"link": 31
|
||||
},
|
||||
{
|
||||
"name": "style_lo",
|
||||
"type": "IMAGE",
|
||||
"link": 25
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "out_img",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
26
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "AesFAStyleBlend"
|
||||
},
|
||||
"widgets_values": [
|
||||
true,
|
||||
1024
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 12,
|
||||
"type": "Reroute",
|
||||
"pos": [
|
||||
539,
|
||||
586
|
||||
],
|
||||
"size": [
|
||||
75,
|
||||
26
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "",
|
||||
"type": "*",
|
||||
"link": 28
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
30,
|
||||
31
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"showOutputText": false,
|
||||
"horizontal": false
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
338,
|
||||
668
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 314
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
28
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"S4.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 13,
|
||||
"type": "Reroute",
|
||||
"pos": [
|
||||
528.2286926786944,
|
||||
402.29128892008794
|
||||
],
|
||||
"size": [
|
||||
75,
|
||||
26
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "",
|
||||
"type": "*",
|
||||
"link": 32
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
33,
|
||||
34
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"showOutputText": false,
|
||||
"horizontal": false
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
174,
|
||||
290
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 314
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
32
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"006.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "SaveImage",
|
||||
"pos": [
|
||||
1039,
|
||||
287
|
||||
],
|
||||
"size": [
|
||||
321.92869267869446,
|
||||
328.691288920088
|
||||
],
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 26
|
||||
}
|
||||
],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"StyleTransfer"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
25,
|
||||
11,
|
||||
0,
|
||||
10,
|
||||
2,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
26,
|
||||
10,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
28,
|
||||
3,
|
||||
0,
|
||||
12,
|
||||
0,
|
||||
"*"
|
||||
],
|
||||
[
|
||||
30,
|
||||
12,
|
||||
0,
|
||||
9,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
31,
|
||||
12,
|
||||
0,
|
||||
10,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
32,
|
||||
2,
|
||||
0,
|
||||
13,
|
||||
0,
|
||||
"*"
|
||||
],
|
||||
[
|
||||
33,
|
||||
13,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
34,
|
||||
13,
|
||||
0,
|
||||
10,
|
||||
0,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.9090909090909091,
|
||||
"offset": [
|
||||
-109.12869267869443,
|
||||
-65.69128892008793
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
Reference in New Issue
Block a user