initial deploy

app.py

@@ -0,0 +1,95 @@
+import gradio as gr
+import torch
+from inference.animegan_inference import load_animegan, run_animegan
+from inference.apdrawing_inference import load_apdrawing, run_apdrawing
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load all models once at startup
+G_texture = load_animegan("models/animegan/weights/GeneratorV2_live_action_cartoon_texture.pt", device)
+G_color   = load_animegan("models/animegan/weights/GeneratorV2_live_action_cartoon_color_2.pt", device)
+G_sketch  = load_apdrawing("models/apdrawing/weights/apdrawing_200.pt", device)
+
+# Inference functions for Gardio
+
+def predict_live_action(image, mode, alpha):
+    if image is None:
+        return None
+    return run_animegan(image, G_texture, G_color, mode=mode, alpha=alpha, device=device)
+
+def predict_sketch(image):
+    if image is None:
+        return None
+    return run_apdrawing(image, G_sketch, device)
+
+# UI
+
+with gr.Blocks(title="Disney Style Transfer") as demo:
+
+    gr.Markdown("""
+    # 🎨 Disney Style Transfer
+    Transform real faces or sketches into Disney-style images using two different GAN models.
+    """)
+
+    with gr.Tabs():
+
+        # Tab 1 — Live Action to Disney
+        with gr.Tab("🎬 Live Action → Disney"):
+            gr.Markdown("Upload a real face photo and convert it to Disney style.")
+            with gr.Row():
+                with gr.Column():
+                    input_img = gr.Image(type="pil", label="Input Image", image_mode="RGB")
+                    mode = gr.Radio(
+                        choices=["ensemble", "texture_only", "color_only"],
+                        value="ensemble",
+                        label="Generation mode",
+                        info="Ensemble combines both models. Texture only keeps stylization. Color only keeps palette."
+                    )
+                    alpha = gr.Slider(
+                        minimum=0.0, maximum=1.0, value=0.8, step=0.05,
+                        label="Color transfer strength (alpha)",
+                        info="Only applies in ensemble mode. Higher = more color from color model.",
+                        visible=True
+                    )
+                    btn1 = gr.Button("Generate", variant="primary")
+                with gr.Column():
+                    output_img1 = gr.Image(type="pil", label="Generated Image")
+
+            # Show/hide alpha slider based on mode
+            mode.change(
+                fn=lambda m: gr.update(visible=(m == "ensemble")),
+                inputs=mode,
+                outputs=alpha
+            )
+
+            btn1.click(
+                fn=predict_live_action,
+                inputs=[input_img, mode, alpha],
+                outputs=output_img1
+            )
+
+            gr.Examples(
+                examples=[["examples/live1.jpg", "ensemble", 0.8]],
+                inputs=[input_img, mode, alpha],
+                outputs=output_img1,
+                fn=predict_live_action,
+                label="Examples"
+            )
+
+        # Tab 2 — Sketch to Disney
+        with gr.Tab("✏️ Sketch → Disney"):
+            gr.Markdown("Upload a sketch or line drawing and convert it to a colored Disney-style image.")
+            with gr.Row():
+                with gr.Column():
+                    input_sketch = gr.Image(type="pil", label="Input Sketch", image_mode="RGB")
+                    btn2 = gr.Button("Generate", variant="primary")
+                with gr.Column():
+                    output_img2 = gr.Image(type="pil", label="Generated Image")
+
+            btn2.click(
+                fn=predict_sketch,
+                inputs=input_sketch,
+                outputs=output_img2
+            )
+
+demo.launch()

inference/animegan_inference.py

