app

BidirectionalTranslation/LICENSE

@@ -0,0 +1,26 @@
+Manga Filling Style Conversion with Screentone Variational Autoencoder
+
+Copyright (c) 2020 The Chinese University of Hong Kong
+
+Copyright and License Information: The source code, the binary executable, and all data files (hereafter, Software) are copyrighted by The Chinese University of Hong Kong and Tien-Tsin Wong (hereafter, Author), Copyright (c) 2021 The Chinese University of Hong Kong. All Rights Reserved.
+
+The Author grants to you ("Licensee") a non-exclusive license to use the Software for academic, research and commercial purposes, without fee. For commercial use, Licensee should submit a WRITTEN NOTICE to the Author. The notice should clearly identify the software package/system/hardware (name, version, and/or model number) using the Software. Licensee may distribute the Software to third parties provided that the copyright notice and this statement appears on all copies. Licensee agrees that the copyright notice and this statement will appear on all copies of the Software, or portions thereof. The Author retains exclusive ownership of the Software.
+
+Licensee may make derivatives of the Software, provided that such derivatives can only be used for the purposes specified in the license grant above.
+
+THE AUTHOR MAKES NO REPRESENTATIONS OR WARRANTIES ABOUT THE SUITABILITY OF THE SOFTWARE, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, OR NON-INFRINGEMENT. THE AUTHOR SHALL NOT BE LIABLE FOR ANY DAMAGES SUFFERED BY LICENSEE AS A RESULT OF USING, MODIFYING OR DISTRIBUTING THE SOFTWARE OR ITS DERIVATIVES.
+
+By using the source code, Licensee agrees to cite the following papers in
+Licensee's publication/work:
+
+  Minshan Xie, Chengze Li, Xueting Liu and Tien-Tsin Wong
+  "Manga Filling Style Conversion with Screentone Variational Autoencoder"
+  ACM Transactions on Graphics (SIGGRAPH Asia 2020 issue), Vol. 39, No. 6, December 2020, pp. 226:1-226:15.
+
+
+By using or copying the Software, Licensee agrees to abide by the intellectual property laws, and all other applicable laws of the U.S., and the terms of this license.
+
+Author shall have the right to terminate this license immediately by written notice upon Licensee's breach of, or non-compliance with, any of its terms.
+Licensee may be held legally responsible for any infringement that is caused or encouraged by Licensee's failure to abide by the terms of this license.
+
+For more information or comments, send mail to: ttwong@acm.org

BidirectionalTranslation/README.md

