eehhhhhhh!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

app.py

@@ -1,580 +1,1374 @@
-import spaces
-import contextlib
-import gc
-import json
-import logging
-import math
-import os
-import random
-import shutil
-import sys
-import time
-import itertools
-from pathlib import Path
-
-import cv2
-import numpy as np
-from PIL import Image, ImageDraw
 import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from torch.optim import lr_scheduler
+import numpy as np
 import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch.utils.data import Dataset
-from torchvision import transforms
-from tqdm.auto import tqdm
-
-import accelerate
-from accelerate import Accelerator
-from accelerate.logging import get_logger
-from accelerate.utils import ProjectConfiguration, set_seed
-
-from datasets import load_dataset
-from huggingface_hub import create_repo, upload_folder
-from packaging import version
-from safetensors.torch import load_model
-from peft import LoraConfig
-import gradio as gr
-import pandas as pd
-
-import transformers
-from transformers import (
-    AutoTokenizer,
-    PretrainedConfig,
-    CLIPVisionModelWithProjection,
-    CLIPImageProcessor,
-    CLIPProcessor,
-)
-
-import diffusers
-from diffusers import (
-    AutoencoderKL,
-    DDPMScheduler,
-    ColorGuiderPixArtModel,
-    ColorGuiderSDModel,
-    UNet2DConditionModel,
-    PixArtTransformer2DModel,
-    ColorFlowPixArtAlphaPipeline,
-    ColorFlowSDPipeline,
-    UniPCMultistepScheduler,
-)
-from colorflow_utils.utils import *
-
-sys.path.append('./BidirectionalTranslation')
-from options.test_options import TestOptions
-from models import create_model
-from util import util
-
-from huggingface_hub import snapshot_download
-
-
-article = r"""
-If ColorFlow is helpful, please help to ⭐ the <a href='https://github.com/TencentARC/ColorFlow' target='_blank'>Github Repo</a>. Thanks! [![GitHub Stars](https://img.shields.io/github/stars/TencentARC/ColorFlow)](https://github.com/TencentARC/ColorFlow)
----
-
-📧 **Contact**
-<br>
-If you have any questions, please feel free to reach me out at <b>zhuangjh23@mails.tsinghua.edu.cn</b>.
-
-📝 **Citation**
-<br>
-If our work is useful for your research, please consider citing:
-```bibtex
-@misc{zhuang2024colorflow,
-      title={ColorFlow: Retrieval-Augmented Image Sequence Colorization}, 
-      author={Junhao Zhuang and Xuan Ju and Zhaoyang Zhang and Yong Liu and Shiyi Zhang and Chun Yuan and Ying Shan},
-      year={2024},
-      eprint={2412.11815},
-      archivePrefix={arXiv},
-      primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2412.11815},
-}
-```
-"""
-
-model_global_path = snapshot_download(repo_id="TencentARC/ColorFlow", cache_dir='./colorflow/', repo_type="model")
-print(model_global_path)
-
-
-transform = transforms.Compose([
-    transforms.ToTensor(),  
-    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  
-])
-weight_dtype = torch.float16
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# line model
-line_model_path = model_global_path + '/LE/erika.pth'
-line_model = res_skip()
-line_model.load_state_dict(torch.load(line_model_path))
-line_model.eval()
-line_model.to(device)
-
-# screen model
-global opt
-
-opt = TestOptions().parse(model_global_path)
-ScreenModel = create_model(opt, model_global_path)
-ScreenModel.setup(opt)
-ScreenModel.eval()
-
-image_processor = CLIPImageProcessor()
-image_encoder = CLIPVisionModelWithProjection.from_pretrained(model_global_path + '/image_encoder/').to(device)
-
-
-examples = [
-    [
-        "./assets/example_5/input.png", 
-        ["./assets/example_5/ref1.png", "./assets/example_5/ref2.png", "./assets/example_5/ref3.png"], 
-        "GrayImage(ScreenStyle)", 
-        "800x512",  
-        0, 
-        10 
-    ],
-    [
-        "./assets/example_4/input.jpg", 
-        ["./assets/example_4/ref1.jpg", "./assets/example_4/ref2.jpg", "./assets/example_4/ref3.jpg"], 
-        "GrayImage(ScreenStyle)", 
-        "640x640",  
-        0, 
-        10 
-    ],
-    [
-        "./assets/example_3/input.png", 
-        ["./assets/example_3/ref1.png", "./assets/example_3/ref2.png", "./assets/example_3/ref3.png"], 
-        "GrayImage(ScreenStyle)", 
-        "800x512", 
-        0, 
-        10 
-    ],
-    [
-        "./assets/example_2/input.png",  
-        ["./assets/example_2/ref1.png", "./assets/example_2/ref2.png", "./assets/example_2/ref3.png"], 
-        "GrayImage(ScreenStyle)",  
-        "800x512",  
-        0,  
-        10  
-    ],
-    [
-        "./assets/example_6/input.png", 
-        ["./assets/example_6/ref1.png", "./assets/example_6/ref2.png", "./assets/example_6/ref3.png"], 
-        "Sketch_Shading", 
-        "512x800",  
-        0, 
-        10 
-    ],
-    [
-        "./assets/example_7/input.jpg", 
-        ["./assets/example_7/ref1.jpg", "./assets/example_7/ref2.jpg", "./assets/example_7/ref3.jpg", "./assets/example_7/ref4.jpg"], 
-        "Sketch_Shading", 
-        "640x640",  
-        2, 
-        10 
-    ],
-    [
-        "./assets/example_1/input.jpg", 
-        ["./assets/example_1/ref1.jpg", "./assets/example_1/ref2.jpg", "./assets/example_1/ref3.jpg"], 
-        "Sketch",  
-        "640x640", 
-        1, 
-        10  
-    ],
-    [
-        "./assets/example_0/input.jpg", 
-        ["./assets/example_0/ref1.jpg"], 
-        "Sketch", 
-        "640x640",  
-        1, 
-        10 
-    ],
-]
-
-global pipeline
-global MultiResNetModel
-
-@spaces.GPU
-def load_ckpt(input_style):
-    global pipeline
-    global MultiResNetModel
-    if input_style == "Sketch" or input_style == "Sketch_Shading":
-        if input_style == "Sketch":
-            ckpt_path = model_global_path + '/sketch/'
-            rank = 128
-        else:
-            ckpt_path = model_global_path + '/shading/'
-            rank = 128
-        pretrained_model_name_or_path = 'PixArt-alpha/PixArt-XL-2-1024-MS'
-        transformer = PixArtTransformer2DModel.from_pretrained(
-            pretrained_model_name_or_path, subfolder="transformer", revision=None, variant=None
-        )
-        pixart_config = get_pixart_config()
-
-        ColorGuider = ColorGuiderPixArtModel.from_pretrained(ckpt_path)
-
-        transformer_lora_config = LoraConfig(
-            r=rank,
-            lora_alpha=rank,
-            init_lora_weights="gaussian",
-            target_modules=["to_k", "to_q", "to_v", "to_out.0", "proj_in", "proj_out", "ff.net.0.proj", "ff.net.2", "proj", "linear", "linear_1", "linear_2"]
-        )
-        transformer.add_adapter(transformer_lora_config)
-        ckpt_key_t = torch.load(ckpt_path + 'transformer_lora.bin', map_location='cpu')
-        transformer.load_state_dict(ckpt_key_t, strict=False)
-
-        transformer.to(device, dtype=weight_dtype)
-        ColorGuider.to(device, dtype=weight_dtype)
-        
-        pipeline = ColorFlowPixArtAlphaPipeline.from_pretrained(
-            pretrained_model_name_or_path,
-            transformer=transformer,
-            colorguider=ColorGuider,
-            safety_checker=None,
-            revision=None,
-            variant=None,
-            torch_dtype=weight_dtype,
-        )
-        pipeline = pipeline.to(device)
-        block_out_channels = [128, 128, 256, 512, 512]
-        
-        MultiResNetModel = MultiHiddenResNetModel(block_out_channels, len(block_out_channels))
-        MultiResNetModel.load_state_dict(torch.load(ckpt_path + 'MultiResNetModel.bin', map_location='cpu'), strict=False)
-        MultiResNetModel.to(device, dtype=weight_dtype)
-
-    elif input_style == "GrayImage(ScreenStyle)":
-        ckpt_path = model_global_path + '/GraySD/'
-        rank = 64
-        pretrained_model_name_or_path = 'stable-diffusion-v1-5/stable-diffusion-v1-5'
-        unet = UNet2DConditionModel.from_pretrained(
-            pretrained_model_name_or_path, subfolder="unet", revision=None, variant=None
-        )
-        ColorGuider = ColorGuiderSDModel.from_pretrained(ckpt_path)
-        ColorGuider.to(device, dtype=weight_dtype)
-        unet.to(device, dtype=weight_dtype)
-        
-        pipeline = ColorFlowSDPipeline.from_pretrained(
-            pretrained_model_name_or_path,
-            unet=unet,
-            colorguider=ColorGuider,
-            safety_checker=None,
-            revision=None,
-            variant=None,
-            torch_dtype=weight_dtype,
-        )
-        pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
-        unet_lora_config = LoraConfig(
-            r=rank,
-            lora_alpha=rank,
-            init_lora_weights="gaussian",
-            target_modules=["to_k", "to_q", "to_v", "to_out.0", "ff.net.0.proj", "ff.net.2"],#ff.net.0.proj ff.net.2
-        )
-        pipeline.unet.add_adapter(unet_lora_config)
-        pipeline.unet.load_state_dict(torch.load(ckpt_path + 'unet_lora.bin', map_location='cpu'), strict=False)
-        pipeline = pipeline.to(device)
-        block_out_channels = [128, 128, 256, 512, 512]
-        
-        MultiResNetModel = MultiHiddenResNetModel(block_out_channels, len(block_out_channels))
-        MultiResNetModel.load_state_dict(torch.load(ckpt_path + 'MultiResNetModel.bin', map_location='cpu'), strict=False)
-        MultiResNetModel.to(device, dtype=weight_dtype)
-
-    
-
-
-
-global cur_input_style
-cur_input_style = "Sketch"
-load_ckpt(cur_input_style)
-cur_input_style = "Sketch_Shading"
-load_ckpt(cur_input_style)
-cur_input_style = "GrayImage(ScreenStyle)"
-load_ckpt(cur_input_style)
-cur_input_style = None
-
-@spaces.GPU
-def fix_random_seeds(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)
-
-def process_multi_images(files):
-    images = [Image.open(file.name) for file in files]
-    imgs = []
-    for i, img in enumerate(images):
-        imgs.append(img)
-    return imgs 
-
-@spaces.GPU
-def extract_lines(image):
-    src = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-
-    rows = int(np.ceil(src.shape[0] / 16)) * 16
-    cols = int(np.ceil(src.shape[1] / 16)) * 16
-
-    patch = np.ones((1, 1, rows, cols), dtype="float32")
-    patch[0, 0, 0:src.shape[0], 0:src.shape[1]] = src
-
-    tensor = torch.from_numpy(patch).to(device)
-
-    with torch.no_grad():
-        y = line_model(tensor)
-
-    yc = y.cpu().numpy()[0, 0, :, :]
-    yc[yc > 255] = 255
-    yc[yc < 0] = 0
-
-    outimg = yc[0:src.shape[0], 0:src.shape[1]]
-    outimg = outimg.astype(np.uint8)
-    outimg = Image.fromarray(outimg)
-    torch.cuda.empty_cache()
-    return outimg
-
-@spaces.GPU
-def to_screen_image(input_image):
-    global opt
-    global ScreenModel
-    input_image = input_image.convert('RGB')
-    input_image = get_ScreenVAE_input(input_image, opt)
-    h = input_image['h']
-    w = input_image['w']
-    ScreenModel.set_input(input_image)
-    fake_B, fake_B2, SCR = ScreenModel.forward(AtoB=True)
-    images=fake_B2[:,:,:h,:w]
-    im = util.tensor2im(images)
-    image_pil = Image.fromarray(im)
-    torch.cuda.empty_cache()
-    return image_pil
-
-@spaces.GPU
-def extract_line_image(query_image_, input_style, resolution):
-    if resolution == "640x640":
-        tar_width = 640
-        tar_height = 640
-    elif resolution == "512x800":
-        tar_width = 512
-        tar_height = 800
-    elif resolution == "800x512":
-        tar_width = 800
-        tar_height = 512
+from torch.nn.modules.normalization import LayerNorm
+import os
+from torch.nn.utils import spectral_norm
+from torchvision import models
+
+###############################################################################
+# Helper functions
+###############################################################################
+
+
+def init_weights(net, init_type='normal', init_gain=0.02):
+    """Initialize network weights.
+    Parameters:
+        net (network)   -- network to be initialized
+        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.
+    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
+    work better for some applications. Feel free to try yourself.
+    """
+    def init_func(m):  # define the initialization function
+        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, init_gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=init_gain)
+            elif init_type == 'kaiming':
+                #init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+                init.kaiming_normal_(m.weight.data, a=0.2, mode='fan_in', nonlinearity='leaky_relu')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=init_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:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
+            init.normal_(m.weight.data, 1.0, init_gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)  # apply the initialization function <init_func>
+
+
+def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[], init=True):
+    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
+    Parameters:
+        net (network)      -- the network to be initialized
+        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        gain (float)       -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+    Return an initialized network.
+    """
+    if len(gpu_ids) > 0 and torch.cuda.is_available():
+        net.to(gpu_ids[0])
+    if init:
+        init_weights(net, init_type, init_gain=init_gain)
+    return net
+
+
+def get_scheduler(optimizer, opt):
+    """Return a learning rate scheduler
+    Parameters:
+        optimizer          -- the optimizer of the network
+        opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　
+                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
+    For 'linear', we keep the same learning rate for the first <opt.niter> epochs
+    and linearly decay the rate to zero over the next <opt.niter_decay> epochs.
+    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
+    See https://pytorch.org/docs/stable/optim.html for more details.
+    """
+    if opt.lr_policy == 'linear':
+        def lambda_rule(epoch):
+            lr_l = 1.0 - max(0, epoch + 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
+
+class LayerNormWarpper(nn.Module):
+    def __init__(self, num_features):
+        super(LayerNormWarpper, self).__init__()
+        self.num_features = int(num_features)
+
+    def forward(self, x):
+        x = nn.LayerNorm([self.num_features, x.size()[2], x.size()[3]], elementwise_affine=False).to(x.device)(x)
+        return x
+
+def get_norm_layer(norm_type='instance'):
+    """Return a normalization layer
+    Parameters:
+        norm_type (str) -- the name of the normalization layer: batch | instance | none
+    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
+    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
+    """
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+    elif norm_type == 'layer':
+        norm_layer = functools.partial(LayerNormWarpper)
+    elif norm_type == 'none':
+        norm_layer = None
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+
+def get_non_linearity(layer_type='relu'):
+    if layer_type == 'relu':
+        nl_layer = functools.partial(nn.ReLU, inplace=True)
+    elif layer_type == 'lrelu':
+        nl_layer = functools.partial(
+            nn.LeakyReLU, negative_slope=0.2, inplace=True)
+    elif layer_type == 'elu':
+        nl_layer = functools.partial(nn.ELU, inplace=True)
+    elif layer_type == 'selu':
+        nl_layer = functools.partial(nn.SELU, inplace=True)
+    elif layer_type == 'prelu':
+        nl_layer = functools.partial(nn.PReLU)
+    else:
+        raise NotImplementedError(
+            'nonlinearity activitation [%s] is not found' % layer_type)
+    return nl_layer
+
+
+def define_G(input_nc, output_nc, nz, ngf, netG='unet_128', norm='batch', nl='relu', use_noise=False, 
+             use_dropout=False, init_type='xavier', init_gain=0.02, gpu_ids=[], where_add='input', upsample='bilinear'):
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+    nl_layer = get_non_linearity(layer_type=nl)
+    # print(norm, norm_layer)
+
+    if nz == 0:
+        where_add = 'input'
+
+    if netG == 'unet_128' and where_add == 'input':
+        net = G_Unet_add_input(input_nc, output_nc, nz, 7, ngf, norm_layer=norm_layer, nl_layer=nl_layer, use_noise=use_noise,
+                               use_dropout=use_dropout, upsample=upsample, device=gpu_ids)
+    elif netG == 'unet_128_G' and where_add == 'input':
+        net = G_Unet_add_input_G(input_nc, output_nc, nz, 7, ngf, norm_layer=norm_layer, nl_layer=nl_layer, use_noise=use_noise,
+                               use_dropout=use_dropout, upsample=upsample, device=gpu_ids)
+    elif netG == 'unet_256' and where_add == 'input':
+        net = G_Unet_add_input(input_nc, output_nc, nz, 8, ngf, norm_layer=norm_layer, nl_layer=nl_layer, use_noise=use_noise,
+                               use_dropout=use_dropout, upsample=upsample, device=gpu_ids)
+    elif netG == 'unet_256_G' and where_add == 'input':
+        net = G_Unet_add_input_G(input_nc, output_nc, nz, 8, ngf, norm_layer=norm_layer, nl_layer=nl_layer, use_noise=use_noise,
+                               use_dropout=use_dropout, upsample=upsample, device=gpu_ids)
+    elif netG == 'unet_128' and where_add == 'all':
+        net = G_Unet_add_all(input_nc, output_nc, nz, 7, ngf, norm_layer=norm_layer, nl_layer=nl_layer, use_noise=use_noise,
+                             use_dropout=use_dropout, upsample=upsample)
+    elif netG == 'unet_256' and where_add == 'all':
+        net = G_Unet_add_all(input_nc, output_nc, nz, 8, ngf, norm_layer=norm_layer, nl_layer=nl_layer, use_noise=use_noise,
+                             use_dropout=use_dropout, upsample=upsample)
+    else:
+        raise NotImplementedError('Generator model name [%s] is not recognized' % net)
+    # print(net)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+def define_C(input_nc, output_nc, nz, ngf, netC='unet_128', norm='instance', nl='relu',
+             use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[], upsample='basic'):
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+    nl_layer = get_non_linearity(layer_type=nl)
+
+    if netC == 'resnet_9blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
+    elif netC == 'resnet_6blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
+    elif netC == 'unet_128':
+        net = G_Unet_add_input_C(input_nc, output_nc, 0, 7, ngf, norm_layer=norm_layer, nl_layer=nl_layer,
+                               use_dropout=use_dropout, upsample=upsample)
+    elif netC == 'unet_256':
+        net = G_Unet_add_input(input_nc, output_nc, 0, 8, ngf, norm_layer=norm_layer, nl_layer=nl_layer,
+                               use_dropout=use_dropout, upsample=upsample)
+    elif netC == 'unet_32':
+        net = G_Unet_add_input(input_nc, output_nc, 0, 5, ngf, norm_layer=norm_layer, nl_layer=nl_layer,
+                               use_dropout=use_dropout, upsample=upsample)
     else:
-        gr.Info("Unsupported resolution")
-
-    query_image = process_image(query_image_, int(tar_width*1.5), int(tar_height*1.5))
-    if input_style == "GrayImage(ScreenStyle)":
-        extracted_line = to_screen_image(query_image)
-        extracted_line = Image.blend(extracted_line.convert('L').convert('RGB'), query_image.convert('L').convert('RGB'), 0.5)
-        input_context = extracted_line
-    elif input_style == "Sketch":
-        query_image = query_image.convert('L').convert('RGB')
-        extracted_line = extract_lines(query_image)
-        extracted_line = extracted_line.convert('L').convert('RGB')
-        input_context = extracted_line
-    elif input_style == "Sketch_Shading":
-        query_image = query_image.convert('L').convert('RGB')
-        extracted_line = extract_lines(query_image)
-        extracted_line = extracted_line.convert('L').convert('RGB')
-        array1 = np.array(query_image)
-        array2 = np.array(extracted_line)
-        array2[array1 < 0.3 * 255.0] = 0
-        gray_rate = 125
-        up_bound = 145
-        array2[(array2 > gray_rate) & (array1 < up_bound) & (array1 > 0.3 * 255.0)] = gray_rate
-        input_context = Image.fromarray(np.uint8(array2))
-    torch.cuda.empty_cache()
-    return input_context, extracted_line, input_context  
-
-@spaces.GPU(duration=180)
-def colorize_image(VAE_input, input_context, reference_images, resolution, seed, input_style, num_inference_steps):
-    if VAE_input is None or input_context is None:
-        gr.Info("Please preprocess the image first")
-        raise ValueError("Please preprocess the image first")
-    global cur_input_style
-    global pipeline
-    global MultiResNetModel
-    if input_style != cur_input_style:
-        gr.Info(f"Loading {input_style} model...")
-        load_ckpt(input_style)
-        cur_input_style = input_style
-        gr.Info(f"{input_style} model loaded")
-    reference_images = process_multi_images(reference_images)
-    fix_random_seeds(seed)
-    if resolution == "640x640":
-        tar_width = 640
-        tar_height = 640
-    elif resolution == "512x800":
-        tar_width = 512
-        tar_height = 800
-    elif resolution == "800x512":
-        tar_width = 800
-        tar_height = 512
+        raise NotImplementedError('Generator model name [%s] is not recognized' % net)
+
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+def define_D(input_nc, ndf, netD, norm='batch', nl='lrelu', init_type='xavier', init_gain=0.02, num_Ds=1, gpu_ids=[]):
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+    nl = 'lrelu'  # use leaky relu for D
+    nl_layer = get_non_linearity(layer_type=nl)
+
+    if netD == 'basic_128':
+        net = D_NLayers(input_nc, ndf, n_layers=2, norm_layer=norm_layer, nl_layer=nl_layer)
+    elif netD == 'basic_256':
+        net = D_NLayers(input_nc, ndf, n_layers=3, norm_layer=norm_layer, nl_layer=nl_layer)
+    elif netD == 'basic_128_multi':
+        net = D_NLayersMulti(input_nc=input_nc, ndf=ndf, n_layers=2, norm_layer=norm_layer, num_D=num_Ds, nl_layer=nl_layer)
+    elif netD == 'basic_256_multi':
+        net = D_NLayersMulti(input_nc=input_nc, ndf=ndf, n_layers=3, norm_layer=norm_layer, num_D=num_Ds, nl_layer=nl_layer)
+    else:
+        raise NotImplementedError('Discriminator model name [%s] is not recognized' % net)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+def define_E(input_nc, output_nc, ndf, netE, norm='batch', nl='lrelu',
+             init_type='xavier', init_gain=0.02, gpu_ids=[], vaeLike=False):
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+    nl = 'lrelu'  # use leaky relu for E
+    nl_layer = get_non_linearity(layer_type=nl)
+    if netE == 'resnet_128':
+        net = E_ResNet(input_nc, output_nc, ndf, n_blocks=4, norm_layer=norm_layer,
+                       nl_layer=nl_layer, vaeLike=vaeLike)
+    elif netE == 'resnet_256':
+        net = E_ResNet(input_nc, output_nc, ndf, n_blocks=5, norm_layer=norm_layer,
+                       nl_layer=nl_layer, vaeLike=vaeLike)
+    elif netE == 'conv_128':
+        net = E_NLayers(input_nc, output_nc, ndf, n_layers=4, norm_layer=norm_layer,
+                        nl_layer=nl_layer, vaeLike=vaeLike)
+    elif netE == 'conv_256':
+        net = E_NLayers(input_nc, output_nc, ndf, n_layers=5, norm_layer=norm_layer,
+                        nl_layer=nl_layer, vaeLike=vaeLike)
     else:
-        gr.Info("Unsupported resolution")
-    validation_mask = Image.open('./assets/mask.png').convert('RGB').resize((tar_width*2, tar_height*2))
-    gr.Info("Image retrieval in progress...")
-    query_image_bw = process_image(input_context, int(tar_width), int(tar_height))
-    query_image = query_image_bw.convert('RGB')
-    query_image_vae = process_image(VAE_input, int(tar_width*1.5), int(tar_height*1.5))
-    reference_images = [process_image(ref_image, tar_width, tar_height) for ref_image in reference_images]
-    query_patches_pil = process_image_Q_varres(query_image, tar_width, tar_height)
-    reference_patches_pil = []
-    for reference_image in reference_images:
-        reference_patches_pil += process_image_ref_varres(reference_image, tar_width, tar_height)
-    combined_image = None
-    with torch.no_grad():
-        clip_img = image_processor(images=query_patches_pil, return_tensors="pt").pixel_values.to(image_encoder.device, dtype=image_encoder.dtype)
-        query_embeddings = image_encoder(clip_img).image_embeds
-        reference_patches_pil_gray = [rimg.convert('RGB').convert('RGB') for rimg in reference_patches_pil]
-        clip_img = image_processor(images=reference_patches_pil_gray, return_tensors="pt").pixel_values.to(image_encoder.device, dtype=image_encoder.dtype)
-        reference_embeddings = image_encoder(clip_img).image_embeds
-        cosine_similarities = F.cosine_similarity(query_embeddings.unsqueeze(1), reference_embeddings.unsqueeze(0), dim=-1)
-        sorted_indices = torch.argsort(cosine_similarities, descending=True, dim=1).tolist()
-        top_k = 3
-        top_k_indices = [cur_sortlist[:top_k] for cur_sortlist in sorted_indices]
-        combined_image = Image.new('RGB', (tar_width * 2, tar_height * 2), 'white')
-        combined_image.paste(query_image_bw.resize((tar_width, tar_height)), (tar_width//2, tar_height//2))
-        idx_table = {0:[(1,0), (0,1), (0,0)], 1:[(1,3), (0,2),(0,3)], 2:[(2,0),(3,1), (3,0)], 3:[(2,3), (3,2),(3,3)]}
-        for i in range(2):
-            for j in range(2):
-                idx_list = idx_table[i * 2 + j]
-                for k in range(top_k):
-                    ref_index = top_k_indices[i * 2 + j][k]
-                    idx_y = idx_list[k][0]
-                    idx_x = idx_list[k][1]
-                    combined_image.paste(reference_patches_pil[ref_index].resize((tar_width//2-2, tar_height//2-2)), (tar_width//2 * idx_x + 1, tar_height//2 * idx_y + 1))
-    gr.Info("Model inference in progress...")
-    generator = torch.Generator(device=device).manual_seed(seed)
-    image = pipeline(
-        "manga", cond_image=combined_image, cond_mask=validation_mask, num_inference_steps=num_inference_steps, generator=generator
-    ).images[0]
-    gr.Info("Post-processing image...")
-    with torch.no_grad():
-        width, height = image.size
-        new_width = width // 2
-        new_height = height // 2
-        left = (width - new_width) // 2
-        top = (height - new_height) // 2
-        right = left + new_width
-        bottom = top + new_height
-        center_crop = image.crop((left, top, right, bottom))
-        up_img = center_crop.resize(query_image_vae.size)
-        test_low_color = transform(up_img).unsqueeze(0).to(device, dtype=weight_dtype)
-        query_image_vae = transform(query_image_vae).unsqueeze(0).to(device, dtype=weight_dtype)
-
-        h_color, hidden_list_color = pipeline.vae._encode(test_low_color,return_dict = False, hidden_flag = True)
-        h_bw, hidden_list_bw = pipeline.vae._encode(query_image_vae, return_dict = False, hidden_flag = True)
-
-        hidden_list_double = [torch.cat((hidden_list_color[hidden_idx], hidden_list_bw[hidden_idx]), dim = 1) for hidden_idx in range(len(hidden_list_color))]
-
-
-        hidden_list = MultiResNetModel(hidden_list_double)
-        output = pipeline.vae._decode(h_color.sample(),return_dict = False, hidden_list = hidden_list)[0]
-
-        output[output > 1] = 1
-        output[output < -1] = -1
-        high_res_image = Image.fromarray(((output[0] * 0.5 + 0.5).permute(1, 2, 0).detach().cpu().numpy() * 255).astype(np.uint8)).convert("RGB")
-    gr.Info("Colorization complete!")
-    torch.cuda.empty_cache()
-    return high_res_image, up_img, image, query_image_bw
-
-with gr.Blocks() as demo:
-    gr.HTML(
+        raise NotImplementedError('Encoder model name [%s] is not recognized' % net)
+
+    return init_net(net, init_type, init_gain, gpu_ids, False)
+
+
+class ResnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, norm_layer=None, use_dropout=False, n_blocks=6, padding_type='replicate'):
+        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:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        model = [nn.ReplicationPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0,
+                           bias=use_bias)]
+        if norm_layer is not None:
+            model += [norm_layer(ngf)]
+        model += [nn.ReLU(True)]
+
+        # n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [nn.ReplicationPad2d(1),nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,
+                                stride=2, padding=0, bias=use_bias)]
+            # model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,
+            #                     stride=2, padding=1, bias=use_bias)]
+            if norm_layer is not None:
+                model += [norm_layer(ngf * mult * 2)]
+            model += [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)]
+            # if norm_layer is not None:
+            #     model += [norm_layer(ngf * mult / 2)]
+            # model += [nn.ReLU(True)]
+            model += upsampleLayer(ngf * mult, int(ngf * mult / 2), upsample='bilinear', padding_type=padding_type)
+            if norm_layer is not None:
+                model += [norm_layer(int(ngf * mult / 2))]
+            model += [nn.ReLU(True)]
+            model +=[nn.ReplicationPad2d(1),
+                     nn.Conv2d(int(ngf * mult / 2), int(ngf * mult / 2), kernel_size=3, padding=0)]
+            if norm_layer is not None:
+                model += [norm_layer(ngf * mult / 2)]
+            model += [nn.ReLU(True)]
+        model += [nn.ReplicationPad2d(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)]
+        if norm_layer is not None:
+            conv_block += [norm_layer(dim)]
+        conv_block += [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)]
+        if norm_layer is not None:
+            conv_block += [norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        out = x + self.conv_block(x)
+        return out
+
+
+class D_NLayersMulti(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3,
+                 norm_layer=nn.BatchNorm2d,  num_D=1, nl_layer=None):
+        super(D_NLayersMulti, self).__init__()
+        # st()
+        self.num_D = num_D
+        self.nl_layer=nl_layer
+        if num_D == 1:
+            layers = self.get_layers(input_nc, ndf, n_layers, norm_layer)
+            self.model = nn.Sequential(*layers)
+        else:
+            layers = self.get_layers(input_nc, ndf, n_layers, norm_layer)
+            self.add_module("model_0", nn.Sequential(*layers))
+            self.down = nn.functional.interpolate
+            for i in range(1, num_D):
+                ndf_i = int(round(ndf / (2**i)))
+                layers = self.get_layers(input_nc, ndf_i, n_layers, norm_layer)
+                self.add_module("model_%d" % i, nn.Sequential(*layers))
+
+    def get_layers(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        kw = 3
+        padw = 1
+        sequence = [spectral_norm(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 += [spectral_norm(nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
+                          kernel_size=kw, stride=2, padding=padw))]
+            if norm_layer:
+                sequence += [norm_layer(ndf * nf_mult)]
+
+            sequence += [self.nl_layer()]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2**n_layers, 8)
+        sequence += [spectral_norm(nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
+                      kernel_size=kw, stride=1, padding=padw))]
+        if norm_layer:
+            sequence += [norm_layer(ndf * nf_mult)]
+        sequence += [self.nl_layer()]
+
+        sequence += [spectral_norm(nn.Conv2d(ndf * nf_mult, 1,
+                               kernel_size=kw, stride=1, padding=padw))]
+
+        return sequence
+
+    def forward(self, input):
+        if self.num_D == 1:
+            return self.model(input)
+        result = []
+        down = input
+        for i in range(self.num_D):
+            model = getattr(self, "model_%d" % i)
+            result.append(model(down))
+            if i != self.num_D - 1:
+                down = self.down(down, scale_factor=0.5, mode='bilinear')
+        return result
+
+class D_NLayers(nn.Module):
+    """Defines a PatchGAN discriminator"""
+
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(D_NLayers, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        kw = 3 
+        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):  # gradually increase the number of filters
+            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)]  # output 1 channel prediction map
+        self.model = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.model(input)
+
+
+class G_Unet_add_input(nn.Module):
+    def __init__(self, input_nc, output_nc, nz, num_downs, ngf=64, 
+                 norm_layer=None, nl_layer=None, use_dropout=False, use_noise=False,
+                 upsample='basic', device=0):
+        super(G_Unet_add_input, self).__init__()
+        self.nz = nz
+        max_nchn = 8
+        noise = []
+        for i in range(num_downs+1):
+            if use_noise:
+                noise.append(True)
+            else:
+                noise.append(False)
+
+        # construct unet structure
+        #print(num_downs)
+        unet_block = UnetBlock_A(ngf * max_nchn, ngf * max_nchn, ngf * max_nchn, noise=noise[num_downs-1], 
+                               innermost=True, norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        for i in range(num_downs - 5):
+            unet_block = UnetBlock_A(ngf * max_nchn, ngf * max_nchn, ngf * max_nchn, unet_block, noise[num_downs-i-3],
+                                   norm_layer=norm_layer, nl_layer=nl_layer, use_dropout=use_dropout, upsample=upsample)
+        unet_block = UnetBlock_A(ngf * 4, ngf * 4, ngf * max_nchn, unet_block, noise[2],
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock_A(ngf * 2, ngf * 2, ngf * 4, unet_block, noise[1],
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock_A(ngf, ngf, ngf * 2, unet_block, noise[0],
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock_A(input_nc + nz, output_nc, ngf, unet_block, None, 
+                               outermost=True, norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+
+        self.model = unet_block
+
+    def forward(self, x, z=None):
+        if self.nz > 0:
+            z_img = z.view(z.size(0), z.size(1), 1, 1).expand(
+                z.size(0), z.size(1), x.size(2), x.size(3))
+            x_with_z = torch.cat([x, z_img], 1)
+        else:
+            x_with_z = x  # no z
+
+
+        return torch.tanh(self.model(x_with_z))
+        # return self.model(x_with_z)
+
+class G_Unet_add_input_G(nn.Module):
+    def __init__(self, input_nc, output_nc, nz, num_downs, ngf=64, 
+                 norm_layer=None, nl_layer=None, use_dropout=False, use_noise=False,
+                 upsample='basic', device=0):
+        super(G_Unet_add_input_G, self).__init__()
+        self.nz = nz
+        max_nchn = 8
+        noise = []
+        for i in range(num_downs+1):
+            if use_noise:
+                noise.append(True)
+            else:
+                noise.append(False)
+        # construct unet structure
+        #print(num_downs)
+        unet_block = UnetBlock_G(ngf * max_nchn, ngf * max_nchn, ngf * max_nchn, noise=False,
+                               innermost=True, norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        for i in range(num_downs - 5):
+            unet_block = UnetBlock_G(ngf * max_nchn, ngf * max_nchn, ngf * max_nchn, unet_block, noise=False,
+                                   norm_layer=norm_layer, nl_layer=nl_layer, use_dropout=use_dropout, upsample=upsample)
+        unet_block = UnetBlock_G(ngf * 4, ngf * 4, ngf * max_nchn, unet_block, noise[2],
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample='basic')
+        unet_block = UnetBlock_G(ngf * 2, ngf * 2, ngf * 4, unet_block, noise[1],
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample='basic')
+        unet_block = UnetBlock_G(ngf, ngf, ngf * 2, unet_block, noise[0],
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample='basic')
+        unet_block = UnetBlock_G(input_nc + nz, output_nc, ngf, unet_block, None,
+                               outermost=True, norm_layer=norm_layer, nl_layer=nl_layer, upsample='basic')
+
+        self.model = unet_block
+
+    def forward(self, x, z=None):
+        if self.nz > 0:
+            z_img = z.view(z.size(0), z.size(1), 1, 1).expand(
+                z.size(0), z.size(1), x.size(2), x.size(3))
+            x_with_z = torch.cat([x, z_img], 1)
+        else:
+            x_with_z = x  # no z
+
+        # return F.tanh(self.model(x_with_z))
+        return self.model(x_with_z)
+
+class G_Unet_add_input_C(nn.Module):
+    def __init__(self, input_nc, output_nc, nz, num_downs, ngf=64, 
+                 norm_layer=None, nl_layer=None, use_dropout=False, use_noise=False,
+                 upsample='basic', device=0):
+        super(G_Unet_add_input_C, self).__init__()
+        self.nz = nz
+        max_nchn = 8
+        # construct unet structure
+        #print(num_downs)
+        unet_block = UnetBlock(ngf * max_nchn, ngf * max_nchn, ngf * max_nchn, noise=False,
+                               innermost=True, norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        for i in range(num_downs - 5):
+            unet_block = UnetBlock(ngf * max_nchn, ngf * max_nchn, ngf * max_nchn, unet_block, noise=False,
+                                   norm_layer=norm_layer, nl_layer=nl_layer, use_dropout=use_dropout, upsample=upsample)
+        unet_block = UnetBlock(ngf * 4, ngf * 4, ngf * max_nchn, unet_block, noise=False,
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock(ngf * 2, ngf * 2, ngf * 4, unet_block, noise=False,
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock(ngf, ngf, ngf * 2, unet_block, noise=False,
+                               norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock(input_nc + nz, output_nc, ngf, unet_block, noise=False,
+                               outermost=True, norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+
+        self.model = unet_block
+
+    def forward(self, x, z=None):
+        if self.nz > 0:
+            z_img = z.view(z.size(0), z.size(1), 1, 1).expand(
+                z.size(0), z.size(1), x.size(2), x.size(3))
+            x_with_z = torch.cat([x, z_img], 1)
+        else:
+            x_with_z = x  # no z
+
+        # return torch.tanh(self.model(x_with_z))
+        return self.model(x_with_z)
+
+def upsampleLayer(inplanes, outplanes, kw=1, upsample='basic', padding_type='replicate'):
+    # padding_type = 'zero'
+    if upsample == 'basic':
+        upconv = [nn.ConvTranspose2d(inplanes, outplanes, kernel_size=4, stride=2, padding=1)]#, padding_mode='replicate'
+    elif upsample == 'bilinear' or upsample == 'nearest' or upsample == 'linear':
+        upconv = [nn.Upsample(scale_factor=2, mode=upsample, align_corners=True),
+                  #nn.ReplicationPad2d(1),
+                  nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=1, padding=0)]
+        # p = kw//2
+        # upconv = [nn.Upsample(scale_factor=2, mode=upsample, align_corners=True),
+        #           nn.Conv2d(inplanes, outplanes, kernel_size=kw, stride=1, padding=p, padding_mode='replicate')]
+    else:
+        raise NotImplementedError(
+            'upsample layer [%s] not implemented' % upsample)
+    return upconv
+
+class UnetBlock_G(nn.Module):
+    def __init__(self, input_nc, outer_nc, inner_nc,
+                 submodule=None, noise=None, outermost=False, innermost=False, 
+                 norm_layer=None, nl_layer=None, use_dropout=False, upsample='basic', padding_type='replicate'):
+        super(UnetBlock_G, self).__init__()
+        self.outermost = outermost
+        p = 0
+        downconv = []
+        if padding_type == 'reflect':
+            downconv += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            downconv += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError(
+                'padding [%s] is not implemented' % padding_type)
+
+        downconv += [nn.Conv2d(input_nc, inner_nc,
+                               kernel_size=3, stride=2, padding=p)]
+        # downsample is different from upsample
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc) if norm_layer is not None else None
+        uprelu = nl_layer()
+        uprelu2 = nl_layer()
+        uppad = nn.ReplicationPad2d(1)
+        upnorm = norm_layer(outer_nc) if norm_layer is not None else None
+        upnorm2 = norm_layer(outer_nc) if norm_layer is not None else None
+        self.noiseblock = ApplyNoise(outer_nc)
+        self.noise = noise
+
+        if outermost:
+            upconv = upsampleLayer(inner_nc * 2, inner_nc, upsample=upsample, padding_type=padding_type)
+            uppad = nn.ReplicationPad2d(3)
+            upconv2 = nn.Conv2d(inner_nc, outer_nc, kernel_size=7, padding=0)
+            down = downconv
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [norm_layer(inner_nc)]
+            # upconv = upsampleLayer(inner_nc * 2, outer_nc, upsample=upsample, padding_type=padding_type)
+            # upconv2 = nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=0)
+            # down = downconv
+            # up = [uprelu] + upconv
+            # if upnorm is not None:
+            #     up += [norm_layer(outer_nc)]
+            up +=[uprelu2, uppad, upconv2] #+ [nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = upsampleLayer(inner_nc, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p)
+            down = [downrelu] + downconv
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [upnorm]
+            up += [uprelu2, uppad, upconv2]
+            if upnorm2 is not None:
+                up += [upnorm2]
+            model = down + up
+        else:
+            upconv = upsampleLayer(inner_nc * 2, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p)
+            down = [downrelu] + downconv
+            if downnorm is not None:
+                down += [downnorm]
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [upnorm]
+            up += [uprelu2, uppad, upconv2]
+            if upnorm2 is not None:
+                up += [upnorm2]
+
+            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:
+            x2 = self.model(x)
+            if self.noise:
+                x2 = self.noiseblock(x2, self.noise)
+            return torch.cat([x2, x], 1)
+
+
+class UnetBlock(nn.Module):
+    def __init__(self, input_nc, outer_nc, inner_nc,
+                 submodule=None, noise=None, outermost=False, innermost=False, 
+                 norm_layer=None, nl_layer=None, use_dropout=False, upsample='basic', padding_type='replicate'):
+        super(UnetBlock, self).__init__()
+        self.outermost = outermost
+        p = 0
+        downconv = []
+        if padding_type == 'reflect':
+            downconv += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            downconv += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError(
+                'padding [%s] is not implemented' % padding_type)
+
+        downconv += [nn.Conv2d(input_nc, inner_nc,
+                               kernel_size=3, stride=2, padding=p)]
+        # downsample is different from upsample
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc) if norm_layer is not None else None
+        uprelu = nl_layer()
+        uprelu2 = nl_layer()
+        uppad = nn.ReplicationPad2d(1)
+        upnorm = norm_layer(outer_nc) if norm_layer is not None else None
+        upnorm2 = norm_layer(outer_nc) if norm_layer is not None else None
+        self.noiseblock = ApplyNoise(outer_nc)
+        self.noise = noise
+
+        if outermost:
+            upconv = upsampleLayer(inner_nc * 2, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p)
+            down = downconv
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [upnorm]
+            up +=[uprelu2, uppad, upconv2] #+ [nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = upsampleLayer(inner_nc, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p)
+            down = [downrelu] + downconv
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [upnorm]
+            up += [uprelu2, uppad, upconv2]
+            if upnorm2 is not None:
+                up += [upnorm2]
+            model = down + up
+        else:
+            upconv = upsampleLayer(inner_nc * 2, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p)
+            down = [downrelu] + downconv
+            if downnorm is not None:
+                down += [downnorm]
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [upnorm]
+            up += [uprelu2, uppad, upconv2]
+            if upnorm2 is not None:
+                up += [upnorm2]
+
+            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:
+            x2 = self.model(x)
+            if self.noise:
+                x2 = self.noiseblock(x2, self.noise)
+            return torch.cat([x2, x], 1)
+
+# Defines the submodule with skip connection.
+# X -------------------identity---------------------- X
+#   |-- downsampling -- |submodule| -- upsampling --|
+class UnetBlock_A(nn.Module):
+    def __init__(self, input_nc, outer_nc, inner_nc,
+                 submodule=None, noise=None, outermost=False, innermost=False, 
+                 norm_layer=None, nl_layer=None, use_dropout=False, upsample='basic', padding_type='replicate'):
+        super(UnetBlock_A, self).__init__()
+        self.outermost = outermost
+        p = 0
+        downconv = []
+        if padding_type == 'reflect':
+            downconv += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            downconv += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError(
+                'padding [%s] is not implemented' % padding_type)
+
+        downconv += [spectral_norm(nn.Conv2d(input_nc, inner_nc,
+                               kernel_size=3, stride=2, padding=p))]
+        # downsample is different from upsample
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc) if norm_layer is not None else None
+        uprelu = nl_layer()
+        uprelu2 = nl_layer()
+        uppad = nn.ReplicationPad2d(1)
+        upnorm = norm_layer(outer_nc) if norm_layer is not None else None
+        upnorm2 = norm_layer(outer_nc) if norm_layer is not None else None
+        self.noiseblock = ApplyNoise(outer_nc)
+        self.noise = noise
+
+        if outermost:
+            upconv = upsampleLayer(inner_nc * 1, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = spectral_norm(nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p))
+            down = downconv
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [upnorm]
+            up +=[uprelu2, uppad, upconv2] #+ [nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = upsampleLayer(inner_nc, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = spectral_norm(nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p))
+            down = [downrelu] + downconv
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [upnorm]
+            up += [uprelu2, uppad, upconv2]
+            if upnorm2 is not None:
+                up += [upnorm2]
+            model = down + up
+        else:
+            upconv = upsampleLayer(inner_nc * 1, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = spectral_norm(nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p))
+            down = [downrelu] + downconv
+            if downnorm is not None:
+                down += [downnorm]
+            up = [uprelu] + upconv
+            if upnorm is not None:
+                up += [upnorm]
+            up += [uprelu2, uppad, upconv2]
+            if upnorm2 is not None:
+                up += [upnorm2]
+
+            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:
+            x2 = self.model(x)
+            if self.noise:
+                x2 = self.noiseblock(x2, self.noise)
+            if x2.shape[-1]==x.shape[-1]:
+                return x2 + x
+            else:
+                x2 = F.interpolate(x2, x.shape[2:])
+                return x2 + x
+
+
+class E_ResNet(nn.Module):
+    def __init__(self, input_nc=3, output_nc=1, ndf=64, n_blocks=4,
+                 norm_layer=None, nl_layer=None, vaeLike=False):
+        super(E_ResNet, self).__init__()
+        self.vaeLike = vaeLike
+        max_ndf = 4
+        conv_layers = [
+            nn.Conv2d(input_nc, ndf, kernel_size=3, stride=2, padding=1, bias=True)]
+        for n in range(1, n_blocks):
+            input_ndf = ndf * min(max_ndf, n)
+            output_ndf = ndf * min(max_ndf, n + 1)
+            conv_layers += [BasicBlock(input_ndf,
+                                       output_ndf, norm_layer, nl_layer)]
+        conv_layers += [nl_layer(), nn.AdaptiveAvgPool2d(4)]
+        if vaeLike:
+            self.fc = nn.Sequential(*[nn.Linear(output_ndf * 16, output_nc)])
+            self.fcVar = nn.Sequential(*[nn.Linear(output_ndf * 16, output_nc)])
+        else:
+            self.fc = nn.Sequential(*[nn.Linear(output_ndf * 16, output_nc)])
+        self.conv = nn.Sequential(*conv_layers)
+
+    def forward(self, x):
+        x_conv = self.conv(x)
+        conv_flat = x_conv.view(x.size(0), -1)
+        output = self.fc(conv_flat)
+        if self.vaeLike:
+            outputVar = self.fcVar(conv_flat)
+            return output, outputVar
+        else:
+            return output
+        return output
+
+
+# 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 G_Unet_add_all(nn.Module):
+    def __init__(self, input_nc, output_nc, nz, num_downs, ngf=64, 
+                 norm_layer=None, nl_layer=None, use_dropout=False, use_noise=False, upsample='basic'):
+        super(G_Unet_add_all, self).__init__()
+        self.nz = nz
+        self.mapping = G_mapping(self.nz, self.nz, 512, normalize_latents=False, lrmul=1)
+        self.truncation_psi = 0
+        self.truncation_cutoff = 0
+
+        # - 2 means we start from feature map with height and width equals 4.
+        # as this example, we get num_layers = 18.
+        num_layers = int(np.log2(512)) * 2 - 2
+        # Noise inputs.
+        self.noise_inputs = []
+        for layer_idx in range(num_layers):
+            res = layer_idx // 2 + 2
+            shape = [1, 1, 2 ** res, 2 ** res]
+            self.noise_inputs.append(torch.randn(*shape).to("cuda" if torch.cuda.is_available() else "cpu"))
+
+        # construct unet structure
+        unet_block = UnetBlock_with_z(ngf * 8, ngf * 8, ngf * 8, nz, submodule=None, innermost=True, 
+                                      norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock_with_z(ngf * 8, ngf * 8, ngf * 8, nz, submodule=unet_block,
+                                      norm_layer=norm_layer, nl_layer=nl_layer, use_dropout=use_dropout, upsample=upsample)
+        for i in range(num_downs - 6):
+            unet_block = UnetBlock_with_z(ngf * 8, ngf * 8, ngf * 8, nz, submodule=unet_block, 
+                                          norm_layer=norm_layer, nl_layer=nl_layer, use_dropout=use_dropout, upsample=upsample)
+        unet_block = UnetBlock_with_z(ngf * 4, ngf * 4, ngf * 8, nz, submodule=unet_block,
+                                      norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock_with_z(ngf * 2, ngf * 2, ngf * 4, nz, submodule=unet_block,
+                                      norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock_with_z(ngf, ngf, ngf * 2, nz, submodule=unet_block, 
+                                      norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        unet_block = UnetBlock_with_z(input_nc, output_nc, ngf, nz, submodule=unet_block,
+                                      outermost=True, norm_layer=norm_layer, nl_layer=nl_layer, upsample=upsample)
+        self.model = unet_block
+
+    def forward(self, x, z):
+
+        dlatents1, num_layers = self.mapping(z)
+        dlatents1 = dlatents1.unsqueeze(1)
+        dlatents1 = dlatents1.expand(-1, int(num_layers), -1)
+
+        # Apply truncation trick.
+        if self.truncation_psi and self.truncation_cutoff:
+            coefs = np.ones([1, num_layers, 1], dtype=np.float32)
+            for i in range(num_layers):
+                if i < self.truncation_cutoff:
+                    coefs[:, i, :] *= self.truncation_psi
+            """Linear interpolation.
+               a + (b - a) * t (a = 0)
+               reduce to
+               b * t
+            """
+            dlatents1 = dlatents1 * torch.Tensor(coefs).to(dlatents1.device)
+
+        return torch.tanh(self.model(x, dlatents1, self.noise_inputs))
+
+
+class ApplyNoise(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.weight = nn.Parameter(torch.randn(channels), requires_grad=True)
+        self.bias = nn.Parameter(torch.zeros(channels), requires_grad=True)
+
+    def forward(self, x, noise):
+        W,_ = torch.split(self.weight.view(1, -1, 1, 1), self.channels // 2, dim=1)
+        B,_ = torch.split(self.bias.view(1, -1, 1, 1), self.channels // 2, dim=1)
+        Z = torch.zeros_like(W)
+        w = torch.cat([W,Z], dim=1).to(x.device)
+        b = torch.cat([B,Z], dim=1).to(x.device)
+        adds = w * torch.randn_like(x) + b
+        return x + adds.type_as(x)
+
+
+class FC(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 gain=2**(0.5),
+                 use_wscale=False,
+                 lrmul=1.0,
+                 bias=True):
+        """
+            The complete conversion of Dense/FC/Linear Layer of original Tensorflow version.
+        """
+        super(FC, self).__init__()
+        he_std = gain * in_channels ** (-0.5)  # He init
+        if use_wscale:
+            init_std = 1.0 / lrmul
+            self.w_lrmul = he_std * lrmul
+        else:
+            init_std = he_std / lrmul
+            self.w_lrmul = lrmul
+
+        self.weight = torch.nn.Parameter(torch.randn(out_channels, in_channels) * init_std)
+        if bias:
+            self.bias = torch.nn.Parameter(torch.zeros(out_channels))
+            self.b_lrmul = lrmul
+        else:
+            self.bias = None
+
+    def forward(self, x):
+        if self.bias is not None:
+            out = F.linear(x, self.weight * self.w_lrmul, self.bias * self.b_lrmul)
+        else:
+            out = F.linear(x, self.weight * self.w_lrmul)
+        out = F.leaky_relu(out, 0.2, inplace=True)
+        return out
+
+
+class ApplyStyle(nn.Module):
     """