@@ -0,0 +1,45 @@
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+
+def load_animegan(weights_path, device):
+    from models.animegan.generator import GeneratorV2
+    checkpoint = torch.load(weights_path, map_location=device)
+    state_dict = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint
+    state_dict = {k[len("module."):] if k.startswith("module.") else k: v for k, v in state_dict.items()}
+    G = GeneratorV2().to(device)
+    G.load_state_dict(state_dict, strict=True)
+    G.eval()
+    return G
+
+def preprocess(pil_img, device, size=(256, 256)):
+    img = np.array(pil_img.convert("RGB").resize(size))
+    tensor = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0)
+    tensor = (tensor / 127.5) - 1.0
+    return tensor.to(device)
+
+def postprocess(tensor):
+    img = ((tensor.squeeze(0).permute(1, 2, 0) + 1) * 127.5).clamp(0, 255).byte().cpu().numpy()
+    return img  # RGB numpy
+
+def transfer_palette(source, target, alpha=1.0):
+    source_lab = cv2.cvtColor(source, cv2.COLOR_RGB2LAB).astype(np.float32)
+    target_lab = cv2.cvtColor(target, cv2.COLOR_RGB2LAB).astype(np.float32)
+    result_lab = target_lab.copy()
+    result_lab[:, :, 1] = source_lab[:, :, 1] * alpha + target_lab[:, :, 1] * (1 - alpha)
+    result_lab[:, :, 2] = source_lab[:, :, 2] * alpha + target_lab[:, :, 2] * (1 - alpha)
+    return cv2.cvtColor(result_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
+
+def run_animegan(pil_img, G_texture, G_color, mode="ensemble", alpha=0.8, device="cpu"):
+    inp = preprocess(pil_img, device)
+    with torch.no_grad():
+        if mode == "texture_only":
+            result = postprocess(G_texture(inp))
+        elif mode == "color_only":
+            result = postprocess(G_color(inp))
+        else:  # ensemble
+            out_texture = postprocess(G_texture(inp))
+            out_color   = postprocess(G_color(inp))
+            result = transfer_palette(source=out_color, target=out_texture, alpha=alpha)
+    return Image.fromarray(result)

inference/apdrawing_inference.py

@@ -0,0 +1,30 @@
+import torch
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+
+def load_apdrawing(weights_path, device):
+    from models.apdrawing.generator import define_G
+    netG = define_G(1, 3, 64, 'unet_256', 'batch',
+                    use_dropout=False, init_type='normal',
+                    init_gain=0.02, gpu_ids=[])
+    ckpt = torch.load(weights_path, map_location='cpu')
+    state_dict = ckpt['G'] if isinstance(ckpt, dict) and 'G' in ckpt else ckpt
+    netG.load_state_dict(state_dict, strict=True)
+    netG.to(device).eval()
+    return netG
+
+def preprocess_sketch(pil_img):
+    img = pil_img.convert('RGB').resize((256, 256), Image.BICUBIC)
+    tensor = transforms.ToTensor()(img)
+    tensor = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(tensor)
+    gray = tensor[0] * 0.299 + tensor[1] * 0.587 + tensor[2] * 0.114
+    return gray.unsqueeze(0).unsqueeze(0)  # (1, 1, 256, 256)
+
+def run_apdrawing(pil_img, netG, device):
+    inp = preprocess_sketch(pil_img).to(device)
+    with torch.no_grad():
+        out = netG(inp)
+    out = out.squeeze(0).permute(1, 2, 0).cpu().numpy()
+    out = ((out + 1) * 127.5).clip(0, 255).astype(np.uint8)
+    return Image.fromarray(out)

models/animegan/generator.py

@@ -0,0 +1,202 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+
+# NOTE!!! This code is an adaptation form the original repo, we are not owners nor the creators behind the design, we justa dapted it to our needs
+# Original: https://github.com/ptran1203/pytorch-animeGAN/tree/master
+
+def initialize_weights(net):
+    for m in net.modules():
+        try:
+            if isinstance(m, nn.Conv2d):
+                # m.weight.data.normal_(0, 0.02)
+                torch.nn.init.xavier_uniform_(m.weight)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.ConvTranspose2d):
+                # m.weight.data.normal_(0, 0.02)
+                torch.nn.init.xavier_uniform_(m.weight)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                # m.weight.data.normal_(0, 0.02)
+                torch.nn.init.xavier_uniform_(m.weight)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+        except Exception as e:
+            # print(f'SKip layer {m}, {e}')
+            pass
+
+class LayerNorm2d(nn.LayerNorm):
+    """ LayerNorm for channels of '2D' spatial NCHW tensors """
+    def __init__(self, num_channels, eps=1e-6, affine=True):
+        super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 3, 1, 2)
+        return x
+
+
+def get_norm(norm_type, channels):
+    if norm_type == "instance":
+        return nn.InstanceNorm2d(channels)
+    elif norm_type == "layer":
+        # return LayerNorm2d
+        return nn.GroupNorm(num_groups=1, num_channels=channels, affine=True)
+        # return partial(nn.GroupNorm, 1, out_ch, 1e-5, True)
+    else:
+        raise ValueError(norm_type)
+
+class ConvBlock(nn.Module):
+    """Stack of Conv2D + Norm + LeakyReLU"""
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        kernel_size=3,
+        stride=1,
+        groups=1,
+        padding=1,
+        bias=False,
+        norm_type="instance"
+    ):
+        super(ConvBlock, self).__init__()
+
+        # if kernel_size == 3 and stride == 1:
+        #     self.pad = nn.ReflectionPad2d((1, 1, 1, 1))
+        # elif kernel_size == 7 and stride == 1:
+        #     self.pad = nn.ReflectionPad2d((3, 3, 3, 3))
+        # elif stride == 2:
+        #     self.pad = nn.ReflectionPad2d((0, 1, 1, 0))
+        # else:
+        #     self.pad = None
+        
+        self.pad = nn.ReflectionPad2d(padding)
+        self.conv = nn.Conv2d(
+            channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            groups=groups,
+            padding=0,
+            bias=bias
+        )
+        self.ins_norm = get_norm(norm_type, out_channels)
+        self.activation = nn.LeakyReLU(0.2, True)
+
+        # initialize_weights(self)
+
+    def forward(self, x):
+        if self.pad is not None:
+            x = self.pad(x)
+        out = self.conv(x)
+        out = self.ins_norm(out)
+        out = self.activation(out)
+        return out
+
+
+class InvertedResBlock(nn.Module):
+    def __init__(
+        self,
+        channels=256,
+        out_channels=256,
+        expand_ratio=2,
+        norm_type="instance",
+    ):
+        super(InvertedResBlock, self).__init__()
+        bottleneck_dim = round(expand_ratio * channels)
+        self.conv_block = ConvBlock(
+            channels,
+            bottleneck_dim,
+            kernel_size=1,
+            padding=0,
+            norm_type=norm_type,
+            bias=False
+        )
+        self.conv_block2 = ConvBlock(
+            bottleneck_dim,
+            bottleneck_dim,
+            groups=bottleneck_dim,
+            norm_type=norm_type,
+            bias=True
+        )
+        self.conv = nn.Conv2d(
+            bottleneck_dim,
+            out_channels,
+            kernel_size=1,
+            padding=0,
+            bias=False
+        )
+        self.norm = get_norm(norm_type, out_channels)
+
+    def forward(self, x):
+        out = self.conv_block(x)
+        out = self.conv_block2(out)
+        # out = self.activation(out)
+        out = self.conv(out)
+        out = self.norm(out)
+
+        if out.shape[1] != x.shape[1]:
+            # Only concate if same shape
+            return out
+        return out + x
+
+class GeneratorV2(nn.Module):
+    def __init__(self, dataset=''):
+        super(GeneratorV2, self).__init__()
+        self.name = f'{self.__class__.__name__}_{dataset}'
+
+        self.conv_block1 = nn.Sequential(
+            ConvBlock(3, 32, kernel_size=7, stride=1, padding=3, norm_type="layer"),
+            ConvBlock(32, 64, kernel_size=3, stride=2, padding=(0, 1, 0, 1), norm_type="layer"),
+            ConvBlock(64, 64, kernel_size=3, stride=1, norm_type="layer"),
+        )
+
+        self.conv_block2 = nn.Sequential(
+            ConvBlock(64, 128, kernel_size=3, stride=2, padding=(0, 1, 0, 1), norm_type="layer"),
+            ConvBlock(128, 128, kernel_size=3, stride=1, norm_type="layer"),
+        )
+
+        self.res_blocks = nn.Sequential(
+            ConvBlock(128, 128, kernel_size=3, stride=1, norm_type="layer"),
+            InvertedResBlock(128, 256, expand_ratio=2, norm_type="layer"),
+            InvertedResBlock(256, 256, expand_ratio=2, norm_type="layer"),
+            InvertedResBlock(256, 256, expand_ratio=2, norm_type="layer"),
+            InvertedResBlock(256, 256, expand_ratio=2, norm_type="layer"),
+            ConvBlock(256, 128, kernel_size=3, stride=1, norm_type="layer"),
+        )
+
+        self.conv_block3 = nn.Sequential(
+            # UpConvLNormLReLU(128, 128, norm_type="layer"),
+            ConvBlock(128, 128, kernel_size=3, stride=1, norm_type="layer"),
+            ConvBlock(128, 128, kernel_size=3, stride=1, norm_type="layer"),
+        )
+
+        self.conv_block4 = nn.Sequential(
+            # UpConvLNormLReLU(128, 64, norm_type="layer"),
+            ConvBlock(128, 64, kernel_size=3, stride=1, norm_type="layer"),
+            ConvBlock(64, 64, kernel_size=3, stride=1, norm_type="layer"),
+            ConvBlock(64, 32, kernel_size=7, padding=3, stride=1, norm_type="layer"),
+        )
+
+        self.decode_blocks = nn.Sequential(
+            nn.Conv2d(32, 3, kernel_size=1, stride=1, padding=0),
+            nn.Tanh(),
+        )
+
+        initialize_weights(self)
+
+    def forward(self, x):
+        out = self.conv_block1(x)
+        out = self.conv_block2(out)
+        out = self.res_blocks(out)
+        out = F.interpolate(out, scale_factor=2, mode="bilinear")
+        out = self.conv_block3(out)
+        out = F.interpolate(out, scale_factor=2, mode="bilinear")
+        out = self.conv_block4(out)
+        img = self.decode_blocks(out)
+
+        return img

