TSSAT and Clean up (#10)
This commit is contained in:
Fuou Marinas
@@ -9,7 +9,7 @@ Click name to jump to workflow
|
||||
4. [**MicroAST**](#microast). Paper: [Towards Super-Fast Ultra-Resolution Arbitrary Style Transfer](https://github.com/EndyWon/MicroAST)
|
||||
5. [**UniST**](#unist). Paper: [Two Birds, One Stone: A Unified Framework for Joint Learning of Image and Video Style Transfers](https://github.com/NevSNev/UniST)
|
||||
6. [**AesPA-Net**](#aespa-net). Paper: [Aesthetic Pattern-Aware Style Transfer Networks](https://github.com/Kibeom-Hong/AesPA-Net)
|
||||
|
||||
7. [**TSSAT**](#tssat). Paper: [Two-Stage Statistics-Aware Transformation for Artistic Style Transfer](https://github.com/HalbertCH/TSSAT)
|
||||
|
||||
[**Extra nodes**](#extra-nodes)
|
||||
|
||||
@@ -26,6 +26,8 @@ All nodes support batched input (i.e video) but is generally not recommended. On
|
||||
|  |  |
|
||||
| **UniST (2023) - Medium - Only squares** | **AesPA-Net (2023) - Medium - Only squares** |
|
||||
|  |  |
|
||||
| **TSSAT (2023) - Slow** | **X** |
|
||||
|  | X |
|
||||
|
||||
### Neural Neighbor
|
||||
|
||||
@@ -138,9 +140,9 @@ Download models from [Usage](https://github.com/Kibeom-Hong/AesPA-Net?tab=readme
|
||||
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
|
||||
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.
|
||||
@@ -149,6 +151,22 @@ The original authors provided a `vgg_normalised_conv5_1.t7` model, which can onl
|
||||
|
||||
|
||||
|
||||
|
||||
### TSSAT
|
||||
|
||||
Download the `TSSAT-model.zip` in the [Model Testing](https://github.com/HalbertCH/TSSAT?tab=readme-ov-file#model-testing) section, and unzip it
|
||||
|
||||
```
|
||||
models/tssat/decoder_iter_160000.pth
|
||||
|
||||
# Outside folder, can be ignored from the zip if already downloaded from above
|
||||
models/vgg_normalised.pth
|
||||
```
|
||||
|
||||
[workflow_tssat.json](workflows/workflow_tssat.json)
|
||||
|
||||
|
||||
|
||||
## Extra nodes
|
||||
|
||||
### Coral Color Transfer
|
||||
|
||||
+3
-1
@@ -1,4 +1,4 @@
|
||||
from .run import NeuralNeighbor, CAST, EFDM, MicroAST, UniST, UniST_Video, AesPA
|
||||
from .run import NeuralNeighbor, CAST, EFDM, MicroAST, UniST, UniST_Video, AesPA, TSSAT
|
||||
from .run_extra import CoralColorTransfer
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
@@ -10,6 +10,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"UniST": UniST,
|
||||
"UniST_Video": UniST_Video,
|
||||
"AesPA": AesPA,
|
||||
"TSSAT": TSSAT,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -21,6 +22,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UniST": "UniST",
|
||||
"UniST_Video": "UniST Video",
|
||||
"AesPA": "AesPA-Net",
|
||||
"TSSAT": "TSSAT",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -21,3 +21,5 @@ 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"
|
||||
|
||||
TSSAT_DEC_PATH = "models/tssat/decoder_iter_160000.pth"
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.6 MiB |
+106
-238
@@ -1,14 +1,116 @@
|
||||
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,
|
||||
)
|
||||
from .utils import adaptive_get_keys, adaptive_gram_weight, feature_wct_simple
|
||||
|
||||
|
||||
class MODEL_AESPA:
|
||||
def __init__(
|
||||
self, vgg_51_path: str, dec_path: str, transformer_path: str, device
|
||||
) -> None:
|
||||
pretrained_vgg = torch.load(vgg_51_path)
|
||||
self.encoder = VGGEncoder(pretrained_vgg)
|
||||
self.decoder = VGGDecoder()
|
||||
self.transformer = AdaptiveMultiAttn_Transformer_v2(
|
||||
in_planes=512,
|
||||
out_planes=512,
|
||||
query_planes=512,
|
||||
key_planes=512 + 256 + 128 + 64,
|
||||
)
|
||||
|
||||
self.decoder.load_state_dict(torch.load(dec_path)["state_dict"])
|
||||
self.transformer.load_state_dict(torch.load(transformer_path)["state_dict"])
|
||||
|
||||
self.encoder.eval().to(device)
|
||||
self.decoder.eval().to(device)
|
||||
self.transformer.eval().to(device)
|
||||
|
||||
self.device = device
|
||||
|
||||
def forward(self, content: torch.Tensor, style: torch.Tensor):
|
||||
gray_content = tf_f.rgb_to_grayscale(content).repeat(1, 3, 1, 1).to(self.device)
|
||||
gray_style = tf_f.rgb_to_grayscale(style).repeat(1, 3, 1, 1).to(self.device)
|
||||
|
||||
style_weight = (
|
||||
adaptive_gram_weight(style, 1, 8, self.encoder)
|
||||
+ adaptive_gram_weight(style, 2, 8, self.encoder)
|
||||
+ adaptive_gram_weight(style, 3, 8, self.encoder)
|
||||
) / 3
|
||||
|
||||
gray_style_weight = (
|
||||
adaptive_gram_weight(gray_style, 1, 8, self.encoder)
|
||||
+ adaptive_gram_weight(gray_style, 2, 8, self.encoder)
|
||||
+ adaptive_gram_weight(gray_style, 3, 8, self.encoder)
|
||||
) / 3
|
||||
|
||||
style_adaptive_alpha = (
|
||||
style_weight.unsqueeze(1).to(self.device)
|
||||
+ gray_style_weight.unsqueeze(1).to(self.device)
|
||||
) / 2
|
||||
|
||||
content_skips = {}
|
||||
style_skips = {}
|
||||
|
||||
_ = self.encoder.encode(content, content_skips) # content_feat
|
||||
_ = self.encoder.encode(style, style_skips) # style_feat
|
||||
|
||||
if gray_content is not None:
|
||||
gray_content_skips = {}
|
||||
_ = self.encoder.encode(
|
||||
gray_content, gray_content_skips
|
||||
) # gray_content_feat
|
||||
if gray_style is not None:
|
||||
gray_style_skips = {}
|
||||
_ = self.encoder.encode(gray_style, gray_style_skips) # gray_style_feat
|
||||
|
||||
(
|
||||
local_transformed_feature,
|
||||
_,
|
||||
_,
|
||||
_,
|
||||
_,
|
||||
) = self.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 = self.decoder.decode(
|
||||
transformed_feature, content_skips, style_skips
|
||||
)
|
||||
|
||||
return stylized_image
|
||||
|
||||
|
||||
def inference_aespa(
|
||||
@@ -17,9 +119,7 @@ def inference_aespa(
|
||||
device,
|
||||
size: int,
|
||||
do_crop: bool,
|
||||
encoder: VGGEncoder,
|
||||
decoder: VGGDecoder,
|
||||
transformer: AdaptiveMultiAttn_Transformer_v2,
|
||||
model: MODEL_AESPA,
|
||||
):
|
||||
content_tf = test_transform(
|
||||
size, do_crop, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
|
||||
@@ -27,245 +127,13 @@ def inference_aespa(
|
||||
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)
|
||||
stylized_image = model.forward(content, style)
|
||||
|
||||
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 = []
|
||||
|
||||
|
||||
+181
-565
@@ -3,76 +3,7 @@ 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
|
||||
from .utils import init_weights, mean_variance_norm
|
||||
|
||||
|
||||
class AdaIN(nn.Module):
|
||||
@@ -158,180 +89,6 @@ class Transform(nn.Module):
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
@@ -448,204 +205,211 @@ class AdaptiveMultiAttn_Transformer_v2(nn.Module):
|
||||
|
||||
|
||||
class VGGEncoder(nn.Module):
|
||||
def __init__(self, vgg_state_dict):
|
||||
super(VGGEncoder, self).__init__()
|
||||
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)
|
||||
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'])
|
||||
###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'])
|
||||
###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.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.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"])
|
||||
|
||||
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'])
|
||||
###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.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'])
|
||||
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"])
|
||||
|
||||
|
||||
###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'])
|
||||
###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.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.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.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.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.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'])
|
||||
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"])
|
||||
|
||||
###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'])
|
||||
###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"])
|
||||
|
||||
def forward(self, x):
|
||||
skips = {}
|
||||
for level in [1, 2, 3, 4]:
|
||||
x = self.encode(x, skips, level)
|
||||
return x
|
||||
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"])
|
||||
|
||||
def encode(self, x, skips):
|
||||
is_maxpool=False
|
||||
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"])
|
||||
|
||||
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)
|
||||
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"])
|
||||
|
||||
out = self.relu(self.conv2_1(self.pad(pooled_feature)))
|
||||
skips['conv2_1'] = out
|
||||
###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"])
|
||||
|
||||
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)
|
||||
def forward(self, x):
|
||||
skips = {}
|
||||
for level in [1, 2, 3, 4]:
|
||||
x = self.encode(x, skips, level)
|
||||
return x
|
||||
|
||||
out = self.relu(self.conv3_1(self.pad(pooled_feature)))
|
||||
skips['conv3_1'] = out
|
||||
def encode(self, x, skips):
|
||||
is_maxpool = False
|
||||
|
||||
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
|
||||
out = self.conv0(x)
|
||||
out = self.relu(self.conv1_1(self.pad(out)))
|
||||
skips["conv1_1"] = out
|
||||
|
||||
##################################
|
||||
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.conv1_2(self.pad(out)))
|
||||
skips["conv1_2"] = out
|
||||
resize_w, resize_h = out.size(2), out.size(3)
|
||||
pooled_feature = self.pool(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
|
||||
# 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)
|
||||
|
||||
return out
|
||||
out = self.relu(self.conv2_1(self.pad(pooled_feature)))
|
||||
skips["conv2_1"] = out
|
||||
|
||||
def get_features(self, x, level):
|
||||
is_maxpool = False
|
||||
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)
|
||||
|
||||
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
|
||||
# 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.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_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)))
|
||||
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
|
||||
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):
|
||||
@@ -730,154 +494,6 @@ class VGGDecoder(nn.Module):
|
||||
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#####
|
||||
#################################
|
||||
|
||||
+133
-537
@@ -1,127 +1,123 @@
|
||||
"""
|
||||
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 random
|
||||
|
||||
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
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
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)
|
||||
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_()
|
||||
|
||||
|
||||
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 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 adaptive_gram_weight(image, level, ratio, encoder):
|
||||
if level == 0:
|
||||
encoded_features = image
|
||||
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
|
||||
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 get_squeeze_feat(feat):
|
||||
_feat = feat.squeeze(0)
|
||||
size = _feat.size(0)
|
||||
return _feat.view(size, -1).clone()
|
||||
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)
|
||||
|
||||
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
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
return targetFeature
|
||||
|
||||
|
||||
def SwitchWhiten2d(x):
|
||||
@@ -148,7 +144,8 @@ def SwitchWhiten2d(x):
|
||||
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)
|
||||
# P = torch.baddbmm(1.5, P, -0.5, torch.matrix_power(P, 3), cov_N)
|
||||
P = torch.baddbmm(P, torch.matrix_power(P, 3), cov_N, beta=1.5, alpha=-0.5)
|
||||
|
||||
wm = P.mul_(rTr.sqrt())
|
||||
x_hat = torch.bmm(wm, in_data - mean)
|
||||
@@ -156,312 +153,22 @@ def SwitchWhiten2d(x):
|
||||
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 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 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 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 gram_matrix(y):
|
||||
@@ -472,124 +179,13 @@ def gram_matrix(y):
|
||||
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 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()
|
||||
|
||||
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),
|
||||
)
|
||||
# return patches.view(b, number_of_patches, c, h, w)
|
||||
return patches.view(b, -1, c, kernel, kernel)
|
||||
|
||||
@@ -1,10 +1,24 @@
|
||||
import torch
|
||||
|
||||
def load_a_ckpt(ckpt_path: str):
|
||||
state_dict = torch.load(ckpt_path)
|
||||
if hasattr(state_dict, "_metadata"):
|
||||
del state_dict._metadata
|
||||
return state_dict
|
||||
from . import net
|
||||
|
||||
|
||||
class MODEL_CAST:
|
||||
def __init__(
|
||||
self, vgg_path: str, net_ae_path: str, net_dec_b_path: str, device
|
||||
) -> None:
|
||||
vgg = net.vgg
|
||||
vgg.load_state_dict(torch.load(vgg_path))
|
||||
vgg = torch.nn.Sequential(*list(vgg.children())[:31])
|
||||
|
||||
self.netAE = net.ADAIN_Encoder(vgg)
|
||||
self.netDec_B = net.Decoder()
|
||||
|
||||
self.netAE.load_state_dict(load_a_ckpt(net_ae_path))
|
||||
self.netDec_B.load_state_dict(load_a_ckpt(net_dec_b_path))
|
||||
|
||||
self.netAE.to(device).eval()
|
||||
self.netDec_B.to(device).eval()
|
||||
|
||||
|
||||
def inference_ucast(
|
||||
@@ -25,3 +39,10 @@ def inference_ucast(
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return fake_B
|
||||
|
||||
|
||||
def load_a_ckpt(ckpt_path: str):
|
||||
state_dict = torch.load(ckpt_path)
|
||||
if hasattr(state_dict, "_metadata"):
|
||||
del state_dict._metadata
|
||||
return state_dict
|
||||
|
||||
@@ -2,6 +2,7 @@ import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision import transforms
|
||||
|
||||
from . import net
|
||||
from .function import (
|
||||
adaptive_instance_normalization,
|
||||
adaptive_mean_normalization,
|
||||
@@ -12,6 +13,20 @@ from .function import (
|
||||
)
|
||||
|
||||
|
||||
class MODEL_EFDM:
|
||||
def __init__(self, dec_path: str, vgg_path: str, device) -> None:
|
||||
self.decoder = net.decoder
|
||||
vgg = net.vgg
|
||||
|
||||
self.decoder.load_state_dict(torch.load(dec_path))
|
||||
vgg.load_state_dict(torch.load(vgg_path))
|
||||
|
||||
self.vgg = torch.nn.Sequential(*list(vgg.children())[:31])
|
||||
|
||||
self.vgg.eval().to(device)
|
||||
self.decoder.eval().to(device)
|
||||
|
||||
|
||||
def inference_efdm(
|
||||
src_img: torch.Tensor,
|
||||
style_img: torch.Tensor,
|
||||
|
||||
@@ -1,8 +1,39 @@
|
||||
import torch
|
||||
|
||||
import torch.nn.functional as F
|
||||
from torchvision import transforms
|
||||
|
||||
from . import net_microAST
|
||||
|
||||
|
||||
class MODEL_MICROAST:
|
||||
def __init__(
|
||||
self,
|
||||
content_dec_path: str,
|
||||
style_enc_path: str,
|
||||
modulator_path: str,
|
||||
decoder_path: str,
|
||||
device,
|
||||
) -> None:
|
||||
content_encoder = net_microAST.Encoder()
|
||||
style_encoder = net_microAST.Encoder()
|
||||
modulator = net_microAST.Modulator()
|
||||
decoder = net_microAST.Decoder()
|
||||
|
||||
content_encoder.eval()
|
||||
style_encoder.eval()
|
||||
modulator.eval()
|
||||
decoder.eval()
|
||||
|
||||
content_encoder.load_state_dict(torch.load(content_dec_path))
|
||||
style_encoder.load_state_dict(torch.load(style_enc_path))
|
||||
modulator.load_state_dict(torch.load(modulator_path))
|
||||
decoder.load_state_dict(torch.load(decoder_path))
|
||||
|
||||
self.network = net_microAST.TestNet(
|
||||
content_encoder, style_encoder, modulator, decoder
|
||||
)
|
||||
self.network.to(device)
|
||||
|
||||
|
||||
def inference_microast(
|
||||
src_img: torch.Tensor,
|
||||
@@ -11,7 +42,7 @@ def inference_microast(
|
||||
do_crop: bool,
|
||||
alpha: float,
|
||||
device,
|
||||
network
|
||||
network,
|
||||
):
|
||||
content_tf = test_transform(size, do_crop)
|
||||
|
||||
|
||||
@@ -0,0 +1,396 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Created on Sat Aug 1 11:35:21 2020
|
||||
|
||||
@author: ZJU
|
||||
"""
|
||||
|
||||
import torch
|
||||
import os
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
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.contiguous().view(N, C, -1).var(dim=2) + eps
|
||||
feat_std = feat_var.sqrt().view(N, C, 1, 1)
|
||||
feat_mean = feat.contiguous().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
|
||||
|
||||
def _calc_feat_flatten_mean_std(feat):
|
||||
# takes 3D feat (C, H, W), return mean and std of array within channels
|
||||
assert (feat.size()[0] == 3)
|
||||
assert (isinstance(feat, torch.FloatTensor))
|
||||
feat_flatten = feat.view(3, -1)
|
||||
mean = feat_flatten.mean(dim=-1, keepdim=True)
|
||||
std = feat_flatten.std(dim=-1, keepdim=True)
|
||||
return feat_flatten, mean, std
|
||||
|
||||
decoder = nn.Sequential(
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 256, (3, 3)),
|
||||
nn.ReLU(),
|
||||
nn.Upsample(scale_factor=2, mode='nearest'),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 128, (3, 3)),
|
||||
nn.ReLU(),
|
||||
nn.Upsample(scale_factor=2, mode='nearest'),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(128, 128, (3, 3)),
|
||||
nn.ReLU(),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(128, 64, (3, 3)),
|
||||
nn.ReLU(),
|
||||
nn.Upsample(scale_factor=2, mode='nearest'),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(64, 64, (3, 3)),
|
||||
nn.ReLU(),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(64, 3, (3, 3)),
|
||||
)
|
||||
|
||||
vgg = nn.Sequential(
|
||||
nn.Conv2d(3, 3, (1, 1)),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(3, 64, (3, 3)),
|
||||
nn.ReLU(), # relu1-1
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(64, 64, (3, 3)),
|
||||
nn.ReLU(), # relu1-2
|
||||
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(64, 128, (3, 3)),
|
||||
nn.ReLU(), # relu2-1
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(128, 128, (3, 3)),
|
||||
nn.ReLU(), # relu2-2
|
||||
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(128, 256, (3, 3)),
|
||||
nn.ReLU(), # relu3-1
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(), # relu3-2
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(), # relu3-3
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 256, (3, 3)),
|
||||
nn.ReLU(), # relu3-4
|
||||
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(256, 512, (3, 3)),
|
||||
nn.ReLU(), # relu4-1, this is the last layer used
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu4-2
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu4-3
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu4-4
|
||||
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu5-1
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu5-2
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU(), # relu5-3
|
||||
nn.ReflectionPad2d((1, 1, 1, 1)),
|
||||
nn.Conv2d(512, 512, (3, 3)),
|
||||
nn.ReLU() # relu5-4
|
||||
)
|
||||
|
||||
|
||||
def adaptive_instance_normalization(content_feat, style_feat):
|
||||
assert (content_feat.size()[:2] == style_feat.size()[:2])
|
||||
size = content_feat.size()
|
||||
style_mean, style_std = calc_mean_std(style_feat)
|
||||
content_mean, content_std = calc_mean_std(content_feat)
|
||||
normalized_feat = (content_feat - content_mean.expand(
|
||||
size)) / content_std.expand(size)
|
||||
return normalized_feat * style_std.expand(size) + style_mean.expand(size)
|
||||
|
||||
|
||||
def TSSAT(cf, sf, patch_size=5, stride=1): # cf,sf Batch_size x C x H x W
|
||||
b, c, h, w = sf.size() # 2 x 256 x 64 x 64
|
||||
kh, kw = patch_size, patch_size
|
||||
sh, sw = stride, stride
|
||||
|
||||
# Create convolutional filters by style features
|
||||
sf_unfold = sf.unfold(2, kh, sh).unfold(3, kw, sw)
|
||||
patches = sf_unfold.permute(0, 2, 3, 1, 4, 5)
|
||||
patches = patches.reshape(b, -1, c, kh, kw)
|
||||
patches_norm = torch.norm(patches.reshape(*patches.shape[:2], -1), dim=2).reshape(b, -1, 1, 1, 1)
|
||||
patches_norm = patches / patches_norm
|
||||
# patches size is 2 x 3844 x 256 x 3 x 3
|
||||
|
||||
cf = adaptive_instance_normalization(cf, sf)
|
||||
for i in range(b):
|
||||
cf_temp = cf[i].unsqueeze(0) # [1 x 256 x 64 x 64]
|
||||
patches_norm_temp = patches_norm[i] # [3844, 256, 3, 3]
|
||||
patches_temp = patches[i]
|
||||
|
||||
_, _, ch, cw = cf.size()
|
||||
for c_i in range(0, ch, patch_size):
|
||||
###################################################
|
||||
if (c_i + patch_size) > ch:
|
||||
break
|
||||
elif (c_i + 2*patch_size) > ch:
|
||||
ckh = ch - c_i
|
||||
else:
|
||||
ckh = patch_size
|
||||
###################################################
|
||||
|
||||
for c_j in range(0, cw, patch_size):
|
||||
###################################################
|
||||
if (c_j + patch_size) > cw:
|
||||
break
|
||||
elif (c_j + 2 * patch_size) > cw:
|
||||
ckw = cw - c_j
|
||||
else:
|
||||
ckw = patch_size
|
||||
###################################################
|
||||
|
||||
temp = cf_temp[:, :, c_i:c_i + ckh, c_j:c_j + ckw]
|
||||
conv_out = F.conv2d(temp, patches_norm_temp, stride=patch_size)
|
||||
index = conv_out.argmax(dim=1).squeeze()
|
||||
style_temp = patches_temp[index].unsqueeze(0)
|
||||
stylized_part = adaptive_instance_normalization(temp, style_temp)
|
||||
|
||||
if c_j == 0:
|
||||
p = stylized_part
|
||||
else:
|
||||
p = torch.cat([p, stylized_part], 3)
|
||||
|
||||
if c_i == 0:
|
||||
q = p
|
||||
else:
|
||||
q = torch.cat([q, p], 2)
|
||||
|
||||
if i == 0:
|
||||
out = q
|
||||
else:
|
||||
out = torch.cat([out, q], 0)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def exponent(input, device):
|
||||
b, h, w = input.size()
|
||||
x = (torch.ones(b, h, w) + 0.02).to(device)
|
||||
y = torch.pow(x, input)
|
||||
out = y / torch.sum(y, dim=-1, keepdim=True)
|
||||
return out
|
||||
|
||||
|
||||
|
||||
class SANet(nn.Module):
|
||||
def __init__(self, in_planes):
|
||||
super(SANet, 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))
|
||||
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)
|
||||
O += content
|
||||
return O
|
||||
|
||||
class Transform(nn.Module):
|
||||
def __init__(self, in_planes):
|
||||
super(Transform, self).__init__()
|
||||
self.sanet4_1 = SANet(in_planes = in_planes)
|
||||
self.sanet5_1 = SANet(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):
|
||||
self.upsample5_1 = nn.Upsample(size=(content4_1.size()[2], content4_1.size()[3]), mode='nearest')
|
||||
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))))
|
||||
|
||||
class Net(nn.Module):
|
||||
def __init__(self, encoder, decoder, start_iter):
|
||||
super(Net, self).__init__()
|
||||
enc_layers = list(encoder.children())
|
||||
self.enc_1 = nn.Sequential(*enc_layers[:4]) # input -> relu1_1
|
||||
self.enc_2 = nn.Sequential(*enc_layers[4:11]) # relu1_1 -> relu2_1
|
||||
self.enc_3 = nn.Sequential(*enc_layers[11:18]) # relu2_1 -> relu3_1
|
||||
self.enc_4 = nn.Sequential(*enc_layers[18:31]) # relu3_1 -> relu4_1
|
||||
self.enc_5 = nn.Sequential(*enc_layers[31:44]) # relu4_1 -> relu5_1
|
||||
#transform
|
||||
self.transform = Transform(in_planes = 512)
|
||||
self.decoder = decoder
|
||||
self.sm = nn.Softmax(dim=-1)
|
||||
|
||||
if(start_iter > 0):
|
||||
self.transform.load_state_dict(torch.load('transformer_iter_' + str(start_iter) + '.pth'))
|
||||
self.decoder.load_state_dict(torch.load('decoder_iter_' + str(start_iter) + '.pth'))
|
||||
self.mse_loss = nn.MSELoss()
|
||||
# fix the encoder
|
||||
for name in ['enc_1', 'enc_2', 'enc_3', 'enc_4', 'enc_5']:
|
||||
for param in getattr(self, name).parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
# extract relu1_1, relu2_1, relu3_1, relu4_1, relu5_1 from input image
|
||||
def encode_with_intermediate(self, input):
|
||||
results = [input]
|
||||
for i in range(5):
|
||||
func = getattr(self, 'enc_{:d}'.format(i + 1))
|
||||
results.append(func(results[-1]))
|
||||
return results[1:]
|
||||
|
||||
def attention_map(self, input1, input2, eps=1e-5):
|
||||
F = mean_variance_norm(input1)
|
||||
G = mean_variance_norm(input2)
|
||||
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)
|
||||
|
||||
n = h * w
|
||||
T = torch.zeros((b, n, n - 1))
|
||||
# REMOVE DIAGONAL
|
||||
for i in range(b):
|
||||
T[i] = S[i].flatten()[1:].view(n - 1, n + 1)[:, :-1].reshape(n, n - 1)
|
||||
T = self.sm(T / 100)
|
||||
|
||||
return T
|
||||
|
||||
def attention_loss(self, content, stylization):
|
||||
attention_map1 = self.attention_map(content, content)
|
||||
attention_map2 = self.attention_map(stylization, stylization)
|
||||
return self.mse_loss(attention_map1, attention_map2)
|
||||
|
||||
def calc_content_loss(self, input, target, norm = False):
|
||||
if not norm:
|
||||
return self.mse_loss(input, target)
|
||||
else:
|
||||
return self.mse_loss(mean_variance_norm(input), mean_variance_norm(target))
|
||||
|
||||
def calc_style_loss(self, input, target):
|
||||
input_mean, input_std = calc_mean_std(input)
|
||||
target_mean, target_std = calc_mean_std(target)
|
||||
return self.mse_loss(input_mean, target_mean) + \
|
||||
self.mse_loss(input_std, target_std)
|
||||
|
||||
def calc_local_style_loss(self, cf, sf, patch_size=5, stride=1):
|
||||
b, c, h, w = sf.size() # 2 x 256 x 64 x 64
|
||||
kh, kw = patch_size, patch_size
|
||||
sh, sw = stride, stride
|
||||
|
||||
# Create convolutional filters by style features
|
||||
sf_unfold = sf.unfold(2, kh, sh).unfold(3, kw, sw)
|
||||
patches = sf_unfold.permute(0, 2, 3, 1, 4, 5)
|
||||
patches = patches.reshape(b, -1, c, kh, kw)
|
||||
patches_norm = torch.norm(patches.reshape(*patches.shape[:2], -1), dim=2).reshape(b, -1, 1, 1, 1)
|
||||
patches_norm = patches / patches_norm
|
||||
# patches size is 2 x 3844 x 256 x 3 x 3
|
||||
|
||||
loss = 0
|
||||
for i in range(b):
|
||||
cf_temp = cf[i].unsqueeze(0) # [1 x 256 x 64 x 64]
|
||||
patches_norm_temp = patches_norm[i] # [3844, 256, 3, 3]
|
||||
patches_temp = patches[i]
|
||||
|
||||
_, _, ch, cw = cf.size()
|
||||
for c_i in range(0, ch, patch_size):
|
||||
###################################################
|
||||
if (c_i + patch_size) > ch:
|
||||
break
|
||||
elif (c_i + 2 * patch_size) > ch:
|
||||
ckh = ch - c_i
|
||||
else:
|
||||
ckh = patch_size
|
||||
###################################################
|
||||
|
||||
for c_j in range(0, cw, patch_size):
|
||||
###################################################
|
||||
if (c_j + patch_size) > cw:
|
||||
break
|
||||
elif (c_j + 2 * patch_size) > cw:
|
||||
ckw = cw - c_j
|
||||
else:
|
||||
ckw = patch_size
|
||||
###################################################
|
||||
|
||||
temp = cf_temp[:, :, c_i:c_i + ckh, c_j:c_j + ckw]
|
||||
conv_out = F.conv2d(temp, patches_norm_temp, stride=patch_size)
|
||||
index = conv_out.argmax(dim=1).squeeze()
|
||||
style_temp = patches_temp[index].unsqueeze(0)
|
||||
|
||||
input_mean, input_std = calc_mean_std(temp)
|
||||
target_mean, target_std = calc_mean_std(style_temp)
|
||||
loss += self.mse_loss(input_mean, target_mean) + self.mse_loss(input_std, target_std)
|
||||
|
||||
return loss
|
||||
|
||||
def forward(self, content, style):
|
||||
style_feats = self.encode_with_intermediate(style)
|
||||
content_feats = self.encode_with_intermediate(content)
|
||||
|
||||
stylized = TSSAT(content_feats[3], style_feats[3])
|
||||
g_t = self.decoder(stylized)
|
||||
g_t_feats = self.encode_with_intermediate(g_t)
|
||||
|
||||
loss_c = self.calc_content_loss(g_t_feats[3], content_feats[3]) + self.calc_content_loss(g_t_feats[4], content_feats[4]) # True; True
|
||||
loss_s = self.calc_style_loss(g_t_feats[0], style_feats[0])
|
||||
for i in range(1, 5):
|
||||
loss_s += self.calc_style_loss(g_t_feats[i], style_feats[i])
|
||||
|
||||
"""LOCAL STYLE LOSSES"""
|
||||
loss_s_local = self.calc_local_style_loss(g_t_feats[3], style_feats[3])
|
||||
|
||||
"""ATTENTION LOSSES"""
|
||||
loss_attention1 = self.attention_loss(content_feats[3], g_t_feats[3])
|
||||
loss_attention2 = self.attention_loss(content_feats[4], g_t_feats[4])
|
||||
loss_attention = loss_attention1 + loss_attention2
|
||||
|
||||
"""IDENTITY LOSSES"""
|
||||
Icc = self.decoder(TSSAT(content_feats[3], content_feats[3]))
|
||||
Iss = self.decoder(TSSAT(style_feats[3], style_feats[3]))
|
||||
l_identity1 = self.calc_content_loss(Icc, content) + self.calc_content_loss(Iss, style)
|
||||
Fcc = self.encode_with_intermediate(Icc)
|
||||
Fss = self.encode_with_intermediate(Iss)
|
||||
l_identity2 = self.calc_content_loss(Fcc[0], content_feats[0]) + self.calc_content_loss(Fss[0], style_feats[0])
|
||||
for i in range(1, 5):
|
||||
l_identity2 += self.calc_content_loss(Fcc[i], content_feats[i]) + self.calc_content_loss(Fss[i], style_feats[i])
|
||||
|
||||
return g_t, loss_c, loss_s, loss_s_local, loss_attention, l_identity1, l_identity2
|
||||
@@ -0,0 +1,33 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Created on Sat Aug 1 11:38:11 2020
|
||||
|
||||
@author: ZJU
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from torch.utils import data
|
||||
|
||||
|
||||
def InfiniteSampler(n):
|
||||
# i = 0
|
||||
i = n - 1
|
||||
order = np.random.permutation(n)
|
||||
while True:
|
||||
yield order[i]
|
||||
i += 1
|
||||
if i >= n:
|
||||
np.random.seed()
|
||||
order = np.random.permutation(n)
|
||||
i = 0
|
||||
|
||||
|
||||
class InfiniteSamplerWrapper(data.sampler.Sampler):
|
||||
def __init__(self, data_source):
|
||||
self.num_samples = len(data_source)
|
||||
|
||||
def __iter__(self):
|
||||
return iter(InfiniteSampler(self.num_samples))
|
||||
|
||||
def __len__(self):
|
||||
return 2**31
|
||||
@@ -0,0 +1,168 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import tqdm
|
||||
from comfy.utils import ProgressBar
|
||||
from torchvision import transforms
|
||||
|
||||
from . import net
|
||||
|
||||
|
||||
class MODEL_TSSAT:
|
||||
def __init__(self, dec_path: str, vgg_path: str, device) -> None:
|
||||
self.decoder = net.decoder
|
||||
vgg = net.vgg
|
||||
|
||||
self.decoder.load_state_dict(torch.load(dec_path))
|
||||
vgg.load_state_dict(torch.load(vgg_path))
|
||||
|
||||
self.enc_1 = nn.Sequential(*list(vgg.children())[:4]) # input -> relu1_1
|
||||
self.enc_2 = nn.Sequential(*list(vgg.children())[4:11]) # relu1_1 -> relu2_1
|
||||
self.enc_3 = nn.Sequential(*list(vgg.children())[11:18]) # relu2_1 -> relu3_1
|
||||
self.enc_4 = nn.Sequential(*list(vgg.children())[18:31]) # relu3_1 -> relu4_1
|
||||
|
||||
self.enc_1.eval().to(device)
|
||||
self.enc_2.eval().to(device)
|
||||
self.enc_3.eval().to(device)
|
||||
self.enc_4.eval().to(device)
|
||||
self.decoder.eval().to(device)
|
||||
|
||||
def encode(self, item):
|
||||
out = self.enc_1(item)
|
||||
out = self.enc_2(out)
|
||||
out = self.enc_3(out)
|
||||
out = self.enc_4(out)
|
||||
return out
|
||||
|
||||
|
||||
def inference_tssat(
|
||||
src_img: torch.Tensor,
|
||||
style_img: torch.Tensor,
|
||||
device,
|
||||
size: int,
|
||||
do_crop: False,
|
||||
max_steps: int,
|
||||
model: MODEL_TSSAT,
|
||||
):
|
||||
content_tf = test_transform(size, do_crop)
|
||||
content = content_tf(src_img).to(device) # [1, C, H, W]
|
||||
style = content_tf(style_img).to(device) # [1, C, H, W]
|
||||
|
||||
for _ in range(max_steps):
|
||||
Content4_1 = model.encode(content)
|
||||
Style4_1 = model.encode(style)
|
||||
content = model.decoder(TSSAT(Content4_1, Style4_1))
|
||||
|
||||
return content
|
||||
|
||||
|
||||
def forward(item, enc_1, enc_2, enc_3, enc_4):
|
||||
out = enc_1(item)
|
||||
out = enc_2(out)
|
||||
out = enc_3(out)
|
||||
out = enc_4(out)
|
||||
return out
|
||||
|
||||
|
||||
def test_transform(size, crop):
|
||||
transform_list = []
|
||||
if size != 0:
|
||||
transform_list.append(transforms.Resize(size))
|
||||
if crop:
|
||||
transform_list.append(transforms.CenterCrop(size))
|
||||
transform = transforms.Compose(transform_list)
|
||||
# [B, C, H, W]
|
||||
return transform
|
||||
|
||||
|
||||
def TSSAT(cf, sf, patch_size=5, stride=1): # cf,sf Batch_size x C x H x W
|
||||
b, c, h, w = sf.size() # 2 x 256 x 64 x 64
|
||||
# print(cf.size())
|
||||
kh, kw = patch_size, patch_size
|
||||
sh, sw = stride, stride
|
||||
|
||||
# Create convolutional filters by style features
|
||||
sf_unfold = sf.unfold(2, kh, sh).unfold(3, kw, sw)
|
||||
patches = sf_unfold.permute(0, 2, 3, 1, 4, 5)
|
||||
patches = patches.reshape(b, -1, c, kh, kw)
|
||||
patches_norm = torch.norm(patches.reshape(*patches.shape[:2], -1), dim=2).reshape(
|
||||
b, -1, 1, 1, 1
|
||||
)
|
||||
patches_norm = patches / patches_norm
|
||||
# patches size is 2 x 3844 x 256 x 3 x 3
|
||||
|
||||
cf = adaptive_instance_normalization(cf, sf)
|
||||
|
||||
for i in range(b):
|
||||
cf_temp = cf[i].unsqueeze(0) # [1 x 256 x 64 x 64]
|
||||
patches_norm_temp = patches_norm[i] # [3844, 256, 3, 3]
|
||||
patches_temp = patches[i]
|
||||
|
||||
_, _, ch, cw = cf.size()
|
||||
pbar = ProgressBar(ch)
|
||||
for c_i in tqdm.tqdm(range(0, ch, patch_size), desc="Optimizing"):
|
||||
###################################################
|
||||
if (c_i + patch_size) > ch:
|
||||
break
|
||||
elif (c_i + 2 * patch_size) > ch:
|
||||
ckh = ch - c_i
|
||||
else:
|
||||
ckh = patch_size
|
||||
###################################################
|
||||
|
||||
for c_j in range(0, cw, patch_size):
|
||||
###################################################
|
||||
if (c_j + patch_size) > cw:
|
||||
break
|
||||
elif (c_j + 2 * patch_size) > cw:
|
||||
ckw = cw - c_j
|
||||
else:
|
||||
ckw = patch_size
|
||||
###################################################
|
||||
|
||||
temp = cf_temp[:, :, c_i : c_i + ckh, c_j : c_j + ckw]
|
||||
conv_out = F.conv2d(temp, patches_norm_temp, stride=patch_size)
|
||||
index = conv_out.argmax(dim=1).squeeze()
|
||||
style_temp = patches_temp[index].unsqueeze(0)
|
||||
stylized_part = adaptive_instance_normalization(temp, style_temp)
|
||||
|
||||
if c_j == 0:
|
||||
p = stylized_part
|
||||
else:
|
||||
p = torch.cat([p, stylized_part], 3)
|
||||
|
||||
if c_i == 0:
|
||||
q = p
|
||||
else:
|
||||
q = torch.cat([q, p], 2)
|
||||
|
||||
pbar.update_absolute(c_i, ch)
|
||||
|
||||
if i == 0:
|
||||
out = q
|
||||
else:
|
||||
out = torch.cat([out, q], 0)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def adaptive_instance_normalization(content_feat, style_feat):
|
||||
assert content_feat.size()[:2] == style_feat.size()[:2]
|
||||
size = content_feat.size()
|
||||
style_mean, style_std = calc_mean_std(style_feat)
|
||||
content_mean, content_std = calc_mean_std(content_feat)
|
||||
normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(
|
||||
size
|
||||
)
|
||||
return normalized_feat * style_std.expand(size) + style_mean.expand(size)
|
||||
|
||||
|
||||
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.contiguous().view(N, C, -1).var(dim=2) + eps
|
||||
feat_std = feat_var.sqrt().view(N, C, 1, 1)
|
||||
feat_mean = feat.contiguous().view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
|
||||
return feat_mean, feat_std
|
||||
@@ -1,37 +0,0 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .models import Transformer
|
||||
|
||||
|
||||
def get_keys(d, name):
|
||||
if "state_dict" in d:
|
||||
d = d["state_dict"]
|
||||
d_filt = {k[len(name) + 1 :]: v for k, v in d.items() if k[: len(name)] == name}
|
||||
return d_filt
|
||||
|
||||
|
||||
class video_Style_transfer(nn.Module):
|
||||
def __init__(self, ckpt_path: str, encoder_path: str, decoder_path: str):
|
||||
super(video_Style_transfer, self).__init__()
|
||||
self.ckpt_path = ckpt_path
|
||||
self.model = Transformer(encoder_path, decoder_path)
|
||||
self.load_weights()
|
||||
|
||||
def load_weights(self):
|
||||
# print("Loading model from checkpoint")
|
||||
ckpt = torch.load(self.ckpt_path, map_location="cpu")
|
||||
self.model.load_state_dict(get_keys(ckpt, "model"), strict=False)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
content_frames,
|
||||
style_images,
|
||||
content_type="image",
|
||||
id_loss="transfer",
|
||||
tab=None,
|
||||
):
|
||||
transfer_result = self.model.forward(
|
||||
content_frames, style_images, content_type, id_loss, tab
|
||||
)
|
||||
return transfer_result
|
||||
@@ -1,7 +1,34 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as transforms
|
||||
|
||||
from .model import video_Style_transfer
|
||||
from .models import Transformer
|
||||
|
||||
|
||||
class video_Style_transfer(nn.Module):
|
||||
def __init__(self, ckpt_path: str, encoder_path: str, decoder_path: str):
|
||||
super(video_Style_transfer, self).__init__()
|
||||
self.ckpt_path = ckpt_path
|
||||
self.model = Transformer(encoder_path, decoder_path)
|
||||
self.load_weights()
|
||||
|
||||
def load_weights(self):
|
||||
# print("Loading model from checkpoint")
|
||||
ckpt = torch.load(self.ckpt_path, map_location="cpu")
|
||||
self.model.load_state_dict(get_keys(ckpt, "model"), strict=False)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
content_frames,
|
||||
style_images,
|
||||
content_type="image",
|
||||
id_loss="transfer",
|
||||
tab=None,
|
||||
):
|
||||
transfer_result = self.model.forward(
|
||||
content_frames, style_images, content_type, id_loss, tab
|
||||
)
|
||||
return transfer_result
|
||||
|
||||
|
||||
def inference_unist(
|
||||
@@ -63,3 +90,10 @@ def match_shape(a: torch.Tensor, b: torch.Tensor):
|
||||
return torch.cat([a.repeat(repeat_factor, 1, 1, 1), a[:remainder]], dim=0)
|
||||
else: # B_a > B_b
|
||||
return a[:B_b]
|
||||
|
||||
|
||||
def get_keys(d, name):
|
||||
if "state_dict" in d:
|
||||
d = d["state_dict"]
|
||||
d_filt = {k[len(name) + 1 :]: v for k, v in d.items() if k[: len(name)] == name}
|
||||
return d_filt
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui-styletransferplus"
|
||||
description = "Nodes: Neural Neighbor, CAST, EFDM, MicroAST, UniST, AesPA, ..."
|
||||
version = "1.0.1"
|
||||
description = "Nodes: Neural Neighbor, CAST, EFDM, MicroAST, UniST, AesPA, TSSAT ..."
|
||||
version = "1.0.2"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["torch", "torchvision", "scikit-image", "einops"]
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import tqdm
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
@@ -27,6 +26,7 @@ from .constants import (
|
||||
MICROAST_DECODER_PATH,
|
||||
MICROAST_MODULATOR_PATH,
|
||||
MICROAST_STYLE_ENCODER_PATH,
|
||||
TSSAT_DEC_PATH,
|
||||
UNIST_DEC_PATH,
|
||||
UNIST_ENC_PATH,
|
||||
UNIST_PATH,
|
||||
@@ -34,35 +34,25 @@ from .constants import (
|
||||
)
|
||||
|
||||
# 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,
|
||||
)
|
||||
from .module_aespa.aespa_model import MODEL_AESPA, inference_aespa
|
||||
|
||||
# CAST
|
||||
from .module_cast import net as net_cast
|
||||
from .module_cast.cast_model import inference_ucast, load_a_ckpt
|
||||
from .module_cast.cast_model import MODEL_CAST, inference_ucast
|
||||
|
||||
# EFDM
|
||||
from .module_efdm import net as net_efdm
|
||||
from .module_efdm.efdm_model import inference_efdm
|
||||
from .module_efdm.efdm_model import MODEL_EFDM, inference_efdm
|
||||
|
||||
# MicroAST
|
||||
from .module_microast import net_microAST
|
||||
from .module_microast.microast_model import inference_microast
|
||||
from .module_microast.microast_model import MODEL_MICROAST, inference_microast
|
||||
|
||||
# Neural Neighbor
|
||||
from .module_neural_neighbor.neural_neighbor_model import inference_neural_neighbor
|
||||
|
||||
# TSSAT
|
||||
from .module_tssat.tssat_model import MODEL_TSSAT, inference_tssat
|
||||
|
||||
# UniST
|
||||
from .module_unist.model import video_Style_transfer
|
||||
from .module_unist.unist_model import inference_unist
|
||||
from .module_unist.unist_model import inference_unist, video_Style_transfer
|
||||
|
||||
base_dir = Path(__file__).resolve().parent
|
||||
|
||||
@@ -172,28 +162,18 @@ class CAST:
|
||||
print(f"{style_img.shape=}")
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
vgg = net_cast.vgg
|
||||
vgg.load_state_dict(torch.load(f"{base_dir}/{CAST_VGG_PATH}"))
|
||||
vgg = torch.nn.Sequential(*list(vgg.children())[:31])
|
||||
|
||||
netAE = net_cast.ADAIN_Encoder(vgg)
|
||||
netDec_B = net_cast.Decoder()
|
||||
|
||||
netAE.load_state_dict(
|
||||
load_a_ckpt(f"{base_dir}/models/{model_arch}_model/{CAST_NET_AE_PATH}")
|
||||
model = MODEL_CAST(
|
||||
vgg_path=f"{base_dir}/{CAST_VGG_PATH}",
|
||||
net_ae_path=f"{base_dir}/models/{model_arch}_model/{CAST_NET_AE_PATH}",
|
||||
net_dec_b_path=f"{base_dir}/models/{model_arch}_model/{CAST_NET_DEC_B_PATH}",
|
||||
device=device,
|
||||
)
|
||||
netDec_B.load_state_dict(
|
||||
load_a_ckpt(f"{base_dir}/models/{model_arch}_model/{CAST_NET_DEC_B_PATH}")
|
||||
)
|
||||
|
||||
netAE = netAE.to(device).eval()
|
||||
netDec_B = netDec_B.to(device).eval()
|
||||
|
||||
params = {
|
||||
"style_img": style_img,
|
||||
"device": device,
|
||||
"netAE": netAE,
|
||||
"netDec_B": netDec_B,
|
||||
"netAE": model.netAE,
|
||||
"netDec_B": model.netDec_B,
|
||||
}
|
||||
|
||||
num_frames = src_img.size(0)
|
||||
@@ -265,25 +245,17 @@ class EFDM:
|
||||
"cuda" if torch.cuda.is_available() and not use_cpu else "cpu"
|
||||
)
|
||||
|
||||
decoder = net_efdm.decoder
|
||||
vgg = net_efdm.vgg
|
||||
|
||||
decoder.eval()
|
||||
vgg.eval()
|
||||
|
||||
decoder.load_state_dict(torch.load(f"{base_dir}/{EFDM_PATH}"))
|
||||
vgg.load_state_dict(torch.load(f"{base_dir}/{CAST_VGG_PATH}"))
|
||||
|
||||
vgg = torch.nn.Sequential(*list(vgg.children())[:31])
|
||||
|
||||
vgg.to(device)
|
||||
decoder.to(device)
|
||||
model = MODEL_EFDM(
|
||||
dec_path=f"{base_dir}/{EFDM_PATH}",
|
||||
vgg_path=f"{base_dir}/{CAST_VGG_PATH}",
|
||||
device=device,
|
||||
)
|
||||
|
||||
params = {
|
||||
"style_img": style_img.permute(0, 3, 1, 2),
|
||||
"device": device,
|
||||
"decoder": decoder,
|
||||
"vgg": vgg,
|
||||
"decoder": model.decoder,
|
||||
"vgg": model.vgg,
|
||||
"alpha": style_strength,
|
||||
"size": size,
|
||||
"style_type": model_arch,
|
||||
@@ -343,29 +315,13 @@ class MicroAST:
|
||||
print(f"{style_img.shape=}")
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
content_encoder = net_microAST.Encoder()
|
||||
style_encoder = net_microAST.Encoder()
|
||||
modulator = net_microAST.Modulator()
|
||||
decoder = net_microAST.Decoder()
|
||||
|
||||
content_encoder.eval()
|
||||
style_encoder.eval()
|
||||
modulator.eval()
|
||||
decoder.eval()
|
||||
|
||||
content_encoder.load_state_dict(
|
||||
torch.load(f"{base_dir}/{MICROAST_CONTENT_ENCODER_PATH}")
|
||||
model = MODEL_MICROAST(
|
||||
content_dec_path=f"{base_dir}/{MICROAST_CONTENT_ENCODER_PATH}",
|
||||
style_enc_path=f"{base_dir}/{MICROAST_STYLE_ENCODER_PATH}",
|
||||
modulator_path=f"{base_dir}/{MICROAST_MODULATOR_PATH}",
|
||||
decoder_path=f"{base_dir}/{MICROAST_DECODER_PATH}",
|
||||
device=device,
|
||||
)
|
||||
style_encoder.load_state_dict(
|
||||
torch.load(f"{base_dir}/{MICROAST_STYLE_ENCODER_PATH}")
|
||||
)
|
||||
modulator.load_state_dict(torch.load(f"{base_dir}/{MICROAST_MODULATOR_PATH}"))
|
||||
decoder.load_state_dict(torch.load(f"{base_dir}/{MICROAST_DECODER_PATH}"))
|
||||
|
||||
network = net_microAST.TestNet(
|
||||
content_encoder, style_encoder, modulator, decoder
|
||||
)
|
||||
network.to(device)
|
||||
|
||||
params = {
|
||||
"style_img": style_img.permute(0, 3, 1, 2),
|
||||
@@ -373,7 +329,7 @@ class MicroAST:
|
||||
"alpha": 1.0,
|
||||
"size": size,
|
||||
"do_crop": do_crop,
|
||||
"network": network,
|
||||
"network": model.network,
|
||||
}
|
||||
|
||||
num_frames = src_img.size(0)
|
||||
@@ -564,35 +520,19 @@ class AesPA:
|
||||
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,
|
||||
model = MODEL_AESPA(
|
||||
vgg_51_path=f"{base_dir}/{VGG_NORM_CONV51_PATH}",
|
||||
dec_path=f"{base_dir}/{AESPA_DEC_PATH}",
|
||||
transformer_path=f"{base_dir}/{AESPA_TRANSFORMER_PATH}",
|
||||
device=device,
|
||||
)
|
||||
|
||||
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,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
num_frames = src_img.size(0)
|
||||
@@ -609,6 +549,72 @@ class AesPA:
|
||||
return (torch.cat(result, dim=0),)
|
||||
|
||||
|
||||
class TSSAT:
|
||||
@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},
|
||||
),
|
||||
"max_iter": (
|
||||
"INT",
|
||||
{"default": 1, "min": 1, "step": 1},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("out_img",)
|
||||
FUNCTION = "todo"
|
||||
CATEGORY = "StyleTransferPlus"
|
||||
|
||||
def todo(
|
||||
self,
|
||||
src_img: torch.Tensor,
|
||||
style_img: torch.Tensor,
|
||||
do_crop: bool,
|
||||
size: int,
|
||||
max_iter: int,
|
||||
):
|
||||
print(f"{src_img.shape=}")
|
||||
print(f"{style_img.shape=}")
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = MODEL_TSSAT(
|
||||
dec_path=f"{base_dir}/{TSSAT_DEC_PATH}",
|
||||
vgg_path=f"{base_dir}/{CAST_VGG_PATH}",
|
||||
device=device,
|
||||
)
|
||||
|
||||
params = {
|
||||
"style_img": style_img.permute(0, 3, 1, 2),
|
||||
"device": device,
|
||||
"size": size,
|
||||
"do_crop": do_crop,
|
||||
"max_steps": max_iter,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
num_frames = src_img.size(0)
|
||||
|
||||
result: list[torch.Tensor] = []
|
||||
with torch.no_grad():
|
||||
for i in range(num_frames):
|
||||
params["src_img"] = src_img[i].unsqueeze(0).permute(0, 3, 1, 2)
|
||||
res_tensor = inference_tssat(**params)
|
||||
result.append(res_tensor.permute(0, 2, 3, 1))
|
||||
|
||||
return (torch.cat(result, dim=0),)
|
||||
|
||||
|
||||
def serialize_floats(lst: list[float]):
|
||||
return ",".join(map(str, lst))
|
||||
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 597 KiB |
@@ -0,0 +1,194 @@
|
||||
{
|
||||
"last_node_id": 9,
|
||||
"last_link_id": 23,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 4,
|
||||
"type": "SaveImage",
|
||||
"pos": [
|
||||
1050,
|
||||
327
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 270
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 23
|
||||
}
|
||||
],
|
||||
"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": [
|
||||
21
|
||||
],
|
||||
"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": [
|
||||
22
|
||||
],
|
||||
"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": 9,
|
||||
"type": "TSSAT",
|
||||
"pos": [
|
||||
704.2608950403187,
|
||||
332.7147398449658
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 126
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "src_img",
|
||||
"type": "IMAGE",
|
||||
"link": 21
|
||||
},
|
||||
{
|
||||
"name": "style_img",
|
||||
"type": "IMAGE",
|
||||
"link": 22
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "out_img",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
23
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "TSSAT"
|
||||
},
|
||||
"widgets_values": [
|
||||
false,
|
||||
1024,
|
||||
1
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
21,
|
||||
2,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
22,
|
||||
3,
|
||||
0,
|
||||
9,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
23,
|
||||
9,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 1.2100000000000002,
|
||||
"offset": [
|
||||
-289.7051144856484,
|
||||
-164.0565891162732
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
Reference in New Issue
Block a user