From 40c8f35dc886ce6bef2467c37acb80c0cdc7843e Mon Sep 17 00:00:00 2001 From: FuouM Date: Wed, 14 Aug 2024 00:32:46 +0700 Subject: [PATCH] unist works, but video only 3 frames --- .gitignore | 1 + README.md | 35 ++- __init__.py | 6 +- constants.py | 6 +- module_microast/microast_model.py | 3 +- module_unist/model.py | 37 +++ module_unist/models.py | 378 ++++++++++++++++++++++++++ module_unist/transformer_layers.py | 50 ++++ module_unist/transformer_modules.py | 36 +++ module_unist/transformer_subLayers.py | 131 +++++++++ module_unist/unist_model.py | 81 ++++++ requirements.txt | 3 +- run.py | 136 ++++++++- 13 files changed, 891 insertions(+), 12 deletions(-) create mode 100644 module_unist/model.py create mode 100644 module_unist/models.py create mode 100644 module_unist/transformer_layers.py create mode 100644 module_unist/transformer_modules.py create mode 100644 module_unist/transformer_subLayers.py create mode 100644 module_unist/unist_model.py diff --git a/.gitignore b/.gitignore index a5fd31f..2edc8a5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ *.pth *.tar +*.pt # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index 3cd1210..61877aa 100644 --- a/README.md +++ b/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} +} +``` \ No newline at end of file diff --git a/__init__.py b/__init__.py index 0ac2141..ae4035b 100644 --- a/__init__.py +++ b/__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", } diff --git a/constants.py b/constants.py index 12daafa..78e6eba 100644 --- a/constants.py +++ b/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" \ No newline at end of file +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" diff --git a/module_microast/microast_model.py b/module_microast/microast_model.py index c8b503f..4cbc85d 100644 --- a/module_microast/microast_model.py +++ b/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, diff --git a/module_unist/model.py b/module_unist/model.py new file mode 100644 index 0000000..73c1712 --- /dev/null +++ b/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 diff --git a/module_unist/models.py b/module_unist/models.py new file mode 100644 index 0000000..38fe54c --- /dev/null +++ b/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 diff --git a/module_unist/transformer_layers.py b/module_unist/transformer_layers.py new file mode 100644 index 0000000..3349b91 --- /dev/null +++ b/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 diff --git a/module_unist/transformer_modules.py b/module_unist/transformer_modules.py new file mode 100644 index 0000000..4514d06 --- /dev/null +++ b/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 diff --git a/module_unist/transformer_subLayers.py b/module_unist/transformer_subLayers.py new file mode 100644 index 0000000..8163957 --- /dev/null +++ b/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 diff --git a/module_unist/unist_model.py b/module_unist/unist_model.py new file mode 100644 index 0000000..f1c8e7a --- /dev/null +++ b/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 diff --git a/requirements.txt b/requirements.txt index 169f15f..9004e51 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ torch torchvision -scikit-image \ No newline at end of file +scikit-image +einops \ No newline at end of file diff --git a/run.py b/run.py index 99b2d0b..67b44be 100644 --- a/run.py +++ b/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))