diff --git a/README.md b/README.md index b4a2f5a..a4b4b8c 100644 --- a/README.md +++ b/README.md @@ -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 | ![efdm](example_outputs/efdm.png) | ![microast](example_outputs/microast.png) | | **UniST (2023) - Medium - Only squares** | **AesPA-Net (2023) - Medium - Only squares** | | ![efdm](example_outputs/unist_image.png) | ![aespa](example_outputs/aespa.png) | +| **TSSAT (2023) - Slow** | **X** | +| ![tssat](example_outputs/tssat.png) | X | ### 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 ![wf_aespa](workflows/wf_aespa.png) + +### 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) + +![wf_tssat](workflows/wf_tssat.png) + ## Extra nodes ### Coral Color Transfer diff --git a/__init__.py b/__init__.py index 685c4b0..cc014f6 100644 --- a/__init__.py +++ b/__init__.py @@ -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", } diff --git a/constants.py b/constants.py index 6c4af4a..4679f9e 100644 --- a/constants.py +++ b/constants.py @@ -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" \ No newline at end of file diff --git a/example_outputs/tssat.png b/example_outputs/tssat.png new file mode 100644 index 0000000..df459f4 Binary files /dev/null and b/example_outputs/tssat.png differ diff --git a/module_aespa/aespa_model.py b/module_aespa/aespa_model.py index 3c54dd2..7265447 100644 --- a/module_aespa/aespa_model.py +++ b/module_aespa/aespa_model.py @@ -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 = [] diff --git a/module_aespa/aespanet_models.py b/module_aespa/aespanet_models.py index 2581e9a..15b9856 100644 --- a/module_aespa/aespanet_models.py +++ b/module_aespa/aespanet_models.py @@ -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##### ################################# diff --git a/module_aespa/utils.py b/module_aespa/utils.py index dfdf9c2..b595680 100644 --- a/module_aespa/utils.py +++ b/module_aespa/utils.py @@ -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) diff --git a/module_cast/cast_model.py b/module_cast/cast_model.py index db73e69..0dbe1b7 100644 --- a/module_cast/cast_model.py +++ b/module_cast/cast_model.py @@ -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 diff --git a/module_efdm/efdm_model.py b/module_efdm/efdm_model.py index 65c0f78..2527423 100644 --- a/module_efdm/efdm_model.py +++ b/module_efdm/efdm_model.py @@ -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, diff --git a/module_microast/microast_model.py b/module_microast/microast_model.py index 4cbc85d..e20578b 100644 --- a/module_microast/microast_model.py +++ b/module_microast/microast_model.py @@ -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) diff --git a/module_tssat/net.py b/module_tssat/net.py new file mode 100644 index 0000000..f8ca6f8 --- /dev/null +++ b/module_tssat/net.py @@ -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 \ No newline at end of file diff --git a/module_tssat/sampler.py b/module_tssat/sampler.py new file mode 100644 index 0000000..62f444e --- /dev/null +++ b/module_tssat/sampler.py @@ -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 diff --git a/module_tssat/tssat_model.py b/module_tssat/tssat_model.py new file mode 100644 index 0000000..1a73528 --- /dev/null +++ b/module_tssat/tssat_model.py @@ -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 diff --git a/module_unist/model.py b/module_unist/model.py deleted file mode 100644 index 73c1712..0000000 --- a/module_unist/model.py +++ /dev/null @@ -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 diff --git a/module_unist/unist_model.py b/module_unist/unist_model.py index 57a4c37..b8331b5 100644 --- a/module_unist/unist_model.py +++ b/module_unist/unist_model.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index f640677..8129d07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/run.py b/run.py index 44443be..eee480c 100644 --- a/run.py +++ b/run.py @@ -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)) diff --git a/workflows/wf_tssat.png b/workflows/wf_tssat.png new file mode 100644 index 0000000..c55ea51 Binary files /dev/null and b/workflows/wf_tssat.png differ diff --git a/workflows/workflow_tssat.json b/workflows/workflow_tssat.json new file mode 100644 index 0000000..f4b157f --- /dev/null +++ b/workflows/workflow_tssat.json @@ -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 +} \ No newline at end of file