add aespa node
This commit is contained in:
FuouM
@@ -1,6 +1,8 @@
|
||||
# Checkpoint files
|
||||
*.pth
|
||||
*.tar
|
||||
*.pt
|
||||
*.t7
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
|
||||
@@ -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 |
|
||||
| :-----------------------------------------------------: | :---------------------------------------: |
|
||||
|  |  |
|
||||
| **EFDM (2022) - Fast** | **MicroAST (2023) - Fast** |
|
||||
|  |  |
|
||||
| **UniST (2023) - Medium - Only squares** | **X** |
|
||||
|  | X |
|
||||
| Neural Neighbor (2020-2022) - Slowest | CAST (2022-2023) - Fast |
|
||||
| :-----------------------------------------------------: | :------------------------------------------: |
|
||||
|  |  |
|
||||
| **EFDM (2022) - Fast** | **MicroAST (2023) - Fast** |
|
||||
|  |  |
|
||||
| **UniST (2023) - Medium - Only squares** | **AesPA-Net (2023) - Medium - Only squares** |
|
||||
|  |  |
|
||||
|
||||
### 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) |
|
||||
|  |  |
|
||||
||(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)
|
||||
|
||||
|
||||
|
||||
## 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}
|
||||
}
|
||||
```
|
||||
+3
-1
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 762 KiB |
@@ -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
|
||||
@@ -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
|
||||
@@ -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),
|
||||
)
|
||||
@@ -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]
|
||||
|
||||
+2
-2
@@ -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"]
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 569 KiB |
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user