diff --git a/.gitignore b/.gitignore index 5ec8bfe..a5fd31f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,5 @@ *.pth - +*.tar # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index 58a0684..0f5ce59 100644 --- a/README.md +++ b/README.md @@ -17,8 +17,8 @@ All nodes support batched input (i.e video) but is generally not recommended. On | Neural Neighbor (2020-2022) - Slowest | CAST (2022-2023) - Fast | |:-:|:-:| |![neural_neighbor](example_outputs/neural_neighbor.png) | ![cast](example_outputs/cast.png) | -| **EFDM (2022) - Fast** | **MicroAST (2023)** | -| ![efdm](example_outputs/efdm.png) | X | +| **EFDM (2022) - Fast** | **MicroAST (2023) - Fast** | +| ![efdm](example_outputs/efdm.png) | ![microast](example_outputs/microast.png) | ### Neural Neighbor @@ -88,7 +88,20 @@ Arguments: ### MicroAST -Download `vgg_normalized.pth` as instructed in [**CAST**](#cast) +Download models from https://github.com/EndyWon/MicroAST/tree/main/models and place them in the models folder + +``` +models/microast/content_encoder_iter_160000.pth.tar +models/microast/decoder_iter_160000.pth.tar +models/microast/modulator_iter_160000.pth.tar +models/microast/style_encoder_iter_160000.pth.tar +``` + +Arguments: Similar to [EFDM](#efdm) + +[workflow_microast.json](workflows/workflow_microast.json) + +![wf_microast](workflows/wf_microast.png) ## Credits diff --git a/__init__.py b/__init__.py index 291e963..887df5b 100644 --- a/__init__.py +++ b/__init__.py @@ -1,15 +1,17 @@ -from .run import NeuralNeighbor, CAST, EFDM +from .run import NeuralNeighbor, CAST, EFDM, MicroAST NODE_CLASS_MAPPINGS = { "NeuralNeighbor": NeuralNeighbor, "CAST": CAST, "EFDM": EFDM, + "MicroAST": MicroAST, } NODE_DISPLAY_NAME_MAPPINGS = { "NeuralNeighbor": "Neural Neighbor", "CAST": "CAST", "EFDM": "EFDM", + "MicroAST": "MicroAST", } diff --git a/constants.py b/constants.py index 7b4b709..12daafa 100644 --- a/constants.py +++ b/constants.py @@ -7,4 +7,9 @@ CAST_VGG_PATH = "models/vgg_normalised.pth" EFDM_STYLE_TYPES = ["adain", "adamean", "adastd", "efdm", "hm"] EFDM_DEFAULT = "efdm" -EFDM_PATH = "models/hm_decoder_iter_160000.pth" \ No newline at end of file +EFDM_PATH = "models/hm_decoder_iter_160000.pth" + +MICROAST_CONTENT_ENCODER_PATH = "models/microast/content_encoder_iter_160000.pth.tar" +MICROAST_STYLE_ENCODER_PATH = "models/microast/style_encoder_iter_160000.pth.tar" +MICROAST_DECODER_PATH = "models/microast/decoder_iter_160000.pth.tar" +MICROAST_MODULATOR_PATH = "models/microast/modulator_iter_160000.pth.tar" \ No newline at end of file diff --git a/example_outputs/microast.png b/example_outputs/microast.png new file mode 100644 index 0000000..bfa0803 Binary files /dev/null and b/example_outputs/microast.png differ diff --git a/module_efdm/efdm_model.py b/module_efdm/efdm_model.py index a8c6925..65c0f78 100644 --- a/module_efdm/efdm_model.py +++ b/module_efdm/efdm_model.py @@ -50,8 +50,8 @@ def inference_efdm( ) content = content.expand_as(style) - print(f"Resized: {content.shape=}") - print(f"Resized: {style.shape=}") + # print(f"Resized: {content.shape=}") + # print(f"Resized: {style.shape=}") if preserve_color: tmp_content = content.squeeze() diff --git a/module_microast/function.py b/module_microast/function.py new file mode 100644 index 0000000..5af9124 --- /dev/null +++ b/module_microast/function.py @@ -0,0 +1,77 @@ +import torch + + +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 normalization(content_feat): + size = content_feat.size() + content_mean, content_std = calc_mean_std(content_feat) + + normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand( + size + ) + return normalized_feat + + +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_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 + + +def _mat_sqrt(x): + U, D, V = torch.svd(x) + return torch.mm(torch.mm(U, D.pow(0.5).diag()), V.t()) + + +def coral(source, target): + # assume both source and target are 3D array (C, H, W) + # Note: flatten -> f + + source_f, source_f_mean, source_f_std = _calc_feat_flatten_mean_std(source) + source_f_norm = ( + source_f - source_f_mean.expand_as(source_f) + ) / source_f_std.expand_as(source_f) + source_f_cov_eye = torch.mm(source_f_norm, source_f_norm.t()) + torch.eye(3) + + target_f, target_f_mean, target_f_std = _calc_feat_flatten_mean_std(target) + target_f_norm = ( + target_f - target_f_mean.expand_as(target_f) + ) / target_f_std.expand_as(target_f) + target_f_cov_eye = torch.mm(target_f_norm, target_f_norm.t()) + torch.eye(3) + + source_f_norm_transfer = torch.mm( + _mat_sqrt(target_f_cov_eye), + torch.mm(torch.inverse(_mat_sqrt(source_f_cov_eye)), source_f_norm), + ) + + source_f_transfer = source_f_norm_transfer * target_f_std.expand_as( + source_f_norm + ) + target_f_mean.expand_as(source_f_norm) + + return source_f_transfer.view(source.size()) diff --git a/module_microast/microast_model.py b/module_microast/microast_model.py new file mode 100644 index 0000000..c8b503f --- /dev/null +++ b/module_microast/microast_model.py @@ -0,0 +1,46 @@ +import torch + +from . import net_microAST +import torch.nn.functional as F +from torchvision import transforms + + +def microast_inference( + src_img: torch.Tensor, + style_img: torch.Tensor, + size: int, + do_crop: bool, + alpha: float, + device, + network +): + content_tf = test_transform(size, do_crop) + + content = content_tf(src_img).to(device) + style = content_tf(style_img).to(device) + style = F.interpolate( + style, + size=(content.shape[2], content.shape[3]), + mode="bilinear", + align_corners=False, + ) + + # print(f"Resized: {content.shape=}") + # print(f"Resized: {style.shape=}") + + torch.cuda.synchronize() + output = network(content, style, alpha) + torch.cuda.synchronize() + + return output + + +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 diff --git a/module_microast/net_microAST.py b/module_microast/net_microAST.py new file mode 100644 index 0000000..0265129 --- /dev/null +++ b/module_microast/net_microAST.py @@ -0,0 +1,426 @@ +import torch.nn as nn + +from .function import adaptive_instance_normalization as featMod +from .function import calc_mean_std + + +class ConvLayer(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, groupnum): + super(ConvLayer, self).__init__() + # Padding Layer + padding_size = kernel_size // 2 + self.reflection_pad = nn.ReflectionPad2d(padding_size) + + # Convolution Layer + self.conv_layer = nn.Conv2d( + in_channels, out_channels, kernel_size, stride, groups=groupnum + ) + + def forward(self, x): + x = self.reflection_pad(x) + x = self.conv_layer(x) + return x + + +class ResidualLayer(nn.Module): + def __init__(self, channels=128, kernel_size=3, groupnum=1): + super(ResidualLayer, self).__init__() + self.conv1 = ConvLayer( + channels, channels, kernel_size, stride=1, groupnum=groupnum + ) + self.relu = nn.ReLU(inplace=True) + self.conv2 = ConvLayer( + channels, channels, kernel_size, stride=1, groupnum=groupnum + ) + + def forward(self, x, weight=None, bias=None, filterMod=False): + if filterMod: + x1 = self.conv1(x) + x2 = weight * x1 + bias * x + + x3 = self.relu(x2) + x4 = self.conv2(x3) + x5 = weight * x4 + bias * x3 + return x + x5 + else: + return x + self.conv2(self.relu(self.conv1(x))) + + +# Control the number of channels +slim_factor = 1 + + +class Encoder(nn.Module): + def __init__(self): + super(Encoder, self).__init__() + self.enc1 = nn.Sequential( + ConvLayer(3, int(16 * slim_factor), kernel_size=9, stride=1, groupnum=1), + nn.ReLU(inplace=True), + ConvLayer( + int(16 * slim_factor), + int(32 * slim_factor), + kernel_size=3, + stride=2, + groupnum=int(16 * slim_factor), + ), + nn.ReLU(inplace=True), + ConvLayer( + int(32 * slim_factor), + int(32 * slim_factor), + kernel_size=1, + stride=1, + groupnum=1, + ), + nn.ReLU(inplace=True), + ConvLayer( + int(32 * slim_factor), + int(64 * slim_factor), + kernel_size=3, + stride=2, + groupnum=int(32 * slim_factor), + ), + nn.ReLU(inplace=True), + ConvLayer( + int(64 * slim_factor), + int(64 * slim_factor), + kernel_size=1, + stride=1, + groupnum=1, + ), + nn.ReLU(inplace=True), + ResidualLayer(int(64 * slim_factor), kernel_size=3), + ) + self.enc2 = nn.Sequential(ResidualLayer(int(64 * slim_factor), kernel_size=3)) + + def forward(self, x): + x1 = self.enc1(x) + x2 = self.enc2(x1) + out = [x1, x2] + return out + + +class Decoder(nn.Module): + def __init__(self): + super(Decoder, self).__init__() + self.dec1 = ResidualLayer(int(64 * slim_factor), kernel_size=3) + self.dec2 = ResidualLayer(int(64 * slim_factor), kernel_size=3) + self.dec3 = nn.Sequential( + nn.Upsample(scale_factor=2, mode="bilinear"), + ConvLayer( + int(64 * slim_factor), + int(32 * slim_factor), + kernel_size=3, + stride=1, + groupnum=int(32 * slim_factor), + ), + nn.ReLU(inplace=True), + ConvLayer( + int(32 * slim_factor), + int(32 * slim_factor), + kernel_size=1, + stride=1, + groupnum=1, + ), + nn.ReLU(inplace=True), + nn.Upsample(scale_factor=2, mode="bilinear"), + ConvLayer( + int(32 * slim_factor), + int(16 * slim_factor), + kernel_size=3, + stride=1, + groupnum=int(16 * slim_factor), + ), + nn.ReLU(inplace=True), + ConvLayer( + int(16 * slim_factor), + int(16 * slim_factor), + kernel_size=1, + stride=1, + groupnum=1, + ), + nn.ReLU(inplace=True), + ConvLayer(int(16 * slim_factor), 3, kernel_size=9, stride=1, groupnum=1), + ) + + def forward(self, x, s, w, b, alpha): + x1 = featMod(x[1], s[1]) + x1 = alpha * x1 + (1 - alpha) * x[1] + + x2 = self.dec1(x1, w[1], b[1], filterMod=True) + + x3 = featMod(x2, s[0]) + x3 = alpha * x3 + (1 - alpha) * x2 + + x4 = self.dec2(x3, w[0], b[0], filterMod=True) + + out = self.dec3(x4) + return out + + +class Modulator(nn.Module): + def __init__(self): + super(Modulator, self).__init__() + self.weight1 = nn.Sequential( + ConvLayer( + int(64 * slim_factor), + int(64 * slim_factor), + kernel_size=3, + stride=1, + groupnum=int(64 * slim_factor), + ), + nn.AdaptiveAvgPool2d((1, 1)), + ) + self.bias1 = nn.Sequential( + ConvLayer( + int(64 * slim_factor), + int(64 * slim_factor), + kernel_size=3, + stride=1, + groupnum=int(64 * slim_factor), + ), + nn.AdaptiveAvgPool2d((1, 1)), + ) + self.weight2 = nn.Sequential( + ConvLayer( + int(64 * slim_factor), + int(64 * slim_factor), + kernel_size=3, + stride=1, + groupnum=int(64 * slim_factor), + ), + nn.AdaptiveAvgPool2d((1, 1)), + ) + self.bias2 = nn.Sequential( + ConvLayer( + int(64 * slim_factor), + int(64 * slim_factor), + kernel_size=3, + stride=1, + groupnum=int(64 * slim_factor), + ), + nn.AdaptiveAvgPool2d((1, 1)), + ) + + def forward(self, x): + w1 = self.weight1(x[0]) + b1 = self.bias1(x[0]) + + w2 = self.weight2(x[1]) + b2 = self.bias2(x[1]) + + return [w1, w2], [b1, b2] + + +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 +) + + +class Net(nn.Module): + def __init__(self, vgg, content_encoder, style_encoder, modulator, decoder): + super(Net, self).__init__() + vgg_enc_layers = list(vgg.children()) + self.vgg_enc_1 = nn.Sequential(*vgg_enc_layers[:4]) # input -> relu1_1 + self.vgg_enc_2 = nn.Sequential(*vgg_enc_layers[4:11]) # relu1_1 -> relu2_1 + self.vgg_enc_3 = nn.Sequential(*vgg_enc_layers[11:18]) # relu2_1 -> relu3_1 + self.vgg_enc_4 = nn.Sequential(*vgg_enc_layers[18:31]) # relu3_1 -> relu4_1 + + self.content_encoder = content_encoder + self.style_encoder = style_encoder + self.modulator = modulator + self.decoder = decoder + self.mse_loss = nn.MSELoss() + + # fix the encoder + for name in ["vgg_enc_1", "vgg_enc_2", "vgg_enc_3", "vgg_enc_4"]: + for param in getattr(self, name).parameters(): + param.requires_grad = False + + # extract relu1_1, relu2_1, relu3_1, relu4_1 from input image + def encode_with_vgg_intermediate(self, input): + results = [input] + for i in range(4): + func = getattr(self, "vgg_enc_{:d}".format(i + 1)) + results.append(func(results[-1])) + return results[1:] + + # extract relu4_1 from input image + def encode_vgg_content(self, input): + for i in range(4): + input = getattr(self, "vgg_enc_{:d}".format(i + 1))(input) + return input + + def calc_content_loss(self, input, target): + assert input.size() == target.size() + return self.mse_loss(input, target) + + def calc_style_loss(self, input, target): + assert input.size() == target.size() + 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 forward(self, content, style, alpha=1.0): + assert 0 <= alpha <= 1 + + # extract style modulation signals + style_feats = self.style_encoder(style) + filter_weights, filter_biases = self.modulator(style_feats) + + # extract content features + content_feats = self.content_encoder(content) + + # generate results + res = self.decoder( + content_feats, style_feats, filter_weights, filter_biases, alpha + ) + + # vgg content and style loss + res_feats_vgg = self.encode_with_vgg_intermediate(res) + + style_feats_vgg = self.encode_with_vgg_intermediate(style) + content_feats_vgg = self.encode_vgg_content(content) + + loss_c = self.calc_content_loss(res_feats_vgg[-1], content_feats_vgg) + loss_s = self.calc_style_loss(res_feats_vgg[0], style_feats_vgg[0]) + for i in range(1, 4): + loss_s = loss_s + self.calc_style_loss(res_feats_vgg[i], style_feats_vgg[i]) + + res_style_feats = self.style_encoder(res) + res_filter_weights, res_filter_biases = self.modulator(res_style_feats) + + # style signal contrastive loss + loss_contrastive = 0.0 + for i in range(int(style.size(0))): + pos_loss = 0.0 + neg_loss = 0.0 + + for j in range(int(style.size(0))): + if j == i: + FeatMod_loss = self.calc_style_loss( + res_style_feats[0][i].unsqueeze(0), + style_feats[0][j].unsqueeze(0), + ) + self.calc_style_loss( + res_style_feats[1][i].unsqueeze(0), + style_feats[1][j].unsqueeze(0), + ) + FilterMod_loss = ( + self.calc_content_loss( + res_filter_weights[0][i], filter_weights[0][j] + ) + + self.calc_content_loss( + res_filter_weights[1][i], filter_weights[1][j] + ) + + self.calc_content_loss( + res_filter_biases[0][i], filter_biases[0][j] + ) + + self.calc_content_loss( + res_filter_biases[1][i], filter_biases[1][j] + ) + ) + pos_loss = FeatMod_loss + FilterMod_loss + else: + FeatMod_loss = self.calc_style_loss( + res_style_feats[0][i].unsqueeze(0), + res_style_feats[0][j].unsqueeze(0), + ) + self.calc_style_loss( + res_style_feats[1][i].unsqueeze(0), + style_feats[1][j].unsqueeze(0), + ) + FilterMod_loss = ( + self.calc_content_loss( + res_filter_weights[0][i], filter_weights[0][j] + ) + + self.calc_content_loss( + res_filter_weights[1][i], filter_weights[1][j] + ) + + self.calc_content_loss( + res_filter_biases[0][i], filter_biases[0][j] + ) + + self.calc_content_loss( + res_filter_biases[1][i], filter_biases[1][j] + ) + ) + neg_loss = neg_loss + FeatMod_loss + FilterMod_loss + + loss_contrastive = loss_contrastive + pos_loss / neg_loss + + return res, loss_c, loss_s, loss_contrastive + + +class TestNet(nn.Module): + def __init__(self, content_encoder, style_encoder, modulator, decoder): + super(TestNet, self).__init__() + + self.content_encoder = content_encoder + self.style_encoder = style_encoder + self.modulator = modulator + self.decoder = decoder + + def forward(self, content, style, alpha=1.0): + assert 0 <= alpha <= 1 + + style_feats = self.style_encoder(style) + filter_weights, filter_biases = self.modulator(style_feats) + + content_feats = self.content_encoder(content) + + res = self.decoder( + content_feats, style_feats, filter_weights, filter_biases, alpha + ) + + return res diff --git a/run.py b/run.py index 3068d45..dcab48c 100644 --- a/run.py +++ b/run.py @@ -21,11 +21,17 @@ from .constants import ( EFDM_DEFAULT, EFDM_PATH, EFDM_STYLE_TYPES, + MICROAST_CONTENT_ENCODER_PATH, + MICROAST_DECODER_PATH, + MICROAST_MODULATOR_PATH, + MICROAST_STYLE_ENCODER_PATH, ) from .module_cast import net as net_cast from .module_cast.cast_model import inference_ucast, load_a_ckpt from .module_efdm import net as net_efdm from .module_efdm.efdm_model import inference_efdm +from .module_microast import net_microAST +from .module_microast.microast_model import microast_inference from .module_neural_neighbor.neural_neighbor_model import inference_neural_neighbor base_dir = Path(__file__).resolve().parent @@ -268,7 +274,87 @@ class EFDM: params["src_img"] = src_img[i].unsqueeze(0).permute(0, 3, 1, 2) res_tensor = inference_efdm(**params) result.append(res_tensor.permute(0, 2, 3, 1)) - print(f"{res_tensor.shape=}") + pbar.update_absolute(i, num_frames) + + return (torch.cat(result, dim=0),) + + +class MicroAST: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "src_img": ("IMAGE",), + "style_img": ("IMAGE",), + "do_crop": ( + "BOOLEAN", + {"default": False}, + ), + "size": ( + "INT", + {"default": 512, "min": 1, "step": 1}, + ), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("res_img",) + FUNCTION = "todo" + CATEGORY = "StyleTransferPlus" + + def todo( + self, + src_img: torch.Tensor, + style_img: torch.Tensor, + do_crop: bool, + size: int, + ): + print(f"{src_img.shape=}") + print(f"{style_img.shape=}") + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + 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}") + ) + 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), + "device": device, + "alpha": 1.0, + "size": size, + "do_crop": do_crop, + "network": network, + } + + num_frames = src_img.size(0) + pbar = ProgressBar(num_frames) + + 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 = microast_inference(**params) + result.append(res_tensor.permute(0, 2, 3, 1)) pbar.update_absolute(i, num_frames) return (torch.cat(result, dim=0),) diff --git a/workflows/wf_microast.png b/workflows/wf_microast.png new file mode 100644 index 0000000..fef0955 Binary files /dev/null and b/workflows/wf_microast.png differ diff --git a/workflows/workflow_microast.json b/workflows/workflow_microast.json new file mode 100644 index 0000000..f9e9a6d --- /dev/null +++ b/workflows/workflow_microast.json @@ -0,0 +1,193 @@ +{ + "last_node_id": 7, + "last_link_id": 15, + "nodes": [ + { + "id": 4, + "type": "SaveImage", + "pos": [ + 1050, + 327 + ], + "size": { + "0": 315, + "1": 270 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 15 + } + ], + "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": [ + 13 + ], + "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": [ + 14 + ], + "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": 7, + "type": "MicroAST", + "pos": [ + 680, + 292 + ], + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "src_img", + "type": "IMAGE", + "link": 13 + }, + { + "name": "style_img", + "type": "IMAGE", + "link": 14 + } + ], + "outputs": [ + { + "name": "res_img", + "type": "IMAGE", + "links": [ + 15 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "MicroAST" + }, + "widgets_values": [ + false, + 512 + ] + } + ], + "links": [ + [ + 13, + 2, + 0, + 7, + 0, + "IMAGE" + ], + [ + 14, + 3, + 0, + 7, + 1, + "IMAGE" + ], + [ + 15, + 7, + 0, + 4, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.1000000000000008, + "offset": [ + -166.99103191074192, + -98.00415725314276 + ] + } + }, + "version": 0.4 +} \ No newline at end of file