diff --git a/README.md b/README.md index a4b4b8c..c11e99b 100644 --- a/README.md +++ b/README.md @@ -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 | ![efdm](example_outputs/efdm.png) | ![microast](example_outputs/microast.png) | | **UniST (2023) - Medium - Only squares** | **AesPA-Net (2023) - Medium - Only squares** | | ![efdm](example_outputs/unist_image.png) | ![aespa](example_outputs/aespa.png) | -| **TSSAT (2023) - Slow** | **X** | -| ![tssat](example_outputs/tssat.png) | X | +| **TSSAT (2023) - Slow** | **AesFA (2024) - Fast - Only squares** | +| ![tssat](example_outputs/tssat.png) | ![aesfa](example_outputs/aesfa.png) | ### Neural Neighbor @@ -167,6 +169,22 @@ models/vgg_normalised.pth ![wf_tssat](workflows/wf_tssat.png) + +### 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) + +![wf_aesfa_w_blend](workflows/wf_aesfa_w_blend.png) + +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 diff --git a/__init__.py b/__init__.py index cc014f6..7a4479c 100644 --- a/__init__.py +++ b/__init__.py @@ -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", } diff --git a/constants.py b/constants.py index 4679f9e..6ff0f7e 100644 --- a/constants.py +++ b/constants.py @@ -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" \ No newline at end of file +TSSAT_DEC_PATH = "models/tssat/decoder_iter_160000.pth" + +AESFA_PATH = "models/aesfa/main.pth" \ No newline at end of file diff --git a/example_outputs/aesfa.png b/example_outputs/aesfa.png new file mode 100644 index 0000000..3d8b1c2 Binary files /dev/null and b/example_outputs/aesfa.png differ diff --git a/module_aesfa/aesfa_model.py b/module_aesfa/aesfa_model.py new file mode 100644 index 0000000..bc12cdf --- /dev/null +++ b/module_aesfa/aesfa_model.py @@ -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 diff --git a/module_aesfa/blocks.py b/module_aesfa/blocks.py new file mode 100644 index 0000000..b8ee7eb --- /dev/null +++ b/module_aesfa/blocks.py @@ -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 \ No newline at end of file diff --git a/module_aesfa/model.py b/module_aesfa/model.py new file mode 100644 index 0000000..f52240c --- /dev/null +++ b/module_aesfa/model.py @@ -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 diff --git a/module_aesfa/networks.py b/module_aesfa/networks.py new file mode 100644 index 0000000..22ca02d --- /dev/null +++ b/module_aesfa/networks.py @@ -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 diff --git a/module_aesfa/vgg19.py b/module_aesfa/vgg19.py new file mode 100644 index 0000000..1f77629 --- /dev/null +++ b/module_aesfa/vgg19.py @@ -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 +) + diff --git a/pyproject.toml b/pyproject.toml index 8129d07..4c403bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/run.py b/run.py index eee480c..80ba259 100644 --- a/run.py +++ b/run.py @@ -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)) diff --git a/workflows/wf_aesfa_w_blend.png b/workflows/wf_aesfa_w_blend.png new file mode 100644 index 0000000..9283da8 Binary files /dev/null and b/workflows/wf_aesfa_w_blend.png differ diff --git a/workflows/workflow_aesfa_w_blend.json b/workflows/workflow_aesfa_w_blend.json new file mode 100644 index 0000000..9501264 --- /dev/null +++ b/workflows/workflow_aesfa_w_blend.json @@ -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 +} \ No newline at end of file