add microast node

2024-08-13 14:07:33 +07:00
parent 8d1ea2600f
commit d0caadfb31
12 changed files with 857 additions and 9 deletions
@@ -1,5 +1,5 @@
 *.pth
-
+*.tar

 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -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

@@ -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",
 }


@@ -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"
+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"
@@ -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()
@@ -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())
@@ -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
@@ -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
@@ -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),)
@@ -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
+}