@@ -0,0 +1,72 @@
+# Bidirectional Translation
+
+Pytorch implementation for multimodal comic-to-manga translation. 
+
+**Note**: The current software works well with PyTorch 1.6.0+. 
+
+## Prerequisites
+- Linux
+- Python 3
+- CPU or NVIDIA GPU + CUDA CuDNN
+
+## Getting Started ###
+### Installation
+- Clone this repo:
+```bash
+git clone https://github.com/msxie/ScreenStyle.git
+cd ScreenStyle/MangaScreening
+```
+- Install PyTorch and dependencies from http://pytorch.org
+- Install python libraries [tensorboardX](https://github.com/lanpa/tensorboardX)
+- Install other libraries
+For pip users:
+```
+pip install -r requirements.txt
+```
+
+## Data praperation
+The training requires paired data (including manga image, western image and their line drawings). 
+The line drawing can be extracted using [MangaLineExtraction](https://github.com/ljsabc/MangaLineExtraction).
+
+  ```
+${DATASET} 
+|-- color2manga 
+|   |-- val 
+|   |   |-- ${FOLDER}
+|   |   |   |-- imgs
+|   |   |   |   |-- 0001.png 
+|   |   |   |   |-- ...
+|   |   |   |-- line
+|   |   |   |   |-- 0001.png 
+|   |   |   |   |-- ...
+  ```
+
+### Use a Pre-trained Model
+- Download the pre-trained [ScreenVAE](https://drive.google.com/file/d/1OBxWHjijMwi9gfTOfDiFiHRZA_CXNSWr/view?usp=sharing) model and place under `checkpoints/ScreenVAE/` folder.
+
+- Download the pre-trained [color2manga](https://drive.google.com/file/d/18-N1W0t3igWLJWFyplNZ5Fa2YHWASCZY/view?usp=sharing) model and place under `checkpoints/color2manga/` folder.
+- Generate results with the model
+```bash
+bash ./scripts/test_western2manga.sh
+```
+
+## Copyright and License
+You are granted with the [LICENSE](LICENSE) for both academic and commercial usages.
+
+## Citation
+If you find the code helpful in your resarch or work, please cite the following papers.
+```
+@article{xie-2020-manga,
+        author   = {Minshan Xie and Chengze Li and Xueting Liu and Tien-Tsin Wong},
+        title    = {Manga Filling Style Conversion with Screentone Variational Autoencoder},
+        journal  = {ACM Transactions on Graphics (SIGGRAPH Asia 2020 issue)},
+        month    = {December},
+        year     = {2020},
+        volume   = {39},
+        number   = {6},
+        pages    = {226:1--226:15}
+    }
+```
+
+### Acknowledgements
+This code borrows heavily from the [pytorch-CycleGAN-and-pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) repository.

BidirectionalTranslation/data/__init__.py

@@ -0,0 +1,100 @@
+"""This package includes all the modules related to data loading and preprocessing
+
+ To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.
+ You need to implement four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point from data loader.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
+
+Now you can use the dataset class by specifying flag '--dataset_mode dummy'.
+See our template dataset class 'template_dataset.py' for more details.
+"""
+import importlib
+import torch.utils.data
+from data.base_dataset import BaseDataset
+
+
+def find_dataset_using_name(dataset_name):
+    """Import the module "data/[dataset_name]_dataset.py".
+
+    In the file, the class called DatasetNameDataset() will
+    be instantiated. It has to be a subclass of BaseDataset,
+    and it is case-insensitive.
+    """
+    dataset_filename = "data." + dataset_name + "_dataset"
+    datasetlib = importlib.import_module(dataset_filename)
+
+    dataset = None
+    target_dataset_name = dataset_name.replace('_', '') + 'dataset'
+    for name, cls in datasetlib.__dict__.items():
+        if name.lower() == target_dataset_name.lower() \
+           and issubclass(cls, BaseDataset):
+            dataset = cls
+
+    if dataset is None:
+        raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))
+
+    return dataset
+
+
+def get_option_setter(dataset_name):
+    """Return the static method <modify_commandline_options> of the dataset class."""
+    dataset_class = find_dataset_using_name(dataset_name)
+    return dataset_class.modify_commandline_options
+
+
+def create_dataset(opt):
+    """Create a dataset given the option.
+
+    This function wraps the class CustomDatasetDataLoader.
+        This is the main interface between this package and 'train.py'/'test.py'
+
+    Example:
+        >>> from data import create_dataset
+        >>> dataset = create_dataset(opt)
+    """
+    data_loader = CustomDatasetDataLoader(opt)
+    dataset = data_loader.load_data()
+    return dataset
+
+
+class CustomDatasetDataLoader():
+    """Wrapper class of Dataset class that performs multi-threaded data loading"""
+
+    def __init__(self, opt):
+        """Initialize this class
+
+        Step 1: create a dataset instance given the name [dataset_mode]
+        Step 2: create a multi-threaded data loader.
+        """
+        self.opt = opt
+        dataset_class = find_dataset_using_name(opt.dataset_mode)
+        self.dataset = dataset_class(opt)
+        print("dataset [%s] was created" % type(self.dataset).__name__)
+
+        train_sampler = None
+        if len(opt.gpu_ids) > 1:
+            train_sampler = torch.utils.data.distributed.DistributedSampler(self.dataset)    
+
+        self.dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=opt.batch_size,
+            #shuffle=not opt.serial_batches,
+            num_workers=int(opt.num_threads),
+            pin_memory=True, sampler=train_sampler
+            )
+
+    def load_data(self):
+        return self
+
+    def __len__(self):
+        """Return the number of data in the dataset"""
+        return min(len(self.dataset), self.opt.max_dataset_size)
+
+    def __iter__(self):
+        """Return a batch of data"""
+        for i, data in enumerate(self.dataloader):
+            if i * self.opt.batch_size >= self.opt.max_dataset_size:
+                break
+            yield data

BidirectionalTranslation/data/aligned_dataset.py

@@ -0,0 +1,60 @@
+import os.path
+from data.base_dataset import BaseDataset, get_params, get_transform
+from data.image_folder import make_dataset
+from PIL import Image
+
+
+class AlignedDataset(BaseDataset):
+    """A dataset class for paired image dataset.
+
+    It assumes that the directory '/path/to/data/train' contains image pairs in the form of {A,B}.
+    During test time, you need to prepare a directory '/path/to/data/test'.
+    """
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseDataset.__init__(self, opt)
+        self.dir_AB = os.path.join(opt.dataroot, opt.phase)  # get the image directory
+        self.AB_paths = sorted(make_dataset(self.dir_AB, opt.max_dataset_size))  # get image paths
+        assert(self.opt.load_size >= self.opt.crop_size)   # crop_size should be smaller than the size of loaded image
+        self.input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
+        self.output_nc = self.opt.input_nc if self.opt.direction == 'BtoA' else self.opt.output_nc
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns a dictionary that contains A, B, A_paths and B_paths
+            A (tensor) - - an image in the input domain
+            B (tensor) - - its corresponding image in the target domain
+            A_paths (str) - - image paths
+            B_paths (str) - - image paths (same as A_paths)
+        """
+        # read a image given a random integer index
+        AB_path = self.AB_paths[index%len(self.AB_paths)]
+        AB = Image.open(AB_path).convert('RGB')
+        # split AB image into A and B
+        w, h = AB.size
+        w2 = int(w / 2)
+        A = AB.crop((0, 0, w2, h))
+        B = AB.crop((w2, 0, w, h))
+
+        # apply the same transform to both A and B
+        transform_params = get_params(self.opt, A.size)
+        A_transform = get_transform(self.opt, transform_params, grayscale=(self.input_nc == 1))
+        B_transform = get_transform(self.opt, transform_params, grayscale=(self.output_nc == 1))
+
+        A = A_transform(A)
+        B = B_transform(B)
+
+        return {'A': A, 'B': B, 'A_paths': AB_path, 'B_paths': AB_path}
+
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return len(self.AB_paths)*100

BidirectionalTranslation/data/base_dataset.py

@@ -0,0 +1,164 @@
+"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
+
+It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
+"""
+import random
+import numpy as np
+import torch.utils.data as data
+from PIL import Image, ImageOps
+import torchvision.transforms as transforms
+from abc import ABC, abstractmethod
+
+
+class BaseDataset(data.Dataset, ABC):
+    """This class is an abstract base class (ABC) for datasets.
+
+    To create a subclass, you need to implement the following four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
+    """
+
+    def __init__(self, opt):
+        """Initialize the class; save the options in the class
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        self.opt = opt
+        self.root = opt.dataroot
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        return parser
+
+    @abstractmethod
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return 0
+
+    @abstractmethod
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns:
+            a dictionary of data with their names. It ususally contains the data itself and its metadata information.
+        """
+        pass
+
+
+def get_params(opt, size):
+    w, h = size
+    new_h = h
+    new_w = w
+    crop = 0
+    if opt.preprocess == 'resize_and_crop':
+        new_h = new_w = opt.load_size
+    elif opt.preprocess == 'scale_width_and_crop':
+        new_w = opt.load_size
+        new_h = opt.load_size * h // w
+
+    # x = random.randint(0, np.maximum(0, new_w - opt.crop_size))
+    # y = random.randint(0, np.maximum(0, new_h - opt.crop_size))
+
+    x = random.randint(crop, np.maximum(0, new_w - opt.crop_size-crop))
+    y = random.randint(crop, np.maximum(0, new_h - opt.crop_size-crop))
+
+    flip = random.random() > 0.5
+
+    return {'crop_pos': (x, y), 'flip': flip}
+
+
+def get_transform(opt, params=None, grayscale=False, method=Image.BICUBIC, convert=True):
+    transform_list = []
+    if grayscale:
+        transform_list.append(transforms.Grayscale(1))
+    if 'resize' in opt.preprocess:
+        osize = [opt.load_size, opt.load_size]
+        transform_list.append(transforms.Resize(osize, method))
+    elif 'scale_width' in opt.preprocess:
+        transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, method)))
+
+    if 'crop' in opt.preprocess:
+        if params is None:
+            # transform_list.append(transforms.RandomCrop(opt.crop_size))
+            transform_list.append(transforms.CenterCrop(opt.crop_size))
+        else:
+            transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size)))
+
+    if opt.preprocess == 'none':
+        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=2**8, method=method)))
+
+    if not opt.no_flip:
+        if params is None:
+            transform_list.append(transforms.RandomHorizontalFlip())
+        elif params['flip']:
+            transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
+            
+    # transform_list += [transforms.ToTensor()]
+    if convert:
+        transform_list += [transforms.ToTensor()]
+        if grayscale:
+            transform_list += [transforms.Normalize((0.5,), (0.5,))]
+        else:
+            transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    return transforms.Compose(transform_list)
+
+
+def __make_power_2(img, base, method=Image.BICUBIC):
+    ow, oh = img.size
+    h = int((oh+base-1) // base * base)
+    w = int((ow+base-1) // base * base)
+    if (h == oh) and (w == ow):
+        return img
+
+    __print_size_warning(ow, oh, w, h)
+    return ImageOps.expand(img, (0, 0, w-ow, h-oh), fill=255)
+
+
+def __scale_width(img, target_width, method=Image.BICUBIC):
+    ow, oh = img.size
+    if (ow == target_width):
+        return img
+    w = target_width
+    h = int(target_width * oh / ow)
+    return img.resize((w, h), method)
+
+
+def __crop(img, pos, size):
+    ow, oh = img.size
+    x1, y1 = pos
+    tw = th = size
+    if (ow > tw or oh > th):
+        return img.crop((x1, y1, x1 + tw, y1 + th))
+    return img
+
+
+def __flip(img, flip):
+    if flip:
+        return img.transpose(Image.FLIP_LEFT_RIGHT)
+    return img
+
+
+def __print_size_warning(ow, oh, w, h):
+    """Print warning information about image size(only print once)"""
+    if not hasattr(__print_size_warning, 'has_printed'):
+        print("The image size needs to be a multiple of 4. "
+              "The loaded image size was (%d, %d), so it was adjusted to "
+              "(%d, %d). This adjustment will be done to all images "
+              "whose sizes are not multiples of 4" % (ow, oh, w, h))
+        __print_size_warning.has_printed = True
+

BidirectionalTranslation/data/image_folder.py

@@ -0,0 +1,66 @@
+"""A modified image folder class
+
+We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
+so that this class can load images from both current directory and its subdirectories.
+"""
+
+import torch.utils.data as data
+
+from PIL import Image
+import os
+import os.path
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG', '.npz', 'npy',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir, max_dataset_size=float("inf")):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+    for root, _, fnames in sorted(os.walk(dir)):
+        for fname in fnames:
+            if is_image_file(fname):
+                path = os.path.join(root, fname)
+                images.append(path)
+    return images[:min(max_dataset_size, len(images))]
+
+
+def default_loader(path):
+    return Image.open(path).convert('RGB')
+
+
+class ImageFolder(data.Dataset):
+
+    def __init__(self, root, transform=None, return_paths=False,
+                 loader=default_loader):
+        imgs = make_dataset(root)
+        if len(imgs) == 0:
+            raise(RuntimeError("Found 0 images in: " + root + "\n"
+                               "Supported image extensions are: " +
+                               ",".join(IMG_EXTENSIONS)))
+
+        self.root = root
+        self.imgs = imgs
+        self.transform = transform
+        self.return_paths = return_paths
+        self.loader = loader
+
+    def __getitem__(self, index):
+        path = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.return_paths:
+            return img, path
+        else:
+            return img
+
+    def __len__(self):
+        return len(self.imgs)

BidirectionalTranslation/data/singleCo_dataset.py

@@ -0,0 +1,85 @@
+import os.path
+from data.base_dataset import BaseDataset, get_params, get_transform
+from data.image_folder import make_dataset
+from PIL import Image, ImageEnhance
+import random
+import numpy as np
+import torch
+import torch.nn.functional as F
+import cv2
+
+
+class SingleCoDataset(BaseDataset):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def __init__(self, opt):
+        self.opt = opt
+        self.root = opt.dataroot
+        self.dir_A = os.path.join(opt.dataroot, opt.phase, opt.folder, 'imgs')
+
+        self.A_paths = make_dataset(self.dir_A)
+
+        self.A_paths = sorted(self.A_paths)
+
+        self.A_size = len(self.A_paths)
+        # self.transform = get_transform(opt)
+
+    def __getitem__(self, index):
+        A_path = self.A_paths[index]
+
+        A_img = Image.open(A_path).convert('RGB')
+        # enhancer = ImageEnhance.Brightness(A_img)
+        # A_img = enhancer.enhance(1.5)
+        if os.path.exists(A_path.replace('imgs','line')[:-4]+'.jpg'):
+            # L_img = Image.open(A_path.replace('imgs','line')[:-4]+'.png')
+            L_img = cv2.imread(A_path.replace('imgs','line')[:-4]+'.jpg')
+            kernel = np.ones((3,3), np.uint8)
+            L_img = cv2.erode(L_img, kernel, iterations=1)
+            L_img = Image.fromarray(L_img)
+        else:
+            L_img = A_img
+        if A_img.size!=L_img.size:
+            # L_img = L_img.resize(A_img.size, Image.ANTIALIAS)
+            A_img = A_img.resize(L_img.size, Image.ANTIALIAS)
+        if A_img.size[1]>2500:
+            A_img = A_img.resize((A_img.size[0]//2, A_img.size[1]//2), Image.ANTIALIAS)
+
+        ow, oh = A_img.size
+        transform_params = get_params(self.opt, A_img.size)
+        A_transform = get_transform(self.opt, transform_params, grayscale=False)
+        L_transform = get_transform(self.opt, transform_params, grayscale=True)
+        A = A_transform(A_img)
+        L = L_transform(L_img)
+
+        # base = 2**9
+        # h = int((oh+base-1) // base * base)
+        # w = int((ow+base-1) // base * base)
+        # A = F.pad(A.unsqueeze(0), (0,w-ow, 0,h-oh), 'replicate').squeeze(0)
+        # L = F.pad(L.unsqueeze(0), (0,w-ow, 0,h-oh), 'replicate').squeeze(0)
+
+        tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114
+        Ai = tmp.unsqueeze(0)
+        
+        return {'A': A, 'Ai': Ai, 'L': L, 
+                'B': torch.zeros(1), 'Bs': torch.zeros(1), 'Bi': torch.zeros(1), 'Bl': torch.zeros(1), 
+                'A_paths': A_path, 'h': oh, 'w': ow}
+
+    def __len__(self):
+        return self.A_size
+
+    def name(self):
+        return 'SingleCoDataset'
+
+
+def M_transform(feat, opt, params=None):
+    outfeat = feat.copy()
+    oh,ow = feat.shape[1:]
+    x1, y1 = params['crop_pos']
+    tw = th = opt.crop_size
+    if (ow > tw or oh > th):
+        outfeat = outfeat[:,y1:y1+th,x1:x1+tw]
+    if params['flip']:
+        outfeat = np.flip(outfeat, 2)#outfeat[:,:,::-1]
+    return torch.from_numpy(outfeat.copy()).float()*2-1.0

BidirectionalTranslation/data/singleSr_dataset.py

@@ -0,0 +1,73 @@
+import os.path
+from data.base_dataset import BaseDataset, get_params, get_transform
+from data.image_folder import make_dataset
+from PIL import Image
+import random
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class SingleSrDataset(BaseDataset):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def __init__(self, opt):
+        self.opt = opt
+        self.root = opt.dataroot
+        self.dir_B = os.path.join(opt.dataroot, opt.phase, opt.folder, 'imgs')
+        # self.dir_B = os.path.join(opt.dataroot, opt.phase, 'test/imgs', opt.folder)
+
+        self.B_paths = make_dataset(self.dir_B)
+
+        self.B_paths = sorted(self.B_paths)
+
+        self.B_size = len(self.B_paths)
+        # self.transform = get_transform(opt)
+        # print(self.B_size)
+
+    def __getitem__(self, index):
+        B_path = self.B_paths[index]
+
+        B_img = Image.open(B_path).convert('RGB')
+        if os.path.exists(B_path.replace('imgs','line').replace('.jpg','.png')):
+            L_img = Image.open(B_path.replace('imgs','line').replace('.jpg','.png'))#.convert('RGB')
+        else:
+            L_img = Image.open(B_path.replace('imgs','line').replace('.png','.jpg'))#.convert('RGB')
+        B_img = B_img.resize(L_img.size, Image.ANTIALIAS)
+
+        ow, oh = B_img.size
+        transform_params = get_params(self.opt, B_img.size)
+        B_transform = get_transform(self.opt, transform_params, grayscale=True)
+        B = B_transform(B_img)
+        L = B_transform(L_img)
+
+        # base = 2**8
+        # h = int((oh+base-1) // base * base)
+        # w = int((ow+base-1) // base * base)
+        # B = F.pad(B.unsqueeze(0), (0,w-ow, 0,h-oh), 'replicate').squeeze(0)
+        # L = F.pad(L.unsqueeze(0), (0,w-ow, 0,h-oh), 'replicate').squeeze(0)
+
+        return {'B': B, 'Bs': B, 'Bi': B, 'Bl': L, 
+                'A': torch.zeros(1), 'Ai': torch.zeros(1), 'L': torch.zeros(1), 
+                'A_paths': B_path, 'h': oh, 'w': ow}
+
+    def __len__(self):
+        return self.B_size
+
+    def name(self):
+        return 'SingleSrDataset'
+
+
+def M_transform(feat, opt, params=None):
+    outfeat = feat.copy()
+    if params is not None:
+        oh,ow = feat.shape[1:]
+        x1, y1 = params['crop_pos']
+        tw = th = opt.crop_size
+        if (ow > tw or oh > th):
+            outfeat = outfeat[:,y1:y1+th,x1:x1+tw]
+        if params['flip']:
+            outfeat = np.flip(outfeat, 2).copy()#outfeat[:,:,::-1]
+    return torch.from_numpy(outfeat).float()*2-1.0

BidirectionalTranslation/models/__init__.py

@@ -0,0 +1,61 @@
+"""This package contains modules related to objective functions, optimizations, and network architectures.
+To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
+You need to implement the following five functions:
+    -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
+    -- <set_input>:                     unpack data from dataset and apply preprocessing.
+    -- <forward>:                       produce intermediate results.
+    -- <optimize_parameters>:           calculate loss, gradients, and update network weights.
+    -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
+In the function <__init__>, you need to define four lists:
+    -- self.loss_names (str list):          specify the training losses that you want to plot and save.
+    -- self.model_names (str list):         specify the images that you want to display and save.
+    -- self.visual_names (str list):        define networks used in our training.
+    -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
+Now you can use the model class by specifying flag '--model dummy'.
+See our template model class 'template_model.py' for an example.
+"""
+
+import importlib
+from models.base_model import BaseModel
+
+
+def find_model_using_name(model_name):
+    """Import the module "models/[model_name]_model.py".
+    In the file, the class called DatasetNameModel() will
+    be instantiated. It has to be a subclass of BaseModel,
+    and it is case-insensitive.
+    """
+    model_filename = "models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+    model = None
+    target_model_name = model_name.replace('_', '') + 'model'
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() \
+           and issubclass(cls, BaseModel):
+            model = cls
+
+    if model is None:
+        print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+def get_option_setter(model_name):
+    """Return the static method <modify_commandline_options> of the model class."""
+    model_class = find_model_using_name(model_name)
+    return model_class.modify_commandline_options
+
+
+def create_model(opt, ckpt_root):
+    """Create a model given the option.
+    This function warps the class CustomDatasetDataLoader.
+    This is the main interface between this package and 'train.py'/'test.py'
+    Example:
+        >>> from models import create_model
+        >>> model = create_model(opt)
+    """
+    model = find_model_using_name(opt.model)
+    instance = model(opt, ckpt_root = ckpt_root)
+    print("model [%s] was created" % type(instance).__name__)
+    return instance

BidirectionalTranslation/models/base_model.py

@@ -0,0 +1,277 @@
+import os
+import torch
+from collections import OrderedDict
+from abc import ABC, abstractmethod
+from . import networks
+import numpy as np
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+class BaseModel(ABC):
+    """This class is an abstract base class (ABC) for models.
+    To create a subclass, you need to implement the following five functions:
+        -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
+        -- <set_input>:                     unpack data from dataset and apply preprocessing.
+        -- <forward>:                       produce intermediate results.
+        -- <optimize_parameters>:           calculate losses, gradients, and update network weights.
+        -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
+    """
+
+    def __init__(self, opt):
+        """Initialize the BaseModel class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+
+        When creating your custom class, you need to implement your own initialization.
+        In this fucntion, you should first call  `BaseModel.__init__(self, opt)`
+        Then, you need to define four lists:
+            -- self.loss_names (str list):          specify the training losses that you want to plot and save.
+            -- self.model_names (str list):         specify the images that you want to display and save.
+            -- self.visual_names (str list):        define networks used in our training.
+            -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
+        """
+        self.opt = opt
+        self.gpu_ids = opt.gpu_ids
+        self.isTrain = opt.isTrain
+        self.iter = 0
+        self.last_iter = 0
+        self.device = torch.device('cuda:{}'.format(
+            self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')  # get device name: CPU or GPU
+        # save all the checkpoints to save_dir
+        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        try:
+            os.mkdir(self.save_dir)
+        except:
+            pass
+        # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
+        if opt.preprocess != 'scale_width':
+            torch.backends.cudnn.benchmark = True
+        self.loss_names = []
+        self.model_names = []
+        self.visual_names = []
+        self.optimizers = []
+        self.image_paths = []
+
+        self.label_colours = np.random.randint(255, size=(100,3))
+
+    def save_suppixel(self,l_inds):
+        im_target_rgb = np.array([self.label_colours[ c % 100 ] for c in l_inds])
+        b,h,w = l_inds.shape
+        im_target_rgb = im_target_rgb.reshape(b,h,w,3).transpose(0,3,1,2)/127.5-1.0
+        return torch.from_numpy(im_target_rgb)
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new model-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        return parser
+
+    @abstractmethod
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input (dict): includes the data itself and its metadata information.
+        """
+        pass
+
+    @abstractmethod
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        pass
+
+    def is_train(self):
+        """check if the current batch is good for training."""
+        return True
+
+    @abstractmethod
+    def optimize_parameters(self):
+        """Calculate losses, gradients, and update network weights; called in every training iteration"""
+        pass
+
+    def setup(self, opt):
+        """Load and print networks; create schedulers
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        if self.isTrain:
+            self.schedulers = [networks.get_scheduler(
+                optimizer, opt) for optimizer in self.optimizers]
+        if not self.isTrain or opt.continue_train:
+            self.load_networks(opt.epoch)
+        self.print_networks(opt.verbose)
+
+    def eval(self):
+        """Make models eval mode during test time"""
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                net.eval()
+
+    def test(self):
+        """Forward function used in test time.
+
+        This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
+        It also calls <compute_visuals> to produce additional visualization results
+        """
+        with torch.no_grad():
+            self.forward()
+            self.compute_visuals()
+
+    def compute_visuals(self):
+        """Calculate additional output images for visdom and HTML visualization"""
+        pass
+
+    def get_image_paths(self):
+        """ Return image paths that are used to load current data"""
+        return self.image_paths
+
+    def update_learning_rate(self):
+        """Update learning rates for all the networks; called at the end of every epoch"""
+        for scheduler in self.schedulers:
+            scheduler.step()
+        lr = self.optimizers[0].param_groups[0]['lr']
+        print('learning rate = %.7f' % lr)
+
+    def get_current_visuals(self):
+        """Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
+        visual_ret = OrderedDict()
+        for name in self.visual_names:
+            if isinstance(name, str):
+                if 'Lab' in name:
+                    labimg = getattr(self, name).cpu()
+                    labimg[:,0,:,:]+=1
+                    labimg[:,0,:,:]*=50
+                    labimg[:,1:,:,:] *= 110
+                    labimg = labimg.permute((0,2,3,1))
+                    for i in range(labimg.shape[0]):
+                        labimg[i,:,:,:]=lab2rgb(labimg[i,:,:,:])
+                    visual_ret[name] = (labimg.permute((0,3,1,2))*2-1.0).to(self.device)
+                elif 'Fm' in name:
+                    visual_ret[name] = self.save_suppixel(getattr(self, name).cpu()).to(self.device)
+                else:
+                    visual_ret[name] = getattr(self, name)
+        return visual_ret
+
+    def get_current_losses(self):
+        """Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
+        errors_ret = OrderedDict()
+        for name in self.loss_names:
+            if isinstance(name, str):
+                # float(...) works for both scalar tensor and float number
+                errors_ret[name] = float(getattr(self, 'loss_' + name))
+        return errors_ret
+
+    def save_networks(self, epoch):
+        """Save all the networks to 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):
+                save_filename = '%s_net_%s.pth' % (epoch, name)
+                save_path = os.path.join(self.save_dir, save_filename)
+                # print(save_path)
+                net = getattr(self, 'net' + name)
+
+                if len(self.gpu_ids) > 0 and torch.cuda.is_available():
+                    torch.save(net.state_dict(), save_path)
+                    # net.cuda(self.gpu_ids[0])
+                else:
+                    torch.save(net.cpu().state_dict(), save_path)
+
+        save_filename = '%s_net_opt.pth' % (epoch)
+        save_path = os.path.join(self.save_dir, save_filename)
+        save_dict = {'iter': str(self.iter // self.opt.print_freq * self.opt.print_freq)}
+        for i, name in enumerate(self.optimizer_names):
+            save_dict.update({name.lower(): self.optimizers[i].state_dict()})
+        torch.save(save_dict, save_path)
+
+
+    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
+        """Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
+        key = keys[i]
+        if i + 1 == len(keys):  # at the end, pointing to a parameter/buffer
+            if module.__class__.__name__.startswith('InstanceNorm') and \
+                    (key == 'running_mean' or key == 'running_var'):
+                if getattr(module, key) is None:
+                    state_dict.pop('.'.join(keys))
+            if module.__class__.__name__.startswith('InstanceNorm') and \
+               (key == 'num_batches_tracked'):
+                state_dict.pop('.'.join(keys))
+        else:
+            self.__patch_instance_norm_state_dict(
+                state_dict, getattr(module, key), keys, i + 1)
+
+    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, 'net' + name)
+                # if isinstance(net, torch.nn.DataParallel):
+                if isinstance(net, DDP):
+                    net = net.module
+                # print(net)
+                print('loading the model from %s' % load_path)
+                # if you are using PyTorch newer than 0.4 (e.g., built from
+                # GitHub source), you can remove str() on self.device
+                state_dict = torch.load(
+                    load_path, map_location=lambda storage, loc: storage.cuda())
+                if hasattr(state_dict, '_metadata'):
+                    del state_dict._metadata
+
+                # patch InstanceNorm checkpoints prior to 0.4
+                # need to copy keys here because we mutate in loop
+                #for key in list(state_dict.keys()):
+                #    self.__patch_instance_norm_state_dict(
+                #        state_dict, net, key.split('.'))
+
+                net.load_state_dict(state_dict)
+                del state_dict
+
+    def print_networks(self, verbose):
+        """Print the total number of parameters in the network and (if verbose) network architecture
+
+        Parameters:
+            verbose (bool) -- if verbose: print the network architecture
+        """
+        print('---------- Networks initialized -------------')
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                num_params = 0
+                for param in net.parameters():
+                    num_params += param.numel()
+                if verbose:
+                    print(net)
+                print('[Network %s] Total number of parameters : %.3f M' %
+                      (name, num_params / 1e6))
+        print('-----------------------------------------------')
+
+    def set_requires_grad(self, nets, requires_grad=False):
+        """Set requires_grad=False for all the networks to avoid unnecessary computations
+        Parameters:
+            nets (network list)   -- a list of networks
+            requires_grad (bool)  -- whether the networks require gradients or not
+        """
+        if not isinstance(nets, list):
+            nets = [nets]
+        for net in nets:
+            if net is not None:
+                for param in net.parameters():
+                    param.requires_grad = requires_grad

BidirectionalTranslation/models/cycle_ganstft_model.py

@@ -0,0 +1,103 @@
+import random
+import torch
+from .base_model import BaseModel
+from . import networks
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+
+class CycleGANSTFTModel(BaseModel):
+
+    def __init__(self, opt, ckpt_root):
+
+        BaseModel.__init__(self, opt)
+
+        use_vae = True
+        self.interchnnls = 4
+        use_noise = False
+        self.half_size = opt.batch_size //2
+        self.device=opt.local_rank
+        self.gpu_ids=[self.device]
+        self.local_rank = opt.local_rank
+        self.cropsize = opt.crop_size
+
+        self.model_names = ['G_INTSCR2RGB','G_RGB2INTSCR','E']
+        self.netG_INTSCR2RGB = networks.define_G(self.interchnnls + 1, 3, opt.nz, opt.ngf, netG='unet_256', 
+                                      norm='layer', nl='lrelu', use_dropout=opt.use_dropout, init_type='kaiming', init_gain=opt.init_gain,
+                                      gpu_ids=self.gpu_ids, where_add='all', upsample='bilinear', use_noise=use_noise)
+        self.netG_RGB2INTSCR = networks.define_G(4, self.interchnnls, 0, opt.ngf, netG='unet_256', 
+                                      norm='layer', nl='lrelu', use_dropout=opt.use_dropout, init_type='kaiming', init_gain=opt.init_gain,
+                                      gpu_ids=self.gpu_ids, where_add='input', upsample='bilinear', use_noise=use_noise)        
+        self.netE = networks.define_E(opt.output_nc, opt.nz, opt.nef, netE=opt.netE, norm='none', nl='lrelu',
+                                      init_type='xavier', init_gain=opt.init_gain, gpu_ids=self.gpu_ids, vaeLike=use_vae)
+        self.nets = [self.netG_INTSCR2RGB, self.netG_RGB2INTSCR, self.netE]
+
+        self.netSVAE = networks.define_SVAE(inc=1, outc=self.interchnnls, outplanes=64, blocks=3, netVAE='SVAE', 
+            save_dir= ckpt_root+'/ScreenStyle/ScreenVAE',init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids)
+
+
+    def set_input(self, input):
+        AtoB = self.opt.direction == 'AtoB'
+        self.real_RGB = input['A'].to(self.device)
+        self.real_Ai = self.grayscale(self.real_RGB)
+        self.real_L = input['L'].to(self.device)
+        self.real_ML = input['Bl'].to(self.device)
+        self.real_M = input['B'].to(self.device)
+
+        self.h = input['h']
+        self.w = input['w']
+
+    def grayscale(self, input_image):
+        rate = torch.Tensor([0.299, 0.587, 0.114]).reshape(1, 3, 1, 1).to(input_image.device)
+        # tmp = input_image[:,0, ...] * 0.299 + input_image[:,1, ...] * 0.587 + input_image[:,2, ...] * 0.114
+        return (input_image*rate).sum(1,keepdims=True)
+
+    def forward(self, AtoB=True, sty=None):
+        if AtoB:
+            real_LRGB = torch.cat([self.real_L, self.real_RGB],1)
+            fake_SCR = self.netG_RGB2INTSCR(real_LRGB)
+            fake_M = self.netSVAE(fake_SCR, line=self.real_L, img_input=False)
+            fake_M = torch.clamp(fake_M, -1,1)
+            fake_M2 = self.norm(torch.mul(self.denorm(fake_M), self.denorm(self.real_L)))#*self.mask2
+            return fake_M[:,:,:self.h, :self.w], fake_M2[:,:,:self.h, :self.w], fake_SCR[:,:,:self.h, :self.w]
+        else:
+            if sty is None:  # use encoded z
+                z0, _ = self.netE(self.real_RGB)
+            else:
+                z0 = sty
+                # z0 = self.get_z_random(self.real_A.size(0), self.opt.nz)
+            real_SCR = self.netSVAE(self.real_M, self.real_ML, output_screen_only=True) #8
+            real_LSCR = torch.cat([self.real_ML, real_SCR], 1)
+            fake_nRGB = self.netG_INTSCR2RGB(real_LSCR, z0)
+            fake_nRGB = torch.clamp(fake_nRGB, -1,1)
+            fake_RGB = self.norm(torch.mul(self.denorm(fake_nRGB), self.denorm(self.real_ML)))
+            return fake_RGB[:,:,:self.h, :self.w], real_SCR[:,:,:self.h, :self.w], self.real_ML[:,:,:self.h, :self.w]
+
+    def norm(self, im):
+        return im * 2.0 - 1
+
+    def denorm(self, im):
+        return (im + 1) / 2.0
+
+    def optimize_parameters(self):
+        pass
+
+    def get_z_random(self, batch_size, nz, random_type='gauss', truncation=False, tvalue=1):
+        z = None
+        if random_type == 'uni':
+            z = torch.rand(batch_size, nz) * 2.0 - 1.0
+        elif random_type == 'gauss':
+            z = torch.randn(batch_size, nz) * tvalue
+            # do the truncation trick
+            if truncation:
+                k = 0
+                while (k < 15 * nz):
+                    if torch.max(z) <= tvalue:
+                        break
+                    zabs = torch.abs(z)
+                    zz = torch.randn(batch_size, nz)
+                    z[zabs > tvalue] = zz[zabs > tvalue]
+                    k += 1
+                z = torch.clamp(z, -tvalue, tvalue)
+
+        return z.detach().to(self.device)

BidirectionalTranslation/models/networks.py

@@ -0,0 +1,1375 @@
+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
+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:
+        assert(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).cuda()(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:
+        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:
+        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"))
+
+        # 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):
+    """
+        @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
+    """
+    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=lambda storage, loc: storage.cuda())
+                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)

BidirectionalTranslation/options/base_options.py

@@ -0,0 +1,142 @@
+import argparse
+import os
+from util import util
+import torch
+import models
+import data
+
+class BaseOptions():
+    def __init__(self):
+        self.initialized = False
+    
+    def initialize(self, parser):
+        """Initialize options used during both training and test time."""
+        # Basic options
+        parser.add_argument('--dataroot', required=False, help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')
+        parser.add_argument('--batch_size', type=int, default=2, help='input batch size')
+        parser.add_argument('--load_size', type=int, default=512, help='scale images to this size')  # Modified default
+        parser.add_argument('--crop_size', type=int, default=1024, help='then crop to this size')    # Modified default
+        parser.add_argument('--input_nc', type=int, default=1, help='# of input image channels')     # Modified default
+        parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')   # Modified default
+        parser.add_argument('--nz', type=int, default=64, help='#latent vector')                     # Modified default
+        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2, -1 for CPU mode')
+        parser.add_argument('--name', type=str, default='color2manga_cycle_ganstft', help='name of the experiment')  # Modified default
+        parser.add_argument('--preprocess', type=str, default='none', help='not implemented')         # Modified default
+        parser.add_argument('--dataset_mode', type=str, default='aligned', help='aligned,single')
+        parser.add_argument('--model', type=str, default='cycle_ganstft', help='chooses which model to use')
+        parser.add_argument('--direction', type=str, default='BtoA', help='AtoB or BtoA')            # Modified default
+        parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+        parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
+        parser.add_argument('--local_rank', default=0, type=int, help='# threads for loading data')
+        parser.add_argument('--checkpoints_dir', type=str, default=self.model_global_path+'/ScreenStyle/color2manga/', help='models are saved here')  # Modified default
+        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
+        parser.add_argument('--use_dropout', action='store_true', help='use dropout for the generator')
+        parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset.')
+        parser.add_argument('--no_flip', action='store_false', help='if specified, do not flip the images for data argumentation')  # Modified default
+
+        # Model parameters
+        parser.add_argument('--level', type=int, default=0, help='level to train')
+        parser.add_argument('--num_Ds', type=int, default=2, help='number of Discriminators')
+        parser.add_argument('--netD', type=str, default='basic_256_multi', help='selects model to use for netD')
+        parser.add_argument('--netD2', type=str, default='basic_256_multi', help='selects model to use for netD2')
+        parser.add_argument('--netG', type=str, default='unet_256', help='selects model to use for netG')
+        parser.add_argument('--netC', type=str, default='unet_128', help='selects model to use for netC')
+        parser.add_argument('--netE', type=str, default='conv_256', help='selects model to use for netE')
+        parser.add_argument('--nef', type=int, default=48, help='# of encoder filters in the first conv layer')  # Modified default
+        parser.add_argument('--ngf', type=int, default=48, help='# of gen filters in the last conv layer')       # Modified default
+        parser.add_argument('--ndf', type=int, default=32, help='# of discrim filters in the first conv layer')  # Modified default
+        parser.add_argument('--norm', type=str, default='layer', help='instance normalization or batch normalization')
+        parser.add_argument('--upsample', type=str, default='bilinear', help='basic | bilinear')                  # Modified default
+        parser.add_argument('--nl', type=str, default='prelu', help='non-linearity activation: relu | lrelu | elu')
+        parser.add_argument('--no_encode', action='store_true', help='if specified, print more debugging information')
+        parser.add_argument('--color2screen', action='store_true', help='continue training: load the latest model including RGB model')  # Modified default
+
+        # Extra parameters
+        parser.add_argument('--where_add', type=str, default='all', help='input|all|middle; where to add z in the network G')
+        parser.add_argument('--conditional_D', action='store_true', help='if use conditional GAN for D')
+        parser.add_argument('--init_type', type=str, default='kaiming', help='network initialization [normal | xavier | kaiming | orthogonal]')
+        parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
+        parser.add_argument('--center_crop', action='store_true', help='if apply for center cropping for the test')  # Modified default
+        parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')
+        parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}')
+        parser.add_argument('--display_winsize', type=int, default=256, help='display window size')
+
+        # Special tasks
+        self.initialized = True
+        return parser
+
+    def gather_options(self):
+        """Initialize our parser with basic options (only once)."""
+        if not self.initialized:
+            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+            parser = self.initialize(parser)
+        
+        # Get the basic options
+        opt, _ = parser.parse_known_args()
+
+        # Modify model-related parser options
+        model_name = opt.model
+        model_option_setter = models.get_option_setter(model_name)
+        parser = model_option_setter(parser, self.isTrain)
+        opt, _ = parser.parse_known_args()  # Parse again with new defaults
+
+        # Modify dataset-related parser options
+        dataset_name = opt.dataset_mode
+        dataset_option_setter = data.get_option_setter(dataset_name)
+        parser = dataset_option_setter(parser, self.isTrain)
+
+        # Save and return the parser
+        self.parser = parser
+        return parser.parse_args()
+
+    def print_options(self, opt):
+        """Print and save options."""
+        message = ''
+        message += '----------------- Options ---------------\n'
+        for k, v in sorted(vars(opt).items()):
+            comment = ''
+            default = self.parser.get_default(k)
+            if v != default:
+                comment = '\t[default: %s]' % str(default)
+            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
+        message += '----------------- End -------------------'
+        print(message)
+
+        # Save to the disk
+        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        if not os.path.exists(expr_dir):
+            try:
+                util.mkdirs(expr_dir)
+            except:
+                pass
+        file_name = os.path.join(expr_dir, 'opt.txt')
+        with open(file_name, 'wt') as opt_file:
+            opt_file.write(message)
+            opt_file.write('\n')
+
+    def parse(self, model_global_path):
+        """Parse options, create checkpoints directory suffix, and set up gpu device."""
+        self.model_global_path = model_global_path
+        opt = self.gather_options()
+        opt.isTrain = self.isTrain  # train or test
+        
+
+        # Process opt.suffix
+        if opt.suffix:
+            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
+            opt.name = opt.name + suffix
+
+        self.print_options(opt)
+
+        # Set gpu ids
+        str_ids = opt.gpu_ids.split(',')
+        opt.gpu_ids = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                opt.gpu_ids.append(id)
+        if len(opt.gpu_ids) > 0:
+            torch.cuda.set_device(opt.gpu_ids[0])
+
+        self.opt = opt
+        return self.opt

BidirectionalTranslation/options/test_options.py

@@ -0,0 +1,19 @@
+from .base_options import BaseOptions
+
+class TestOptions(BaseOptions):
+    def initialize(self, parser):
+        BaseOptions.initialize(self, parser)
+        
+        
+        # Additional test-specific arguments
+        parser.add_argument('--results_dir', type=str, default='../results/', help='saves results here.')
+        parser.add_argument('--phase', type=str, default='val', help='train, val, test, etc')
+        parser.add_argument('--num_test', type=int, default=30, help='how many test images to run')
+        parser.add_argument('--n_samples', type=int, default=1, help='#samples')
+        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio for the results')
+        parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
+        parser.add_argument('--folder', type=str, default='intra', help='saves results here.')
+        parser.add_argument('--sync', action='store_true', help='use the same latent code for different input images')
+
+        self.isTrain = False
+        return parser

BidirectionalTranslation/requirements.txt

@@ -0,0 +1,8 @@
+torch~=1.6.0
+torchvision~=0.4.0
+tensorboardx~=1.9
+scipy==1.1
+dominate~=2.3.1
+scikit-image~=0.16.2
+opencv-python~=3.4.2
+lpips

BidirectionalTranslation/scripts/test_western2manga.sh

@@ -0,0 +1,49 @@
+set -ex
+# models
+RESULTS_DIR='./results/test/western2manga'
+
+# dataset
+CLASS='color2manga'
+MODEL='cycle_ganstft'
+DIRECTION='BtoA' # from domain A to domain B
+PREPROCESS='none'
+LOAD_SIZE=512 # scale images to this size
+CROP_SIZE=1024 # then crop to this size
+INPUT_NC=1  # number of channels in the input image
+OUTPUT_NC=3  # number of channels in the input image
+NGF=48
+NEF=48
+NDF=32
+NZ=64
+
+# misc
+GPU_ID=0   # gpu id
+NUM_TEST=30 # number of input images duirng test
+NUM_SAMPLES=1 # number of samples per input images
+NAME=${CLASS}_${MODEL}
+
+# command
+CUDA_VISIBLE_DEVICES=${GPU_ID} \
+python3 ./test.py \
+  --dataroot ./datasets/${CLASS} \
+  --results_dir ${RESULTS_DIR} \
+  --checkpoints_dir ./checkpoints/${CLASS}/ \
+  --name ${NAME} \
+  --model ${MODEL} \
+  --direction ${DIRECTION} \
+  --preprocess ${PREPROCESS} \
+  --load_size ${LOAD_SIZE} \
+  --crop_size ${CROP_SIZE} \
+  --input_nc ${INPUT_NC} \
+  --output_nc ${OUTPUT_NC} \
+  --nz ${NZ} \
+  --netE conv_256 \
+  --num_test ${NUM_TEST} \
+  --n_samples ${NUM_SAMPLES} \
+  --upsample bilinear \
+  --ngf ${NGF} \
+  --nef ${NEF} \
+  --ndf ${NDF} \
+  --center_crop \
+  --color2screen \
+  --no_flip

BidirectionalTranslation/test.py

@@ -0,0 +1,71 @@
+import os
+from options.test_options import TestOptions
+from data import create_dataset
+from models import create_model
+from util.visualizer import save_images
+from itertools import islice
+from util import html
+import cv2
+
+seed = 10
+import torch
+import numpy as np
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+np.random.seed(seed)
+
+# options
+opt = TestOptions().parse()
+opt.num_threads = 1   # test code only supports num_threads=1
+opt.batch_size = 1   # test code only supports batch_size=1
+opt.serial_batches = True  # no shuffle
+
+model = create_model(opt)
+model.setup(opt)
+model.eval() 
+print('Loading model %s' % opt.model)
+
+testdata = ['manga_paper']
+# fake_sty = model.get_z_random(1, 64, truncation=True)
+
+opt.dataset_mode = 'singleSr'
+for folder in testdata:
+    opt.folder = folder
+    # create dataset
+    dataset = create_dataset(opt)
+    web_dir = os.path.join(opt.results_dir, opt.folder + '_Sr2Co')
+    webpage = html.HTML(web_dir, 'Training = %s, Phase = %s, Class =%s' % (opt.name, opt.phase, opt.name))
+    # fake_sty = model.get_z_random(1, 64, truncation=True)
+    for i, data in enumerate(islice(dataset, opt.num_test)):
+        h = data['h']
+        w = data['w']
+        model.set_input(data)
+        fake_sty = model.get_z_random(1, 64, truncation=True, tvalue=1.25)
+        fake_B, SCR, line = model.forward(AtoB=False, sty=fake_sty)
+        images=[fake_B[:,:,:h,:w]]
+        names=['color']
+
+        img_path = 'input_%3.3d' % i
+        save_images(webpage, images, names, img_path, aspect_ratio=opt.aspect_ratio, width=opt.crop_size)
+    webpage.save()
+
+testdata = ['western_paper']
+
+opt.dataset_mode = 'singleCo'
+for folder in testdata:
+    opt.folder = folder
+    # create dataset
+    dataset = create_dataset(opt)
+    web_dir = os.path.join(opt.results_dir, opt.folder + '_Sr2Co')
+    webpage = html.HTML(web_dir, 'Training = %s, Phase = %s, Class =%s' % (opt.name, opt.phase, opt.name))
+    for i, data in enumerate(islice(dataset, opt.num_test)):
+        h = data['h']
+        w = data['w']
+        model.set_input(data)
+        fake_B, fake_B2, SCR = model.forward(AtoB=True)
+        images=[fake_B2[:,:,:h,:w]]
+        names=['manga']
+
+        img_path = 'input_%3.3d' % i
+        save_images(webpage, images, names, img_path, aspect_ratio=opt.aspect_ratio, width=opt.crop_size)
+    webpage.save()

BidirectionalTranslation/util/html.py

@@ -0,0 +1,86 @@
+import dominate
+from dominate.tags import meta, h3, table, tr, td, p, a, img, br
+import os
+
+
+class HTML:
+    """This HTML class allows us to save images and write texts into a single HTML file.
+
+     It consists of functions such as <add_header> (add a text header to the HTML file),
+     <add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).
+     It is based on Python library 'dominate', a Python library for creating and manipulating HTML documents using a DOM API.
+    """
+
+    def __init__(self, web_dir, title, refresh=0):
+        """Initialize the HTML classes
+
+        Parameters:
+            web_dir (str) -- a directory that stores the webpage. HTML file will be created at <web_dir>/index.html; images will be saved at <web_dir/images/
+            title (str)   -- the webpage name
+            reflect (int) -- how often the website refresh itself; if 0; no refreshing
+        """
+        self.title = title
+        self.web_dir = web_dir
+        self.img_dir = os.path.join(self.web_dir, 'images')
+        if not os.path.exists(self.web_dir):
+            os.makedirs(self.web_dir)
+        if not os.path.exists(self.img_dir):
+            os.makedirs(self.img_dir)
+
+        self.doc = dominate.document(title=title)
+        if refresh > 0:
+            with self.doc.head:
+                meta(http_equiv="refresh", content=str(refresh))
+
+    def get_image_dir(self):
+        """Return the directory that stores images"""
+        return self.img_dir
+
+    def add_header(self, text):
+        """Insert a header to the HTML file
+
+        Parameters:
+            text (str) -- the header text
+        """
+        with self.doc:
+            h3(text)
+
+    def add_images(self, ims, txts, links, width=400):
+        """add images to the HTML file
+
+        Parameters:
+            ims (str list)   -- a list of image paths
+            txts (str list)  -- a list of image names shown on the website
+            links (str list) --  a list of hyperref links; when you click an image, it will redirect you to a new page
+        """
+        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
+        self.doc.add(self.t)
+        with self.t:
+            with tr():
+                for im, txt, link in zip(ims, txts, links):
+                    with td(style="word-wrap: break-word;", halign="center", valign="top"):
+                        with p():
+                            with a(href=os.path.join('images', link)):
+                                img(style="width:%dpx" % width, src=os.path.join('images', im))
+                            br()
+                            p(txt)
+
+    def save(self):
+        """save the current content to the HMTL file"""
+        html_file = '%s/index.html' % self.web_dir
+        f = open(html_file, 'wt')
+        f.write(self.doc.render())
+        f.close()
+
+
+if __name__ == '__main__':  # we show an example usage here.
+    html = HTML('web/', 'test_html')
+    html.add_header('hello world')
+
+    ims, txts, links = [], [], []
+    for n in range(4):
+        ims.append('image_%d.png' % n)
+        txts.append('text_%d' % n)
+        links.append('image_%d.png' % n)
+    html.add_images(ims, txts, links)
+    html.save()

BidirectionalTranslation/util/util.py

@@ -0,0 +1,136 @@
+from __future__ import print_function
+import torch
+import numpy as np
+from PIL import Image
+import os
+import pickle
+
+
+def tensor2im(input_image, imtype=np.uint8):
+    """"Convert a Tensor array into a numpy image array.
+    Parameters:
+        input_image (tensor) --  the input image tensor array
+        imtype (type)        --  the desired type of the converted numpy array
+    """
+    if not isinstance(input_image, np.ndarray):
+        if isinstance(input_image, torch.Tensor):  # get the data from a variable
+            image_tensor = input_image.data
+        else:
+            return input_image
+        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
+        if image_numpy.shape[0] == 1:  # grayscale to RGB
+            image_numpy = np.tile(image_numpy, (3, 1, 1))
+        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
+    else:  # if it is a numpy array, do nothing
+        image_numpy = input_image
+    return image_numpy.astype(imtype)
+
+
+def tensor2vec(vector_tensor):
+    numpy_vec = vector_tensor.data.cpu().numpy()
+    if numpy_vec.ndim == 4:
+        return numpy_vec[:, :, 0, 0]
+    else:
+        return numpy_vec
+
+
+def pickle_load(file_name):
+    data = None
+    with open(file_name, 'rb') as f:
+        data = pickle.load(f)
+    return data
+
+
+def pickle_save(file_name, data):
+    with open(file_name, 'wb') as f:
+        pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+def diagnose_network(net, name='network'):
+    """Calculate and print the mean of average absolute(gradients)
+    Parameters:
+        net (torch network) -- Torch network
+        name (str) -- the name of the network
+    """
+    mean = 0.0
+    count = 0
+    for param in net.parameters():
+        if param.grad is not None:
+            mean += torch.mean(torch.abs(param.grad.data))
+            count += 1
+    if count > 0:
+        mean = mean / count
+    print(name)
+    print(mean)
+
+
+def interp_z(z0, z1, num_frames, interp_mode='linear'):
+    zs = []
+    if interp_mode == 'linear':
+        for n in range(num_frames):
+            ratio = n / float(num_frames - 1)
+            z_t = (1 - ratio) * z0 + ratio * z1
+            zs.append(z_t[np.newaxis, :])
+        zs = np.concatenate(zs, axis=0).astype(np.float32)
+
+    if interp_mode == 'slerp':
+        z0_n = z0 / (np.linalg.norm(z0) + 1e-10)
+        z1_n = z1 / (np.linalg.norm(z1) + 1e-10)
+        omega = np.arccos(np.dot(z0_n, z1_n))
+        sin_omega = np.sin(omega)
+        if sin_omega < 1e-10 and sin_omega > -1e-10:
+            zs = interp_z(z0, z1, num_frames, interp_mode='linear')
+        else:
+            for n in range(num_frames):
+                ratio = n / float(num_frames - 1)
+                z_t = np.sin((1 - ratio) * omega) / sin_omega * z0 + np.sin(ratio * omega) / sin_omega * z1
+                zs.append(z_t[np.newaxis, :])
+        zs = np.concatenate(zs, axis=0).astype(np.float32)
+
+    return zs
+
+
+def save_image(image_numpy, image_path):
+    """Save a numpy image to the disk
+    Parameters:
+        image_numpy (numpy array) -- input numpy array
+        image_path (str)          -- the path of the image
+    """
+    image_pil = Image.fromarray(image_numpy)
+    image_pil.save(image_path)
+
+
+def print_numpy(x, val=True, shp=False):
+    """Print the mean, min, max, median, std, and size of a numpy array
+    Parameters:
+        val (bool) -- if print the values of the numpy array
+        shp (bool) -- if print the shape of the numpy array
+    """
+    x = x.astype(np.float64)
+    if shp:
+        print('shape,', x.shape)
+    if val:
+        x = x.flatten()
+        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
+            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
+
+
+def mkdirs(paths):
+    """create empty directories if they don't exist
+    Parameters:
+        paths (str list) -- a list of directory paths
+    """
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+
+def mkdir(path):
+    """create a single empty directory if it didn't exist
+    Parameters:
+        path (str) -- a single directory path
+    """
+    if not os.path.exists(path):
+        os.makedirs(path, exist_ok=True)

BidirectionalTranslation/util/visualizer.py

@@ -0,0 +1,221 @@
+import numpy as np
+import os
+import sys
+import ntpath
+import time
+from . import util
+from . import html
+from subprocess import Popen, PIPE
+import cv2
+
+
+# if sys.version_info[0] == 2:
+#     VisdomExceptionBase = Exception
+# else:
+#     VisdomExceptionBase = ConnectionError
+
+
+def save_images(webpage, images, names, image_path, aspect_ratio=1.0, width=256):
+    """Save images to the disk.
+    Parameters:
+        webpage (the HTML class)  -- the HTML webpage class that stores these imaegs (see html.py for more details)
+        images (numpy array list) -- a list of numpy array that stores images
+        names (str list)          -- a str list stores the names of the images above
+        image_path (str)         -- the string is used to create image paths
+        aspect_ratio (float)     -- the aspect ratio of saved images
+        width (int)              -- the images will be resized to width x width
+    This function will save images stored in 'visuals' to the HTML file specified by 'webpage'.
+    """
+    image_dir = webpage.get_image_dir()
+    name = ntpath.basename(image_path)
+
+    webpage.add_header(name)
+    ims, txts, links = [], [], []
+
+    for label, im_data in zip(names, images):
+        im = util.tensor2im(im_data)
+        image_name = '%s_%s.jpg' % (name, label)
+        save_path = os.path.join(image_dir, image_name)
+        h, w, _ = im.shape
+        if aspect_ratio > 1.0:
+            im = cv2.resize(im, (h, int(w * aspect_ratio)), interpolation=cv2.INTER_CUBIC)
+        if aspect_ratio < 1.0:
+            im = cv2.resize(im, (int(h / aspect_ratio), w), interpolation=cv2.INTER_CUBIC)
+        util.save_image(im, save_path)
+
+        ims.append(image_name)
+        txts.append(label)
+        links.append(image_name)
+    webpage.add_images(ims, txts, links, width=width)
+
+
+class Visualizer():
+    """This class includes several functions that can display/save images and print/save logging information.
+    It uses a Python library 'visdom' for display, and a Python library 'dominate' (wrapped in 'HTML') for creating HTML files with images.
+    """
+
+    def __init__(self, opt):
+        """Initialize the Visualizer class
+        Parameters:
+            opt -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        Step 1: Cache the training/test options
+        Step 2: connect to a visdom server
+        Step 3: create an HTML object for saveing HTML filters
+        Step 4: create a logging file to store training losses
+        """
+        self.opt = opt  # cache the option
+        self.display_id = opt.display_id
+        self.use_html = opt.isTrain and not opt.no_html
+        self.win_size = opt.display_winsize
+        self.name = opt.name
+        self.port = opt.display_port
+        self.saved = False
+        # if self.display_id > 0:  # connect to a visdom server given <display_port> and <display_server>
+        #     import visdom
+        #     self.ncols = opt.display_ncols
+        #     self.vis = visdom.Visdom(server=opt.display_server, port=opt.display_port, env=opt.display_env)
+        #     if not self.vis.check_connection():
+        #         self.create_visdom_connections()
+        if self.use_html:  # create an HTML object at <checkpoints_dir>/web/; images will be saved under <checkpoints_dir>/web/images/
+            self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
+            self.img_dir = os.path.join(self.web_dir, 'images')
+            print('create web directory %s...' % self.web_dir)
+            util.mkdirs([self.web_dir, self.img_dir])
+        # create a logging file to store training losses
+        self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
+        with open(self.log_name, "a") as log_file:
+            now = time.strftime("%c")
+            log_file.write('================ Training Loss (%s) ================\n' % now)
+
+    def reset(self):
+        """Reset the self.saved status"""
+        self.saved = False
+
+    def create_visdom_connections(self):
+        """If the program could not connect to Visdom server, this function will start a new server at port < self.port > """
+        cmd = sys.executable + ' -m visdom.server -p %d &>/dev/null &' % self.port
+        print('\n\nCould not connect to Visdom server. \n Trying to start a server....')
+        print('Command: %s' % cmd)
+        Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE)
+
+    def display_current_results(self, visuals, epoch, save_result):
+        """Display current results on visdom; save current results to an HTML file.
+        Parameters:
+            visuals (OrderedDict) - - dictionary of images to display or save
+            epoch (int) - - the current epoch
+            save_result (bool) - - if save the current results to an HTML file
+        """
+        # if self.display_id > 0:  # show images in the browser using visdom
+        #     ncols = self.ncols
+        #     if ncols > 0:        # show all the images in one visdom panel
+        #         ncols = min(ncols, len(visuals))
+        #         h, w = next(iter(visuals.values())).shape[:2]
+        #         table_css = """<style>
+        #                 table {border-collapse: separate; border-spacing: 4px; white-space: nowrap; text-align: center}
+        #                 table td {width: % dpx; height: % dpx; padding: 4px; outline: 4px solid black}
+        #                 </style>""" % (w, h)  # create a table css
+        #         # create a table of images.
+        #         title = self.name
+        #         label_html = ''
+        #         label_html_row = ''
+        #         images = []
+        #         idx = 0
+        #         for label, image in visuals.items():
+        #             image_numpy = util.tensor2im(image)
+        #             label_html_row += '<td>%s</td>' % label
+        #             images.append(image_numpy.transpose([2, 0, 1]))
+        #             idx += 1
+        #             if idx % ncols == 0:
+        #                 label_html += '<tr>%s</tr>' % label_html_row
+        #                 label_html_row = ''
+        #         white_image = np.ones_like(image_numpy.transpose([2, 0, 1])) * 255
+        #         while idx % ncols != 0:
+        #             images.append(white_image)
+        #             label_html_row += '<td></td>'
+        #             idx += 1
+        #         if label_html_row != '':
+        #             label_html += '<tr>%s</tr>' % label_html_row
+        #         try:
+        #             self.vis.images(images, nrow=ncols, win=self.display_id + 1,
+        #                             padding=2, opts=dict(title=title + ' images'))
+        #             label_html = '<table>%s</table>' % label_html
+        #             self.vis.text(table_css + label_html, win=self.display_id + 2,
+        #                           opts=dict(title=title + ' labels'))
+        #         except VisdomExceptionBase:
+        #             self.create_visdom_connections()
+
+        #     else:     # show each image in a separate visdom panel;
+        #         idx = 1
+        #         try:
+        #             for label, image in visuals.items():
+        #                 image_numpy = util.tensor2im(image)
+        #                 self.vis.image(image_numpy.transpose([2, 0, 1]), opts=dict(title=label),
+        #                                win=self.display_id + idx)
+        #                 idx += 1
+        #         except VisdomExceptionBase:
+        #             self.create_visdom_connections()
+
+        if self.use_html and (save_result or not self.saved):  # save images to an HTML file if they haven't been saved.
+            self.saved = True
+            # save images to the disk
+            for label, image in visuals.items():
+                image_numpy = util.tensor2im(image)
+                img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.png' % (epoch, label))
+                util.save_image(image_numpy, img_path)
+
+            # update website
+            webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, refresh=1)
+            for n in range(epoch, 0, -1):
+                webpage.add_header('epoch [%d]' % n)
+                ims, txts, links = [], [], []
+
+                for label, image_numpy in visuals.items():
+                    image_numpy = util.tensor2im(image)
+                    img_path = 'epoch%.3d_%s.png' % (n, label)
+                    ims.append(img_path)
+                    txts.append(label)
+                    links.append(img_path)
+                webpage.add_images(ims, txts, links, width=self.win_size)
+            webpage.save()
+
+    def plot_current_losses(self, epoch, counter_ratio, losses):
+        """display the current losses on visdom display: dictionary of error labels and values
+        Parameters:
+            epoch (int)           -- current epoch
+            counter_ratio (float) -- progress (percentage) in the current epoch, between 0 to 1
+            losses (OrderedDict)  -- training losses stored in the format of (name, float) pairs
+        """
+        if not hasattr(self, 'plot_data'):
+            self.plot_data = {'X': [], 'Y': [], 'legend': list(losses.keys())}
+        self.plot_data['X'].append(epoch + counter_ratio)
+        self.plot_data['Y'].append([losses[k] for k in self.plot_data['legend']])
+        # try:
+        #     self.vis.line(
+        #         X=np.stack([np.array(self.plot_data['X'])] * len(self.plot_data['legend']), 1),
+        #         Y=np.array(self.plot_data['Y']),
+        #         opts={
+        #             'title': self.name + ' loss over time',
+        #             'legend': self.plot_data['legend'],
+        #             'xlabel': 'epoch',
+        #             'ylabel': 'loss'},
+        #         win=self.display_id)
+        # except VisdomExceptionBase:
+        #     self.create_visdom_connections()
+
+    # losses: same format as |losses| of plot_current_losses
+    def print_current_losses(self, epoch, iters, losses, t_comp, t_data):
+        """print current losses on console; also save the losses to the disk
+        Parameters:
+            epoch (int) -- current epoch
+            iters (int) -- current training iteration during this epoch (reset to 0 at the end of every epoch)
+            losses (OrderedDict) -- training losses stored in the format of (name, float) pairs
+            t_comp (float) -- computational time per data point (normalized by batch_size)
+            t_data (float) -- data loading time per data point (normalized by batch_size)
+        """
+        message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, iters, t_comp, t_data)
+        for k, v in losses.items():
+            message += '%s: %.3f ' % (k, v)
+
+        print(message)  # print the message
+        with open(self.log_name, "a") as log_file:
+            log_file.write('%s\n' % message)  # save the message

app.py

@@ -0,0 +1,507 @@
+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.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 util_colorflow.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
+
+model_global_path = snapshot_download(repo_id="JunhaoZhuang/ColorFlow", cache_dir='./colorflow/')
+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
+
+# 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.cuda()
+
+# 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('cuda')
+
+
+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_1/input.jpg", 
+        ["./assets/example_1/ref1.jpg", "./assets/example_1/ref2.jpg", "./assets/example_1/ref3.jpg"], 
+        "Sketch",  
+        "640x640", 
+        0, 
+        10  
+    ],
+    [
+        "./assets/example_0/input.jpg", 
+        ["./assets/example_0/ref1.jpg"], 
+        "Sketch", 
+        "640x640",  
+        0, 
+        10 
+    ],
+]
+
+global pipeline
+global MultiResNetModel
+
+def load_ckpt(input_style):
+    global pipeline
+    global MultiResNetModel
+    if input_style == "Sketch":
+        ckpt_path = model_global_path + '/sketch/'
+        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('cuda', dtype=weight_dtype)
+        ColorGuider.to('cuda', 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("cuda")
+        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('cuda', 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('cuda', dtype=weight_dtype)
+        unet.to('cuda', 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("cuda")
+        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('cuda', dtype=weight_dtype)
+
+    
+
+
+
+global cur_input_style
+cur_input_style = "Sketch"
+load_ckpt(cur_input_style)
+cur_input_style = "GrayImage(ScreenStyle)"
+load_ckpt(cur_input_style)
+
+
+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 
+
+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).cuda()
+
+    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
+
+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
+
+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
+    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
+    torch.cuda.empty_cache()
+    return input_context, extracted_line, input_context  
+
+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
+    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='cuda').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('cuda', dtype=weight_dtype)
+        query_image_vae = transform(query_image_vae).unsqueeze(0).to('cuda', 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(
+    """
+<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/">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) 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>
+    """
+)
+    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"], 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=6,
+        )
+demo.launch(server_name="0.0.0.0", server_port=22348)

diffusers/.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,110 @@
+name: "\U0001F41B Bug Report"
+description: Report a bug on Diffusers
+labels: [ "bug" ]
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks a lot for taking the time to file this issue 🤗.
+        Issues do not only help to improve the library, but also publicly document common problems, questions, workflows for the whole community!
+        Thus, issues are of the same importance as pull requests when contributing to this library ❤️.
+        In order to make your issue as **useful for the community as possible**, let's try to stick to some simple guidelines:
+        - 1. Please try to be as precise and concise as possible.
+             *Give your issue a fitting title. Assume that someone which very limited knowledge of Diffusers can understand your issue. Add links to the source code, documentation other issues, pull requests etc...*
+        - 2. If your issue is about something not working, **always** provide a reproducible code snippet. The reader should be able to reproduce your issue by **only copy-pasting your code snippet into a Python shell**.
+             *The community cannot solve your issue if it cannot reproduce it. If your bug is related to training, add your training script and make everything needed to train public. Otherwise, just add a simple Python code snippet.*
+        - 3. Add the **minimum** amount of code / context that is needed to understand, reproduce your issue.
+             *Make the life of maintainers easy. `diffusers` is getting many issues every day. Make sure your issue is about one bug and one bug only. Make sure you add only the context, code needed to understand your issues - nothing more. Generally, every issue is a way of documenting this library, try to make it a good documentation entry.*
+        - 4. For issues related to community pipelines (i.e., the pipelines located in the `examples/community` folder), please tag the author of the pipeline in your issue thread as those pipelines are not maintained.
+  - type: markdown
+    attributes:
+      value: |
+        For more in-detail information on how to write good issues you can have a look [here](https://huggingface.co/course/chapter8/5?fw=pt).
+  - type: textarea
+    id: bug-description
+    attributes:
+      label: Describe the bug
+      description: A clear and concise description of what the bug is. If you intend to submit a pull request for this issue, tell us in the description. Thanks!
+      placeholder: Bug description
+    validations:
+      required: true
+  - type: textarea
+    id: reproduction
+    attributes:
+      label: Reproduction
+      description: Please provide a minimal reproducible code which we can copy/paste and reproduce the issue.
+      placeholder: Reproduction
+    validations:
+      required: true
+  - type: textarea
+    id: logs
+    attributes:
+      label: Logs
+      description: "Please include the Python logs if you can."
+      render: shell
+  - type: textarea
+    id: system-info
+    attributes:
+      label: System Info
+      description: Please share your system info with us. You can run the command `diffusers-cli env` and copy-paste its output below.
+      placeholder: Diffusers version, platform, Python version, ...
+    validations:
+      required: true
+  - type: textarea
+    id: who-can-help
+    attributes:
+      label: Who can help?
+      description: |
+        Your issue will be replied to more quickly if you can figure out the right person to tag with @.
+        If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
+
+        All issues are read by one of the core maintainers, so if you don't know who to tag, just leave this blank and
+        a core maintainer will ping the right person.
+
+        Please tag a maximum of 2 people.
+
+        Questions on DiffusionPipeline (Saving, Loading, From pretrained, ...): @sayakpaul @DN6
+
+        Questions on pipelines:
+        - Stable Diffusion @yiyixuxu @asomoza
+        - Stable Diffusion XL @yiyixuxu @sayakpaul @DN6
+        - Stable Diffusion 3: @yiyixuxu @sayakpaul @DN6 @asomoza
+        - Kandinsky @yiyixuxu
+        - ControlNet @sayakpaul @yiyixuxu @DN6
+        - T2I Adapter @sayakpaul @yiyixuxu @DN6
+        - IF @DN6
+        - Text-to-Video / Video-to-Video @DN6 @a-r-r-o-w
+        - Wuerstchen @DN6
+        - Other: @yiyixuxu @DN6
+        - Improving generation quality: @asomoza
+
+        Questions on models:
+        - UNet @DN6 @yiyixuxu @sayakpaul
+        - VAE @sayakpaul @DN6 @yiyixuxu
+        - Transformers/Attention @DN6 @yiyixuxu @sayakpaul
+
+        Questions on single file checkpoints: @DN6
+
+        Questions on Schedulers: @yiyixuxu
+
+        Questions on LoRA: @sayakpaul
+
+        Questions on Textual Inversion: @sayakpaul
+
+        Questions on Training:
+        - DreamBooth @sayakpaul
+        - Text-to-Image Fine-tuning @sayakpaul
+        - Textual Inversion @sayakpaul
+        - ControlNet @sayakpaul
+
+        Questions on Tests: @DN6 @sayakpaul @yiyixuxu
+
+        Questions on Documentation: @stevhliu
+
+        Questions on JAX- and MPS-related things: @pcuenca
+
+        Questions on audio pipelines: @sanchit-gandhi
+
+
+
+      placeholder: "@Username ..."

diffusers/.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,4 @@
+contact_links:
+  - name: Questions / Discussions
+    url: https://github.com/huggingface/diffusers/discussions
+    about: General usage questions and community discussions

diffusers/.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: "\U0001F680 Feature Request"
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...].
+
+**Describe the solution you'd like.**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered.**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context.**
+Add any other context or screenshots about the feature request here.

diffusers/.github/ISSUE_TEMPLATE/feedback.md

@@ -0,0 +1,12 @@
+---
+name: "💬 Feedback about API Design"
+about: Give feedback about the current API design
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**What API design would you like to have changed or added to the library? Why?**
+
+**What use case would this enable or better enable? Can you give us a code example?**

diffusers/.github/ISSUE_TEMPLATE/new-model-addition.yml

@@ -0,0 +1,31 @@
+name: "\U0001F31F New Model/Pipeline/Scheduler Addition"
+description: Submit a proposal/request to implement a new diffusion model/pipeline/scheduler
+labels: [ "New model/pipeline/scheduler" ]
+
+body:
+  - type: textarea
+    id: description-request
+    validations:
+      required: true
+    attributes:
+      label: Model/Pipeline/Scheduler description
+      description: |
+        Put any and all important information relative to the model/pipeline/scheduler
+
+  - type: checkboxes
+    id: information-tasks
+    attributes:
+      label: Open source status
+      description: |
+          Please note that if the model implementation isn't available or if the weights aren't open-source, we are less likely to implement it in `diffusers`.
+      options:
+        - label: "The model implementation is available."
+        - label: "The model weights are available (Only relevant if addition is not a scheduler)."
+
+  - type: textarea
+    id: additional-info
+    attributes:
+      label: Provide useful links for the implementation
+      description: |
+        Please provide information regarding the implementation, the weights, and the authors.
+        Please mention the authors by @gh-username if you're aware of their usernames.

diffusers/.github/ISSUE_TEMPLATE/translate.md

@@ -0,0 +1,29 @@
+---
+name: 🌐 Translating a New Language?
+about: Start a new translation effort in your language
+title: '[<languageCode>] Translating docs to <languageName>'
+labels: WIP
+assignees: ''
+
+---
+
+<!--
+Note: Please search to see if an issue already exists for the language you are trying to translate.
+-->
+
+Hi!
+
+Let's bring the documentation to all the <languageName>-speaking community 🌐.
+
+Who would want to translate? Please follow the 🤗 [TRANSLATING guide](https://github.com/huggingface/diffusers/blob/main/docs/TRANSLATING.md). Here is a list of the files ready for translation. Let us know in this issue if you'd like to translate any, and we'll add your name to the list.
+
+Some notes:
+
+* Please translate using an informal tone (imagine you are talking with a friend about Diffusers 🤗).
+* Please translate in a gender-neutral way.
+* Add your translations to the folder called `<languageCode>` inside the [source folder](https://github.com/huggingface/diffusers/tree/main/docs/source).
+* Register your translation in `<languageCode>/_toctree.yml`; please follow the order of the [English version](https://github.com/huggingface/diffusers/blob/main/docs/source/en/_toctree.yml).
+* Once you're finished, open a pull request and tag this issue by including #issue-number in the description, where issue-number is the number of this issue. Please ping @stevhliu for review.
+* 🙋 If you'd like others to help you with the translation, you can also post in the 🤗 [forums](https://discuss.huggingface.co/c/discussion-related-to-httpsgithubcomhuggingfacediffusers/63).
+
+Thank you so much for your help! 🤗