diff --git a/animegan_style_transfer.py b/animegan_style_transfer.py new file mode 100644 index 0000000..d0e7772 --- /dev/null +++ b/animegan_style_transfer.py @@ -0,0 +1,36 @@ +import os +import numpy +from pathlib import Path +from typing import NamedTuple +from torchvision import transforms + +from towhee.operator import Operator +from towhee.utils.pil_utils import to_pil +from towhee.types.image import Image + +import warnings +warnings.filterwarnings("ignore") + +class AnimeganStyleTransfer(Operator): + """ + PyTorch model for image embedding. + """ + def __init__(self, model_name: str, framework: str = 'pytorch') -> None: + super().__init__() + if framework == 'pytorch': + import importlib.util + path = os.path.join(str(Path(__file__).parent), 'pytorch', 'model.py') + opname = os.path.basename(str(Path(__file__))).split('.')[0] + spec = importlib.util.spec_from_file_location(opname, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + self.model = module.Model(model_name) + self.tfms = transforms.Compose([ + transforms.ToTensor() + ]) + + def __call__(self, image: 'towhee.types.Image') -> NamedTuple('Outputs', [('styled_image', numpy.ndarray)]): + img = self.tfms(to_pil(image)).unsqueeze(0) + styled_image = self.model(img) + Outputs = NamedTuple('Outputs', [('styled_image', numpy.ndarray)]) + return Outputs(styled_image) \ No newline at end of file diff --git a/pytorch/__init__.py b/pytorch/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pytorch/model.py b/pytorch/model.py new file mode 100644 index 0000000..361ecc7 --- /dev/null +++ b/pytorch/model.py @@ -0,0 +1,133 @@ +from torch import nn, load, Tensor +import os +from pathlib import Path + + +class ConvNormLReLU(nn.Sequential): + def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1, pad_mode="reflect", groups=1, bias=False): + + pad_layer = { + "zero": nn.ZeroPad2d, + "same": nn.ReplicationPad2d, + "reflect": nn.ReflectionPad2d, + } + if pad_mode not in pad_layer: + raise NotImplementedError + + super(ConvNormLReLU, self).__init__( + pad_layer[pad_mode](padding), + nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, stride=stride, padding=0, groups=groups, bias=bias), + nn.GroupNorm(num_groups=1, num_channels=out_ch, affine=True), + nn.LeakyReLU(0.2, inplace=True) + ) + + +class InvertedResBlock(nn.Module): + def __init__(self, in_ch, out_ch, expansion_ratio=2): + super(InvertedResBlock, self).__init__() + + self.use_res_connect = in_ch == out_ch + bottleneck = int(round(in_ch*expansion_ratio)) + layers = [] + if expansion_ratio != 1: + layers.append(ConvNormLReLU(in_ch, bottleneck, kernel_size=1, padding=0)) + + # dw + layers.append(ConvNormLReLU(bottleneck, bottleneck, groups=bottleneck, bias=True)) + # pw + layers.append(nn.Conv2d(bottleneck, out_ch, kernel_size=1, padding=0, bias=False)) + layers.append(nn.GroupNorm(num_groups=1, num_channels=out_ch, affine=True)) + + self.layers = nn.Sequential(*layers) + + def forward(self, input): + out = self.layers(input) + if self.use_res_connect: + out = input + out + return out + + +class Generator(nn.Module): + def __init__(self, ): + super().__init__() + + self.block_a = nn.Sequential( + ConvNormLReLU(3, 32, kernel_size=7, padding=3), + ConvNormLReLU(32, 64, stride=2, padding=(0,1,0,1)), + ConvNormLReLU(64, 64) + ) + + self.block_b = nn.Sequential( + ConvNormLReLU(64, 128, stride=2, padding=(0,1,0,1)), + ConvNormLReLU(128, 128) + ) + + self.block_c = nn.Sequential( + ConvNormLReLU(128, 128), + InvertedResBlock(128, 256, 2), + InvertedResBlock(256, 256, 2), + InvertedResBlock(256, 256, 2), + InvertedResBlock(256, 256, 2), + ConvNormLReLU(256, 128), + ) + + self.block_d = nn.Sequential( + ConvNormLReLU(128, 128), + ConvNormLReLU(128, 128) + ) + + self.block_e = nn.Sequential( + ConvNormLReLU(128, 64), + ConvNormLReLU(64, 64), + ConvNormLReLU(64, 32, kernel_size=7, padding=3) + ) + + self.out_layer = nn.Sequential( + nn.Conv2d(32, 3, kernel_size=1, stride=1, padding=0, bias=False), + nn.Tanh() + ) + + def forward(self, input, align_corners=True): + out = self.block_a(input) + half_size = out.size()[-2:] + out = self.block_b(out) + out = self.block_c(out) + + if align_corners: + out = nn.functional.interpolate(out, half_size, mode="bilinear", align_corners=True) + else: + out = nn.functional.interpolate(out, scale_factor=2, mode="bilinear", align_corners=False) + out = self.block_d(out) + + if align_corners: + out = nn.functional.interpolate(out, input.size()[-2:], mode="bilinear", align_corners=True) + else: + out = nn.functional.interpolate(out, scale_factor=2, mode="bilinear", align_corners=False) + out = self.block_e(out) + + out = self.out_layer(out) + return out + +class Model(): + def __init__(self, model_name) -> None: + self._model = Generator() + path = os.path.join(str(Path(__file__).parent), 'weights', model_name + '.pt') + ckpt = load(path) + self._model.load_state_dict(ckpt) + self._model.eval() + + + def __call__(self, img_tensor: Tensor): + img_tensor = img_tensor * 2 - 1 + out = self._model(img_tensor).detach() + out = out.squeeze(0).clip(-1, 1) * 0.5 + 0.5 + return out.numpy() + + def train(self): + """ + For training model + """ + pass + + + \ No newline at end of file diff --git a/pytorch/weights/CelebA.pt b/pytorch/weights/CelebA.pt new file mode 100644 index 0000000..b269a2e --- /dev/null +++ b/pytorch/weights/CelebA.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a3740d98f99efe2ee6c332de2b800f542ddbb2d15e835c07e9bf667c29cef8a7 +size 8603556 diff --git a/pytorch/weights/FacePaintV1.pt b/pytorch/weights/FacePaintV1.pt new file mode 100644 index 0000000..ff16e49 --- /dev/null +++ b/pytorch/weights/FacePaintV1.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f27b45d17c6f4d027753026aeb7e0a558bb95cff5d03a207ac06ca0a372d5316 +size 8603556 diff --git a/pytorch/weights/FacePaintV2.pt b/pytorch/weights/FacePaintV2.pt new file mode 100644 index 0000000..76f61ce --- /dev/null +++ b/pytorch/weights/FacePaintV2.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:06b88a204eb230889444ad868ee5608f4fce5d4ff7b7738acaa4209c2b8fdca7 +size 8601086 diff --git a/pytorch/weights/Hayao.pt b/pytorch/weights/Hayao.pt new file mode 100644 index 0000000..7312811 --- /dev/null +++ b/pytorch/weights/Hayao.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:96e5c586c944fbca18a108a698c02011a96baa22803b77aab8d0b49f5d0b204d +size 8601086 diff --git a/pytorch/weights/Paprika.pt b/pytorch/weights/Paprika.pt new file mode 100644 index 0000000..997eded --- /dev/null +++ b/pytorch/weights/Paprika.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3eaba1b6d01e88ea32b16a6006b04eb3f60327c1dcc841de640d3196898da344 +size 8603556 diff --git a/pytorch/weights/Shinkai.pt b/pytorch/weights/Shinkai.pt new file mode 100644 index 0000000..d0ed068 --- /dev/null +++ b/pytorch/weights/Shinkai.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b2c8512f263f2ccf112a6cbe3ec1cef3bb6e17c9076211e37e6e1f2324fe2b1e +size 8601086