diff --git a/.gitignore b/.gitignore index 2edc8a5..0b6b44e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,8 @@ +# Checkpoint files *.pth *.tar *.pt +*.t7 # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index 237c840..b4a2f5a 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ # ComfyUI-StyleTransferPlus -Advanced Non-Diffusion-based Style Transfer in ComfyUI +Advance Non-Diffusion-based Style Transfer in ComfyUI Click name to jump to workflow @@ -7,7 +7,9 @@ Click name to jump to workflow 2. [**CAST**](#cast). Paper: [Domain Enhanced Arbitrary Image Style Transfer via Contrastive Learning](https://github.com/zyxElsa/CAST_pytorch) 3. [**EFDM**](#efdm). Paper: [Exact Feature Distribution Matching for Arbitrary Style Transfer and Domain Generalization](https://github.com/YBZh/EFDM) 4. [**MicroAST**](#microast). Paper: [Towards Super-Fast Ultra-Resolution Arbitrary Style Transfer](https://github.com/EndyWon/MicroAST) -5. [**UniST**](#unist) : [Two Birds, One Stone: A Unified Framework for Joint Learning of Image and Video Style Transfers](https://github.com/NevSNev/UniST) +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) + [**Extra nodes**](#extra-nodes) @@ -17,13 +19,13 @@ More will come soon All nodes support batched input (i.e video) but is generally not recommended. One should generate 1 or 2 style frames (start and end), then use [ComfyUI-EbSynth](https://github.com/FuouM/ComfyUI-EbSynth) to propagate the style to the entire video. -| Neural Neighbor (2020-2022) - Slowest | CAST (2022-2023) - Fast | -| :-----------------------------------------------------: | :---------------------------------------: | -| ![neural_neighbor](example_outputs/neural_neighbor.png) | ![cast](example_outputs/cast.png) | -| **EFDM (2022) - Fast** | **MicroAST (2023) - Fast** | -| ![efdm](example_outputs/efdm.png) | ![microast](example_outputs/microast.png) | -| **UniST (2023) - Medium - Only squares** | **X** | -| ![efdm](example_outputs/unist_image.png) | X | +| Neural Neighbor (2020-2022) - Slowest | CAST (2022-2023) - Fast | +| :-----------------------------------------------------: | :------------------------------------------: | +| ![neural_neighbor](example_outputs/neural_neighbor.png) | ![cast](example_outputs/cast.png) | +| **EFDM (2022) - Fast** | **MicroAST (2023) - Fast** | +| ![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) | ### Neural Neighbor @@ -47,7 +49,7 @@ See more information on the original repository. No model is needed. Download the CAST and UCAST models in the [Test](https://github.com/zyxElsa/CAST_pytorch?tab=readme-ov-file#test) section (Google Drive) (Unzip). -Download the `vgg_normalized` model in the [Train](https://github.com/zyxElsa/CAST_pytorch?tab=readme-ov-file#train) section +Download the `vgg_normalized.pth` model in the [Train](https://github.com/zyxElsa/CAST_pytorch?tab=readme-ov-file#train) section ``` models/ @@ -127,8 +129,25 @@ The Video node is "more native" than the Image node for batched images (video) i | :-----------------------------------------------------------: | :--------------------------------------------------------------: | | [workflow_unist_image.json](workflows/workflow_microast.json) | [workflow_unist_video.json](workflows/workflow_unist_video.json) | | ![wf_unist_image](workflows/wf_unist_image.png) | ![wf_microast](example_outputs/example_unist_video.png) | -||(Not a workflow-embedded image)| +| | (Not a workflow-embedded image) | +### AesPA-Net + +Download models from [Usage](https://github.com/Kibeom-Hong/AesPA-Net?tab=readme-ov-file#usage) section. + +For the VGG model, download `vgg_normalised_conv5_1.pth` from https://github.com/pietrocarbo/deep-transfer/tree/master/models/autoencoder_vgg19/vgg19_5 + +``` +models\aespa\dec_model.pth +models\aespa\transformer_model.pth +models\aespa\vgg_normalised_conv5_1.pth +``` + +The original authors provided a `vgg_normalised_conv5_1.t7` model, which can only be opened with torchfile, and pyTorch has dropped native support for it. It also doesn't work reliably on Windows. + +[workflow_aespa.json](workflows/workflow_aespa.json) + +![wf_aespa](workflows/wf_aespa.png) ## Extra nodes @@ -140,17 +159,18 @@ The Video node is "more native" than the Image node for batched images (video) i ``` @inproceedings{zhang2020cast, -author = {Zhang, Yuxin and Tang, Fan and Dong, Weiming and Huang, Haibin and Ma, Chongyang and Lee, Tong-Yee and Xu, Changsheng}, -title = {Domain Enhanced Arbitrary Image Style Transfer via Contrastive Learning}, -booktitle = {ACM SIGGRAPH}, -year = {2022}} + author = {Zhang, Yuxin and Tang, Fan and Dong, Weiming and Huang, Haibin and Ma, Chongyang and Lee, Tong-Yee and Xu, Changsheng}, + title = {Domain Enhanced Arbitrary Image Style Transfer via Contrastive Learning}, + booktitle = {ACM SIGGRAPH}, + year = {2022} +} @article{zhang2023unified, - title={A Unified Arbitrary Style Transfer Framework via Adaptive Contrastive Learning}, - author={Zhang, Yuxin and Tang, Fan and Dong, Weiming and Huang, Haibin and Ma, Chongyang and Lee, Tong-Yee and Xu, Changsheng}, - journal={ACM Transactions on Graphics}, - year={2023}, - publisher={ACM New York, NY} + title={A Unified Arbitrary Style Transfer Framework via Adaptive Contrastive Learning}, + author={Zhang, Yuxin and Tang, Fan and Dong, Weiming and Huang, Haibin and Ma, Chongyang and Lee, Tong-Yee and Xu, Changsheng}, + journal={ACM Transactions on Graphics}, + year={2023}, + publisher={ACM New York, NY} } ``` @@ -181,4 +201,15 @@ year = {2022}} year = {2023}, pages = {23545-23554} } +``` + +``` +@article{Hong2023AesPANetAP, + title={AesPA-Net: Aesthetic Pattern-Aware Style Transfer Networks}, + author={Kibeom Hong and Seogkyu Jeon and Junsoo Lee and Namhyuk Ahn and Kunhee Kim and Pilhyeon Lee and Daesik Kim and Youngjung Uh and Hyeran Byun}, + journal={ArXiv}, + year={2023}, + volume={abs/2307.09724}, + url={https://api.semanticscholar.org/CorpusID:259982728} +} ``` \ No newline at end of file diff --git a/__init__.py b/__init__.py index ae4035b..685c4b0 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,4 @@ -from .run import NeuralNeighbor, CAST, EFDM, MicroAST, UniST, UniST_Video +from .run import NeuralNeighbor, CAST, EFDM, MicroAST, UniST, UniST_Video, AesPA from .run_extra import CoralColorTransfer NODE_CLASS_MAPPINGS = { @@ -9,6 +9,7 @@ NODE_CLASS_MAPPINGS = { "CoralColorTransfer": CoralColorTransfer, "UniST": UniST, "UniST_Video": UniST_Video, + "AesPA": AesPA, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -19,6 +20,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CoralColorTransfer": "Coral Color Transfer", "UniST": "UniST", "UniST_Video": "UniST Video", + "AesPA": "AesPA-Net", } diff --git a/constants.py b/constants.py index 78e6eba..6c4af4a 100644 --- a/constants.py +++ b/constants.py @@ -17,3 +17,7 @@ MICROAST_MODULATOR_PATH = "models/microast/modulator_iter_160000.pth.tar" UNIST_PATH = "models/unist/UniST_model.pt" UNIST_DEC_PATH = "models/unist/dec_r41.pth" UNIST_ENC_PATH = "models/unist/vgg_r41.pth" + +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" diff --git a/example_outputs/aespa.png b/example_outputs/aespa.png new file mode 100644 index 0000000..9da76c9 Binary files /dev/null and b/example_outputs/aespa.png differ diff --git a/module_aespa/aespa_model.py b/module_aespa/aespa_model.py new file mode 100644 index 0000000..3c54dd2 --- /dev/null +++ b/module_aespa/aespa_model.py @@ -0,0 +1,295 @@ +import torch +import torchvision.transforms as transforms +import torchvision.transforms.functional as tf_f +import random +import torch.nn.functional as F + +from .aespanet_models import ( + AdaptiveMultiAttn_Transformer_v2, + VGGDecoder, + VGGEncoder, +) + + +def inference_aespa( + src_img: torch.Tensor, + style_img: torch.Tensor, + device, + size: int, + do_crop: bool, + encoder: VGGEncoder, + decoder: VGGDecoder, + transformer: AdaptiveMultiAttn_Transformer_v2, +): + content_tf = test_transform( + size, do_crop, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ) + content = content_tf(src_img).to(device) # [B, C, H, W] + style = content_tf(style_img).to(device) + + gray_content = tf_f.rgb_to_grayscale(content).repeat(1, 3, 1, 1).to(device) + gray_style = tf_f.rgb_to_grayscale(style).repeat(1, 3, 1, 1).to(device) + + style_weight = ( + adaptive_gram_weight(style, 1, 8, encoder) + + adaptive_gram_weight(style, 2, 8, encoder) + + adaptive_gram_weight(style, 3, 8, encoder) + ) / 3 + + gray_style_weight = ( + adaptive_gram_weight(gray_style, 1, 8, encoder) + + adaptive_gram_weight(gray_style, 2, 8, encoder) + + adaptive_gram_weight(gray_style, 3, 8, encoder) + ) / 3 + + style_adaptive_alpha = ( + style_weight.unsqueeze(1).cuda() + gray_style_weight.unsqueeze(1).cuda() + ) / 2 + + content_skips = {} + style_skips = {} + + _ = encoder.encode(content, content_skips) # content_feat + _ = encoder.encode(style, style_skips) # style_feat + + if gray_content is not None: + gray_content_skips = {} + _ = encoder.encode(gray_content, gray_content_skips) # gray_content_feat + if gray_style is not None: + gray_style_skips = {} + _ = encoder.encode(gray_style, gray_style_skips) # gray_style_feat + + ( + local_transformed_feature, + _, + _, + _, + _, + ) = transformer.forward( + content_skips["conv4_1"], + style_skips["conv4_1"], + content_skips["conv5_1"], + style_skips["conv5_1"], + adaptive_get_keys(content_skips, 4, 4, target_feat=content_skips["conv4_1"]), + adaptive_get_keys(style_skips, 1, 4, target_feat=style_skips["conv4_1"]), + adaptive_get_keys(content_skips, 5, 5, target_feat=content_skips["conv5_1"]), + adaptive_get_keys(style_skips, 1, 5, target_feat=style_skips["conv5_1"]), + ) + + if gray_content is not None: + global_transformed_feat = feature_wct_simple( + gray_content_skips["conv4_1"], gray_style_skips["conv4_1"] + ) + else: + global_transformed_feat = feature_wct_simple( + content_skips["conv4_1"], style_skips["conv4_1"] + ) + + transformed_feature = ( + global_transformed_feat * (1 - style_adaptive_alpha.unsqueeze(-1).unsqueeze(-1)) + + style_adaptive_alpha.unsqueeze(-1).unsqueeze(-1) * local_transformed_feature + ) + + stylized_image = decoder.decode(transformed_feature, content_skips, style_skips) + + output = un_normalize_batch(stylized_image) + + return output + + +def adaptive_gram_weight(image, level, ratio, encoder): + if level == 0: + encoded_features = image + else: + encoded_features = encoder.get_features(image, level) # B x C x W x H + global_gram = gram_matrix(encoded_features) + + B, C, w, h = encoded_features.size() + target_w, target_h = w // ratio, h // ratio + + patches = extract_image_patches(encoded_features, target_w, target_h) + _, patches_num, _, _, _ = patches.size() + cos = torch.nn.CosineSimilarity(eps=1e-6) + + intra_gram_statistic = [] + inter_gram_statistic = [] + comb = torch.combinations(torch.arange(patches_num), r=2) + if patches_num >= 10: + sampling_num = int(comb.size(0) * 0.05) + else: + sampling_num = comb.size(0) + for idx in range(B): + if patches_num < 2: + continue + cos_gram = [] + + for patch in range(0, patches_num): + cos_gram.append( + cos(global_gram, gram_matrix(patches[idx][patch].unsqueeze(0))) + .mean() + .item() + ) + + intra_gram_statistic.append(torch.tensor(cos_gram)) + + cos_gram = [] + for idxes in random.choices(list(comb), k=sampling_num): + cos_gram.append( + cos( + gram_matrix(patches[idx][idxes[0]].unsqueeze(0)), + gram_matrix(patches[idx][idxes[1]].unsqueeze(0)), + ) + .mean() + .item() + ) + + inter_gram_statistic.append(torch.tensor(cos_gram)) + + intra_gram_statistic = torch.stack(intra_gram_statistic).mean(dim=1) + inter_gram_statistic = torch.stack(inter_gram_statistic).mean(dim=1) + results = (intra_gram_statistic + inter_gram_statistic) / 2 + + ##For boosting value + results = 1 / (1 + torch.exp(-10 * (results - 0.6))) + + return results + + +def gram_matrix(y): + (b, ch, h, w) = y.size() + features = y.view(b, ch, w * h) + features_t = features.transpose(1, 2) + gram = features.bmm(features_t) / (ch * h * w) + return gram + + +def extract_image_patches(x, kernel, stride=1): + b, c, h, w = x.shape + + # Extract patches + patches = x.unfold(2, kernel, stride).unfold(3, kernel, stride) + patches = patches.contiguous().view(b, c, -1, kernel, kernel) + patches = patches.permute(0, 2, 1, 3, 4).contiguous() + + # return patches.view(b, number_of_patches, c, h, w) + return patches.view(b, -1, c, kernel, kernel) + + +def adaptive_get_keys(feat_skips, start_layer_idx, last_layer_idx, target_feat): + B, C, th, tw = target_feat.shape + results = [] + target_conv_layer = "conv" + str(last_layer_idx) + "_1" + _, _, h, w = feat_skips[target_conv_layer].shape + for i in range(start_layer_idx, last_layer_idx + 1): + target_conv_layer = "conv" + str(i) + "_1" + if i == last_layer_idx: + results.append(mean_variance_norm(feat_skips[target_conv_layer])) + else: + results.append( + mean_variance_norm(F.interpolate(feat_skips[target_conv_layer], (h, w))) + ) + + return F.interpolate(torch.cat(results, dim=1), (th, tw)) + + +def mean_variance_norm(feat): + size = feat.size() + mean, std = calc_mean_std(feat) + normalized_feat = (feat - mean.expand(size)) / std.expand(size) + return normalized_feat + + +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 feature_wct_simple(content_feat, style_feat, alpha=1): + target_feature = Bw_wct_core(content_feat, style_feat) + + target_feature = target_feature.view_as(content_feat) + target_feature = alpha * target_feature + (1 - alpha) * content_feat + return target_feature + + +def Bw_wct_core(content_feat, style_feat, weight=1, registers=None, device="cpu"): + N, C, H, W = content_feat.size() + cont_min = content_feat.min().item() + cont_max = content_feat.max().item() + + whiten_cF, _, _ = SwitchWhiten2d(content_feat) + _, wm_s, s_mean = SwitchWhiten2d(style_feat) + + targetFeature = torch.bmm(torch.inverse(wm_s), whiten_cF) + targetFeature = targetFeature.view(N, C, H, W) + targetFeature = targetFeature + s_mean.unsqueeze(2).expand_as(targetFeature) + targetFeature.clamp_(cont_min, cont_max) + + return targetFeature + + +def SwitchWhiten2d(x): + N, C, H, W = x.size() + + in_data = x.view(N, C, -1) + + eye = in_data.data.new().resize_(C, C) + eye = torch.nn.init.eye_(eye).view(1, C, C).expand(N, C, C) + + # calculate other statistics + mean_in = in_data.mean(-1, keepdim=True) + x_in = in_data - mean_in + # (N x g) x C x C + cov_in = torch.bmm(x_in, torch.transpose(x_in, 1, 2)).div(H * W) + + mean = mean_in + cov = cov_in + 1e-5 * eye + + # perform whitening using Newton's iteration + Ng, c, _ = cov.size() + P = torch.eye(c).to(cov).expand(Ng, c, c) + + rTr = (cov * P).sum((1, 2), keepdim=True).reciprocal_() + cov_N = cov * rTr + for k in range(5): + P = torch.baddbmm(1.5, P, -0.5, torch.matrix_power(P, 3), cov_N) + + wm = P.mul_(rTr.sqrt()) + x_hat = torch.bmm(wm, in_data - mean) + + return x_hat, wm, mean + + +def test_transform(size, crop, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]): + transform_list = [] + + if size != 0: + # Ensure size is even + size = (size // 2) * 2 + + # Limit size to 1024 + if size > 1024: + size = 1024 + + if crop: + transform_list.extend( + [transforms.Resize(size), transforms.CenterCrop(size)] + ) + else: + 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_aespa/aespanet_models.py b/module_aespa/aespanet_models.py new file mode 100644 index 0000000..2581e9a --- /dev/null +++ b/module_aespa/aespanet_models.py @@ -0,0 +1,950 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils import spectral_norm + +from .utils import feature_wct_simple, init_weights + + +def size_arrange(x): + x_w, x_h = x.size(2), x.size(3) + + if (x_w % 2) != 0: + x_w = (x_w // 2) * 2 + if (x_h % 2) != 0: + x_h = (x_h // 2) * 2 + + if (x_h > 1024) or (x_w > 1024): + old_x_w = x_w + x_w = x_w // 2 + x_h = int(x_h * x_w / old_x_w) + + return F.interpolate(x, size=(x_w, x_h)) + + +def get_LL_HH(x): + pooled = torch.nn.functional.avg_pool2d(x, 2) + # up_pooled = torch.nn.functional.interpolate(pooled, scale_factor=2, mode='nearest') + up_pooled = torch.nn.functional.interpolate( + pooled, size=[x.size(2), x.size(3)], mode="nearest" + ) + HH = x - up_pooled + LL = up_pooled + return HH, LL + + +def split_image(x, scale=1): + origin_w, origin_h = x.size(2), x.size(3) + down_sample = torch.nn.functional.avg_pool2d(x, 2**scale) + up_sample = torch.nn.functional.interpolate( + down_sample, (origin_w, origin_h), mode="bilinear" + ) + return get_LL_HH(up_sample) + + +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 mean_variance_norm(feat): + size = feat.size() + mean, std = calc_mean_std(feat) + normalized_feat = (feat - mean.expand(size)) / std.expand(size) + return normalized_feat + + +class Noise(nn.Module): + def __init__(self, use_noise, sigma=0.2): + super(Noise, self).__init__() + self.use_noise = use_noise + self.sigma = sigma + + def forward(self, x): + if self.use_noise: + return x + self.sigma * torch.autograd.Variable( + torch.cuda.FloatTensor(x.size()).normal_(), requires_grad=False + ) + else: + return x + + +class AdaIN(nn.Module): + def __init__(self): + super(AdaIN, self).__init__() + + def forward(self, content, style, style_strength=1.0, eps=1e-5): + b, c, h, w = content.size() + + content_std, content_mean = torch.std_mean( + content.view(b, c, -1), dim=2, keepdim=True + ) + style_std, style_mean = torch.std_mean( + style.view(b, c, -1), dim=2, keepdim=True + ) + + normalized_content = (content.view(b, c, -1) - content_mean) / ( + content_std + eps + ) + + stylized_content = (normalized_content * style_std) + style_mean + + output = ( + 1 - style_strength + ) * content + style_strength * stylized_content.view(b, c, h, w) + return output + + +class StylizedSANet(nn.Module): + def __init__(self, in_planes): + super(StylizedSANet, self).__init__() + self.f = nn.Conv2d(in_planes, in_planes, (1, 1)) + self.g = nn.Conv2d(in_planes, in_planes, (1, 1)) + self.h = nn.Conv2d(in_planes, in_planes, (1, 1)) + self.sm = nn.Softmax(dim=-1) + self.out_conv = nn.Conv2d(in_planes, in_planes, (1, 1)) + self.style_transform = AdaIN() + + def forward(self, content, style): + F = self.f(mean_variance_norm(content)) + G = self.g(mean_variance_norm(style)) + H = self.h(style) + b, c, h, w = F.size() + F = F.view(b, -1, w * h).permute(0, 2, 1) + b, c, h, w = G.size() + G = G.view(b, -1, w * h) + S = torch.bmm(F, G) + S = self.sm(S) + b, c, h, w = H.size() + H = H.view(b, -1, w * h) + O = torch.bmm(H, S.permute(0, 2, 1)) + b, c, h, w = content.size() + O = O.view(b, c, h, w) + O = self.out_conv(O) + # stylized_content = self.style_transform(content, style) + # O += stylized_content + O += content + return O + # 아니면 여기 adaptive하게 설정하는데 일단은 0.5 0.5 로 해볼까? + + +class Transform(nn.Module): + def __init__(self, in_planes): + super(Transform, self).__init__() + self.sanet4_1 = StylizedSANet(in_planes=in_planes) + self.sanet5_1 = StylizedSANet(in_planes=in_planes) + self.merge_conv_pad = nn.ReflectionPad2d((1, 1, 1, 1)) + self.merge_conv = nn.Conv2d(in_planes, in_planes, (3, 3)) + + def forward(self, content4_1, style4_1, content5_1, style5_1): + # return self.merge_conv(self.merge_conv_pad(self.sanet4_1(content4_1, style4_1) + self.upsample5_1(self.sanet5_1(content5_1, style5_1)))) + return self.merge_conv( + self.merge_conv_pad( + self.sanet4_1(content4_1, style4_1) + + nn.functional.interpolate( + self.sanet5_1(content5_1, style5_1), + size=( + self.sanet4_1(content4_1, style4_1).size(2), + self.sanet4_1(content4_1, style4_1).size(3), + ), + ) + ) + ) + + +class AdaAttN(nn.Module): + def __init__(self, in_planes, max_sample=256 * 256, key_planes=None): + super(AdaAttN, self).__init__() + if key_planes is None: + key_planes = in_planes + self.f = nn.Conv2d(key_planes, key_planes, (1, 1)) + self.g = nn.Conv2d(key_planes, key_planes, (1, 1)) + self.h = nn.Conv2d(in_planes, in_planes, (1, 1)) + self.sm = nn.Softmax(dim=-1) + self.max_sample = max_sample + + def forward(self, content, style, content_key, style_key, seed=None): + F = self.f(content_key) + G = self.g(style_key) + H = self.h(style) + b, _, h_g, w_g = G.size() + G = G.view(b, -1, w_g * h_g).contiguous() + if w_g * h_g > self.max_sample: + if seed is not None: + torch.manual_seed(seed) + index = torch.randperm(w_g * h_g).to(content.device)[: self.max_sample] + G = G[:, :, index] + style_flat = ( + H.view(b, -1, w_g * h_g)[:, :, index].transpose(1, 2).contiguous() + ) + else: + style_flat = H.view(b, -1, w_g * h_g).transpose(1, 2).contiguous() + b, _, h, w = F.size() + F = F.view(b, -1, w * h).permute(0, 2, 1) + S = torch.bmm(F, G) + # S: b, n_c, n_s + S = self.sm(S) + # mean: b, n_c, c + mean = torch.bmm(S, style_flat) + # std: b, n_c, c + std = torch.sqrt(torch.relu(torch.bmm(S, style_flat**2) - mean**2)) + # mean, std: b, c, h, w + mean = mean.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous() + std = std.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous() + return std * mean_variance_norm(content) + mean, S + + +class AdaAttn_Transformer(nn.Module): + def __init__(self, in_planes, key_planes=None, shallow_layer=False): + super(AdaAttn_Transformer, self).__init__() + self.attn_adain_4_1 = AdaAttN(in_planes=in_planes, key_planes=key_planes) + self.attn_adain_5_1 = AdaAttN( + in_planes=in_planes, + key_planes=key_planes + 512 if shallow_layer else key_planes, + ) + self.upsample5_1 = nn.Upsample(scale_factor=2, mode="nearest") + self.merge_conv_pad = nn.ReflectionPad2d((1, 1, 1, 1)) + self.merge_conv = nn.Conv2d(in_planes, in_planes, (3, 3)) + + def forward( + self, + content4_1, + style4_1, + content5_1, + style5_1, + content4_1_key, + style4_1_key, + content5_1_key, + style5_1_key, + seed=None, + ): + feature_4_1, attn_4_1 = self.attn_adain_4_1( + content4_1, style4_1, content4_1_key, style4_1_key, seed=seed + ) + feature_5_1, attn_5_1 = self.attn_adain_5_1( + content5_1, style5_1, content5_1_key, style5_1_key, seed=seed + ) + # stylized_results = self.merge_conv(self.merge_conv_pad(feature_4_1 + self.upsample5_1(feature_5_1))) + stylized_results = self.merge_conv( + self.merge_conv_pad( + feature_4_1 + + nn.functional.interpolate( + feature_5_1, size=(feature_4_1.size(2), feature_4_1.size(3)) + ) + ) + ) + return stylized_results, feature_4_1, feature_5_1 + + +class AdaptiveMultiAdaAttN(nn.Module): + def __init__( + self, in_planes, max_sample=256 * 256, query_planes=None, key_planes=None + ): + super(AdaptiveMultiAdaAttN, self).__init__() + if key_planes is None: + key_planes = in_planes + self.f = nn.Conv2d(query_planes, key_planes, (1, 1)) + self.g = nn.Conv2d(key_planes, key_planes, (1, 1)) + self.h = nn.Conv2d(in_planes, in_planes, (1, 1)) + self.sm = nn.Softmax(dim=-1) + self.out_conv = nn.Conv2d(in_planes, in_planes, (1, 1)) + self.max_sample = max_sample + + def forward(self, content, style, content_key, style_key, seed=None): + F = self.f(content_key) + G = self.g(style_key) + H = self.h(style) + b, _, h_g, w_g = G.size() + G = G.view(b, -1, w_g * h_g).contiguous() + if w_g * h_g > self.max_sample: + if seed is not None: + torch.manual_seed(seed) + index = torch.randperm(w_g * h_g).to(content.device)[: self.max_sample] + G = G[:, :, index] + style_flat = ( + H.view(b, -1, w_g * h_g)[:, :, index].transpose(1, 2).contiguous() + ) + else: + style_flat = H.view(b, -1, w_g * h_g).transpose(1, 2).contiguous() + b, _, h, w = F.size() + F = F.view(b, -1, w * h).permute(0, 2, 1) + S = torch.bmm(F, G) + # S: b, n_c, n_s + S = self.sm(S) + # mean: b, n_c, c + mean = torch.bmm(S, style_flat) + # std: b, n_c, c + std = torch.sqrt(torch.relu(torch.bmm(S, style_flat**2) - mean**2)) + # mean, std: b, c, h, w + mean = mean.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous() + std = std.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous() + return std * mean_variance_norm(content) + mean, S + + +class AdaptiveMultiAttn_Transformer(nn.Module): + def __init__( + self, in_planes, query_planes=None, key_planes=None, shallow_layer=False + ): + super(AdaptiveMultiAttn_Transformer, self).__init__() + self.attn_adain_4_1 = AdaptiveMultiAdaAttN( + in_planes=in_planes, query_planes=query_planes, key_planes=key_planes + ) + self.attn_adain_5_1 = AdaptiveMultiAdaAttN( + in_planes=in_planes, query_planes=query_planes, key_planes=key_planes + 512 + ) + self.upsample5_1 = nn.Upsample(scale_factor=2, mode="nearest") + self.merge_conv_pad = nn.ReflectionPad2d((1, 1, 1, 1)) + self.merge_conv = nn.Conv2d(in_planes, in_planes, (3, 3)) + + def forward( + self, + content4_1, + style4_1, + content5_1, + style5_1, + content4_1_key, + style4_1_key, + content5_1_key, + style5_1_key, + seed=None, + ): + feature_4_1, attn_4_1 = self.attn_adain_4_1( + content4_1, style4_1, content4_1_key, style4_1_key, seed=seed + ) + feature_5_1, attn_5_1 = self.attn_adain_5_1( + content5_1, style5_1, content5_1_key, style5_1_key, seed=seed + ) + # stylized_results = self.merge_conv(self.merge_conv_pad(feature_4_1 + self.upsample5_1(feature_5_1))) + stylized_results = self.merge_conv( + self.merge_conv_pad( + feature_4_1 + + nn.functional.interpolate( + feature_5_1, size=(feature_4_1.size(2), feature_4_1.size(3)) + ) + ) + ) + return stylized_results, feature_4_1, feature_5_1, attn_4_1, attn_5_1 + + +class AdaptiveMultiAdaAttN_v2(nn.Module): + def __init__( + self, + in_planes, + out_planes, + max_sample=256 * 256, + query_planes=None, + key_planes=None, + ): + super(AdaptiveMultiAdaAttN_v2, self).__init__() + if key_planes is None: + key_planes = in_planes + self.f = nn.Conv2d(query_planes, key_planes, (1, 1)) + self.g = nn.Conv2d(key_planes, key_planes, (1, 1)) + self.h = nn.Conv2d(in_planes, out_planes, (1, 1)) + self.sm = nn.Softmax(dim=-1) + self.out_conv = nn.Conv2d(in_planes, out_planes, (1, 1)) + self.max_sample = max_sample + + def forward(self, content, style, content_key, style_key, seed=None): + F = self.f(content_key) + G = self.g(style_key) + H = self.h(style) + b, _, h_g, w_g = G.size() + G = G.view(b, -1, w_g * h_g).contiguous() + b, _, h, w = F.size() + F = F.view(b, -1, w * h).permute(0, 2, 1) + S = torch.bmm(F, G) + # S: b, n_c, n_s + S = self.sm(S) + + b, style_c, style_h, style_w = H.size() + H = torch.nn.functional.interpolate(H, (h_g, w_g), mode="bicubic") + if h_g * w_g > self.max_sample: + if seed is not None: + torch.manual_seed(seed) + index = torch.randperm(h_g * w_g).to(content.device)[: self.max_sample] + G = G[:, :, index] + style_flat = ( + H.view(b, -1, h_g * w_g)[:, :, index].transpose(1, 2).contiguous() + ) + else: + style_flat = H.view(b, -1, h_g * w_g).transpose(1, 2).contiguous() + + # mean: b, n_c, c + mean = torch.bmm(S, style_flat) + # std: b, n_c, c + std = torch.sqrt(torch.relu(torch.bmm(S, style_flat**2) - mean**2)) + # mean, std: b, c, h, w + _, _, ch, cw = content.size() + # mean = torch.nn.functional.interpolate(mean.view(b, style_h, style_w, style_c).permute(0, 3, 1, 2).contiguous(), (ch, cw)) + # std = torch.nn.functional.interpolate(std.view(b, style_h, style_w, style_c).permute(0, 3, 1, 2).contiguous(), (ch, cw)) + mean = mean.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous() + std = std.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous() + return std * mean_variance_norm(content) + mean, S + + +class AdaptiveMultiAttn_Transformer_v2(nn.Module): + def __init__( + self, + in_planes, + out_planes, + query_planes=None, + key_planes=None, + shallow_layer=False, + ): + super(AdaptiveMultiAttn_Transformer_v2, self).__init__() + self.attn_adain_4_1 = AdaptiveMultiAdaAttN_v2( + in_planes=in_planes, + out_planes=out_planes, + query_planes=query_planes, + key_planes=key_planes, + ) + self.attn_adain_5_1 = AdaptiveMultiAdaAttN_v2( + in_planes=in_planes, + out_planes=out_planes, + query_planes=query_planes, + key_planes=key_planes + 512, + ) + # self.attn_adain_5_1 = AdaptiveMultiAdaAttN_v2(in_planes=in_planes, out_planes=out_planes, query_planes=query_planes+512, key_planes=key_planes+512) + self.upsample5_1 = nn.Upsample(scale_factor=2, mode="nearest") + self.merge_conv_pad = nn.ReflectionPad2d((1, 1, 1, 1)) + self.merge_conv = nn.Conv2d(out_planes, out_planes, (3, 3)) + + def forward( + self, + content4_1, + style4_1, + content5_1, + style5_1, + content4_1_key, + style4_1_key, + content5_1_key, + style5_1_key, + seed=None, + ): + feature_4_1, attn_4_1 = self.attn_adain_4_1( + content4_1, style4_1, content4_1_key, style4_1_key, seed=seed + ) + feature_5_1, attn_5_1 = self.attn_adain_5_1( + content5_1, style5_1, content5_1_key, style5_1_key, seed=seed + ) + # stylized_results = self.merge_conv(self.merge_conv_pad(feature_4_1 + self.upsample5_1(feature_5_1))) + + stylized_results = self.merge_conv( + self.merge_conv_pad( + feature_4_1 + + nn.functional.interpolate( + feature_5_1, size=(feature_4_1.size(2), feature_4_1.size(3)) + ) + ) + ) + return stylized_results, feature_4_1, feature_5_1, attn_4_1, attn_5_1 + + +class VGGEncoder(nn.Module): + def __init__(self, vgg_state_dict): + super(VGGEncoder, self).__init__() + + self.pad = nn.ReflectionPad2d(1) + self.relu = nn.ReLU(inplace=True) + self.pool = nn.AvgPool2d(2) + self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices = False) + self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices = False) + self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices = False) + self.maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices = False) + + ###Level0### + self.conv0 = nn.Conv2d(3, 3, 1, 1, 0) + self.conv0.weight = nn.Parameter(vgg_state_dict['0.weight']) + self.conv0.bias = nn.Parameter(vgg_state_dict['0.bias']) + + ###Level1### + self.conv1_1 = nn.Conv2d(3, 64, 3, 1, 0) + self.conv1_1.weight = nn.Parameter(vgg_state_dict['2.weight']) + self.conv1_1.bias = nn.Parameter(vgg_state_dict['2.bias']) + + self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 0) + self.conv1_2.weight = nn.Parameter(vgg_state_dict['5.weight']) + self.conv1_2.bias = nn.Parameter(vgg_state_dict['5.bias']) + + ###Level2### + self.conv2_1 = nn.Conv2d(64, 128, 3, 1, 0) + self.conv2_1.weight = nn.Parameter(vgg_state_dict['9.weight']) + self.conv2_1.bias = nn.Parameter(vgg_state_dict['9.bias']) + + self.conv2_2 = nn.Conv2d(128, 128, 3, 1, 0) + self.conv2_2.weight = nn.Parameter(vgg_state_dict['12.weight']) + self.conv2_2.bias = nn.Parameter(vgg_state_dict['12.bias']) + + ###Level3### + self.conv3_1 = nn.Conv2d(128, 256, 3, 1, 0) + self.conv3_1.weight = nn.Parameter(vgg_state_dict['16.weight']) + self.conv3_1.bias = nn.Parameter(vgg_state_dict['16.bias']) + + self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 0) + self.conv3_2.weight = nn.Parameter(vgg_state_dict['19.weight']) + self.conv3_2.bias = nn.Parameter(vgg_state_dict['19.bias']) + + self.conv3_3 = nn.Conv2d(256, 256, 3, 1, 0) + self.conv3_3.weight = nn.Parameter(vgg_state_dict['22.weight']) + self.conv3_3.bias = nn.Parameter(vgg_state_dict['22.bias']) + + self.conv3_4 = nn.Conv2d(256, 256, 3, 1, 0) + self.conv3_4.weight = nn.Parameter(vgg_state_dict['25.weight']) + self.conv3_4.bias = nn.Parameter(vgg_state_dict['25.bias']) + + + ###Level4### + self.conv4_1 = nn.Conv2d(256, 512, 3, 1, 0) + self.conv4_1.weight = nn.Parameter(vgg_state_dict['29.weight']) + self.conv4_1.bias = nn.Parameter(vgg_state_dict['29.bias']) + + self.conv4_2 = nn.Conv2d(512, 512, 3, 1, 0) + self.conv4_2.weight = nn.Parameter(vgg_state_dict['32.weight']) + self.conv4_2.bias = nn.Parameter(vgg_state_dict['32.bias']) + + self.conv4_3 = nn.Conv2d(512, 512, 3, 1, 0) + self.conv4_3.weight = nn.Parameter(vgg_state_dict['35.weight']) + self.conv4_3.bias = nn.Parameter(vgg_state_dict['35.bias']) + + self.conv4_4 = nn.Conv2d(512, 512, 3, 1, 0) + self.conv4_4.weight = nn.Parameter(vgg_state_dict['38.weight']) + self.conv4_4.bias = nn.Parameter(vgg_state_dict['38.bias']) + + ###Level5### + self.conv5_1 = nn.Conv2d(512, 512, 3, 1, 0) + self.conv5_1.weight = nn.Parameter(vgg_state_dict['42.weight']) + self.conv5_1.bias = nn.Parameter(vgg_state_dict['42.bias']) + + def forward(self, x): + skips = {} + for level in [1, 2, 3, 4]: + x = self.encode(x, skips, level) + return x + + def encode(self, x, skips): + is_maxpool=False + + out = self.conv0(x) + out = self.relu(self.conv1_1(self.pad(out))) + skips['conv1_1'] = out + + out = self.relu(self.conv1_2(self.pad(out))) + skips['conv1_2'] = out + resize_w, resize_h = out.size(2), out.size(3) + pooled_feature = self.pool(out) + + #HH = out - F.interpolate(pooled_feature, scale_factor=2, mode='nearest') + HH = out - F.interpolate(pooled_feature, size=[resize_w, resize_h], mode='nearest') + skips['pool1'] = HH + ################################## + if is_maxpool: + pooled_feature = self.maxpool1(out) + + out = self.relu(self.conv2_1(self.pad(pooled_feature))) + skips['conv2_1'] = out + + out = self.relu(self.conv2_2(self.pad(out))) + skips['conv2_2'] = out + resize_w, resize_h = out.size(2), out.size(3) + pooled_feature = self.pool(out) + + #HH = out - F.interpolate(pooled_feature, scale_factor=2, mode='nearest') + HH = out - F.interpolate(pooled_feature, size=[resize_w, resize_h], mode='nearest') + skips['pool2'] = HH + ################################## + if is_maxpool: + pooled_feature = self.maxpool2(out) + + out = self.relu(self.conv3_1(self.pad(pooled_feature))) + skips['conv3_1'] = out + + out = self.relu(self.conv3_2(self.pad(out))) + out = self.relu(self.conv3_3(self.pad(out))) + out = self.relu(self.conv3_4(self.pad(out))) + skips['conv3_4'] = out + resize_w, resize_h = out.size(2), out.size(3) + pooled_feature = self.pool(out) + #HH = out - F.interpolate(pooled_feature, scale_factor=2, mode='nearest') + HH = out - F.interpolate(pooled_feature, size=[resize_w, resize_h], mode='nearest') + skips['pool3'] = HH + + ################################## + if is_maxpool: + pooled_feature = self.maxpool3(out) + + out = self.relu(self.conv4_1(self.pad(pooled_feature))) + skips['conv4_1'] = out + + out = self.relu(self.conv4_2(self.pad(out))) + out = self.relu(self.conv4_3(self.pad(out))) + out = self.relu(self.conv4_4(self.pad(out))) + skips['conv4_4'] = out + resize_w, resize_h = out.size(2), out.size(3) + pooled_feature = self.pool(out) + HH = out - F.interpolate(pooled_feature, size=[resize_w, resize_h], mode='nearest') + skips['pool4'] = HH + + #################################### + if is_maxpool: + pooled_feature = self.maxpool4(out) + out = self.relu(self.conv5_1(self.pad(pooled_feature))) + skips['conv5_1'] = out + + return out + + def get_features(self, x, level): + is_maxpool = False + + out = self.conv0(x) + out = self.relu(self.conv1_1(self.pad(out))) + if level == 1: + return out + + out = self.relu(self.conv1_2(self.pad(out))) + pooled_feature = self.pool(out) + ################################## + if is_maxpool: + pooled_feature = self.maxpool1(out) + out = self.relu(self.conv2_1(self.pad(pooled_feature))) + if level == 2: + return out + + out = self.relu(self.conv2_2(self.pad(out))) + pooled_feature = self.pool(out) + ################################## + if is_maxpool: + pooled_feature = self.maxpool2(out) + out = self.relu(self.conv3_1(self.pad(pooled_feature))) + if level == 3: + return out + + out = self.relu(self.conv3_2(self.pad(out))) + out = self.relu(self.conv3_3(self.pad(out))) + out = self.relu(self.conv3_4(self.pad(out))) + pooled_feature = self.pool(out) + ################################## + if is_maxpool: + pooled_feature = self.maxpool3(out) + out = self.relu(self.conv4_1(self.pad(pooled_feature))) + if level == 4: + return out + + out = self.relu(self.conv4_2(self.pad(out))) + out = self.relu(self.conv4_3(self.pad(out))) + out = self.relu(self.conv4_4(self.pad(out))) + pooled_feature = self.pool(out) + #################################### + if is_maxpool: + pooled_feature = self.maxpool4(out) + out = self.relu(self.conv5_1(self.pad(pooled_feature))) + if level == 5: + return out + + +class VGGDecoder(nn.Module): + def __init__(self): + super(VGGDecoder, self).__init__() + + self.pad = nn.ReflectionPad2d(1) + self.relu = nn.ReLU(inplace=False) + self.adain = AdaIN() + + # self.conv5_1 = nn.Conv2d(512, 512, 3, 1, 0) + # self.conv4_4 = nn.Conv2d(512, 512, 3, 1, 0) + # self.conv4_3 = nn.Conv2d(512, 512, 3, 1, 0) + # self.conv4_2 = nn.Conv2d(512, 512, 3, 1, 0) + self.conv4_1 = nn.Conv2d(512, 256, 3, 1, 0) + self.conv3_4 = nn.Conv2d(256, 256, 3, 1, 0) + self.conv3_3 = nn.Conv2d(256, 256, 3, 1, 0) + self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 0) + self.conv3_1 = nn.Conv2d(256, 128, 3, 1, 0) + self.conv2_2 = nn.Conv2d(128, 128, 3, 1, 0) + self.conv2_1 = nn.Conv2d(128, 64, 3, 1, 0) + self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 0) + self.conv1_1 = nn.Conv2d(64, 3, 3, 1, 0) + + init_weights(self) + + def forward(self, x, skips): + x = self.decode(x, skips) + + def decode(self, stylized_feat, content_skips, style_skips): + out = self.relu(self.conv4_1(self.pad(stylized_feat))) + resize_w, resize_h = ( + content_skips["conv3_4"].size(2), + content_skips["conv3_4"].size(3), + ) + unpooled_feat = F.interpolate(out, size=[resize_w, resize_h], mode="nearest") + out = unpooled_feat + + ##여기에 skip_conneciton 넣어서 stylization boost 하기 (local pattern) + out = self.relu(self.conv3_4(self.pad(out))) + out = self.relu(self.conv3_3(self.pad(out))) + out = self.relu(self.conv3_2(self.pad(out))) + # out = feature_wct_simple(out, style_skips['conv3_1']) + + out = self.relu(self.conv3_1(self.pad(out))) + resize_w, resize_h = ( + content_skips["conv2_2"].size(2), + content_skips["conv2_2"].size(3), + ) + unpooled_feat = F.interpolate(out, size=[resize_w, resize_h], mode="nearest") + out = unpooled_feat + + out = self.relu(self.conv2_2(self.pad(out))) + # out = feature_wct_simple(out, style_skips['conv2_1']) + + out = self.relu(self.conv2_1(self.pad(out))) + resize_w, resize_h = ( + content_skips["conv1_2"].size(2), + content_skips["conv1_2"].size(3), + ) + unpooled_feat = F.interpolate(out, size=[resize_w, resize_h], mode="nearest") + out = unpooled_feat + + out = self.relu(self.conv1_2(self.pad(out))) + # out = feature_wct_simple(out, style_skips['conv1_1']) + out = self.conv1_1(self.pad(out)) + return out + + def reconstruct(self, x): + out = self.relu(self.conv4_1(self.pad(x))) + out = F.interpolate(out, size=[out.size(2), out.size(3)], mode="nearest") + out = self.relu(self.conv3_4(self.pad(out))) + out = self.relu(self.conv3_3(self.pad(out))) + out = self.relu(self.conv3_2(self.pad(out))) + out = self.relu(self.conv3_1(self.pad(out))) + out = F.interpolate(out, size=[out.size(2), out.size(3)], mode="nearest") + out = self.relu(self.conv2_2(self.pad(out))) + out = self.relu(self.conv2_1(self.pad(out))) + out = F.interpolate(out, size=[out.size(2), out.size(3)], mode="nearest") + out = self.relu(self.conv1_2(self.pad(out))) + out = self.conv1_1(self.pad(out)) + return out + + +class Baseline_net(nn.Module): + def __init__(self, pretrained_vgg=None): + super(Baseline_net, self).__init__() + + self.encoder = VGGEncoder(pretrained_vgg) + self.decoder = VGGDecoder() + self.adain = AdaIN() + + # for v2 transformer + query_channels = 512 # +256+128+64 + key_channels = 512 + 256 + 128 + 64 + + self.transformer = AdaptiveMultiAttn_Transformer_v2( + in_planes=512, + out_planes=512, + query_planes=query_channels, + key_planes=key_channels, + ) + + def get_keys(self, feat_skips, last_layer_idx, type="content", need_shallow=True): + results = [] + target_conv_layer = "conv" + str(last_layer_idx) + "_1" + _, _, h, w = feat_skips[target_conv_layer].shape + for i in range(1, last_layer_idx + 1): + target_conv_layer = "conv" + str(i) + "_1" + if i == last_layer_idx: + results.append(mean_variance_norm(feat_skips[target_conv_layer])) + else: + results.append( + mean_variance_norm( + nn.functional.interpolate(feat_skips[target_conv_layer], (h, w)) + ) + ) + + return torch.cat(results, dim=1) + + def adaptive_get_keys( + self, feat_skips, start_layer_idx, last_layer_idx, target_feat + ): + B, C, th, tw = target_feat.shape + results = [] + target_conv_layer = "conv" + str(last_layer_idx) + "_1" + _, _, h, w = feat_skips[target_conv_layer].shape + for i in range(start_layer_idx, last_layer_idx + 1): + target_conv_layer = "conv" + str(i) + "_1" + if i == last_layer_idx: + results.append(mean_variance_norm(feat_skips[target_conv_layer])) + else: + results.append( + mean_variance_norm( + nn.functional.interpolate(feat_skips[target_conv_layer], (h, w)) + ) + ) + + return nn.functional.interpolate(torch.cat(results, dim=1), (th, tw)) + + def encode(self, x, skips): + return self.encoder.encode(x, skips) + + def decode(self, stylized_feat, content_skips, style_skips): + return self.decoder.decode(stylized_feat, content_skips, style_skips) + + def forward( + self, content, style, adaptive_alpha, gray_content=None, gray_style=None + ): + content_ = size_arrange(content) + style_ = size_arrange(style) + + ####Input#### + content_feat, content_skips = content_, {} + style_feat, style_skips = style_, {} + + ####Encode#### + content_feat = self.encode(content_feat, content_skips) + style_feat = self.encode(style_feat, style_skips) + + if gray_content is not None: + gray_content_ = size_arrange(gray_content) + gray_content_feat, gray_content_skips = gray_content_, {} + gray_content_feat = self.encode(gray_content_feat, gray_content_skips) + if gray_style is not None: + gray_style_ = size_arrange(gray_style) + gray_style_feat, gray_style_skips = gray_style_, {} + gray_style_feat = self.encode(gray_style_feat, gray_style_skips) + + ####Style transfer#### + # transformed_feature = self.transformer(content_skips['conv4_1'], style_skips['conv4_1'], content_skips['conv5_1'], style_skips['conv5_1']) + + # Test_Multi_style_level_1 + ( + local_transformed_feature, + attn_style_4_1, + attn_style_5_1, + attn_map_4_1, + attn_map_5_1, + ) = self.transformer( + content_skips["conv4_1"], + style_skips["conv4_1"], + content_skips["conv5_1"], + style_skips["conv5_1"], + self.adaptive_get_keys( + content_skips, 4, 4, target_feat=content_skips["conv4_1"] + ), + self.adaptive_get_keys( + style_skips, 1, 4, target_feat=style_skips["conv4_1"] + ), + self.adaptive_get_keys( + content_skips, 5, 5, target_feat=content_skips["conv5_1"] + ), + self.adaptive_get_keys( + style_skips, 1, 5, target_feat=style_skips["conv5_1"] + ), + ) + + if gray_content is not None: + global_transformed_feat = feature_wct_simple( + gray_content_skips["conv4_1"], gray_style_skips["conv4_1"] + ) + # global_transformed_feat = self.adain(gray_content_skips['conv4_1'], gray_style_skips['conv4_1']) + # global_transformed_feat = self.decorator(gray_content_skips['conv4_1'], gray_style_skips['conv4_1']) + else: + global_transformed_feat = feature_wct_simple( + content_skips["conv4_1"], style_skips["conv4_1"] + ) + # global_transformed_feat = self.adain(content_skips['conv4_1'], style_skips['conv4_1']) + # global_transformed_feat = self.decorator(content_skips['conv4_1'], style_skips['conv4_1']) + + transformed_feature = ( + global_transformed_feat * (1 - adaptive_alpha.unsqueeze(-1).unsqueeze(-1)) + + adaptive_alpha.unsqueeze(-1).unsqueeze(-1) * local_transformed_feature + ) + # transformed_feature = global_transformed_feat + local_transformed_feature + + ## Ablation study ## + # transformed_feature = local_transformed_feature + + ####Decode#### + stylized_image = self.decode(transformed_feature, content_skips, style_skips) + + return ( + stylized_image, + attn_style_4_1, + attn_style_5_1, + attn_map_4_1, + attn_map_5_1, + ) + + +################################# +#####MultiScaleDiscriminator##### +################################# +class MultiScaleImageDiscriminator(nn.Module): + def __init__(self, nc=3, ndf=64): + super(MultiScaleImageDiscriminator, self).__init__() + self.nc = nc + self.output_dim = 1 + + self.conv1 = nn.Sequential( + # 256->128 + spectral_norm(nn.Conv2d(self.nc, ndf, 4, 2, 1, bias=False)), + nn.LeakyReLU(0.2, inplace=True), + ) + + self.skip_out1 = nn.Sequential( + nn.Conv2d(ndf, 32, 3, 1, 1, bias=False), + nn.LeakyReLU(0.2, inplace=True), + nn.AdaptiveAvgPool2d(8), + nn.Conv2d(32, self.output_dim, 3, 1, 1, bias=False), + ) + + self.conv2 = nn.Sequential( + # 128->64 + spectral_norm(nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False)), + # nn.BatchNorm2d(ndf*2), + nn.InstanceNorm2d(ndf * 2), + nn.LeakyReLU(0.2, inplace=True), + ) + + self.conv3 = nn.Sequential( + # 64->32 + spectral_norm(nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False)), + # nn.BatchNorm2d(ndf*4), + nn.InstanceNorm2d(ndf * 4), + nn.LeakyReLU(0.2, inplace=True), + ) + + self.skip_out2 = nn.Sequential( + nn.Conv2d(ndf * 4, 32, 3, 1, 1, bias=False), + nn.LeakyReLU(0.2, inplace=True), + nn.AdaptiveAvgPool2d(8), + nn.Conv2d(32, self.output_dim, 3, 1, 1, bias=False), + ) + + self.conv4 = nn.Sequential( + # 32->16 + spectral_norm(nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False)), + # nn.BatchNorm2d(ndf*8), + nn.InstanceNorm2d(ndf * 8), + nn.LeakyReLU(0.2, inplace=True), + ) + + self.conv5 = nn.Sequential( + # 16->8 + spectral_norm(nn.Conv2d(ndf * 8, self.output_dim, 4, 2, 1, bias=False)), + ) + + init_weights(self) + + def forward(self, input): + out = self.conv1(input) + skip_out1 = self.skip_out1(out) + out = self.conv3(self.conv2(out)) + skip_out2 = self.skip_out2(out) + out = self.conv5(self.conv4(out)) + out = ((out + skip_out1 + skip_out2) * 1 / 3).squeeze() + # out = ((out+skip_out1+skip_out2)).squeeze() + + return out diff --git a/module_aespa/utils.py b/module_aespa/utils.py new file mode 100644 index 0000000..dfdf9c2 --- /dev/null +++ b/module_aespa/utils.py @@ -0,0 +1,595 @@ +""" +from photo_wct.py of https://github.com/NVIDIA/FastPhotoStyle +Copyright (C) 2018 NVIDIA Corporation. +Licensed under the CC BY-NC-SA 4.0 +""" + +import datetime +import os + +import imageio +import numpy as np +import torch +import torch.nn as nn +from PIL import Image +from torchvision import transforms + + +def svd(feat, iden=False, device="cuda:0"): + size = feat.size() + mean = torch.mean(feat, 1) + mean = mean.unsqueeze(1).expand_as(feat) + _feat = feat.clone() + _feat -= mean + if size[1] > 1: + conv = torch.mm(_feat, _feat.t()).div(size[1] - 1) + else: + conv = torch.mm(_feat, _feat.t()) + if iden: + conv += torch.eye(size[0]).to(device) + # pdb.set_trace() + u, e, v = torch.svd(conv, some=False) + return u, e, v + + +def get_squeeze_feat(feat): + _feat = feat.squeeze(0) + size = _feat.size(0) + return _feat.view(size, -1).clone() + + +def get_rank(singular_values, dim, eps=0.00001): + r = dim + for i in range(dim - 1, -1, -1): + if singular_values[i] >= eps: + r = i + 1 + break + return r + + +def covsqrt_mean(feature, inverse=False, tolerance=1e-14): + # I referenced the default svd tolerance value in matlab. + + b, c, h, w = feature.size() + + mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True) + zeromean = feature.view(b, c, -1) - mean + cov = torch.bmm(zeromean, zeromean.transpose(1, 2)) + + eps_matrix_ = (torch.ones_like(cov) * 1e-8).cuda() + # evals, evects = torch.symeig(cov, eigenvectors=True) + evals, evects = torch.symeig(cov + eps_matrix_, eigenvectors=True) + + p = 0.5 + if inverse: + p *= -1 + + covsqrt = [] + for i in range(b): + k = 0 + for j in range(c): + if evals[i][j] > tolerance: + k = j + break + covsqrt.append( + torch.mm( + evects[i][:, k:], + torch.mm(evals[i][k:].pow(p).diag_embed(), evects[i][:, k:].t()), + ).unsqueeze(0) + ) + covsqrt = torch.cat(covsqrt, dim=0) + + return covsqrt, mean + + +def whitening(feature): + b, c, h, w = feature.size() + + inv_covsqrt, mean = covsqrt_mean(feature, inverse=True) + + normalized_feature = torch.matmul(inv_covsqrt, feature.view(b, c, -1) - mean) + + return normalized_feature.view(b, c, h, w) + + +def whitening_edit(feature): + b, c, h, w = feature.size() + cont_feat = get_squeeze_feat(feature) + # cont_min = cont_feat.min() + # cont_max = cont_feat.max() + cont_mean = torch.mean(cont_feat, 1).unsqueeze(1).expand_as(cont_feat) + cont_feat -= cont_mean + + _, c_e, c_v = svd(cont_feat, iden=True) + k_c = get_rank(c_e, cont_feat.size()[0]) + c_d = (c_e[0:k_c]).pow(-0.5) + step1 = torch.mm(c_v[:, 0:k_c], torch.diag(c_d)) + step2 = torch.mm(step1, (c_v[:, 0:k_c].t())) + normalized_feature = torch.mm(step2, cont_feat) + + # inv_covsqrt, mean = covsqrt_mean(feature, inverse=True) + + # normalized_feature = torch.matmul(inv_covsqrt, feature.view(b, c, -1)-mean) + + return normalized_feature.view(b, c, h, w) + + +def coloring(feature, target): + b, c, h, w = feature.size() + + covsqrt, mean = covsqrt_mean(target) + + colored_feature = torch.matmul(covsqrt, feature.view(b, c, -1)) + mean + + return colored_feature.view(b, c, h, w) + + +def SwitchWhiten2d(x): + N, C, H, W = x.size() + + in_data = x.view(N, C, -1) + + eye = in_data.data.new().resize_(C, C) + eye = torch.nn.init.eye_(eye).view(1, C, C).expand(N, C, C) + + # calculate other statistics + mean_in = in_data.mean(-1, keepdim=True) + x_in = in_data - mean_in + # (N x g) x C x C + cov_in = torch.bmm(x_in, torch.transpose(x_in, 1, 2)).div(H * W) + + mean = mean_in + cov = cov_in + 1e-5 * eye + + # perform whitening using Newton's iteration + Ng, c, _ = cov.size() + P = torch.eye(c).to(cov).expand(Ng, c, c) + + rTr = (cov * P).sum((1, 2), keepdim=True).reciprocal_() + cov_N = cov * rTr + for k in range(5): + P = torch.baddbmm(1.5, P, -0.5, torch.matrix_power(P, 3), cov_N) + + wm = P.mul_(rTr.sqrt()) + x_hat = torch.bmm(wm, in_data - mean) + + return x_hat, wm, mean + + +def wct_core(cont_feat, styl_feat, weight=1, registers=None, device="cuda:0"): + cont_feat = get_squeeze_feat(cont_feat) + cont_min = cont_feat.min() + cont_max = cont_feat.max() + cont_mean = torch.mean(cont_feat, 1).unsqueeze(1).expand_as(cont_feat) + cont_feat -= cont_mean + + if not registers: + _, c_e, c_v = svd(cont_feat, iden=True, device=device) + + styl_feat = get_squeeze_feat(styl_feat) + s_mean = torch.mean(styl_feat, 1) + _, s_e, s_v = svd(styl_feat, iden=True, device=device) + k_s = get_rank(s_e, styl_feat.size()[0]) + s_d = (s_e[0:k_s]).pow(0.5) + EDE = torch.mm( + torch.mm(s_v[:, 0:k_s], torch.diag(s_d) * weight), (s_v[:, 0:k_s].t()) + ) + + if registers is not None: + registers["EDE"] = EDE + registers["s_mean"] = s_mean + registers["c_v"] = c_v + registers["c_e"] = c_e + else: + EDE = registers["EDE"] + s_mean = registers["s_mean"] + _, c_e, c_v = svd(cont_feat, iden=True, device=device) + + k_c = get_rank(c_e, cont_feat.size()[0]) + c_d = (c_e[0:k_c]).pow(-0.5) + # TODO could be more fast + step1 = torch.mm(c_v[:, 0:k_c], torch.diag(c_d)) + step2 = torch.mm(step1, (c_v[:, 0:k_c].t())) + whiten_cF = torch.mm(step2, cont_feat) + + targetFeature = torch.mm(EDE, whiten_cF) + targetFeature = targetFeature + s_mean.unsqueeze(1).expand_as(targetFeature) + targetFeature.clamp_(cont_min, cont_max) + + return targetFeature + + +def Bw_wct_core(content_feat, style_feat, weight=1, registers=None, device="cpu"): + N, C, H, W = content_feat.size() + cont_min = content_feat.min().item() + cont_max = content_feat.max().item() + + whiten_cF, _, _ = SwitchWhiten2d(content_feat) + _, wm_s, s_mean = SwitchWhiten2d(style_feat) + + targetFeature = torch.bmm(torch.inverse(wm_s), whiten_cF) + targetFeature = targetFeature.view(N, C, H, W) + targetFeature = targetFeature + s_mean.unsqueeze(2).expand_as(targetFeature) + targetFeature.clamp_(cont_min, cont_max) + + return targetFeature + + +def wct_core_segment( + content_feat, + style_feat, + content_segment, + style_segment, + label_set, + label_indicator, + weight=1, + registers=None, + device="cpu", +): + def resize(feat, target): + size = (target.size(2), target.size(1)) + if len(feat.shape) == 2: + return np.asarray(Image.fromarray(feat).resize(size, Image.NEAREST)) + else: + return np.asarray( + Image.fromarray(feat, mode="RGB").resize(size, Image.NEAREST) + ) + + def get_index(feat, label): + mask = np.where(feat.reshape(feat.shape[0] * feat.shape[1]) == label) + if mask[0].size <= 0: + return None + return torch.LongTensor(mask[0]).cuda() + + squeeze_content_feat = content_feat.squeeze(0) + squeeze_style_feat = style_feat.squeeze(0) + + content_feat_view = squeeze_content_feat.view( + squeeze_content_feat.size(0), -1 + ).clone() + style_feat_view = squeeze_style_feat.view(squeeze_style_feat.size(0), -1).clone() + + resized_content_segment = resize(content_segment, squeeze_content_feat) + resized_style_segment = resize(style_segment, squeeze_style_feat) + + target_feature = content_feat_view.clone() + for label in label_set: + if not label_indicator[label]: + continue + content_index = get_index(resized_content_segment, label) + style_index = get_index(resized_style_segment, label) + if content_index is None or style_index is None: + continue + masked_content_feat = torch.index_select(content_feat_view, 1, content_index) + masked_style_feat = torch.index_select(style_feat_view, 1, style_index) + # _target_feature = Bw_wct_core(masked_content_feat, masked_style_feat, device=device) + _target_feature = wct_core( + masked_content_feat, masked_style_feat, weight, registers, device=device + ) + if torch.__version__ >= "0.4.0": + # XXX reported bug in the original repository + new_target_feature = torch.transpose(target_feature, 1, 0) + new_target_feature.index_copy_( + 0, content_index, torch.transpose(_target_feature, 1, 0) + ) + target_feature = torch.transpose(new_target_feature, 1, 0) + else: + target_feature.index_copy_(1, content_index, _target_feature) + return target_feature + + +def feature_wct( + content_feat, + style_feat, + content_segment=None, + style_segment=None, + label_set=None, + label_indicator=None, + weight=1, + registers=None, + alpha=1, + device="cuda:0", +): + if label_set is not None: + target_feature = wct_core_segment( + content_feat, + style_feat, + content_segment, + style_segment, + label_set, + label_indicator, + weight, + registers, + device=device, + ) + else: + target_feature = Bw_wct_core(content_feat, style_feat, device=device) + + target_feature = target_feature.view_as(content_feat) + target_feature = alpha * target_feature + (1 - alpha) * content_feat + return target_feature + + +def feature_wct_simple(content_feat, style_feat, alpha=1): + target_feature = Bw_wct_core(content_feat, style_feat) + + target_feature = target_feature.view_as(content_feat) + target_feature = alpha * target_feature + (1 - alpha) * content_feat + return target_feature + + + + + +def init_weights(net): + for m in net.modules(): + if isinstance(m, nn.Conv2d): + m.weight.data.normal_(0, 0.02) + if m.bias is not None: + m.bias.data.zero_() + elif isinstance(m, nn.ConvTranspose2d): + m.weight.data.normal_(0, 0.02) + if m.bias is not None: + m.bias.data.zero_() + elif isinstance(m, nn.Linear): + m.weight.data.normal_(0, 0.02) + if m.bias is not None: + m.bias.data.zero_() + + +class Timer: + def __init__(self, msg="Elapsed time: {}", verbose=True): + self.msg = msg + self.start_time = None + self.verbose = verbose + + def __enter__(self): + self.start_time = datetime.datetime.now() + + def __exit__(self, exc_type, exc_value, exc_tb): + if self.verbose: + print(self.msg.format(datetime.datetime.now() - self.start_time)) + + +def _normalizer(denormalize=False): + # set Mean and Std of RGB channels of IMAGENET to use pre-trained VGG net + MEAN = [0.485, 0.456, 0.406] + STD = [0.229, 0.224, 0.225] + + if denormalize: + MEAN = [-mean / std for mean, std in zip(MEAN, STD)] + STD = [1 / std for std in STD] + + return transforms.Normalize(mean=MEAN, std=STD) + + +def open_image(image_path, image_size=None): + normalize = _normalizer() + image = Image.open(image_path) + _transforms = [] + if image_size is not None: + image = transforms.Resize(image_size)(image) + # _transforms.append(transforms.Resize(image_size)) + w, h = image.size + _transforms.append(transforms.CenterCrop((h // 16 * 16, w // 16 * 16))) + _transforms.append(transforms.ToTensor()) + _transforms.append(normalize) + transform = transforms.Compose(_transforms) + return transform(image).unsqueeze(0) + + +def change_seg(seg): + color_dict = { + (0, 0, 255): 3, # blue + (0, 255, 0): 2, # green + (0, 0, 0): 0, # black + (255, 255, 255): 1, # white + (255, 0, 0): 4, # red + (255, 255, 0): 5, # yellow + (128, 128, 128): 6, # grey + (0, 255, 255): 7, # lightblue + (255, 0, 255): 8, # purple + } + arr_seg = np.asarray(seg) + new_seg = np.zeros(arr_seg.shape[:-1]) + for x in range(arr_seg.shape[0]): + for y in range(arr_seg.shape[1]): + if tuple(arr_seg[x, y, :]) in color_dict: + new_seg[x, y] = color_dict[tuple(arr_seg[x, y, :])] + else: + min_dist_index = 0 + min_dist = 99999 + for key in color_dict: + dist = np.sum(np.abs(np.asarray(key) - arr_seg[x, y, :])) + if dist < min_dist: + min_dist = dist + min_dist_index = color_dict[key] + elif dist == min_dist: + try: + min_dist_index = new_seg[x, y - 1, :] + except Exception: + pass + new_seg[x, y] = min_dist_index + return new_seg.astype(np.uint8) + + +def load_segment(image_path, image_size=None): + if not image_path: + return np.asarray([]) + image = Image.open(image_path) + if image_size is not None: + transform = transforms.Resize(image_size, interpolation=Image.NEAREST) + image = transform(image) + w, h = image.size + transform = transforms.CenterCrop((h // 16 * 16, w // 16 * 16)) + image = transform(image) + if len(np.asarray(image).shape) == 3: + image = change_seg(image) + return np.asarray(image) + + +def compute_label_info(content_segment, style_segment): + if not content_segment.size or not style_segment.size: + return None, None + max_label = np.max(content_segment) + 1 + label_set = np.unique(content_segment) + label_indicator = np.zeros(max_label) + for l in label_set: # noqa: E741 + content_mask = np.where( + content_segment.reshape(content_segment.shape[0] * content_segment.shape[1]) + == l + ) + style_mask = np.where( + style_segment.reshape(style_segment.shape[0] * style_segment.shape[1]) == l + ) + + c_size = content_mask[0].size + s_size = style_mask[0].size + if ( + c_size > 10 + and s_size > 10 + and c_size / s_size < 100 + and s_size / c_size < 100 + ): + label_indicator[l] = True + else: + label_indicator[l] = False + return label_set, label_indicator + + +def mkdir(dname): + if not os.path.exists(dname): + os.makedirs(dname) + else: + assert os.path.isdir(dname), "alread exists filename {}".format(dname) + + +def gram_matrix(y): + (b, ch, h, w) = y.size() + features = y.view(b, ch, w * h) + features_t = features.transpose(1, 2) + gram = features.bmm(features_t) / (ch * h * w) + return gram + + +def TVloss(img, tv_weight): + """ + Compute total variation loss. + Inputs: + - img: PyTorch Variable of shape (1, 3, H, W) holding an input image. + - tv_weight: Scalar giving the weight w_t to use for the TV loss. + Returns: + - loss: PyTorch Variable holding a scalar giving the total variation loss + for img weighted by tv_weight. + """ + # w_variance = torch.sum(torch.pow(img[:, :, :, :-1] - img[:, :, :, 1:], 2)) + # h_variance = torch.sum(torch.pow(img[:, :, :-1, :] - img[:, :, 1:, :], 2)) + # loss = tv_weight * (h_variance + w_variance) + w_variance = torch.mean(torch.abs(img[:, :, :, :-1] - img[:, :, :, 1:])) + h_variance = torch.mean(torch.abs(img[:, :, :-1, :] - img[:, :, 1:, :])) + loss = tv_weight * (h_variance + w_variance) + return loss + + +def zeros_like(x): + return torch.autograd.Variable(torch.zeros_like(x).cuda()) + + +def ones_like(x): + return torch.autograd.Variable(torch.ones_like(x).cuda()) + + +def denorm_2(x): + out = (x + 1) / 2 + return out.clamp_(0, 1) + + +def save_video(video, save_path, type="photo"): + video = denorm_2(video) + """ + vid_lst=[] + for i in range(0, num_samples*num_samples, num_samples): + temp_vid = list(video[i:i+num_samples]) + temp_vid = torch.cat(temp_vid, dim=-1) + vid_lst.append(temp_vid) + + save_videos = torch.cat(vid_lst, dim=2) + """ + + save_videos = video.data.cpu().numpy().transpose(0, 2, 3, 1) + outputdata = save_videos * 255 + # outputdata = ((save_videos+1)/2) * 255 + outputdata = outputdata.astype(np.uint8) + dir_path = save_path + if not os.path.exists(dir_path): + os.makedirs(dir_path) + if type == "photo": + gif_file_path = os.path.join(dir_path, "Photo_StylizedVideo.gif") + elif type == "art": + gif_file_path = os.path.join(dir_path, "Art_StylizedVideo.gif") + else: + gif_file_path = os.path.join(dir_path, "content_StylizedVideo.gif") + imageio.mimsave(gif_file_path, outputdata, fps=25) + + +class GuidedFilter(nn.Module): + def box_filter(self, x: torch.Tensor, r): + ch = x.shape[1] + k = 2 * r + 1 + weight = 1 / ((k) ** 2) # 1/9 + # [c,1,3,3] * 1/9 + box_kernel = torch.ones( + (ch, 1, k, k), dtype=torch.float32, device=x.device + ).fill_(weight) + # same padding + return torch.nn.functional.conv2d(x, box_kernel, padding=r, groups=ch) + + def forward(self, x: torch.Tensor, y: torch.Tensor, r, eps=1e-2): + b, c, h, w = x.shape + device = x.device + + N = self.box_filter(torch.ones((1, 1, h, w), dtype=x.dtype, device=device), r) + + mean_x = self.box_filter(x, r) / N + mean_y = self.box_filter(y, r) / N + cov_xy = self.box_filter(x * y, r) / N - mean_x * mean_y + var_x = self.box_filter(x * x, r) / N - mean_x * mean_x + + A = cov_xy / (var_x + eps) + b = mean_y - A * mean_x + + mean_A = self.box_filter(A, r) / N + mean_b = self.box_filter(b, r) / N + + output = mean_A * x + mean_b + return output + + +def attn_visualization_all(Ic, Is, Fc, Fs, Attn, slide_size=4): + if Ic.size(3) != Is.size(3): + Is = torch.nn.functional.interpolate(Is, size=(Ic.size(2), Ic.size(3))) + B, C, S = Attn.size() + B, c, c_w, c_h = Ic.size() + B, c, s_w, s_h = Is.size() + _, _, f_w, f_h = Fs.size() + attn = torch.zeros(1, Fs.size(2) * Fs.size(3)).cuda() + mask = torch.zeros(1, Fc.size(2) * Fc.size(3)).cuda() + for index in range(0, Fc.size(2) * Fc.size(3), slide_size): + for idx in range(slide_size): + start_idx = index + Fc.size(3) * idx + # start_idx_ref = index + Fs.size(3) * idx + mask[0][start_idx : start_idx + slide_size] = 1 + attn += torch.sum(Attn[0][start_idx : start_idx + slide_size], dim=0) + masked = torch.nn.functional.interpolate( + mask.view(1, 1, Fc.size(2), Fc.size(3)), size=(c_w, c_h), mode="nearest" + ) + attn = torch.nn.functional.interpolate( + attn.view(1, 1, Fs.size(2), Fs.size(3)), size=(s_w, s_h), mode="nearest" + ) + # imsave_no_norm(Is, attn.repeat(1,3,1,1),'style.png') + return ( + attn.repeat(1, 3, 1, 1), + Ic * (masked + 0.2), + Is * (attn * (1 / attn.max()) + 0.2), + Is * (attn * 0.05 + 0.2), + ) diff --git a/module_unist/unist_model.py b/module_unist/unist_model.py index ae5f1e7..57a4c37 100644 --- a/module_unist/unist_model.py +++ b/module_unist/unist_model.py @@ -27,7 +27,7 @@ def inference_unist( return output -def test_transform(size, crop): +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: @@ -36,7 +36,7 @@ def test_transform(size, crop): elif size != 0: transform_list.append(transforms.Resize((size, size))) - transform_list.append(transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])) + transform_list.append(transforms.Normalize(mean, std)) transform = transforms.Compose(transform_list) # [B, C, H, W] diff --git a/pyproject.toml b/pyproject.toml index 7e867d0..f640677 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-styletransferplus" -description = "Nodes:Neural Neighbor, CAST, EFDM, MicroAST, Coral Color Transfer." -version = "1.0.0" +description = "Nodes: Neural Neighbor, CAST, EFDM, MicroAST, UniST, AesPA, ..." +version = "1.0.1" license = {file = "LICENSE"} dependencies = ["torch", "torchvision", "scikit-image", "einops"] diff --git a/run.py b/run.py index 94c482f..44443be 100644 --- a/run.py +++ b/run.py @@ -13,6 +13,8 @@ import tqdm from comfy.utils import ProgressBar from .constants import ( + AESPA_DEC_PATH, + AESPA_TRANSFORMER_PATH, CAST_DEFAULT, CAST_NET_AE_PATH, CAST_NET_DEC_B_PATH, @@ -28,14 +30,37 @@ from .constants import ( UNIST_DEC_PATH, UNIST_ENC_PATH, UNIST_PATH, + VGG_NORM_CONV51_PATH, ) + +# AesPA +from .module_aespa.aespa_model import inference_aespa +from .module_aespa.aespanet_models import ( + AdaptiveMultiAttn_Transformer_v2, +) +from .module_aespa.aespanet_models import ( + VGGDecoder as AESPA_DEC, +) +from .module_aespa.aespanet_models import ( + VGGEncoder as AESPA_VGG, +) + +# CAST from .module_cast import net as net_cast from .module_cast.cast_model import inference_ucast, load_a_ckpt + +# EFDM from .module_efdm import net as net_efdm from .module_efdm.efdm_model import inference_efdm + +# MicroAST from .module_microast import net_microAST from .module_microast.microast_model import inference_microast + +# Neural Neighbor from .module_neural_neighbor.neural_neighbor_model import inference_neural_neighbor + +# UniST from .module_unist.model import video_Style_transfer from .module_unist.unist_model import inference_unist @@ -505,6 +530,85 @@ class UniST_Video: return (final_result,) +class AesPA: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "src_img": ("IMAGE",), + "style_img": ("IMAGE",), + "do_crop": ( + "BOOLEAN", + {"default": False}, + ), + "size": ( + "INT", + {"default": 512, "min": 1, "max": 1024, "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") + + pretrained_vgg = torch.load(f"{base_dir}/{VGG_NORM_CONV51_PATH}") + encoder = AESPA_VGG(pretrained_vgg) + decoder = AESPA_DEC() + transformer = AdaptiveMultiAttn_Transformer_v2( + in_planes=512, + out_planes=512, + query_planes=512, + key_planes=512 + 256 + 128 + 64, + ) + + decoder.load_state_dict( + torch.load(f"{base_dir}/{AESPA_DEC_PATH}")["state_dict"] + ) + transformer.load_state_dict( + torch.load(f"{base_dir}/{AESPA_TRANSFORMER_PATH}")["state_dict"] + ) + + encoder.eval().to(device) + decoder.eval().to(device) + transformer.eval().to(device) + + params = { + "style_img": style_img.permute(0, 3, 1, 2), + "device": device, + "size": size, + "do_crop": do_crop, + "encoder": encoder, + "decoder": decoder, + "transformer": transformer, + } + + num_frames = src_img.size(0) + pbar = ProgressBar(num_frames) + + result: list[torch.Tensor] = [] + with torch.no_grad(): + for i in tqdm.tqdm(range(num_frames), desc="Generating"): + params["src_img"] = src_img[i].unsqueeze(0).permute(0, 3, 1, 2) + res_tensor = inference_aespa(**params) + result.append(res_tensor.permute(0, 2, 3, 1)) + pbar.update_absolute(i, num_frames) + + return (torch.cat(result, dim=0),) + + def serialize_floats(lst: list[float]): return ",".join(map(str, lst)) diff --git a/workflows/wf_aespa.png b/workflows/wf_aespa.png new file mode 100644 index 0000000..36e82fd Binary files /dev/null and b/workflows/wf_aespa.png differ diff --git a/workflows/workflow_aespa.json b/workflows/workflow_aespa.json new file mode 100644 index 0000000..a7d7dd6 --- /dev/null +++ b/workflows/workflow_aespa.json @@ -0,0 +1,193 @@ +{ + "last_node_id": 8, + "last_link_id": 19, + "nodes": [ + { + "id": 4, + "type": "SaveImage", + "pos": [ + 1050, + 327 + ], + "size": { + "0": 315, + "1": 270 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 19 + } + ], + "properties": {}, + "widgets_values": [ + "StyleTransfer" + ] + }, + { + "id": 2, + "type": "LoadImage", + "pos": [ + 309, + 187 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 17 + ], + "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": 3, + "type": "LoadImage", + "pos": [ + 315, + 570 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 18 + ], + "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": 8, + "type": "AesPA", + "pos": [ + 674, + 349 + ], + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "src_img", + "type": "IMAGE", + "link": 17 + }, + { + "name": "style_img", + "type": "IMAGE", + "link": 18 + } + ], + "outputs": [ + { + "name": "out_img", + "type": "IMAGE", + "links": [ + 19 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "AesPA" + }, + "widgets_values": [ + true, + 1024 + ] + } + ], + "links": [ + [ + 17, + 2, + 0, + 8, + 0, + "IMAGE" + ], + [ + 18, + 3, + 0, + 8, + 1, + "IMAGE" + ], + [ + 19, + 8, + 0, + 4, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.1000000000000016, + "offset": [ + -196.07907685850128, + -119.07837620860246 + ] + } + }, + "version": 0.4 +} \ No newline at end of file