unist works, but video only 3 frames

.gitignore

@@ -1,5 +1,6 @@
 *.pth
 *.tar
+*.pt
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -7,6 +7,7 @@ Click name to jump to workflow
 2. [**CAST**](#cast). Paper: [Domain Enhanced Arbitrary Image Style Transfer via Contrastive Learning](https://github.com/zyxElsa/CAST_pytorch)
 3. [**EFDM**](#efdm). Paper: [Exact Feature Distribution Matching for Arbitrary Style Transfer and Domain Generalization](https://github.com/YBZh/EFDM)
 4. [**MicroAST**](#microast). Paper: [Towards Super-Fast Ultra-Resolution Arbitrary Style Transfer](https://github.com/EndyWon/MicroAST)
+5. [**UniST**](#unist) : [Two Birds, One Stone: A Unified Framework for Joint Learning of Image and Video Style Transfers](https://github.com/NevSNev/UniST)
 
 [**Extra nodes**](#extra-nodes)
 
@@ -16,11 +17,13 @@ More will come soon
 
 All nodes support batched input (i.e video) but is generally not recommended. One should generate 1 or 2 style frames (start and end), then use [ComfyUI-EbSynth](https://github.com/FuouM/ComfyUI-EbSynth) to propagate the style to the entire video.
 
-| Neural Neighbor (2020-2022) - Slowest | CAST (2022-2023) - Fast |
-|:-:|:-:|
-|![neural_neighbor](example_outputs/neural_neighbor.png) | ![cast](example_outputs/cast.png) |
-| **EFDM (2022) - Fast** | **MicroAST (2023) - Fast** | 
-| ![efdm](example_outputs/efdm.png) | ![microast](example_outputs/microast.png) |
+|          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) - Fast**         |
+|            ![efdm](example_outputs/efdm.png)            | ![microast](example_outputs/microast.png) |
+|                    **UniST (2023)**                     |                   **X**                   |
+|                            X                            |                     X                     |
 
 ### Neural Neighbor
 
@@ -106,6 +109,17 @@ Arguments: Similar to [EFDM](#efdm)
 ![wf_microast](workflows/wf_microast.png)
 
 
+### UniST
+
+Download `UniST_model.pt` in [Testing](https://github.com/NevSNev/UniST?tab=readme-ov-file#testing) and `vgg_r41.pth, dec_r41.pth` in [Training](https://github.com/NevSNev/UniST?tab=readme-ov-file#training) section and place them in
+
+```
+models/unist/UniST_model.pt
+models/unist/dec_r41.pth
+models/unist/vgg_r41.pth
+```
+
+
 ## Extra nodes
 
 ### Coral Color Transfer
@@ -147,3 +161,14 @@ year = {2022}}
   year={2023}
 }
 ```
+
+```
+@InProceedings{Gu_2023_ICCV,
+    author    = {Gu, Bohai and Fan, Heng and Zhang, Libo},
+    title     = {Two Birds, One Stone: A Unified Framework for Joint Learning of Image and Video Style Transfers},
+    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
+    month     = {October},
+    year      = {2023},
+    pages     = {23545-23554}
+}
+```

__init__.py

@@ -1,4 +1,4 @@
-from .run import NeuralNeighbor, CAST, EFDM, MicroAST
+from .run import NeuralNeighbor, CAST, EFDM, MicroAST, UniST, UniST_Video
 from .run_extra import CoralColorTransfer
 
 NODE_CLASS_MAPPINGS = {
@@ -7,6 +7,8 @@ NODE_CLASS_MAPPINGS = {
     "EFDM": EFDM,
     "MicroAST": MicroAST,
     "CoralColorTransfer": CoralColorTransfer,
+    "UniST": UniST,
+    "UniST_Video": UniST_Video,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -15,6 +17,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "EFDM": "EFDM",
     "MicroAST": "MicroAST",
     "CoralColorTransfer": "Coral Color Transfer",
+    "UniST": "UniST",
+    "UniST_Video": "UniST Video",
 }
 
 

constants.py

@@ -12,4 +12,8 @@ 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"
+MICROAST_MODULATOR_PATH = "models/microast/modulator_iter_160000.pth.tar"
+
+UNIST_PATH = "models/unist/UniST_model.pt"
+UNIST_DEC_PATH = "models/unist/dec_r41.pth"
+UNIST_ENC_PATH = "models/unist/vgg_r41.pth"

module_microast/microast_model.py

@@ -1,11 +1,10 @@
 import torch
 
-from . import net_microAST
 import torch.nn.functional as F
 from torchvision import transforms
 
 
-def microast_inference(
+def inference_microast(
     src_img: torch.Tensor,
     style_img: torch.Tensor,
     size: int,

module_unist/model.py

@@ -0,0 +1,37 @@
+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

module_unist/models.py

@@ -0,0 +1,378 @@
+import math
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from .transformer_layers import DecoderLayer_style, EncoderLayer, EncoderLayer_cross
+
+
+class Encoder_vgg(nn.Module):
+    def __init__(self):
+        super(Encoder_vgg, self).__init__()
+        # vgg
+        # 256 x 256
+        self.conv1 = nn.Conv2d(3, 3, 1, 1, 0)
+        self.reflecPad1 = nn.ReflectionPad2d((1, 1, 1, 1))
+        # 256 x 256
+
+        self.conv2 = nn.Conv2d(3, 64, 3, 1, 0)
+        self.relu2 = nn.ReLU(inplace=True)
+        # 256 x 256
+
+        self.reflecPad3 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv3 = nn.Conv2d(64, 64, 3, 1, 0)
+        self.relu3 = nn.ReLU(inplace=True)
+        # 256 x 256
+
+        self.maxPool = nn.MaxPool2d(kernel_size=2, stride=2)
+        # 128 x 128
+
+        self.reflecPad4 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv4 = nn.Conv2d(64, 128, 3, 1, 0)
+        self.relu4 = nn.ReLU(inplace=True)
+        # 128 x 128
+
+        self.reflecPad5 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv5 = nn.Conv2d(128, 128, 3, 1, 0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 128 x 128 x 128
+
+        self.maxPool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        # 64 x 64
+
+        self.reflecPad6 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv6 = nn.Conv2d(128, 256, 3, 1, 0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 64 x 64
+
+        self.reflecPad7 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv7 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.relu7 = nn.ReLU(inplace=True)
+        # 64 x 64
+
+        self.reflecPad8 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv8 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.relu8 = nn.ReLU(inplace=True)
+        # 64 x 64
+
+        self.reflecPad9 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv9 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.relu9 = nn.ReLU(inplace=True)
+        # 64 x 64
+
+        self.maxPool3 = nn.MaxPool2d(kernel_size=2, stride=2)
+        # 32 x 32
+
+        self.reflecPad10 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv10 = nn.Conv2d(256, 512, 3, 1, 0)
+        self.relu10 = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.reflecPad1(x)
+        x = self.relu2(self.conv2(x))
+        x = self.reflecPad3(x)
+        x = self.relu3(self.conv3(x))
+
+        x = self.maxPool(x)
+
+        x = self.reflecPad4(x)
+        x = self.relu4(self.conv4(x))
+        x = self.reflecPad5(x)
+        x = self.relu5(self.conv5(x))
+
+        x = self.maxPool2(x)
+
+        x = self.reflecPad6(x)
+        x = self.relu6(self.conv6(x))
+        x = self.reflecPad7(x)
+        x = self.relu7(self.conv7(x))
+        x = self.reflecPad8(x)
+        x = self.relu8(self.conv8(x))
+        x = self.reflecPad9(x)
+        x = self.relu9(self.conv9(x))
+
+        x = self.maxPool3(x)
+
+        x = self.reflecPad10(x)
+        x = self.relu10(self.conv10(x))
+
+        return x
+
+
+class Decoder_vgg(nn.Module):
+    def __init__(self):
+        super(Decoder_vgg, self).__init__()
+        # decoder
+        self.reflecPad11 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv11 = nn.Conv2d(512, 256, 3, 1, 0)
+        self.relu11 = nn.ReLU(inplace=True)
+        # 32 x 32
+
+        self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
+        # 64 x 64
+
+        self.reflecPad12 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv12 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.relu12 = nn.ReLU(inplace=True)
+        # 64 x 64
+
+        self.reflecPad13 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv13 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.relu13 = nn.ReLU(inplace=True)
+        # 64 x 64
+
+        self.reflecPad14 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv14 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.relu14 = nn.ReLU(inplace=True)
+        # 64 x 64
+
+        self.reflecPad15 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv15 = nn.Conv2d(256, 128, 3, 1, 0)
+        self.relu15 = nn.ReLU(inplace=True)
+        # 64 x 64
+
+        self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 128 x 128
+
+        self.reflecPad16 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv16 = nn.Conv2d(128, 128, 3, 1, 0)
+        self.relu16 = nn.ReLU(inplace=True)
+        # 128 x 128
+
+        self.reflecPad17 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv17 = nn.Conv2d(128, 64, 3, 1, 0)
+        self.relu17 = nn.ReLU(inplace=True)
+        # 128 x 128
+
+        self.unpool3 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 256 x 256
+
+        self.reflecPad18 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv18 = nn.Conv2d(64, 64, 3, 1, 0)
+        self.relu18 = nn.ReLU(inplace=True)
+        # 256 x 256
+
+        self.reflecPad19 = nn.ReflectionPad2d((1, 1, 1, 1))
+        self.conv19 = nn.Conv2d(64, 3, 3, 1, 0)
+
+    def forward(self, x):
+        # 128x128x128
+        out = self.reflecPad11(x)
+        out = self.conv11(out)
+        out = self.relu11(out)
+
+        out = self.unpool(out)
+
+        out = self.reflecPad12(out)
+        out = self.conv12(out)
+        out = self.relu12(out)
+        out = self.reflecPad13(out)
+        out = self.conv13(out)
+        out = self.relu13(out)
+        out = self.reflecPad14(out)
+        out = self.conv14(out)
+        out = self.relu14(out)
+        out = self.reflecPad15(out)
+        out = self.conv15(out)
+        out = self.relu15(out)
+        out = self.unpool2(out)
+        out = self.reflecPad16(out)
+        out = self.conv16(out)
+        out = self.relu16(out)
+        out = self.reflecPad17(out)
+        out = self.conv17(out)
+        out = self.relu17(out)
+        out = self.unpool3(out)
+        out = self.reflecPad18(out)
+        out = self.conv18(out)
+        out = self.relu18(out)
+        out = self.reflecPad19(out)
+        out = self.conv19(out)
+        return out
+
+
+class Encoder_video(nn.Module):
+    """For Content Enc."""
+
+    def __init__(self, n_layers, n_head, d_k, d_v, d_model, d_inner, dropout=0.1):
+        super().__init__()
+        self.layer_stack = nn.ModuleList(
+            [
+                EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
+                for _ in range(n_layers)
+            ]
+        )
+
+    def forward(self, enc_output):
+        for enc_layer in self.layer_stack:
+            enc_output = enc_layer(enc_output)
+        return enc_output
+
+
+class Encoder_image(nn.Module):
+    """For Style Enc."""
+
+    def __init__(self, n_layers, n_head, d_k, d_v, d_model, d_inner, dropout=0.1):
+        super().__init__()
+        self.layer_stack = nn.ModuleList(
+            [
+                EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
+                for _ in range(n_layers)
+            ]
+        )
+
+    def forward(self, enc_output):
+        for enc_layer in self.layer_stack:
+            enc_output = enc_layer(enc_output)
+        return enc_output
+
+
+class Encoder_cross(nn.Module):
+    def __init__(self, n_layers, n_head, d_k, d_v, d_model, d_inner, dropout=0.1):
+        super().__init__()
+        self.layer_stack = nn.ModuleList(
+            [
+                EncoderLayer_cross(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
+                for _ in range(n_layers)
+            ]
+        )
+
+    def forward(self, dec_output, enc_output):
+        print(f"{dec_output.shape=}")
+        print(f"{enc_output.shape=}")
+        for enc_layer in self.layer_stack:
+            enc_output = enc_layer(dec_output, enc_output)
+        return enc_output
+
+
+class Decoder_style_tranfer(nn.Module):
+    def __init__(self, n_layers, n_head, d_k, d_v, d_model, d_inner, dropout=0.1):
+        super().__init__()
+
+        self.layer_stack = nn.ModuleList(
+            [
+                DecoderLayer_style(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
+                for _ in range(n_layers)
+            ]
+        )
+
+    def forward(self, dec_output, enc_output):
+        print(f"{dec_output.shape=}")
+        print(f"{enc_output.shape=}")
+        for dec_layer in self.layer_stack:
+            # print(dec_layer)
+            dec_output = dec_layer(dec_output, enc_output)
+
+        _, l, _ = dec_output.size()
+        dec_output = rearrange(
+            dec_output, "b (m1 m2) c -> b c m1 m2", m1=int(math.sqrt(l))
+        )
+
+        return dec_output
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        encoder_path: str,
+        decoder_path: str,
+        d_model=512,
+        d_inner=4096,
+        n_layers=3,
+        n_head=8,
+        d_k=64,
+        d_v=64,
+        dropout=0.1,
+    ):
+        super().__init__()
+
+        self.encoder_frames = Encoder_video(
+            d_model=d_model,
+            d_inner=d_inner,
+            n_layers=2,
+            n_head=n_head,
+            d_k=d_k,
+            d_v=d_v,
+            dropout=dropout,
+        )
+
+        self.encoder_style = Encoder_image(
+            d_model=d_model,
+            d_inner=d_inner,
+            n_layers=1,
+            n_head=n_head,
+            d_k=d_k,
+            d_v=d_v,
+            dropout=dropout,
+        )
+
+        self.image_cross_video = Encoder_cross(
+            d_model=d_model,
+            d_inner=d_inner,
+            n_layers=1,
+            n_head=n_head,
+            d_k=d_k,
+            d_v=d_v,
+            dropout=dropout,
+        )
+
+        self.video_cross_image = Encoder_cross(
+            d_model=d_model,
+            d_inner=d_inner,
+            n_layers=1,
+            n_head=n_head,
+            d_k=d_k,
+            d_v=d_v,
+            dropout=dropout,
+        )
+
+        self.style_decoder = Decoder_style_tranfer(
+            d_model=d_model,
+            d_inner=d_inner,
+            n_layers=n_layers,
+            n_head=n_head,
+            d_k=d_k,
+            d_v=d_v,
+            dropout=dropout,
+        )
+
+        self.encoder = Encoder_vgg()
+
+        self.decoder = Decoder_vgg()
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+        print("Loading pretrained AutoEncoder")
+        encoder_ckpt = torch.load(encoder_path, map_location="cpu")
+        decoder_ckpt = torch.load(decoder_path, map_location="cpu")
+        self.encoder.load_state_dict(encoder_ckpt, strict=True)
+        self.decoder.load_state_dict(decoder_ckpt, strict=True)
+
+    def forward(
+        self,
+        content,
+        style_images,
+        content_type="image",
+        id_loss="transfer",
+        tab=None,
+    ):
+        if id_loss == "transfer":
+            b, t, c, h, w = content.size()
+            content = content.view(b * t, c, h, w)
+            enc_content = self.encoder(content)
+            enc_content = rearrange(enc_content, "bt c h w ->bt (h w) c")
+            enc_content = self.encoder_frames(enc_content)
+            if content_type == "image":
+                content_self_attn = self.image_cross_video(enc_content, enc_content)
+            elif content_type == "video":
+                content_self_attn = self.video_cross_image(enc_content, enc_content)
+            enc_style = self.encoder.forward(style_images)
+            enc_style = rearrange(enc_style, "b c h w ->b (h w) c")
+            enc_style_image = self.encoder_style.forward(enc_style)
+            dec_output = self.style_decoder.forward(content_self_attn, enc_style_image)
+            style_result = self.decoder.forward(dec_output)
+            return style_result

module_unist/transformer_layers.py

@@ -0,0 +1,50 @@
+import torch.nn as nn
+
+from .transformer_subLayers import MultiHeadAttention, PositionwiseFeedForward
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
+        super(EncoderLayer, self).__init__()
+        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
+        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)
+
+    def forward(self, enc_input, slf_attn_mask=None):
+        enc_output, enc_slf_attn = self.slf_attn(
+            enc_input, enc_input, enc_input, mask=slf_attn_mask
+        )
+        enc_output = self.pos_ffn(enc_output)
+        return enc_output
+
+
+class EncoderLayer_cross(nn.Module):
+    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
+        super(EncoderLayer_cross, self).__init__()
+        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
+        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)
+
+    def forward(self, dec_input, enc_input, slf_attn_mask=None):
+        enc_output, enc_slf_attn = self.slf_attn(
+            dec_input, enc_input, enc_input, mask=slf_attn_mask
+        )
+        enc_output = self.pos_ffn(enc_output)
+        return enc_output
+
+
+class DecoderLayer_style(nn.Module):
+    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
+        super(DecoderLayer_style, self).__init__()
+        self.enc_attn1 = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
+        self.enc_attn2 = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
+        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)
+
+    # style decoder no need mask
+    def forward(self, dec_input, enc_input):
+        dec_output, dec_slf_attn = self.enc_attn1(
+            dec_input, enc_input, enc_input, mask=None
+        )
+        dec_output, dec_enc_attn = self.enc_attn2(
+            dec_output, enc_input, enc_input, mask=None
+        )
+        dec_output = self.pos_ffn(dec_output)
+        return dec_output

module_unist/transformer_modules.py

@@ -0,0 +1,36 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__author__ = "Yu-Hsiang Huang"
+
+
+class ScaledDotProductAttention(nn.Module):
+    """Scaled Dot-Product Attention"""
+
+    def __init__(self, temperature, attn_dropout=0.1):
+        super().__init__()
+        self.temperature = temperature
+        self.dropout = nn.Dropout(attn_dropout)
+
+    def forward(self, q, k, v, mask=None):
+        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
+
+        if mask is not None:
+            attn = attn.masked_fill(mask == 0, -1e9)
+
+        attn = self.dropout(F.softmax(attn, dim=-1))
+        output = torch.matmul(attn, v)
+
+        return output, attn
+
+
+class Attention(nn.Module):
+    def forward(self, query, key, value):
+        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query.size(-1))
+
+        p_attn = F.softmax(scores, dim=-1)
+        p_val = torch.matmul(p_attn, value)
+        return p_val, p_attn

module_unist/transformer_subLayers.py

@@ -0,0 +1,131 @@
+import math
+
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from .transformer_modules import ScaledDotProductAttention
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1, granular=True):
+        super().__init__()
+        self.n_head = n_head
+        self.d_k = d_k
+        self.d_v = d_v
+        self.query_embedding = nn.Conv2d(512, 512, kernel_size=1, padding=0)
+        self.value_embedding = nn.Conv2d(512, 512, kernel_size=1, padding=0)
+        self.key_embedding = nn.Conv2d(512, 512, kernel_size=1, padding=0)
+        
+        self.granular = granular
+
+        # option b
+        self.w_qs_h = nn.Linear(d_model, n_head * d_k, bias=False)
+        self.w_ks_h = nn.Linear(d_model, n_head * d_k, bias=False)
+        self.w_vs_h = nn.Linear(d_model, n_head * d_v, bias=False)
+        self.fc_h = nn.Linear(n_head * d_v, d_model, bias=False)
+
+        self.attention_h = ScaledDotProductAttention(temperature=d_k**0.5)
+
+        self.w_qs_w = nn.Linear(d_model, n_head * d_k, bias=False)
+        self.w_ks_w = nn.Linear(d_model, n_head * d_k, bias=False)
+        self.w_vs_w = nn.Linear(d_model, n_head * d_v, bias=False)
+        self.fc_w = nn.Linear(n_head * d_v, d_model, bias=False)
+
+        self.attention_w = ScaledDotProductAttention(temperature=d_k**0.5)
+
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.layer_norm1 = nn.LayerNorm(d_model, eps=1e-6)
+        self.layer_norm2 = nn.LayerNorm(d_model, eps=1e-6)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, q, k, v, mask=None):
+        residual = q
+        bt, l, _ = q.size()  # noqa: E741
+
+        q = rearrange(q, "b (m1 m2) c -> b c m1 m2", m1=int(math.sqrt(l)))
+        k = rearrange(k, "b (m1 m2) c -> b c m1 m2", m1=int(math.sqrt(l)))
+        v = rearrange(v, "b (m1 m2) c -> b c m1 m2", m1=int(math.sqrt(l)))
+
+        q = self.query_embedding(q)
+        k_prev = self.key_embedding(k)
+        v = self.value_embedding(v)
+
+        q = rearrange(q, "b c m1 m2 -> (b m1) m2 c")
+        k = rearrange(k_prev, "b c m1 m2 -> (b m1) m2 c")
+        v = rearrange(v, "b c m1 m2 -> (b m1) m2 c")
+
+        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
+        q = self.w_qs_h(q)
+        k = self.w_ks_h(k)
+        v = self.w_vs_h(v)
+
+        q = rearrange(q, "b len (head d_k) -> b len head d_k", d_k=d_k, head=n_head)
+        k = rearrange(k, "b len (head d_k) -> b len head d_k", d_k=d_k, head=n_head)
+        v = rearrange(v, "b len (head d_v) -> b len head d_v", d_v=d_v, head=n_head)
+
+        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
+
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+
+        q, attn = self.attention_h(q, k, v, mask=mask)
+
+        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
+        q = self.dropout1(self.fc_h(q))
+        q = rearrange(q, "(bt h) w c ->bt (h w) c", h=int(math.sqrt(l)))
+        q_mid = self.layer_norm1(q)
+
+        q = rearrange(q_mid, "bt (h w) c ->(bt w) h c", h=int(math.sqrt(l)))
+        if self.granular:
+            k = rearrange(q_mid,'bt (h w) c ->(bt w) h c',h =int(math.sqrt(l)))
+        else:
+            k = rearrange(k_prev, "b c m1 m2 -> (b m2) m1 c")
+        v = rearrange(q_mid, "bt (h w) c ->(bt w) h c", h=int(math.sqrt(l)))
+
+        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
+        q = self.w_qs_w(q)
+        k = self.w_ks_w(k)
+        v = self.w_vs_w(v)
+
+        q = rearrange(q, "b len (head d_k) -> b len head d_k", d_k=d_k, head=n_head)
+        k = rearrange(k, "b len (head d_k) -> b len head d_k", d_k=d_k, head=n_head)
+        v = rearrange(v, "b len (head d_v) -> b len head d_v", d_v=d_v, head=n_head)
+
+        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)  # noqa: F841
+
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+
+        q, attn = self.attention_w(q, k, v, mask=mask)
+        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
+        q = self.dropout2(self.fc_w(q))
+        q = rearrange(q, "(bt w) h c -> bt (h w) c", w=int(math.sqrt(l)))
+        q += residual
+        q = self.layer_norm2(q)
+
+        return q, attn
+
+
+class PositionwiseFeedForward(nn.Module):
+    def __init__(self, d_in, d_hid, dropout=0.1):
+        super().__init__()
+        self.w_1 = nn.Linear(d_in, d_hid)
+        self.w_2 = nn.Linear(d_hid, d_in)
+        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        residual = x
+        x = self.w_2(F.relu(self.w_1(x)))
+        x = self.dropout(x)
+        x += residual
+
+        x = self.layer_norm(x)
+
+        return x

module_unist/unist_model.py

@@ -0,0 +1,81 @@
+import numpy as np
+import torch
+import torchvision.transforms as transforms
+
+from .model import video_Style_transfer
+
+
+def inference_unist(
+    src: torch.Tensor,
+    src_style: torch.Tensor,
+    device,
+    size: int,
+    do_crop: bool,
+    network: video_Style_transfer,
+    content_type="image",
+):
+    content_tf = test_transform(size, do_crop)
+    content = content_tf(src).unsqueeze(0).to(device)  # [1, B, C, H, W]
+
+    style = content_tf(src_style).to(device)  # [B, C, H, W]
+
+    if content_type == "video":
+        style = match_shape(style, content[0])
+
+    print(f"{content.shape=}")
+    print(f"{style.shape=}")
+
+    y_hat = network.forward(
+        content, style, content_type, tab="inference"
+    )  # [B, C, H, W]
+    print(f"{y_hat.shape=}")
+    output = un_normalize_batch(y_hat)  # [B, C, H, W]
+
+    return output
+
+
+def test_transform(size, crop):
+    transform_list = []
+
+    if crop and size != 0:
+        transform_list.append(transforms.Resize(size))
+        transform_list.append(transforms.CenterCrop(size))
+    elif size != 0:
+        transform_list.append(transforms.Resize((size, size)))
+
+    transform_list.append(transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]))
+
+    transform = transforms.Compose(transform_list)
+    # [B, C, H, W]
+    return transform
+
+
+def un_normalize_batch(tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
+    for t, m, s in zip(tensor, mean, std):
+        t.mul_(s).add_(m)
+    return tensor
+
+
+def match_shape(a: torch.Tensor, b: torch.Tensor):
+    B_a, C, H, W = a.shape
+    B_b, _, _, _ = b.shape
+
+    if B_a == 1 and B_b > 1:
+        # Expand a along the batch dimension to match the shape of b
+        expanded_a = a.expand(B_b, -1, -1, -1)
+    elif B_a > 1 and B_a < B_b:
+        # Repeat a along the batch dimension to match the shape of b
+        repeated_a = a.repeat(B_b // B_a, 1, 1, 1)
+        if B_b % B_a != 0:
+            remaining = B_b % B_a
+            repeated_a = torch.cat([repeated_a, a[-remaining:]], dim=0)
+        expanded_a = repeated_a
+    elif B_a > B_b:
+        # Trim a along the batch dimension to match the shape of b
+        trimmed_a = a[:B_b]
+        expanded_a = trimmed_a
+    else:
+        # If B_a == B_b, no change is needed
+        expanded_a = a
+
+    return expanded_a

requirements.txt

@@ -1,3 +1,4 @@
 torch
 torchvision
-scikit-image
+scikit-image
+einops

run.py

@@ -25,14 +25,19 @@ from .constants import (
     MICROAST_DECODER_PATH,
     MICROAST_MODULATOR_PATH,
     MICROAST_STYLE_ENCODER_PATH,
+    UNIST_DEC_PATH,
+    UNIST_ENC_PATH,
+    UNIST_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_microast.microast_model import inference_microast
 from .module_neural_neighbor.neural_neighbor_model import inference_neural_neighbor
+from .module_unist.model import video_Style_transfer
+from .module_unist.unist_model import inference_unist
 
 base_dir = Path(__file__).resolve().parent
 
@@ -353,13 +358,140 @@ class MicroAST:
         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)
+                res_tensor = inference_microast(**params)
                 result.append(res_tensor.permute(0, 2, 3, 1))
                 pbar.update_absolute(i, num_frames)
 
         return (torch.cat(result, dim=0),)
 
 
+class UniST:
+    @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 = ("out_img",)
+    FUNCTION = "todo"
+    CATEGORY = "StyleTransferPlus"
+
+    def todo(
+        self,
+        src_img: torch.Tensor,
+        style_img: torch.Tensor,
+        do_crop: bool,
+        size: int,
+        src_images: torch.Tensor | None = None,
+    ):
+        print(f"{src_img.shape=}")
+        print(f"{style_img.shape=}")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        network = video_Style_transfer(
+            ckpt_path=f"{base_dir}/{UNIST_PATH}",
+            encoder_path=f"{base_dir}/{UNIST_ENC_PATH}",
+            decoder_path=f"{base_dir}/{UNIST_DEC_PATH}",
+        )
+
+        network.eval()
+        network.to(device)
+
+        params = {
+            "src_style": style_img.permute(0, 3, 1, 2),
+            "device": device,
+            "size": size,
+            "do_crop": do_crop,
+            "network": network,
+            "content_type": "image",
+        }
+
+        num_frames = src_img.size(0)
+        pbar = ProgressBar(num_frames)
+
+        result: list[torch.Tensor] = []
+        with torch.no_grad():
+            for i in tqdm.tqdm(range(num_frames), desc="Generating"):
+                params["src"] = src_img[i].unsqueeze(0).permute(0, 3, 1, 2)
+                res_tensor = inference_unist(**params)
+                result.append(res_tensor.permute(0, 2, 3, 1))
+                pbar.update_absolute(i, num_frames)
+
+        return (torch.cat(result, dim=0),)
+
+
+class UniST_Video:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_video": ("IMAGE",),
+                "style_video": ("IMAGE",),
+                "do_crop": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+                "size": (
+                    "INT",
+                    {"default": 512, "min": 1, "step": 1},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("out_img",)
+    FUNCTION = "todo"
+    CATEGORY = "StyleTransferPlus"
+
+    def todo(
+        self,
+        src_video: torch.Tensor,
+        style_video: torch.Tensor,
+        do_crop: bool,
+        size: int,
+        src_images: torch.Tensor | None = None,
+    ):
+        print(f"{src_video.shape=}")
+        print(f"{style_video.shape=}")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        network = video_Style_transfer(
+            ckpt_path=f"{base_dir}/{UNIST_PATH}",
+            encoder_path=f"{base_dir}/{UNIST_ENC_PATH}",
+            decoder_path=f"{base_dir}/{UNIST_DEC_PATH}",
+        )
+
+        network.eval()
+        network.to(device)
+
+        params = {
+            "src": src_video.permute(0, 3, 1, 2),
+            "src_style": style_video.permute(0, 3, 1, 2),
+            "device": device,
+            "size": size,
+            "do_crop": do_crop,
+            "network": network,
+            "content_type": "video",
+        }
+        
+        with torch.no_grad():
+            res_tensor = inference_unist(**params).permute(0, 2, 3, 1)
+
+        return (res_tensor, )
+
+
 def serialize_floats(lst: list[float]):
     return ",".join(map(str, lst))