-<div style="text-align: center;">
-    <h1 style="text-align: center; font-size: 3em;">🎨 ColorFlow:</h1>
-    <h3 style="text-align: center; font-size: 1.8em;">Retrieval-Augmented Image Sequence Colorization</h3>
-    <p style="text-align: center; font-weight: bold;">
-        <a href="https://zhuang2002.github.io/ColorFlow/">Project Page</a> | 
-        <a href="https://arxiv.org/abs/2412.11815">ArXiv Preprint</a> | 
-        <a href="https://github.com/TencentARC/ColorFlow">GitHub Repository</a>
-    </p>
-    <p style="text-align: center; font-weight: bold;">
-        NOTE: Each time you switch the input style, the corresponding model will be reloaded, which may take some time. Please be patient.
-    </p>
-    <p style="text-align: left; font-size: 1.1em;">
-        Welcome to the demo of <strong>ColorFlow</strong>. Follow the steps below to explore the capabilities of our model:
-    </p>
-</div>
-<div style="text-align: left; margin: 0 auto;">
-    <ol style="font-size: 1.1em;">
-        <li>Choose input style: GrayImage(ScreenStyle)、Sketch with Shading or Sketch.</li>
-        <li>Upload your image: Use the 'Upload' button to select the image you want to colorize.</li>
-        <li>Preprocess the image: Click the 'Preprocess' button to decolorize the image.</li>
-        <li>Upload reference images: Upload multiple reference images to guide the colorization.</li>
-        <li>Set sampling parameters (optional): Adjust the settings and click the <b>Colorize</b> button.</li>
-    </ol>
-    <p>
-        ⏱️ <b>ZeroGPU Time Limit</b>: Hugging Face ZeroGPU has an inference time limit of 180 seconds. You may need to log in with a free account to use this demo. Large sampling steps might lead to timeout (GPU Abort). In that case, please consider logging in with a Pro account or running it on your local machine.
-    </p>
-</div>
-<div style="text-align: center;">
-    <p style="text-align: center; font-weight: bold;">
-        注意：每次切换输入样式时，相应的模型将被重新加载，可能需要一些时间。请耐心等待。
-    </p>
-    <p style="text-align: left; font-size: 1.1em;">
-        欢迎使用 <strong>ColorFlow</strong> 演示。请按照以下步骤探索我们模型的能力：
-    </p>
-</div>
-<div style="text-align: left; margin: 0 auto;">
-    <ol style="font-size: 1.1em;">
-        <li>选择输入样式：灰度图(ScreenStyle)、线稿+阴影、线稿。</li>
-        <li>上传您的图像：使用“上传”按钮选择要上色的图像。</li>
-        <li>预处理图像：点击“预处理”按钮以去色图像。</li>
-        <li>上传参考图像：上传多张参考图像以指导上色。</li>
-        <li>设置采样参数（可选）：调整设置并点击 <b>上色</b> 按钮。</li>
-    </ol>
-    <p>
-        ⏱️ <b>ZeroGPU时间限制</b>：Hugging Face ZeroGPU 的推理时间限制为 180 秒。您可能需要使用免费帐户登录以使用此演示。大采样步骤可能会导致超时（GPU 中止）。在这种情况下，请考虑使用专业帐户登录或在本地计算机上运行。
-    </p>
-</div>
+        @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
     """
-)
-    VAE_input = gr.State()
-    input_context = gr.State()
-    # example_loading = gr.State(value=None)
-    
-    with gr.Column():
-        with gr.Row():
-            input_style = gr.Radio(["GrayImage(ScreenStyle)", "Sketch_Shading", "Sketch"], label="Input Style", value="GrayImage(ScreenStyle)")
-        with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(type="pil", label="Image to Colorize")
-                resolution = gr.Radio(["640x640", "512x800", "800x512"], label="Select Resolution(Width*Height)", value="640x640")
-                extract_button = gr.Button("Preprocess (Decolorize)")
-            extracted_image = gr.Image(type="pil", label="Decolorized Result")
-        with gr.Row():
-            reference_images = gr.Files(label="Reference Images (Upload multiple)", file_count="multiple")
-            with gr.Column():
-                output_gallery = gr.Gallery(label="Colorization Results", type="pil")
-                seed = gr.Slider(label="Random Seed", minimum=0, maximum=100000, value=0, step=1)
-                num_inference_steps = gr.Slider(label="Inference Steps", minimum=4, maximum=100, value=10, step=1)
-                colorize_button = gr.Button("Colorize")
-    
-    # progress_text = gr.Textbox(label="Progress", interactive=False)
-    
-    
-    extract_button.click(
-        extract_line_image, 
-        inputs=[input_image, input_style, resolution], 
-        outputs=[extracted_image, VAE_input, input_context]
-    )
-    colorize_button.click(
-        colorize_image, 
-        inputs=[VAE_input, input_context, reference_images, resolution, seed, input_style, num_inference_steps], 
-        outputs=output_gallery
-    )
-
-    with gr.Column():
-        gr.Markdown("### Quick Examples")
-        gr.Examples(
-            examples=examples,
-            inputs=[input_image, reference_images, input_style, resolution, seed, num_inference_steps],
-            label="Examples",
-            examples_per_page=8,
-        )
-    gr.HTML('<a href="https://github.com/TencentARC/ColorFlow"><img src="https://img.shields.io/github/stars/TencentARC/ColorFlow" alt="GitHub Stars"></a>')
-    gr.Markdown(article)
-    # gr.HTML(
-    #     '<a href="https://github.com/TencentARC/ColorFlow"><img src="https://img.shields.io/github/stars/TencentARC/ColorFlow" alt="GitHub Stars"></a>'
-    # )
-
-    
-demo.launch()
+    def __init__(self, latent_size, channels, use_wscale, nl_layer):
+        super(ApplyStyle, self).__init__()
+        modules = [nn.Linear(latent_size, channels*2)]
+        if nl_layer:
+            modules += [nl_layer()]
+        self.linear = nn.Sequential(*modules)
+
+    def forward(self, x, latent):
+        style = self.linear(latent)  # style => [batch_size, n_channels*2]
+        shape = [-1, 2, x.size(1), 1, 1]
+        style = style.view(shape)    # [batch_size, 2, n_channels, ...]
+        x = x * (style[:, 0] + 1.) + style[:, 1]
+        return x
+
+class PixelNorm(nn.Module):
+    def __init__(self, epsilon=1e-8):
+        """
+            @notice: avoid in-place ops.
+            https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+        """
+        super(PixelNorm, self).__init__()
+        self.epsilon = epsilon
+
+    def forward(self, x):
+        tmp  = torch.mul(x, x) # or x ** 2
+        tmp1 = torch.rsqrt(torch.mean(tmp, dim=1, keepdim=True) + self.epsilon)
+
+        return x * tmp1
+
+
+class InstanceNorm(nn.Module):
+    def __init__(self, epsilon=1e-8):
+        """
+            @notice: avoid in-place ops.
+            https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+        """
+        super(InstanceNorm, self).__init__()
+        self.epsilon = epsilon
+
+    def forward(self, x):
+        x   = x - torch.mean(x, (2, 3), True)
+        tmp = torch.mul(x, x) # or x ** 2
+        tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon)
+        return x * tmp
+
+
+class LayerEpilogue(nn.Module):
+    def __init__(self, channels, dlatent_size, use_wscale, use_noise,
+                 use_pixel_norm, use_instance_norm, use_styles, nl_layer=None):
+        super(LayerEpilogue, self).__init__()
+        self.use_noise = use_noise
+        if use_noise:
+            self.noise = ApplyNoise(channels)
+        self.act = nn.LeakyReLU(negative_slope=0.2)
+
+        if use_pixel_norm:
+            self.pixel_norm = PixelNorm()
+        else:
+            self.pixel_norm = None
+
+        if use_instance_norm:
+            self.instance_norm = InstanceNorm()
+        else:
+            self.instance_norm = None
+
+        if use_styles:
+            self.style_mod = ApplyStyle(dlatent_size, channels, use_wscale=use_wscale, nl_layer=nl_layer)
+        else:
+            self.style_mod = None
+
+    def forward(self, x, noise, dlatents_in_slice=None):
+        # if noise is not None:
+        if self.use_noise:
+            x = self.noise(x, noise)
+        x = self.act(x)
+        if self.pixel_norm is not None:
+            x = self.pixel_norm(x)
+        if self.instance_norm is not None:
+            x = self.instance_norm(x)
+        if self.style_mod is not None:
+            x = self.style_mod(x, dlatents_in_slice)
+
+        return x
+
+class G_mapping(nn.Module):
+    def __init__(self,
+                 mapping_fmaps=512,
+                 dlatent_size=512,
+                 resolution=512,
+                 normalize_latents=True,  # Normalize latent vectors (Z) before feeding them to the mapping layers?
+                 use_wscale=True,         # Enable equalized learning rate?
+                 lrmul=0.01,              # Learning rate multiplier for the mapping layers.
+                 gain=2**(0.5),            # original gain in tensorflow.
+                 nl_layer=None
+                 ):
+        super(G_mapping, self).__init__()
+        self.mapping_fmaps = mapping_fmaps
+        func = [
+            nn.Linear(self.mapping_fmaps, dlatent_size)
+        ]
+        if nl_layer:
+            func += [nl_layer()]
+
+        for j in range(0,4):
+            func += [
+                nn.Linear(dlatent_size, dlatent_size)
+            ]
+            if nl_layer:
+                func += [nl_layer()]
+
+        self.func = nn.Sequential(*func)
+            #FC(self.mapping_fmaps, dlatent_size, gain, lrmul=lrmul, use_wscale=use_wscale),
+            #FC(dlatent_size, dlatent_size, gain, lrmul=lrmul, use_wscale=use_wscale),
+
+        self.normalize_latents = normalize_latents
+        self.resolution_log2 = int(np.log2(resolution))
+        self.num_layers = self.resolution_log2 * 2 - 2
+        self.pixel_norm = PixelNorm()
+        # - 2 means we start from feature map with height and width equals 4.
+        # as this example, we get num_layers = 18.
+
+    def forward(self, x):
+        if self.normalize_latents:
+            x = self.pixel_norm(x)
+        out = self.func(x)
+        return out, self.num_layers
+
+class UnetBlock_with_z(nn.Module):
+    def __init__(self, input_nc, outer_nc, inner_nc, nz=0, 
+                 submodule=None, outermost=False, innermost=False, 
+                 norm_layer=None, nl_layer=None, use_dropout=False, 
+                 upsample='basic', padding_type='replicate'):
+        super(UnetBlock_with_z, self).__init__()
+        p = 0
+        downconv = []
+        if padding_type == 'reflect':
+            downconv += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            downconv += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError(
+                'padding [%s] is not implemented' % padding_type)
+
+        self.outermost = outermost
+        self.innermost = innermost
+        self.nz = nz
+
+        # input_nc = input_nc + nz
+        downconv += [spectral_norm(nn.Conv2d(input_nc, inner_nc,
+                               kernel_size=3, stride=2, padding=p))]
+        # downsample is different from upsample
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc) if norm_layer is not None else None
+        uprelu = nl_layer()
+        uprelu2 = nl_layer()
+        uppad = nn.ReplicationPad2d(1)
+        upnorm = norm_layer(outer_nc) if norm_layer is not None else None
+        upnorm2 = norm_layer(outer_nc) if norm_layer is not None else None
+
+        use_styles=False
+        uprelu = nl_layer()
+        if self.nz >0:
+            use_styles=True
+
+        if outermost:
+            self.adaIn = LayerEpilogue(inner_nc, self.nz, use_wscale=True, use_noise=False,
+                                        use_pixel_norm=True, use_instance_norm=True, use_styles=use_styles, nl_layer=nl_layer)
+            upconv = upsampleLayer(
+                inner_nc , outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = spectral_norm(nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p))
+            down = downconv
+            up = [uprelu] + upconv 
+            if upnorm is not None:
+                up += [upnorm]
+            up +=[uprelu2, uppad, upconv2] #+ [nn.Tanh()]
+        elif innermost:
+            self.adaIn = LayerEpilogue(inner_nc, self.nz, use_wscale=True, use_noise=True,
+                                        use_pixel_norm=True, use_instance_norm=True, use_styles=use_styles, nl_layer=nl_layer)
+            upconv = upsampleLayer(
+                inner_nc, outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = spectral_norm(nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p))
+            down = [downrelu] + downconv
+            up = [uprelu] + upconv
+            if norm_layer is not None:
+                up += [norm_layer(outer_nc)]
+            up += [uprelu2, uppad, upconv2]
+            if upnorm2 is not None:
+                up += [upnorm2]
+        else:
+            self.adaIn = LayerEpilogue(inner_nc, self.nz, use_wscale=True, use_noise=False,
+                                        use_pixel_norm=True, use_instance_norm=True, use_styles=use_styles, nl_layer=nl_layer)
+            upconv = upsampleLayer(
+                inner_nc , outer_nc, upsample=upsample, padding_type=padding_type)
+            upconv2 = spectral_norm(nn.Conv2d(outer_nc, outer_nc, kernel_size=3, padding=p))
+            down = [downrelu] + downconv
+            if norm_layer is not None:
+                down += [norm_layer(inner_nc)]
+            up = [uprelu] + upconv
+
+            if norm_layer is not None:
+                up += [norm_layer(outer_nc)]
+            up += [uprelu2, uppad, upconv2]
+            if upnorm2 is not None:
+                up += [upnorm2]
+
+            if use_dropout:
+                up += [nn.Dropout(0.5)]
+        self.down = nn.Sequential(*down)
+        self.submodule = submodule
+        self.up = nn.Sequential(*up)
+
+
+    def forward(self, x, z, noise):
+        if self.outermost:
+            x1 = self.down(x)
+            x2 = self.submodule(x1, z[:,2:], noise[2:])
+            return self.up(x2)
+
+        elif self.innermost:
+            x1 = self.down(x)
+            x_and_z = self.adaIn(x1, noise[0], z[:,0])
+            x2 = self.up(x_and_z)
+            x2 = F.interpolate(x2, x.shape[2:])
+            return x2 + x
+
+        else:
+            x1 = self.down(x)
+            x2 = self.submodule(x1, z[:,2:], noise[2:])
+            x_and_z = self.adaIn(x2, noise[0], z[:,0])
+            return self.up(x_and_z) + x
+
+
+class E_NLayers(nn.Module):
+    def __init__(self, input_nc, output_nc=1, ndf=64, n_layers=4,
+                 norm_layer=None, nl_layer=None, vaeLike=False):
+        super(E_NLayers, self).__init__()
+        self.vaeLike = vaeLike
+
+        kw, padw = 3, 1
+        sequence = [spectral_norm(nn.Conv2d(input_nc, ndf, kernel_size=kw,
+                              stride=2, padding=padw, padding_mode='replicate')), nl_layer()]
+
+        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 += [spectral_norm(nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
+                         kernel_size=kw, stride=2, padding=padw, padding_mode='replicate'))]
+            if norm_layer is not None:
+                sequence += [norm_layer(ndf * nf_mult)]
+            sequence += [nl_layer()]
+        sequence += [nn.AdaptiveAvgPool2d(4)]
+        self.conv = nn.Sequential(*sequence)
+        self.fc = nn.Sequential(*[spectral_norm(nn.Linear(ndf * nf_mult * 16, output_nc))])
+        if vaeLike:
+            self.fcVar = nn.Sequential(*[spectral_norm(nn.Linear(ndf * nf_mult * 16, output_nc))])
+
+    def forward(self, x):
+        x_conv = self.conv(x)
+        conv_flat = x_conv.view(x.size(0), -1)
+        output = self.fc(conv_flat)
+        if self.vaeLike:
+            outputVar = self.fcVar(conv_flat)
+            return output, outputVar
+        return output
+
+class BasicBlock(nn.Module):
+    def __init__(self, inplanes, outplanes):
+        super(BasicBlock, self).__init__()
+        layers = []
+        norm_layer=get_norm_layer(norm_type='layer') #functools.partial(LayerNorm)
+        # norm_layer = None
+        nl_layer=nn.ReLU()
+        if norm_layer is not None:
+            layers += [norm_layer(inplanes)]
+        layers += [nl_layer]
+        layers += [nn.ReplicationPad2d(1),
+                   nn.Conv2d(inplanes, outplanes, kernel_size=3, stride=1,
+                     padding=0, bias=True)]
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+def define_SVAE(inc=96, outc=3, outplanes=64, blocks=1, netVAE='SVAE', model_name='', load_ext=True, save_dir='',
+    init_type="normal", init_gain=0.02, gpu_ids=[]):
+    if netVAE == 'SVAE':
+        net = ScreenVAE(inc=inc, outc=outc, outplanes=outplanes, blocks=blocks, save_dir=save_dir, 
+            init_type=init_type, init_gain=init_gain, gpu_ids=gpu_ids)
+    else:
+        raise NotImplementedError('Encoder model name [%s] is not recognized' % net)
+    init_net(net, init_type=init_type, init_gain=init_gain, gpu_ids=gpu_ids)
+    net.load_networks('latest')
+    return net
+
+
+class ScreenVAE(nn.Module):
+    def __init__(self,inc=1,outc=4, outplanes=64, downs=5, blocks=2,load_ext=True, save_dir='',init_type="normal", init_gain=0.02, gpu_ids=[]):
+        super(ScreenVAE, self).__init__()
+        self.inc = inc
+        self.outc = outc
+        self.save_dir = save_dir
+        norm_layer=functools.partial(LayerNormWarpper)
+        nl_layer=nn.LeakyReLU
+
+        self.model_names=['enc','dec']
+        self.enc=define_C(inc+1, outc*2, 0, 24, netC='resnet_6blocks', 
+                                      norm='layer', nl='lrelu', use_dropout=True, init_type='kaiming', 
+                                      gpu_ids=gpu_ids, upsample='bilinear')
+        self.dec=define_G(outc, inc, 0, 48, netG='unet_128_G', 
+                                      norm='layer', nl='lrelu', use_dropout=True, init_type='kaiming', 
+                                      gpu_ids=gpu_ids, where_add='input', upsample='bilinear', use_noise=True)
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_networks(self, epoch):
+        """Load all the networks from the disk.
+
+        Parameters:
+            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
+        """
+        for name in self.model_names:
+            if isinstance(name, str):
+                load_filename = '%s_net_%s.pth' % (epoch, name)
+                load_path = os.path.join(self.save_dir, load_filename)
+                net = getattr(self, name)
+                if isinstance(net, torch.nn.DataParallel):
+                    net = net.module
+                print('loading the model from %s' % load_path)
+                state_dict = torch.load(
+                    load_path, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+                if hasattr(state_dict, '_metadata'):
+                    del state_dict._metadata
+
+                net.load_state_dict(state_dict)
+                del state_dict
+
+    def npad(self, im, pad=128):
+        h,w = im.shape[-2:]
+        hp = h //pad*pad+pad
+        wp = w //pad*pad+pad
+        return F.pad(im, (0, wp-w, 0, hp-h), mode='replicate')
+
+    def forward(self, x, line=None, img_input=True, output_screen_only=True):
+        if img_input:
+            if line is None:
+                line = torch.ones_like(x)
+            else:
+                line = torch.sign(line)
+                x = torch.clamp(x + (1-line),-1,1)
+            h,w = x.shape[-2:]
+            input = torch.cat([x, line], 1)
+            input = self.npad(input)
+            inter = self.enc(input)[:,:,:h,:w]
+            scr, logvar = torch.split(inter, (self.outc, self.outc), dim=1)
+            if output_screen_only:
+                return scr
+            recons = self.dec(scr)
+            return recons, scr, logvar
+        else:
+            h,w = x.shape[-2:]
+            x = self.npad(x)
+            recons = self.dec(x)[:,:,:h,:w]
+            recons = (recons+1)*(line+1)/2-1
+            return torch.clamp(recons,-1,1)