models/animegan/weights/GeneratorV2_live_action_cartoon_color.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:409e71f59fa121ac447968b79daa910ba1347de5de8510e73e9e784424a71b6e
+size 25827551

models/animegan/weights/GeneratorV2_live_action_cartoon_color_2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a16e6ab6697009852dff912a9aa798fef5e1b5e936192544409be26ff871e76
+size 25827551

models/animegan/weights/GeneratorV2_live_action_cartoon_texture.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67cd02cd47208c05e1209459e2830cc42e6b7c6aeba698c6c8c8afc9a13d4661
+size 25827551

models/apdrawing/generator.py

@@ -0,0 +1,453 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from torch.optim import lr_scheduler
+
+# NOTE!!! This code does not originally belong to us, it´s just an adaptation for our fine tunning
+# Original repo: https://github.com/yiranran/APDrawingGAN.git
+
+# Helper Functions
+
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=True)
+    elif norm_type == 'none':
+        norm_layer = None
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+
+def get_scheduler(optimizer, opt):
+    if opt.lr_policy == 'lambda':
+        def lambda_rule(epoch):
+            lr_l = 1.0 - max(0, epoch + 1 + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1)
+            return lr_l
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif opt.lr_policy == 'step':
+        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
+    elif opt.lr_policy == 'plateau':
+        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
+    elif opt.lr_policy == 'cosine':
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.niter, eta_min=0)
+    else:
+        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
+    return scheduler
+
+
+def init_weights(net, init_type='normal', gain=0.02):
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)
+
+
+def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    if len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        net.to(gpu_ids[0])
+        net = torch.nn.DataParallel(net, gpu_ids)
+    init_weights(net, init_type, gain=init_gain)
+    return net
+
+
+def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[], nnG=9):
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netG == 'resnet_9blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
+    elif netG == 'resnet_6blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
+    elif netG == 'resnet_nblocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=nnG)
+    elif netG == 'unet_128':
+        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_256':#default for pix2pix
+        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_512':
+        net = UnetGenerator(input_nc, output_nc, 9, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_ndown':
+        net = UnetGenerator(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'partunet':
+        net = PartUnet(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'partunet2':
+        net = PartUnet2(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'combiner':
+        net = Combiner(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=2)
+    else:
+        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+def define_D(input_nc, ndf, netD,
+             n_layers_D=3, norm='batch', use_sigmoid=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netD == 'basic':
+        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
+    elif netD == 'n_layers':
+        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
+    elif netD == 'pixel':
+        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
+    else:
+        raise NotImplementedError('Discriminator model name [%s] is not recognized' % net)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+# Classes
+
+# Defines the GAN loss which uses either LSGAN or the regular GAN.
+# When LSGAN is used, it is basically same as MSELoss,
+# but it abstracts away the need to create the target label tensor
+# that has the same size as the input
+class GANLoss(nn.Module):
+    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0):
+        super(GANLoss, self).__init__()
+        self.register_buffer('real_label', torch.tensor(target_real_label))
+        self.register_buffer('fake_label', torch.tensor(target_fake_label))
+        if use_lsgan:
+            self.loss = nn.MSELoss()
+        else:#no_lsgan
+            self.loss = nn.BCELoss()
+
+    def get_target_tensor(self, input, target_is_real):
+        if target_is_real:
+            target_tensor = self.real_label
+        else:
+            target_tensor = self.fake_label
+        return target_tensor.expand_as(input)
+
+    def __call__(self, input, target_is_real):
+        target_tensor = self.get_target_tensor(input, target_is_real)
+        return self.loss(input, target_tensor)
+
+# Defines the generator that consists of Resnet blocks between a few
+# downsampling/upsampling operations.
+# Code and idea originally from Justin Johnson's architecture.
+# https://github.com/jcjohnson/fast-neural-style/
+class ResnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        assert(n_blocks >= 0)
+        super(ResnetGenerator, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0,
+                           bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,
+                                stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2),
+                      nn.ReLU(True)]
+
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                        kernel_size=3, stride=2,
+                                        padding=1, output_padding=1,
+                                        bias=use_bias),
+                    norm_layer(int(ngf * mult / 2)),
+                    nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        return self.model(input)
+
+class Combiner(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        assert(n_blocks >= 0)
+        super(Combiner, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0,
+                           bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        for i in range(n_blocks):
+            model += [ResnetBlock(ngf, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        return self.model(input)
+
+# Define a resnet block
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
+                       norm_layer(dim),
+                       nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
+                       norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        out = x + self.conv_block(x)
+        return out
+
+
+# Defines the Unet generator.
+# |num_downs|: number of downsamplings in UNet. For example,
+# if |num_downs| == 7, image of size 128x128 will become of size 1x1
+# at the bottleneck
+class UnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetGenerator, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+class PartUnet(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(PartUnet, self).__init__()
+
+        # construct unet structure
+        # 3 downs 
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+class PartUnet2(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(PartUnet2, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 2, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 3):
+            unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+
+# Defines the submodule with skip connection.
+# X -------------------identity---------------------- X
+#   |-- downsampling -- |submodule| -- upsampling --|
+class UnetSkipConnectionBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)
+
+
+# Defines the PatchGAN discriminator with the specified arguments.
+class NLayerDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
+        super(NLayerDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [
+            nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):
+            nf_mult_prev = nf_mult
+            nf_mult = min(2**n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
+                          kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2**n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
+                      kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
+
+        if use_sigmoid:#no_lsgan, use sigmoid before calculating bceloss(binary cross entropy)
+            sequence += [nn.Sigmoid()]
+
+        self.model = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        return self.model(input)
+
+
+class PixelDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
+        super(PixelDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        self.net = [
+            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
+            norm_layer(ndf * 2),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
+
+        if use_sigmoid:
+            self.net.append(nn.Sigmoid())
+
+        self.net = nn.Sequential(*self.net)
+
+    def forward(self, input):
+        return self.net(input)

models/apdrawing/weights/apdrawing_200.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c662146e101b06a6cf8a70de79ae86f9bfd58425ecb0a7c9d5251efb05932c4
+size 217718925

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+torch
+torchvision
+opencv-python-headless
+Pillow
+numpy
+scikit-image