378 lines
11 KiB
Python
378 lines
11 KiB
Python
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:
|
|
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
|