add basic files

configs/__init__.py

@@ -0,0 +1,30 @@
+import importlib
+from configs.base_config import BaseConfig
+
+def find_config_by_name(config_name):
+    # load config lib by config name
+    config_file = "configs." + config_name + '_config'
+    config_lib = importlib.import_module(config_file)
+    print(config_lib)
+
+    # find the subclass of BaseConfig
+    config = None
+    target_config_name = config_name.replace('_', '') + 'config'
+    target_config_name = target_config_name.lower()
+    for name, cls in config_lib.__dict__.items():
+        if name.lower() == target_config_name and issubclass(cls, BaseConfig):
+            config = cls
+    
+    if config is None:
+        raise Exception('No valid config found.')
+
+    return config
+
+def parse_config(cfg_file):
+    # parse config using BaseConfig
+    cfg = BaseConfig().parse_config(cfg_file)
+    model_name = cfg['common']['model']
+
+    # re-parse using specified Config class
+    config = find_config_by_name(model_name)
+    return config().parse_config(cfg_file)

configs/base_config.py

@@ -0,0 +1,160 @@
+import yaml
+import copy
+from typing import Union
+
+class BaseConfig():
+
+    def __init__(self):
+        self.__config_dict = {}
+        self.__check_func_dict = {}
+
+        is_greater_than_0 = lambda x: x > 0
+
+        # common config
+        self._add_option('common', 'name', str, 'style_master')
+        self._add_option('common', 'model', str, 'cycle_gan')
+        self._add_option('common', 'phase', str, 'train', check_func=lambda x: x in ['train', 'test'])
+        self._add_option('common', 'gpu_ids', list, [0])
+        self._add_option('common', 'verbose', bool, False)
+
+        # model config
+        self._add_option('model', 'input_nc', int, 3, check_func=is_greater_than_0)
+        self._add_option('model', 'output_nc', int, 3, check_func=is_greater_than_0)
+
+        # dataset config
+        # common dataset options
+        self._add_option('dataset', 'use_absolute_datafile', bool, True)
+        self._add_option('dataset', 'batch_size', int, 1, check_func=is_greater_than_0)
+        self._add_option('dataset', 'n_threads', int, 4, check_func=is_greater_than_0)
+        self._add_option('dataset', 'dataroot', str, './')
+        self._add_option('dataset', 'drop_last', bool, False)
+        self._add_option('dataset', 'landmark_scale', list, None)
+        self._add_option('dataset', 'check_all_data', bool, False)
+        self._add_option('dataset', 'accept_data_error', bool, True)  # Upon loading a bad data, if this is true,
+                                                                        # dataloader will throw an exception and
+                                                                        # load the next good data.
+                                                                        # If this is false, process will crash.
+
+        self._add_option('dataset', 'train_data', dict, {})
+        self._add_option('dataset', 'val_data', dict, {})
+
+        # paired data config
+        self._add_option('dataset', 'paired_trainA_folder', str, '')
+        self._add_option('dataset', 'paired_trainB_folder', str, '')
+        self._add_option('dataset', 'paired_train_filelist', str, '')
+        self._add_option('dataset', 'paired_valA_folder', str, '')
+        self._add_option('dataset', 'paired_valB_folder', str, '')
+        self._add_option('dataset', 'paired_val_filelist', str, '')
+
+        # unpaired data config
+        self._add_option('dataset', 'unpaired_trainA_folder', str, '')
+        self._add_option('dataset', 'unpaired_trainB_folder', str, '')
+        self._add_option('dataset', 'unpaired_trainA_filelist', str, '')
+        self._add_option('dataset', 'unpaired_trainB_filelist', str, '')
+        self._add_option('dataset', 'unpaired_valA_folder', str, '')
+        self._add_option('dataset', 'unpaired_valB_folder', str, '')
+        self._add_option('dataset', 'unpaired_valA_filelist', str, '')
+        self._add_option('dataset', 'unpaired_valB_filelist', str, '')
+
+        # custom data
+        self._add_option('dataset', 'custom_train_data', dict, {})
+        self._add_option('dataset', 'custom_val_data', dict, {})
+
+        # training config
+        self._add_option('training', 'checkpoints_dir', str, './checkpoints')
+        self._add_option('training', 'log_dir', str, './logs')
+        self._add_option('training', 'use_new_log', bool, False)
+        self._add_option('training', 'continue_train', bool, False)
+        self._add_option('training', 'which_epoch', str, 'latest')
+        self._add_option('training', 'n_epochs', int, 100, check_func=is_greater_than_0)
+        self._add_option('training', 'n_epochs_decay', int, 100, check_func=is_greater_than_0)
+        self._add_option('training', 'save_latest_freq', int, 5000, check_func=is_greater_than_0)
+        self._add_option('training', 'print_freq', int, 200, check_func=is_greater_than_0)
+        self._add_option('training', 'save_epoch_freq', int, 5, check_func=is_greater_than_0)
+        self._add_option('training', 'epoch_as_iter', bool, False)
+        self._add_option('training', 'lr', float, 2e-4, check_func=is_greater_than_0)
+        self._add_option('training', 'lr_policy', str, 'linear',
+            check_func=lambda x: x in ['linear', 'step', 'plateau', 'cosine'])
+        self._add_option('training', 'lr_decay_iters', int, 50, check_func=is_greater_than_0)
+        self._add_option('training', 'DDP', bool, False)
+        self._add_option('training', 'num_nodes', int, 1, check_func=is_greater_than_0)
+        self._add_option('training', 'DDP_address', str, '127.0.0.1')
+        self._add_option('training', 'DDP_port', str, '29700')
+        self._add_option('training', 'find_unused_parameters', bool, False) # a DDP option that allows backward on a subgraph of the model
+        self._add_option('training', 'val_percent', float, 5.0, check_func=is_greater_than_0)  # Uses x% of training data to validate
+        self._add_option('training', 'val', bool, True)  # perform validation every epoch
+        self._add_option('training', 'save_training_progress', bool, False)  # save images to create a training progression video
+
+        # testing config
+        self._add_option('testing', 'results_dir', str, './results')
+        self._add_option('testing', 'load_size', int, 512, check_func=is_greater_than_0)
+        self._add_option('testing', 'crop_size', int, 512, check_func=is_greater_than_0)
+        self._add_option('testing', 'preprocess', list, ['scale_width'])
+        self._add_option('testing', 'visual_names', list, [])
+        self._add_option('testing', 'num_test', int, 999999, check_func=is_greater_than_0)
+        self._add_option('testing', 'image_format', str, 'jpg', check_func=lambda x: x in ['input', 'jpg', 'jpeg', 'png'])
+
+    def _add_option(self, group_name, option_name, value_type, default_value, check_func=None):
+        # check name type
+        if not type(group_name) is str or not type(option_name) is str:
+            raise Exception('Type of {} and {} must be str.'.format(group_name, option_name))
+
+        # add group
+        if not group_name in self.__config_dict:
+            self.__config_dict[group_name] = {}
+            self.__check_func_dict[group_name] = {}
+
+        # check type & default value
+        if not type(value_type) is type:
+            try:
+                if value_type.__origin__ is not Union:
+                    raise Exception('{} is not a type.'.format(value_type))
+            except Exception as e:
+                print(e)
+        if not type(default_value) is value_type:
+            try:
+                if value_type.__origin__ is not Union:
+                    raise Exception('Type of {} must be {}.'.format(default_value, value_type))
+            except Exception as e:
+                print(e)
+
+        # add option to dict
+        if not option_name in self.__config_dict[group_name]:
+            if not check_func is None and not check_func(default_value):
+                raise Exception('Checking {}/{} failed.'.format(group_name, option_name))
+            self.__config_dict[group_name][option_name] = default_value
+            self.__check_func_dict[group_name][option_name] = check_func
+        else:
+            raise Exception('{} has been already added.'.format(option_name))
+
+    def parse_config(self, cfg_file):
+        # load config from yaml file
+        with open(cfg_file, 'r') as f:
+            yaml_config = yaml.safe_load(f)
+        if not type(yaml_config) is dict:
+            raise Exception('Loading yaml file failed.')
+
+        # replace default options
+        config_dict = copy.deepcopy(self.__config_dict)
+        for group in config_dict:
+            if group in yaml_config:
+                for option in config_dict[group]:
+                    if option in yaml_config[group]:
+                        value = yaml_config[group][option]
+                        if not type(value) is type(config_dict[group][option]):
+                            try: # if <config_dict[group][option]> is not union, it won't have __origin__ attribute. So will throw an error.
+                                # The line below is necessary because we check if <config_dict[group][option]> has __origin__ attribute.
+                                if config_dict[group][option].__origin__ is Union:
+                                    # check to see if type of <value> belongs to a type in the union.
+                                    if not isinstance(value, config_dict[group][option].__args__):
+                                        raise Exception('Type of {}/{} must be {}.'.format(group, option,
+                                                                config_dict[group][option].__args__))
+                            except Exception as e: # if the error was thrown, we know there's a type error.
+                                print(e)
+                        else:
+                            check_func = self.__check_func_dict[group][option]
+                            if not check_func is None and not check_func(value):
+                                raise Exception('Checking {}/{} failed.'.format(group, option))
+                            config_dict[group][option] = value
+        return config_dict
+

configs/style_based_pix2pixII_config.py

@@ -0,0 +1,42 @@
+from .base_config import BaseConfig
+from typing import Union as Union
+
+class StyleBasedPix2PixIIConfig(BaseConfig):
+
+    def __init__(self):
+        super(StyleBasedPix2PixIIConfig, self).__init__()
+
+        is_greater_than_0 = lambda x: x > 0
+
+        # model config
+        self._add_option('model', 'ngf', int, 64, check_func=is_greater_than_0)
+        self._add_option('model', 'min_feats_size', list, [4, 4])
+
+        # dataset config
+        self._add_option('dataset', 'data_type', list, ['unpaired'])
+        self._add_option('dataset', 'direction', str, 'AtoB')
+        self._add_option('dataset', 'serial_batches', bool, False)
+        self._add_option('dataset', 'load_size', int, 512, check_func=is_greater_than_0)
+        self._add_option('dataset', 'crop_size', int, 512, check_func=is_greater_than_0)
+        self._add_option('dataset', 'preprocess', Union[list, str], ['resize'])
+        self._add_option('dataset', 'no_flip', bool, True)
+
+        # training config
+        self._add_option('training', 'beta1', float, 0.1, check_func=is_greater_than_0)
+        self._add_option('training', 'data_aug_prob', float, 0.0, check_func=lambda x: x >= 0.0)
+        self._add_option('training', 'style_mixing_prob', float, 0.0, check_func=lambda x: x >= 0.0)
+        self._add_option('training', 'phase', int, 1, check_func=lambda x: x in [1, 2, 3, 4])
+        self._add_option('training', 'pretrained_model', str, 'model.pth')
+        self._add_option('training', 'src_text_prompt', str, 'photo')
+        self._add_option('training', 'text_prompt', str, 'a portrait in style of sketch')
+        self._add_option('training', 'image_prompt', str, 'style.png')
+        self._add_option('training', 'lambda_L1', float, 1.0)
+        self._add_option('training', 'lambda_Feat', float, 4.0)
+        self._add_option('training', 'lambda_ST', float, 1.0)
+        self._add_option('training', 'lambda_GAN', float, 1.0)
+        self._add_option('training', 'lambda_CLIP', float, 1.0)
+        self._add_option('training', 'lambda_PROJ', float, 1.0)
+        self._add_option('training', 'ema', float, 0.999)
+
+        # testing config
+        self._add_option('testing', 'aspect_ratio', float, 1.0, check_func=is_greater_than_0)

data/__init__.py

@@ -0,0 +1,58 @@
+"""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 torch.utils.data.distributed import DistributedSampler
+
+class CustomDataLoader():
+    """Wrapper class of Dataset class that performs multi-threaded data loading"""
+
+    def __init__(self, config, dataset, DDP_gpu=None, drop_last=False):
+        """Initialize this class
+
+        Step 1: create a dataset instance given the name [dataset_mode]
+        Step 2: create a multi-threaded data loader.
+        """
+        self.config = config
+        self.dataset = dataset
+
+        if DDP_gpu is None:
+            self.dataloader = torch.utils.data.DataLoader(
+                self.dataset,
+                batch_size=config['dataset']['batch_size'],
+                shuffle=not config['dataset']['serial_batches'],
+                num_workers=int(config['dataset']['n_threads']), drop_last=drop_last)
+        else:
+            sampler = DistributedSampler(self.dataset, num_replicas=self.config['training']['world_size'],
+                                         rank=DDP_gpu)
+            self.dataloader = torch.utils.data.DataLoader(
+                self.dataset,
+                batch_size=config['dataset']['batch_size'],
+                shuffle=False,
+                num_workers=int(config['dataset']['n_threads']),
+                sampler=sampler,
+                drop_last=drop_last)
+
+    def load_data(self):
+        return self
+
+    def __len__(self):
+        """Return the number of data in the dataset"""
+        return min(len(self.dataset), 1e9)
+
+    def __iter__(self):
+        """Return a batch of data"""
+        for i, data in enumerate(self.dataloader):
+            if i * self.config['dataset']['batch_size'] >= 1e9:
+                break
+            yield data

data/deprecated/custom_data.py

@@ -0,0 +1,121 @@
+import os
+import random
+import numpy as np
+from utils.augmentation import ImagePathToImage
+from utils.data_utils import Transforms, check_img_loaded, check_numpy_loaded
+
+
+class CustomData(object):
+
+    def __init__(self, config, shuffle=False):
+        self.paired_file_groups = []
+        self.paired_type_groups = []
+        self.len_of_groups = []
+        self.landmark_scale = config['dataset']['landmark_scale']
+        self.shuffle = shuffle
+        self.config = config
+
+        data_dict = config['dataset']['custom_' + config['common']['phase'] + '_data']
+        if len(data_dict) == 0:
+            self.len_of_groups.append(0)
+            return
+
+        for i, group in enumerate(data_dict.values()):  # one example: (0, group_1),  (1, group_2)
+            data_types = group['data_types']  # one example: 'image', 'patch'
+            data_names = group['data_names']  # one example: 'real_A', 'patch_A'
+            file_list = group['file_list']  # one example: "lmt/data/trainA.txt"
+            assert(len(data_types) == len(data_names))
+
+            self.paired_file_groups.append({})
+            self.paired_type_groups.append({})
+            for data_name, data_type in zip(data_names, data_types):
+                self.paired_file_groups[i][data_name] = []
+                self.paired_type_groups[i][data_name] = data_type
+
+            paired_file = open(file_list, 'rt')
+            lines = paired_file.readlines()
+            if self.shuffle:
+                random.shuffle(lines)
+            for line in lines:
+                items = line.strip().split(' ')
+                if len(items) == len(data_names):
+                    ok = True
+                    for item in items:
+                        ok = ok and os.path.exists(item) and os.path.getsize(item) > 0
+                    if ok:
+                        for data_name, item in zip(data_names, items):
+                            self.paired_file_groups[i][data_name].append(item)
+            paired_file.close()
+
+            self.len_of_groups.append(len(self.paired_file_groups[i][data_names[0]]))
+
+        self.transform = Transforms(config)
+        self.transform.get_transform_from_config()
+        self.transform.get_transforms().insert(0, ImagePathToImage())
+        self.transform = self.transform.compose_transforms()
+
+    def get_len(self):
+        return max(self.len_of_groups)
+
+    def get_item(self, idx):
+        return_dict = {}
+        for i in range(len(self.paired_file_groups)):
+            inner_idx = idx if idx < self.len_of_groups[i] else random.randint(0, self.len_of_groups[i] - 1)
+            img_list = []
+            img_k_list = []
+            for k, v in self.paired_file_groups[i].items():
+                if self.paired_type_groups[i][k] == 'image':
+                    # gather images for processing later
+                    img_k_list.append(k)
+                    img_list.append(v[inner_idx])
+                elif self.paired_type_groups[i][k] == 'landmark':
+                    # different from images, landmark doesn't use data augmentation. So process them directly here.
+                    lmk = np.load(v[inner_idx])
+                    lmk[:, 0] *= self.landmark_scale[0]
+                    lmk[:, 1] *= self.landmark_scale[1]
+                    return_dict[k] = lmk
+                return_dict[k + '_path'] = v[inner_idx]
+
+            # transform all images
+            if len(img_list) == 1:
+                return_dict[img_k_list[0]], _ = self.transform(img_list[0], None)
+            elif len(img_list) > 1:
+                input1, input2 = img_list[0], img_list[1:]
+                output1, output2 = self.transform(input1, input2) # output1 is one image. output2 is a list of images.
+                return_dict[img_k_list[0]] = output1
+                for j in range(1, len(img_list)):
+                    return_dict[img_k_list[j]] = output2[j-1]
+
+        return return_dict
+
+    def split_data_into_bins(self, num_bins):
+        bins = []
+        for i in range(0, num_bins):
+            bins.append([])
+        for i in range(0, len(self.paired_file_groups)):
+            for b in range(0, num_bins):
+                bins[b].append({})
+            for dataname, item_list in self.paired_file_groups[i].items():
+                if len(item_list) < self.config['dataset']['n_threads']:
+                    bins[0][i][dataname] = item_list
+                else:
+                    num_items_in_bin = len(item_list) // num_bins
+                    for j in range(0, len(item_list)):
+                        which_bin = min(j // num_items_in_bin, num_bins - 1)
+                        if dataname not in bins[which_bin][i]:
+                            bins[which_bin][i][dataname] = []
+                        else:
+                            bins[which_bin][i][dataname].append(item_list[j])
+        return bins
+
+    def check_data_helper(self, data):
+        all_pass = True
+        for paired_file_group in data:
+            for k, v in paired_file_group.items():
+                if len(v) > 0:
+                    for v1 in v:
+                        if '.npy' in v1:  # case: numpy array or landmark
+                            all_pass = all_pass and check_numpy_loaded(v1)
+                        else:  # case: image
+                            all_pass = all_pass and check_img_loaded(v1)
+        return all_pass

data/deprecated/landmark_data.py

@@ -0,0 +1,89 @@
+import os
+from PIL import Image
+import numpy as np
+from utils.data_utils import check_create_shuffled_order, check_equal_length
+
+def landmark_path_to_numpy(lmk_path, image_path, image_tensor):
+    """Convert an landmark path to the actual landmarks in a numpy array. Also applies scaling to the landmarks
+    according to final image' size.
+
+    Parameters:
+        lmk_path  --  the landmark file path.
+        image_path  --  the original image file path.
+        image_tensor  --  the image tensor after all transformations.
+    """
+    lmk = np.load(lmk_path)
+    ow, oh = Image.open(image_path).size
+    h, w = image_tensor.size()[1:]
+    lmk[:, 0] *= w / ow
+    lmk[:, 1] *= h / oh
+    return lmk
+
+def add_landmark_data(data, config, paired_data_order):
+    A_lmk_paths = []
+    B_lmk_paths = []
+
+    if config['dataset']['paired_' + config['common']['phase'] + '_filelist'] != '':
+        paired_data_file = open(config['dataset']['paired_' + config['common']['phase'] + '_filelist'], 'r')
+        Lines = paired_data_file.readlines()
+        paired_data_order = check_create_shuffled_order(Lines, paired_data_order)
+        check_equal_length(Lines, paired_data_order, data)
+        for i in paired_data_order:
+            line = Lines[i]
+            if not config['dataset']['use_absolute_datafile']:
+                file3 = os.path.join(config['dataset']['dataroot'], line.split(" ")[2]).strip()
+                file4 = os.path.join(config['dataset']['dataroot'], line.split(" ")[3]).strip()
+            else:
+                file3 = line.split(" ")[2].strip()
+                file4 = line.split(" ")[3].strip()
+            if os.path.exists(file3) and os.path.exists(file4):
+                A_lmk_paths.append(file3)
+                B_lmk_paths.append(file4)
+        paired_data_file.close()
+    elif config['dataset']['paired_' + config['common']['phase'] + 'A_folder'] != '' and \
+            config['dataset']['paired_' + config['common']['phase'] + 'B_folder'] != '' and \
+            os.path.exists(config['dataset']['paired_' + config['common']['phase'] + 'A_lmk_folder']) and \
+            os.path.exists(config['dataset']['paired_' + config['common']['phase'] + 'B_lmk_folder']):
+        dir_A = config['dataset']['paired_' + config['common']['phase'] + 'A_folder']
+        dir_A_lmk = config['dataset']['paired_' + config['common']['phase'] + 'A_lmk_folder']
+        dir_B_lmk = config['dataset']['paired_' + config['common']['phase'] + 'B_lmk_folder']
+        filenames = os.listdir(dir_A)
+        paired_data_order = check_create_shuffled_order(filenames, paired_data_order)
+        check_equal_length(filenames, paired_data_order, data)
+        for i in paired_data_order:
+            filename = filenames[i]
+            A_lmk_path = os.path.join(dir_A_lmk, os.path.splitext(filename)[0] + '.npy')
+            B_lmk_path = os.path.join(dir_B_lmk, os.path.splitext(filename)[0] + '.npy')
+            if os.path.exists(A_lmk_path) and os.path.exists(B_lmk_path):
+                A_lmk_paths.append(A_lmk_path)
+                B_lmk_paths.append(B_lmk_path)
+    else:
+        dir_A = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'pairedA')
+        dir_A_lmk = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'pairedA_lmk')
+        dir_B_lmk = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'pairedB_lmk')
+        if os.path.exists(dir_A_lmk) and os.path.exists(dir_B_lmk):
+            filenames = os.listdir(dir_A)
+            paired_data_order = check_create_shuffled_order(filenames, paired_data_order)
+            check_equal_length(filenames, paired_data_order, data)
+            for i in paired_data_order:
+                filename = filenames[i]
+                A_lmk_path = os.path.join(dir_A_lmk, os.path.splitext(filename)[0] + '.npy')
+                B_lmk_path = os.path.join(dir_B_lmk, os.path.splitext(filename)[0] + '.npy')
+                if os.path.exists(A_lmk_path) and os.path.exists(B_lmk_path):
+                    A_lmk_paths.append(A_lmk_path)
+                    B_lmk_paths.append(B_lmk_path)
+        else:
+            print(dir_A_lmk + " or " + dir_B_lmk + " doesn't exist. Skipping landmark data.")
+
+    data['A_lmk_path'] = A_lmk_paths
+    data['B_lmk_path'] = B_lmk_paths
+
+    return paired_data_order
+
+
+def apply_landmark_transforms(index, data, return_dict):
+    if len(data['A_lmk_path']) > 0:
+        return_dict['A_lmk'] = landmark_path_to_numpy(data['A_lmk_path'][index], data['paired_A_path'][index], return_dict['paired_A'])
+        return_dict['B_lmk'] = landmark_path_to_numpy(data['B_lmk_path'][index], data['paired_B_path'][index], return_dict['paired_B'])
+        return_dict['A_lmk_path'] = data['A_lmk_path'][index]
+        return_dict['B_lmk_path'] = data['B_lmk_path'][index]

data/deprecated/numpy_paired_data.py

@@ -0,0 +1,81 @@
+import os
+from utils.util import check_path_is_img
+from utils.data_utils import Transforms, check_create_shuffled_order, check_equal_length
+from utils.augmentation import NumpyToTensor
+
+
+def add_numpy_paired_data(data, transforms, config, paired_data_order):
+    A_paths = []
+    B_paths = []
+
+    if config['dataset']['paired_' + config['common']['phase'] + '_filelist'] != '':
+        paired_data_file = open(config['dataset']['paired_' + config['common']['phase'] + '_filelist'], 'r')
+        Lines = paired_data_file.readlines()
+        paired_data_order = check_create_shuffled_order(Lines, paired_data_order)
+        check_equal_length(Lines, paired_data_order, data)
+        for i in paired_data_order:
+            line = Lines[i]
+            if not config['dataset']['use_absolute_datafile']:
+                file1 = os.path.join(config['dataset']['dataroot'], line.split(" ")[0]).strip()
+                file2 = os.path.join(config['dataset']['dataroot'], line.split(" ")[1]).strip()
+            else:
+                file1 = line.split(" ")[0].strip()
+                file2 = line.split(" ")[1].strip()
+            if os.path.exists(file1) and os.path.exists(file2):
+                A_paths.append(file1)
+                B_paths.append(file2)
+        paired_data_file.close()
+    elif config['dataset']['paired_' + config['common']['phase'] + 'A_folder'] != '' and \
+            config['dataset']['paired_' + config['common']['phase'] + 'B_folder'] != '':
+        dir_A = config['dataset']['paired_' + config['common']['phase'] + 'A_folder']
+        dir_B = config['dataset']['paired_' + config['common']['phase'] + 'B_folder']
+        filenames = os.listdir(dir_A)
+        paired_data_order = check_create_shuffled_order(filenames, paired_data_order)
+        check_equal_length(filenames, paired_data_order, data)
+        for i in paired_data_order:
+            filename = filenames[i]
+            if not check_path_is_img(filename):
+                continue
+            A_path = os.path.join(dir_A, filename)
+            B_path = os.path.join(dir_B, filename)
+            if os.path.exists(A_path) and os.path.exists(B_path):
+                A_paths.append(A_path)
+                B_paths.append(B_path)
+    else:
+        dir_A = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'numpypairedA')
+        dir_B = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'numpypairedB')
+        if os.path.exists(dir_A) and os.path.exists(dir_B):
+            filenames = os.listdir(dir_A)
+            paired_data_order = check_create_shuffled_order(filenames, paired_data_order)
+            check_equal_length(filenames, paired_data_order, data)
+            for i in paired_data_order:
+                filename = filenames[i]
+                if not check_path_is_img(filename):
+                    continue
+                A_path = os.path.join(dir_A, filename)
+                B_path = os.path.join(dir_B, filename)
+                if os.path.exists(A_path) and os.path.exists(B_path):
+                    A_paths.append(A_path)
+                    B_paths.append(B_path)
+
+    btoA = config['dataset']['direction'] == 'BtoA'
+    # get the number of channels of input image
+    input_nc = config['model']['output_nc'] if btoA else config['model']['input_nc']
+    output_nc = config['model']['input_nc'] if btoA else config['model']['output_nc']
+
+    transform = Transforms(config, input_grayscale_flag=(input_nc == 1), output_grayscale_flag=(output_nc == 1))
+    transform.transform_list.append(NumpyToTensor())
+    transform = transform.compose_transforms()
+
+    data['paired_A_path'] = A_paths
+    data['paired_B_path'] = B_paths
+    transforms['paired'] = transform
+    return paired_data_order
+
+
+def apply_numpy_paired_transforms(index, data, transforms, return_dict):
+    if len(data['paired_A_path']) > 0:
+        return_dict['paired_A'], return_dict['paired_B'] = transforms['paired'] \
+            (data['paired_A_path'][index], data['paired_B_path'][index])
+        return_dict['paired_A_path'] = data['paired_A_path'][index]
+        return_dict['paired_B_path'] = data['paired_B_path'][index]

data/deprecated/numpy_unpaired_data.py

@@ -0,0 +1,100 @@
+import os
+from utils.util import check_path_is_img
+from utils.data_utils import Transforms
+from utils.augmentation import NumpyToTensor
+import random
+
+def add_numpy_unpaired_data(data, transforms, config, shuffle=False):
+    A_paths = []
+    B_paths = []
+    if config['dataset']['unpaired_' + config['common']['phase'] + 'A_filelist'] != '':
+        unpaired_data_file1 = open(config['dataset']['unpaired_' + config['common']['phase'] + 'A_filelist'], 'r')
+        Lines = unpaired_data_file1.readlines()
+        if shuffle:
+            random.shuffle(Lines)
+        for line in Lines:
+            if not config['dataset']['use_absolute_datafile']:
+                file = os.path.join(config['dataset']['dataroot'], line.strip())
+            else:
+                file = line.strip()
+            if os.path.exists(file):
+                A_paths.append(file)
+        unpaired_data_file1.close()
+
+        unpaired_data_file2 = open(config['dataset']['unpaired_' + config['common']['phase'] + 'B_filelist'], 'r')
+        Lines = unpaired_data_file2.readlines()
+        if shuffle:
+            random.shuffle(Lines)
+        for line in Lines:
+            if not config['dataset']['use_absolute_datafile']:
+                file = os.path.join(config['dataset']['dataroot'], line.strip())
+            else:
+                file = line.strip()
+            if os.path.exists(file):
+                B_paths.append(file)
+        unpaired_data_file2.close()
+    elif config['dataset']['unpaired_' + config['common']['phase'] + 'A_folder'] != '' and \
+         config['dataset']['unpaired_' + config['common']['phase'] + 'B_folder'] != '':
+        dir_A = config['dataset']['unpaired_' + config['common']['phase'] + 'A_folder']
+        filenames = os.listdir(dir_A)
+        if shuffle:
+            random.shuffle(filenames)
+        for filename in filenames:
+            if not check_path_is_img(filename):
+                continue
+            A_path = os.path.join(dir_A, filename)
+            if os.path.exists(A_path):
+                A_paths.append(A_path)
+
+        dir_B = config['dataset']['unpaired_' + config['common']['phase'] + 'B_folder']
+        filenames = os.listdir(dir_B)
+        if shuffle:
+            random.shuffle(filenames)
+        for filename in filenames:
+            if not check_path_is_img(filename):
+                continue
+            B_path = os.path.join(dir_B, filename)
+            if os.path.exists(B_path):
+                B_paths.append(B_path)
+
+    else:
+        dir_A = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'numpyunpairedA')
+        dir_B = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'numpyunpairedB')
+        if os.path.exists(dir_A) and os.path.exists(dir_B):
+            filenames = os.listdir(dir_A)
+            if shuffle:
+                random.shuffle(filenames)
+            for filename in filenames:
+                if not check_path_is_img(filename):
+                    continue
+                A_path = os.path.join(dir_A, filename)
+                A_paths.append(A_path)
+            filenames = os.listdir(dir_B)
+            if shuffle:
+                random.shuffle(filenames)
+            for filename in filenames:
+                if not check_path_is_img(filename):
+                    continue
+                B_path = os.path.join(dir_B, filename)
+                B_paths.append(B_path)
+
+
+    btoA = config['dataset']['direction'] == 'BtoA'
+    input_nc = config['model']['output_nc'] if btoA else config['model']['input_nc']
+    output_nc = config['model']['input_nc'] if btoA else config['model']['output_nc']
+
+    transform = Transforms(config, input_grayscale_flag=(input_nc == 1), output_grayscale_flag=(output_nc == 1))
+    transform.transform_list.append(NumpyToTensor())
+    transform = transform.compose_transforms()
+
+    data['unpaired_A_path'] = A_paths
+    data['unpaired_B_path'] = B_paths
+    transforms['unpaired'] = transform
+
+def apply_numpy_unpaired_transforms(index, data, transforms, return_dict):
+    if len(data['unpaired_A_path']) > 0 and len(data['unpaired_B_path']) > 0:
+        index_B = random.randint(0, len(data['unpaired_B_path']) - 1)
+        return_dict['unpaired_A'], return_dict['unpaired_B'] = transforms['unpaired'] \
+            (data['unpaired_A_path'][index], data['unpaired_B_path'][index_B])
+        return_dict['unpaired_A_path'] = data['unpaired_A_path'][index]
+        return_dict['unpaired_B_path'] = data['unpaired_B_path'][index_B]

data/deprecated/paired_data.py

@@ -0,0 +1,80 @@
+import os
+from utils.util import check_path_is_img
+from utils.data_utils import Transforms, check_create_shuffled_order
+from utils.augmentation import ImagePathToImage
+
+
+def add_paired_data(data, transforms, config, paired_data_order):
+    A_paths = []
+    B_paths = []
+
+    if config['dataset']['paired_' + config['common']['phase'] + '_filelist'] != '':
+        paired_data_file = open(config['dataset']['paired_' + config['common']['phase'] + '_filelist'], 'r')
+        Lines = paired_data_file.readlines()
+        paired_data_order = check_create_shuffled_order(Lines, paired_data_order)
+        for i in paired_data_order:
+            line = Lines[i]
+            if not config['dataset']['use_absolute_datafile']:
+                file1 = os.path.join(config['dataset']['dataroot'], line.split(" ")[0]).strip()
+                file2 = os.path.join(config['dataset']['dataroot'], line.split(" ")[1]).strip()
+            else:
+                file1 = line.split(" ")[0].strip()
+                file2 = line.split(" ")[1].strip()
+            if os.path.exists(file1) and os.path.exists(file2):
+                A_paths.append(file1)
+                B_paths.append(file2)
+        paired_data_file.close()
+    elif config['dataset']['paired_' + config['common']['phase'] + 'A_folder'] != '' and \
+            config['dataset']['paired_' + config['common']['phase'] + 'B_folder'] != '':
+        dir_A = config['dataset']['paired_' + config['common']['phase'] + 'A_folder']
+        dir_B = config['dataset']['paired_' + config['common']['phase'] + 'B_folder']
+        filenames = os.listdir(dir_A)
+        paired_data_order = check_create_shuffled_order(filenames, paired_data_order)
+        for i in paired_data_order:
+            filename = filenames[i]
+            if not check_path_is_img(filename):
+                continue
+            A_path = os.path.join(dir_A, filename)
+            B_path = os.path.join(dir_B, filename)
+            if os.path.exists(A_path) and os.path.exists(B_path):
+                A_paths.append(A_path)
+                B_paths.append(B_path)
+    else:
+        dir_A = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'pairedA')
+        dir_B = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'pairedB')
+        if os.path.exists(dir_A) and os.path.exists(dir_B):
+            filenames = os.listdir(dir_A)
+            paired_data_order = check_create_shuffled_order(filenames, paired_data_order)
+            for i in paired_data_order:
+                filename = filenames[i]
+                if not check_path_is_img(filename):
+                    continue
+                A_path = os.path.join(dir_A, filename)
+                B_path = os.path.join(dir_B, filename)
+                if os.path.exists(A_path) and os.path.exists(B_path):
+                    A_paths.append(A_path)
+                    B_paths.append(B_path)
+
+    btoA = config['dataset']['direction'] == 'BtoA'
+    # get the number of channels of input image
+    input_nc = config['model']['output_nc'] if btoA else config['model']['input_nc']
+    output_nc = config['model']['input_nc'] if btoA else config['model']['output_nc']
+
+    transform = Transforms(config, input_grayscale_flag=(input_nc == 1), output_grayscale_flag=(output_nc == 1))
+    transform.get_transform_from_config()
+    transform.get_transforms().insert(0, ImagePathToImage())
+    transform = transform.compose_transforms()
+
+    data['paired_A_path'] = A_paths
+    data['paired_B_path'] = B_paths
+
+    transforms['paired'] = transform
+    return paired_data_order
+
+
+def apply_paired_transforms(index, data, transforms, return_dict):
+    if len(data['paired_A_path']) > 0:
+        return_dict['paired_A'], return_dict['paired_B'] = transforms['paired'] \
+            (data['paired_A_path'][index], data['paired_B_path'][index])
+        return_dict['paired_A_path'] = data['paired_A_path'][index]
+        return_dict['paired_B_path'] = data['paired_B_path'][index]

data/deprecated/patch_data.py

@@ -0,0 +1,44 @@
+import random
+import torch
+
+def load_patches(patch_batch_size, batch_size, patch_size, num_patch, diff_patch, index, data, transforms, return_dict):
+    if patch_size > 0:
+        assert (patch_batch_size % batch_size == 0), \
+            "patch_batch_size is not divisible by batch_size."
+        if 'paired_A' in return_dict or 'paired_B' in return_dict:
+            if not diff_patch:
+                # load patch from current image
+                patchA = return_dict['paired_A'].clone()
+                patchB = return_dict['paired_B'].clone()
+            else:
+                # load patch from a different image
+                pathA = data['paired_A_path'][(index + 1) % len(data['paired_A_path'])]
+                pathB = data['paired_B_path'][(index + 1) % len(data['paired_B_path'])]
+                patchA, patchB = transforms['paired'](pathA, pathB)
+        else:
+            if not diff_patch:
+                # load patch from current image
+                patchA = return_dict['unpaired_A'].clone()
+                patchB = return_dict['unpaired_B'].clone()
+            else:
+                # load patch from a different image
+                pathA = data['unpaired_A_path'][(index + 1) % len(data['unpaired_A_path'])]
+                pathB = data['unpaired_B_path'][(index + 1) % len(data['unpaired_B_path'])]
+                patchA, patchB = transforms['unpaired'](pathA, pathB)
+
+        # crop patch
+        patchAs = []
+        patchBs = []
+        _, h, w = patchA.size()
+
+        for _ in range(num_patch):
+            r = random.randint(0, h - patch_size - 1)
+            c = random.randint(0, w - patch_size - 1)
+            patchAs.append(patchA[:, r:r + patch_size, c:c + patch_size])
+            patchBs.append(patchB[:, r:r + patch_size, c:c + patch_size])
+
+        patchAs = torch.cat(patchAs, 0)
+        patchBs = torch.cat(patchBs, 0)
+
+        return_dict['patch_A'] = patchAs
+        return_dict['patch_B'] = patchBs

data/deprecated/unpaired_data.py

@@ -0,0 +1,101 @@
+import os
+from utils.util import check_path_is_img
+from utils.data_utils import Transforms
+from utils.augmentation import ImagePathToImage
+import random
+
+def add_unpaired_data(data, transforms, config, shuffle=False):
+    A_paths = []
+    B_paths = []
+    if config['dataset']['unpaired_' + config['common']['phase'] + 'A_filelist'] != '':
+        unpaired_data_file1 = open(config['dataset']['unpaired_' + config['common']['phase'] + 'A_filelist'], 'r')
+        Lines = unpaired_data_file1.readlines()
+        if shuffle:
+            random.shuffle(Lines)
+        for line in Lines:
+            if not config['dataset']['use_absolute_datafile']:
+                file = os.path.join(config['dataset']['dataroot'], line.strip())
+            else:
+                file = line.strip()
+            if os.path.exists(file):
+                A_paths.append(file)
+        unpaired_data_file1.close()
+
+        unpaired_data_file2 = open(config['dataset']['unpaired_' + config['common']['phase'] + 'B_filelist'], 'r')
+        Lines = unpaired_data_file2.readlines()
+        if shuffle:
+            random.shuffle(Lines)
+        for line in Lines:
+            if not config['dataset']['use_absolute_datafile']:
+                file = os.path.join(config['dataset']['dataroot'], line.strip())
+            else:
+                file = line.strip()
+            if os.path.exists(file):
+                B_paths.append(file)
+        unpaired_data_file2.close()
+    elif config['dataset']['unpaired_' + config['common']['phase'] + 'A_folder'] != '' and \
+         config['dataset']['unpaired_' + config['common']['phase'] + 'B_folder'] != '':
+        dir_A = config['dataset']['unpaired_' + config['common']['phase'] + 'A_folder']
+        filenames = os.listdir(dir_A)
+        if shuffle:
+            random.shuffle(filenames)
+        for filename in filenames:
+            if not check_path_is_img(filename):
+                continue
+            A_path = os.path.join(dir_A, filename)
+            if os.path.exists(A_path):
+                A_paths.append(A_path)
+
+        dir_B = config['dataset']['unpaired_' + config['common']['phase'] + 'B_folder']
+        filenames = os.listdir(dir_B)
+        if shuffle:
+            random.shuffle(filenames)
+        for filename in filenames:
+            if not check_path_is_img(filename):
+                continue
+            B_path = os.path.join(dir_B, filename)
+            if os.path.exists(B_path):
+                B_paths.append(B_path)
+
+    else:
+        dir_A = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'unpairedA')
+        dir_B = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'unpairedB')
+        if os.path.exists(dir_A) and os.path.exists(dir_B):
+            filenames = os.listdir(dir_A)
+            if shuffle:
+                random.shuffle(filenames)
+            for filename in filenames:
+                if not check_path_is_img(filename):
+                    continue
+                A_path = os.path.join(dir_A, filename)
+                A_paths.append(A_path)
+            filenames = os.listdir(dir_B)
+            if shuffle:
+                random.shuffle(filenames)
+            for filename in filenames:
+                if not check_path_is_img(filename):
+                    continue
+                B_path = os.path.join(dir_B, filename)
+                B_paths.append(B_path)
+
+
+    btoA = config['dataset']['direction'] == 'BtoA'
+    input_nc = config['model']['output_nc'] if btoA else config['model']['input_nc']
+    output_nc = config['model']['input_nc'] if btoA else config['model']['output_nc']
+
+    transform = Transforms(config, input_grayscale_flag=(input_nc == 1), output_grayscale_flag=(output_nc == 1))
+    transform.get_transform_from_config()
+    transform.get_transforms().insert(0, ImagePathToImage())
+    transform = transform.compose_transforms()
+
+    data['unpaired_A_path'] = A_paths
+    data['unpaired_B_path'] = B_paths
+    transforms['unpaired'] = transform
+
+def apply_unpaired_transforms(index, data, transforms, return_dict):
+    if len(data['unpaired_A_path']) > 0 and len(data['unpaired_B_path']) > 0:
+        index_B = random.randint(0, len(data['unpaired_B_path']) - 1)
+        return_dict['unpaired_A'], return_dict['unpaired_B'] = transforms['unpaired'] \
+            (data['unpaired_A_path'][index], data['unpaired_B_path'][index_B])
+        return_dict['unpaired_A_path'] = data['unpaired_A_path'][index]
+        return_dict['unpaired_B_path'] = data['unpaired_B_path'][index_B]

data/static_data.py

@@ -0,0 +1,457 @@
+import os, sys
+import random
+import numpy as np
+from utils.augmentation import ImagePathToImage, NumpyToTensor
+from utils.data_utils import Transforms
+from utils.util import check_path_is_static_data
+import torch
+from PIL import Image
+
+
+def check_dataname_folder_correspondence(data_names, group, group_name):
+    for data_name in data_names:
+        if data_name + '_folder' not in group:
+            print("%s not found in config file. Going to use dataroot mode to load group %s." % (data_name + '_folder', group_name))
+            return False
+    return True
+
+
+def custom_check_path_exists(str1):
+    return True if (str1 == "None" or os.path.exists(str1)) else False
+
+
+def custom_getsize(str1):
+    return 1 if str1 == "None" else os.path.getsize(str1)
+
+
+def check_different_extension_path_exists(str1):
+    acceptable_extensions = ['png', 'jpg', 'jpeg', 'npy', 'npz', 'PNG', 'JPG', 'JPEG']
+    curr_extension = str1.split('.')[-1]
+    for extension in acceptable_extensions:
+        str2 = str1.replace(curr_extension, extension)
+        if os.path.exists(str2):
+            return str2
+    return None
+
+
+class StaticData(object):
+
+    def __init__(self, config, shuffle=False):
+        # private variables
+        self.file_groups = []
+        self.type_groups = []
+        self.group_names = []
+        self.pair_type_groups = []
+        self.len_of_groups = []
+        self.transforms = {}
+        # parameters
+        self.shuffle = shuffle
+        self.config = config
+
+
+    def load_static_data(self):
+        data_dict = self.config['dataset'][self.config['common']['phase'] + '_data']
+        print("----------------loading %s static data.---------------------" % self.config['common']['phase'])
+
+        if len(data_dict) == 0:
+            self.len_of_groups.append(0)
+            return
+
+        self.group_names = list(data_dict.keys())
+        for i, group in enumerate(data_dict.values()):  # examples: (0, group_1),  (1, group_2)
+            data_types = group['data_types']  # examples: 'image', 'patch'
+            data_names = group['data_names']  # examples: 'real_A', 'patch_A'
+            self.file_groups.append({})
+            self.type_groups.append({})
+            self.len_of_groups.append(0)
+            self.pair_type_groups.append(group['paired'])
+
+            # exclude patch data since they are not stored on disk. They will be handled later.
+            data_types, data_names = self.exclude_patch_data(data_types, data_names)
+            assert(len(data_types) == len(data_names))
+
+            if len(data_names) == 0:
+                continue
+
+            for data_name, data_type in zip(data_names, data_types):
+                self.file_groups[i][data_name] = []
+                self.type_groups[i][data_name] = data_type
+
+
+            # paired data
+            if group['paired']:
+                # First way to load data: load a file list
+                if 'file_list' in group:
+                    file_list = group['file_list']
+                    paired_file = open(file_list, 'rt')
+                    lines = paired_file.readlines()
+                    if self.shuffle:
+                        random.shuffle(lines)
+                    for line in lines:
+                        items = line.strip().split(' ')
+                        if len(items) == len(data_names):
+                            ok = True
+                            for item in items:
+                                ok = ok and os.path.exists(item) and os.path.getsize(item) > 0
+                            if ok:
+                                for data_name, item in zip(data_names, items):
+                                    self.file_groups[i][data_name].append(item)
+                    paired_file.close()
+                # second and third way to load data: specify one folder for each dataname, or specify a dataroot folder
+                elif check_dataname_folder_correspondence(data_names, group, self.group_names[i]) or 'dataroot' in group:
+                    dataname_to_dir_dict = {}
+                    for data_name, data_type in zip(data_names, data_types):
+                        if 'dataroot' in group:
+                            # In new data config format, data is stored in dataroot_name/mode/dataname. e.g. FFHQ/train/pairedA
+                            # In old format, data is stored in dataroot_name/mode_dataname. e.g. FFHQ/train_pairedA
+                            # So we need to check both.
+                            dir = os.path.join(group['dataroot'], self.config['common']['phase'], data_name)
+                            if not os.path.exists(dir):
+                                old_dir = os.path.join(group['dataroot'], self.config['common']['phase'] + data_name.replace('_', ''))
+                                if 'numpy' in data_type:
+                                    old_dir += 'numpy'
+                                if not os.path.exists(old_dir):
+                                    print("Both %s and %s does not exist. Please check." % (dir, old_dir))
+                                    sys.exit()
+                                else:
+                                    dir = old_dir
+                        else:
+                            dir = group[data_name + '_folder']
+                            if not os.path.exists(dir):
+                                print("directory %s does not exist. Please check." % dir)
+                                sys.exit()
+                        dataname_to_dir_dict[data_name] = dir
+
+                    filenames = os.listdir(dataname_to_dir_dict[data_names[0]])
+                    if self.shuffle:
+                        random.shuffle(filenames)
+                    for filename in filenames:
+                        if not check_path_is_static_data(filename):
+                            continue
+                        file_paths = []
+                        for data_name in data_names:
+                            file_path = os.path.join(dataname_to_dir_dict[data_name], filename)
+                            checked_extension = check_different_extension_path_exists(file_path)
+                            if checked_extension is not None:
+                                file_paths.append(checked_extension)
+
+                        if len(file_paths) != len(data_names):
+                            print("for file %s , looks like some of the other pair data is missing. Ignoring and proceeding." % filename)
+                            continue
+                        else:
+                            for j in range(len(data_names)):
+                                data_name = data_names[j]
+                                self.file_groups[i][data_name].append(file_paths[j])
+                else:
+                    print("method for loading data is incorrect/unspecified for data group %s." % self.group_names)
+                    sys.exit()
+
+                self.len_of_groups[i] = len(self.file_groups[i][data_names[0]])
+
+            # unpaired data
+            else:
+                # First way to load data: load a file list
+                if 'file_list' in group:
+                    file_list = group['file_list']
+                    unpaired_file = open(file_list, 'rt')
+                    lines = unpaired_file.readlines()
+                    if self.shuffle:
+                        random.shuffle(lines)
+                    item_count = 0
+                    for line in lines:
+                        items = line.strip().split(' ')
+                        if len(items) == len(data_names):
+                            ok = True
+                            for item in items:
+                                ok = ok and custom_check_path_exists(item) and custom_getsize(item) > 0
+                            if ok:
+                                has_data = False
+                                for data_name, item in zip(data_names, items):
+                                    if item != 'None':
+                                        self.file_groups[i][data_name].append(item)
+                                        has_data = True
+                                if has_data:
+                                    item_count += 1
+                    unpaired_file.close()
+                    self.len_of_groups[i] = item_count
+                # second and third way to load data: specify one folder for each dataname, or specify a dataroot folder
+                elif check_dataname_folder_correspondence(data_names, group, self.group_names[i]) or 'dataroot' in group:
+                    max_length = 0
+                    for data_name, data_type in zip(data_names, data_types):
+                        if 'dataroot' in group:
+                            # In new data config format, data is stored in dataroot_name/mode/dataname. e.g. FFHQ/train/pairedA
+                            # In old format, data is stored in dataroot_name/mode_dataname. e.g. FFHQ/train_pairedA
+                            # So we need to check both.
+                            dir = os.path.join(group['dataroot'], self.config['common']['phase'], data_name)
+                            if not os.path.exists(dir):
+                                old_dir = os.path.join(group['dataroot'], self.config['common']['phase'] + data_name.replace('_', ''))
+                                if 'numpy' in data_type:
+                                    old_dir += 'numpy'
+                                if not os.path.exists(old_dir):
+                                    print("Both %s and %s does not exist. Please check." % (dir, old_dir))
+                                    sys.exit()
+                                else:
+                                    dir = old_dir
+                        else:
+                            dir = group[data_name + '_folder']
+                            if not os.path.exists(dir):
+                                print("directory %s does not exist. Please check." % dir)
+                                sys.exit()
+                        filenames = os.listdir(dir)
+                        if self.shuffle:
+                            random.shuffle(filenames)
+
+                        item_count = 0
+                        for filename in filenames:
+                            if not check_path_is_static_data(filename):
+                                continue
+                            fullpath = os.path.join(dir, filename)
+                            if os.path.exists(fullpath):
+                                self.file_groups[i][data_name].append(fullpath)
+                                item_count += 1
+                        max_length = max(item_count, max_length)
+                    self.len_of_groups[i] = max_length
+                else:
+                    print("method for loading data is incorrect/unspecified for data group %s." % self.group_names)
+                    sys.exit()
+
+
+    def create_transforms(self):
+        btoA = self.config['dataset']['direction'] == 'BtoA'
+        input_nc = self.config['model']['output_nc'] if btoA else self.config['model']['input_nc']
+        output_nc = self.config['model']['input_nc'] if btoA else self.config['model']['output_nc']
+        input_grayscale_flag = (input_nc == 1)
+        output_grayscale_flag = (output_nc == 1)
+
+        data_dict = self.config['dataset'][self.config['common']['phase'] + '_data']
+        for i, group in enumerate(data_dict.values()):  # examples: (0, group_1),  (1, group_2)
+
+            if i not in self.transforms:
+                self.transforms[i] = {}
+
+            data_types = group['data_types']  # examples: 'image', 'patch'
+            data_names = group['data_names']  # examples: 'real_A', 'patch_A'
+            data_types, data_names = self.exclude_patch_data(data_types, data_names)
+            for data_name, data_type in zip(data_names, data_types):
+                if data_type in self.transforms[i]:
+                    continue
+                self.transforms[i][data_type] = Transforms(self.config, input_grayscale_flag=input_grayscale_flag,
+                                                                    output_grayscale_flag=output_grayscale_flag)
+                self.transforms[i][data_type].create_transforms_from_list(group['preprocess'])
+                if '.png' in self.file_groups[i][data_name][0] or '.jpg' in self.file_groups[i][data_name][0] or \
+                    '.jpeg' in self.file_groups[i][data_name][0]:
+                    self.transforms[i][data_type].get_transforms().insert(0, ImagePathToImage())
+                elif '.npy' in self.file_groups[i][data_name][0] or '.npz' in self.file_groups[i][data_name][0]:
+                    self.transforms[i][data_type].get_transforms().insert(0, NumpyToTensor())
+                self.transforms[i][data_type] = self.transforms[i][data_type].compose_transforms()
+
+
+    def apply_transformations_to_images(self, img_list, img_dataname_list, transform, return_dict,
+                                        next_img_paths_bucket, next_img_dataname_list):
+
+        if len(img_list) == 1:
+            return_dict[img_dataname_list[0]], _ = transform(img_list[0], None)
+        elif len(img_list) > 1:
+            next_data_count = len(next_img_paths_bucket)
+            img_list += next_img_paths_bucket
+            img_dataname_list += next_img_dataname_list
+
+            input1, input2 = img_list[0], img_list[1:]
+            output1, output2 = transform(input1, input2)  # output1 is one image. output2 is a list of images.
+
+            if next_data_count != 0:
+                output2, next_outputs = output2[:-next_data_count], output2[-next_data_count:]
+                for i in range(next_data_count):
+                    return_dict[img_dataname_list[-next_data_count+i] + '_next'] = next_outputs[i]
+
+            return_dict[img_dataname_list[0]] = output1
+            for j in range(0, len(output2)):
+                return_dict[img_dataname_list[j+1]] = output2[j]
+
+        return return_dict
+
+
+    def calculate_landmark_scale(self, data_path, data_type, i):
+        if data_type == 'image':
+            original_image = Image.open(data_path)
+            original_width, original_height = original_image.size
+        else:
+            original_image = np.load(data_path)
+            original_height, original_width = original_image.shape[0], original_image.shape[1]
+        transformed_image, _ = self.transforms[i][data_type](data_path, None)
+        transformed_height, transformed_width = transformed_image.size()[1:]
+        landmark_scale = (transformed_width / original_width, transformed_height / original_height)
+        return landmark_scale
+
+
+    def get_item(self, idx):
+
+        return_dict = {}
+        data_dict = self.config['dataset'][self.config['common']['phase'] + '_data']
+
+        for i, group in enumerate(data_dict.values()):
+            if self.file_groups[i] == {}:
+                continue
+
+            paired_type = self.pair_type_groups[i]
+            inner_idx = idx if idx < self.len_of_groups[i] else random.randint(0, self.len_of_groups[i] - 1)
+
+            landmark_scale = None
+
+            # for patches since they might need to be loaded from different images.
+            next_img_paths_bucket = []
+            next_img_dataname_list = []
+            next_numpy_paths_bucket = []
+            next_numpy_dataname_list = []
+
+            # First, handle all non-patch data
+            if paired_type:
+                img_paths_bucket = []
+                img_dataname_list = []
+                numpy_paths_bucket = []
+                numpy_dataname_list = []
+
+            for data_name, data_list in self.file_groups[i].items():
+                data_type = self.type_groups[i][data_name]
+                if data_type in ['image', 'numpy']:
+                    if paired_type:
+                        # augmentation will be applied to all images in paired group all at once so need to gather the images here.
+                        if data_type == 'image':
+                            img_paths_bucket.append(data_list[inner_idx])
+                            img_dataname_list.append(data_name)
+                        else:
+                            numpy_paths_bucket.append(data_list[inner_idx])
+                            numpy_dataname_list.append(data_name)
+                        return_dict[data_name + '_path'] = data_list[inner_idx]
+                        if landmark_scale is None:
+                            landmark_scale = self.calculate_landmark_scale(data_list[inner_idx], data_type, i)
+                        if 'diff_patch' in self.config['dataset'] and self.config['dataset']['diff_patch'] and \
+                                data_name in group['patch_sources']:
+                            next_idx = (inner_idx + 1) % (len(data_list) - 1)
+                            if data_type == 'image':
+                                next_img_paths_bucket.append(data_list[next_idx])
+                                next_img_dataname_list.append(data_name)
+                            else:
+                                next_numpy_paths_bucket.append(data_list[next_idx])
+                                next_numpy_dataname_list.append(data_name)
+                    else:
+                        unpaired_inner_idx = random.randint(0, len(data_list) - 1)
+                        return_dict[data_name], _ = self.transforms[i][data_type](data_list[unpaired_inner_idx], None)
+                        if landmark_scale is None:
+                            landmark_scale = self.calculate_landmark_scale(data_list[unpaired_inner_idx], data_type, i)
+                        if 'diff_patch' in self.config['dataset'] and self.config['dataset']['diff_patch'] and \
+                                data_name in group['patch_sources']:
+                            next_idx = (unpaired_inner_idx + 1) % (len(data_list) - 1)
+                            return_dict[data_name + '_next'], _ = self.transforms[i][data_type](data_list[next_idx], None)
+                        return_dict[data_name + '_path'] = data_list[unpaired_inner_idx]
+                elif self.type_groups[i][data_name] == 'landmark':
+                    # We do not apply transformations on landmarks. Only scales landmarks to transformed image's size.
+                    # Also numpy data is passed into network as numpy array and not tensor.
+                    lmk = np.load(data_list[inner_idx])
+                    if self.config['dataset']['landmark_scale'] is not None:
+                        lmk[:, 0] *= self.config['dataset']['landmark_scale'][0]
+                        lmk[:, 1] *= self.config['dataset']['landmark_scale'][1]
+                    else:
+                        if landmark_scale is None:
+                            print("landmark_scale is None. If you have not defined it in config file, please specify "
+                                  "image and numpy data before landmark data and the proper scale will be automatically calculated.")
+                        else:
+                            lmk[:, 0] *= landmark_scale[0]
+                            lmk[:, 1] *= landmark_scale[1]
+                    return_dict[data_name] = lmk
+                    return_dict[data_name + '_path'] = data_list[inner_idx]
+
+
+            if paired_type:
+                # apply augmentations to all images and numpy arrays
+                if 'image' in self.transforms[i]:
+                    return_dict = self.apply_transformations_to_images(img_paths_bucket, img_dataname_list,
+                                                              self.transforms[i]['image'], return_dict,
+                                                                       next_img_paths_bucket,
+                                                                       next_img_dataname_list)
+                if 'numpy' in self.transforms[i]:
+                    return_dict = self.apply_transformations_to_images(numpy_paths_bucket, numpy_dataname_list,
+                                                              self.transforms[i]['numpy'], return_dict,
+                                                                       next_numpy_paths_bucket,
+                                                                       next_numpy_dataname_list)
+
+            # Handle patch data
+            data_types = group['data_types']  # examples: 'image', 'patch'
+            data_names = group['data_names']  # examples: 'real_A', 'patch_A'
+            data_types, data_names = self.filter_patch_data(data_types, data_names)
+
+            if 'patch_sources' in group:
+                patch_sources = group['patch_sources']
+                return_dict = self.load_patches(
+                    data_names,
+                    self.config['dataset']['patch_batch_size'],
+                    self.config['dataset']['batch_size'],
+                    self.config['dataset']['patch_size'],
+                    self.config['dataset']['patch_batch_size'] // self.config['dataset']['batch_size'],
+                    self.config['dataset']['diff_patch'],
+                    patch_sources,
+                    return_dict,
+                )
+
+        return return_dict
+
+
+    def get_len(self):
+        if len(self.len_of_groups) == 0:
+            return 0
+        else:
+            return max(self.len_of_groups)
+
+
+    def exclude_patch_data(self, data_types, data_names):
+        data_types_patch_excluded = []
+        data_names_patch_excluded = []
+        for data_name, data_type in zip(data_names, data_types):
+            if data_type != 'patch':
+                data_types_patch_excluded.append(data_type)
+                data_names_patch_excluded.append(data_name)
+        return data_types_patch_excluded, data_names_patch_excluded
+
+
+    def filter_patch_data(self, data_types, data_names):
+        data_types_patch = []
+        data_names_patch = []
+        for data_name, data_type in zip(data_names, data_types):
+            if data_type == 'patch':
+                data_types_patch.append(data_type)
+                data_names_patch.append(data_name)
+        return data_types_patch, data_names_patch
+
+
+    def load_patches(self, data_names, patch_batch_size, batch_size, patch_size,
+                     num_patch, diff_patch, patch_sources, return_dict):
+
+        if patch_size > 0:
+            assert (patch_batch_size % batch_size == 0), \
+                "patch_batch_size is not divisible by batch_size."
+            assert (len(patch_sources) == len(data_names)), \
+                "length of patch_sources is not the same as number of patch data specified. Please check again in config file."
+
+            rlist = []  # used for cropping patches
+            clist = []  # used for cropping patches
+            for _ in range(num_patch):
+                r = random.randint(0, self.config['dataset']['crop_size'] - patch_size - 1)
+                c = random.randint(0, self.config['dataset']['crop_size'] - patch_size - 1)
+                rlist.append(r)
+                clist.append(c)
+
+            for i in range(len(data_names)):
+                # load transformed image
+                patch = return_dict[patch_sources[i]] if not diff_patch else return_dict[patch_sources[i] + '_next']
+
+                # crop patch
+                patchs = []
+                _, h, w = patch.size()
+
+                for j in range(num_patch):
+                    patchs.append(patch[:, rlist[j]:rlist[j] + patch_size, clist[j]:clist[j] + patch_size])
+                patchs = torch.cat(patchs, 0)
+
+                return_dict[data_names[i]] = patchs
+
+        return return_dict

data/super_dataset.py

@@ -0,0 +1,321 @@
+import copy
+import torch.utils.data as data
+from utils.data_utils import check_img_loaded, check_numpy_loaded
+
+from data.test_data import add_test_data, apply_test_transforms
+from data.test_video_data import TestVideoData
+from data.static_data import StaticData
+
+from multiprocessing import Pool
+import sys
+
+
+class DataBin(object):
+    def __init__(self, filegroups):
+        self.filegroups = filegroups
+
+
+class SuperDataset(data.Dataset):
+    def __init__(self, config, shuffle=False, check_all_data=False, DDP_device=None):
+        self.config = config
+
+        self.check_all_data = check_all_data
+        self.DDP_device = DDP_device
+
+        self.data = {}  # Will be dictionary. Keys are data names, e.g. paired_A, patch_A. Values are lists containing associated data.
+        self.transforms = {}
+
+        if self.config['common']['phase'] == 'test':
+            if not self.config['testing']['test_video'] is None:
+                self.test_video_data = TestVideoData(self.config)
+            else:
+                add_test_data(self.data, self.transforms, self.config)
+            return
+
+        self.static_data = StaticData(self.config, shuffle)
+
+
+    def convert_old_config_to_new(self):
+        data_types = self.config['dataset']['data_type']
+        if len(data_types) == 1 and data_types[0] == 'custom':
+            # convert custom data configuration to new data configuration
+            old_dict = self.config['dataset']['custom_' + self.config['common']['phase'] + '_data']
+            preprocess_list = self.config['dataset']['preprocess']
+            new_datadict = self.config['dataset'][self.config['common']['phase'] + '_data'] = old_dict
+            for i, group in enumerate(new_datadict.values()):  # examples: (0, group_1),  (1, group_2)
+                group['paired'] = True
+                group['preprocess'] = preprocess_list
+                # custom data does not support patch so we skip patch logic.
+        else:
+            new_datadict = self.config['dataset'][self.config['common']['phase'] + '_data'] = {}
+            preprocess_list = self.config['dataset']['preprocess']
+            new_datadict['paired_group'] = {}
+            new_datadict['paired_group']['paired'] = True
+            new_datadict['paired_group']['data_types'] = []
+            new_datadict['paired_group']['data_names'] = []
+            new_datadict['paired_group']['preprocess'] = preprocess_list
+            new_datadict['unpaired_group'] = {}
+            new_datadict['unpaired_group']['paired'] = False
+            new_datadict['unpaired_group']['data_types'] = []
+            new_datadict['unpaired_group']['data_names'] = []
+            new_datadict['unpaired_group']['preprocess'] = preprocess_list
+
+            for i in range(len(self.config['dataset']['data_type'])):
+                data_type = self.config['dataset']['data_type'][i]
+                if data_type == 'paired' or data_type == 'paired_numpy':
+                    if self.config['dataset']['paired_' + self.config['common']['phase'] + '_filelist'] != '':
+                        new_datadict['paired_group']['file_list'] = self.config['dataset'][
+                            'paired_' + self.config['common']['phase'] + '_filelist']
+                    elif self.config['dataset']['paired_' + self.config['common']['phase'] + 'A_folder'] != '' and \
+                            self.config['dataset']['paired_' + self.config['common']['phase'] + 'B_folder'] != '':
+                        new_datadict['paired_group']['paired_A_folder'] = self.config['dataset']['paired_' + self.config['common']['phase'] + 'A_folder']
+                        new_datadict['paired_group']['paired_B_folder'] = self.config['dataset']['paired_' + self.config['common']['phase'] + 'B_folder']
+                    else:
+                        new_datadict['paired_group']['dataroot'] = self.config['dataset']['dataroot']
+
+                    new_datadict['paired_group']['data_names'].append('paired_A')
+                    new_datadict['paired_group']['data_names'].append('paired_B')
+                    if data_type == 'paired':
+                        new_datadict['paired_group']['data_types'].append('image')
+                        new_datadict['paired_group']['data_types'].append('image')
+                    else:
+                        new_datadict['paired_group']['data_types'].append('numpy')
+                        new_datadict['paired_group']['data_types'].append('numpy')
+
+                elif data_type == 'unpaired' or data_type == 'unpaired_numpy':
+                    if self.config['dataset']['unpaired_' + self.config['common']['phase'] + 'A_filelist'] != ''\
+                            and self.config['dataset']['unpaired_' + self.config['common']['phase'] + 'B_filelist'] != '':
+                        # combine those two filelists into one filelist
+                        self.combine_two_filelists_into_one(
+                            self.config['dataset']['unpaired_' + self.config['common']['phase'] + 'A_filelist'],
+                            self.config['dataset']['unpaired_' + self.config['common']['phase'] + 'B_filelist']
+                        )
+                        new_datadict['unpaired_group']['file_list'] = './tmp_filelist.txt'
+                    elif self.config['dataset']['unpaired_' + self.config['common']['phase'] + 'A_folder'] != '' and \
+                            self.config['dataset']['unpaired_' + self.config['common']['phase'] + 'B_folder'] != '':
+                        new_datadict['unpaired_group']['unpaired_A_folder'] = self.config['dataset']['unpaired_' + self.config['common']['phase'] + 'A_folder']
+                        new_datadict['unpaired_group']['unpaired_B_folder'] = self.config['dataset']['unpaired_' + self.config['common']['phase'] + 'B_folder']
+                    else:
+                        new_datadict['unpaired_group']['dataroot'] = self.config['dataset']['dataroot']
+
+                    new_datadict['unpaired_group']['data_names'].append('unpaired_A')
+                    new_datadict['unpaired_group']['data_names'].append('unpaired_B')
+                    if data_type == 'unpaired':
+                        new_datadict['unpaired_group']['data_types'].append('image')
+                        new_datadict['unpaired_group']['data_types'].append('image')
+                    else:
+                        new_datadict['unpaired_group']['data_types'].append('numpy')
+                        new_datadict['unpaired_group']['data_types'].append('numpy')
+
+                elif data_type == 'landmark':
+                    if self.config['dataset']['paired_' + self.config['common']['phase'] + '_filelist'] != '':
+                        new_datadict['paired_group']['file_list'] = self.config['dataset'][
+                            'paired_' + self.config['common']['phase'] + '_filelist']
+                    elif 'paired_' + self.config['common']['phase'] + 'A_lmk_folder' in self.config['dataset'] and \
+                            'paired_' + self.config['common']['phase'] + 'B_lmk_folder' in self.config['dataset'] and \
+                            self.config['dataset']['paired_' + self.config['common']['phase'] + 'A_lmk_folder'] != '' and \
+                            self.config['dataset']['paired_' + self.config['common']['phase'] + 'B_lmk_folder'] != '':
+                        new_datadict['paired_group']['lmk_A_folder'] = self.config['dataset']['paired_' + self.config['common']['phase'] + 'A_lmk_folder']
+                        new_datadict['paired_group']['lmk_B_folder'] = self.config['dataset']['paired_' + self.config['common']['phase'] + 'B_lmk_folder']
+                    else:
+                        new_datadict['paired_group']['dataroot'] = self.config['dataset']['dataroot']
+
+                    new_datadict['paired_group']['data_names'].append('lmk_A')
+                    new_datadict['paired_group']['data_names'].append('lmk_B')
+                    new_datadict['paired_group']['data_types'].append('landmark')
+                    new_datadict['paired_group']['data_types'].append('landmark')
+
+            # Handle patches. This needs to happen after all non-patch data are added first.
+            if 'patch' in self.config['dataset']['data_type']:
+                # determine if patch comes from paired or unpaired image
+                if 'paired_A' in new_datadict['paired_group']['data_names']:
+                    new_datadict['paired_group']['data_types'].append('patch')
+                    new_datadict['paired_group']['data_names'].append('patch_A')
+                    new_datadict['paired_group']['data_types'].append('patch')
+                    new_datadict['paired_group']['data_names'].append('patch_B')
+
+                    if 'patch_sources' not in new_datadict['paired_group']:
+                        new_datadict['paired_group']['patch_sources'] = []
+                    new_datadict['paired_group']['patch_sources'].append('paired_A')
+                    new_datadict['paired_group']['patch_sources'].append('paired_B')
+                else:
+                    new_datadict['unpaired_group']['data_types'].append('patch')
+                    new_datadict['unpaired_group']['data_names'].append('patch_A')
+                    new_datadict['unpaired_group']['data_types'].append('patch')
+                    new_datadict['unpaired_group']['data_names'].append('patch_B')
+
+                    if 'patch_sources' not in new_datadict['unpaired_group']:
+                        new_datadict['unpaired_group']['patch_sources'] = []
+                    new_datadict['unpaired_group']['patch_sources'].append('unpaired_A')
+                    new_datadict['unpaired_group']['patch_sources'].append('unpaired_B')
+
+                if 'diff_patch' not in self.config['dataset']:
+                    self.config['dataset']['diff_patch'] = False
+
+        new_datadict = {key: value for key, value in new_datadict.items() if len(value['data_names']) > 0}
+
+        print('-----------------------------------------------------------------------')
+        print("converted %s data configuration: " % self.config['common']['phase'])
+        for key, value in new_datadict.items():
+            print(key + ': ', value)
+        print('-----------------------------------------------------------------------')
+
+        return self.config
+
+
+    def combine_two_filelists_into_one(self, filelist1, filelist2):
+        tmp_file = open('./tmp_filelist.txt', 'w+')
+        f1 = open(filelist1, 'r')
+        f2 = open(filelist2, 'r')
+        f1_lines = f1.readlines()
+        f2_lines = f2.readlines()
+        min_index = min(len(f1_lines), len(f2_lines))
+        for i in range(min_index):
+            tmp_file.write(f1_lines[i].strip() + ' ' + f2_lines[i].strip() + '\n')
+        if min_index == len(f1_lines):
+            for i in range(min_index, len(f2_lines)):
+                tmp_file.write('None ' + f2_lines[i].strip() + '\n')
+        else:
+            for i in range(min_index, len(f1_lines)):
+                tmp_file.write(f1_lines[i].strip() + ' None\n')
+
+        tmp_file.close()
+        f1.close()
+        f2.close()
+
+
+    def __len__(self):
+        if self.config['common']['phase'] == 'test':
+                if self.config['testing']['test_video'] is not None:
+                    return self.test_video_data.get_len()
+                else:
+                    if len(self.data.keys()) == 0:
+                        return 0
+                    else:
+                        min_len = 999999
+                        for k, v in self.data.items():
+                            length = len(v)
+                            if length < min_len:
+                                min_len = length
+                        return min_len
+        else:
+            return self.static_data.get_len()
+
+
+
+    def get_item_logic(self, index):
+        return_dict = {}
+
+        if self.config['common']['phase'] == 'test':
+            if not self.config['testing']['test_video'] is None:
+                return self.test_video_data.get_item()
+            else:
+                apply_test_transforms(index, self.data, self.transforms, return_dict)
+                return return_dict
+
+        return_dict = self.static_data.get_item(index)
+
+        return return_dict
+
+
+    def __getitem__(self, index):
+        if self.config['dataset']['accept_data_error']:
+            while True:
+                try:
+                    return self.get_item_logic(index)
+                except Exception as e:
+                    print("Exception encountered in super_dataset's getitem function: ", e)
+                    index = (index + 1) % self.__len__()
+        else:
+            return self.get_item_logic(index)
+
+
+    def split_data(self, value_mode, value, mode='split'):
+        new_dataset = copy.deepcopy(self)
+        ret1, new_dataset.static_data = self.split_data_helper(self.static_data, new_dataset.static_data, value_mode, value, mode=mode)
+        if ret1 is not None:
+            self.static_data = ret1
+        return self, new_dataset
+
+
+    def split_data_helper(self, dataset, new_dataset, value_mode, value, mode='split'):
+        for i in range(len(dataset.file_groups)):
+            max_split_index = 0
+            for k in dataset.file_groups[i].keys():
+                length = len(dataset.file_groups[i][k])
+                if value_mode == 'count':
+                    split_index = min(length, value)
+                else:
+                    split_index = int((1 - value) * length)
+                max_split_index = max(max_split_index, split_index)
+                new_dataset.file_groups[i][k] = new_dataset.file_groups[i][k][split_index:]
+                if mode == 'split':
+                    dataset.file_groups[i][k] = dataset.file_groups[i][k][:split_index]
+            new_dataset.len_of_groups[i] -= max_split_index
+            if mode == 'split':
+                dataset.len_of_groups[i] = max_split_index
+        if mode == 'split':
+            return dataset, new_dataset
+        else:
+            return None, new_dataset
+
+
+    def check_data_helper(self, databin):
+        all_pass = True
+        for group in databin.filegroups:
+            for data_name, data_list in group.items():
+                for data in data_list:
+                    if '.npy' in data:  # case: numpy array or landmark
+                        all_pass = all_pass and check_numpy_loaded(data)
+                    else:  # case: image
+                        all_pass = all_pass and check_img_loaded(data)
+        return all_pass
+
+
+    def check_data(self):
+        if self.DDP_device is None or self.DDP_device == 0:
+            print("-----------------------Checking all data-------------------------------")
+            data_ok = True
+            if self.config['dataset']['n_threads'] == 0:
+                data_ok = data_ok and self.check_data_helper(self.static_data)
+            else:
+                # start n_threads number of workers to perform data checking
+                with Pool(processes=self.config['dataset']['n_threads']) as pool:
+                    checks = pool.map(self.check_data_helper,
+                                      self.split_data_into_bins(self.config['dataset']['n_threads']))
+                for check in checks:
+                    data_ok = data_ok and check
+            if data_ok:
+                print("---------------------all data passed check.-----------------------")
+            else:
+                print("---------------------The above data have failed in data checking. "
+                      "Please fix first.---------------------------")
+                sys.exit()
+
+
+    def split_data_into_bins(self, num_bins):
+        bins = []
+        for i in range(num_bins):
+            bins.append(DataBin(filegroups=[]))
+
+        # handle static data
+        bins = self.split_data_into_bins_helper(bins, self.static_data)
+        return bins
+
+
+    def split_data_into_bins_helper(self, bins, dataset):
+        num_bins = len(bins)
+        for bin in bins:
+            for group_idx in range(len(dataset.file_groups)):
+                bin.filegroups.append({})
+
+        for group_idx in range(len(dataset.file_groups)):
+            file_group = dataset.file_groups[group_idx]
+            for data_name, data_list in file_group.items():
+                num_items_in_bin = len(data_list) // num_bins
+                for data_index in range(len(data_list)):
+                    which_bin = min(data_index // num_items_in_bin, num_bins - 1)
+                    if data_name not in bins[which_bin].filegroups[group_idx]:
+                        bins[which_bin].filegroups[group_idx][data_name] = []
+                    bins[which_bin].filegroups[group_idx][data_name].append(data_list[data_index])
+        return bins

data/test_data.py

@@ -0,0 +1,51 @@
+import os
+from utils.util import check_path_is_static_data
+from utils.data_utils import Transforms
+from utils.augmentation import ImagePathToImage, NumpyToTensor
+
+def add_test_data(data, transforms, config):
+    A_paths = []
+    B_paths = []
+
+    if not config['testing']['test_img'] is None:
+        A_paths.append(config['testing']['test_img'])
+        B_paths.append(config['testing']['test_img'])
+    else:
+        files = os.listdir(config['testing']['test_folder'])
+        for fn in files:
+            if not check_path_is_static_data(fn):
+                continue
+            full_path = os.path.join(config['testing']['test_folder'], fn)
+            A_paths.append(full_path)
+            B_paths.append(full_path)
+
+    btoA = config['dataset']['direction'] == 'BtoA'
+    # get the number of channels of input image
+    input_nc = config['model']['output_nc'] if btoA else config['model']['input_nc']
+    output_nc = config['model']['input_nc'] if btoA else config['model']['output_nc']
+
+    transform = Transforms(config, input_grayscale_flag=(input_nc == 1), output_grayscale_flag=(output_nc == 1))
+    transform.create_transforms_from_list(config['testing']['preprocess'])
+    transform.get_transforms().insert(0, ImagePathToImage())
+    transform = transform.compose_transforms()
+
+    transform_np = Transforms(config, input_grayscale_flag=(input_nc == 1), output_grayscale_flag=(output_nc == 1))
+    transform_np.transform_list.append(NumpyToTensor())
+    transform_np = transform_np.compose_transforms()
+
+    data['test_A_path'] = A_paths
+    data['test_B_path'] = B_paths
+    transforms['test'] = transform
+    transforms['test_np'] = transform_np
+
+def apply_test_transforms(index, data, transforms, return_dict):
+    if len(data['test_A_path']) > 0:
+        ext_name = os.path.splitext(data['test_A_path'][index])[1]
+        if not ext_name.lower() in ['.npy', '.npz']:
+            return_dict['test_A'], return_dict['test_B'] = transforms['test'] \
+                (data['test_A_path'][index], data['test_B_path'][index])
+        else:
+            return_dict['test_A'], return_dict['test_B'] = transforms['test_np'] \
+                (data['test_A_path'][index], data['test_B_path'][index])
+        return_dict['test_A_path'] = data['test_A_path'][index]
+        return_dict['test_B_path'] = data['test_B_path'][index]

data/test_video_data.py

@@ -0,0 +1,28 @@
+import imp
+import cv2
+from PIL import Image
+from utils.data_utils import Transforms
+
+
+class TestVideoData(object):
+
+    def __init__(self, config):
+
+        self.vcap = cv2.VideoCapture(config['testing']['test_video'])
+        self.transform = Transforms(config)
+        self.transform.create_transforms_from_list(config['testing']['preprocess'])
+        self.transform = self.transform.compose_transforms()
+
+    def __del__(self):
+        self.vcap.release()
+
+    def get_len(self):
+        return int(self.vcap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    def get_item(self):
+        return_dict = {}
+        _, frame = self.vcap.read()
+        frame = Image.fromarray(frame[:,:,::-1]).convert('RGB')
+        return_dict['test_A'], return_dict['test_B'] = self.transform(frame, frame)
+        return_dict['test_A_path'], return_dict['test_B_path'] = 'A.jpg', 'B.jpg'
+        return return_dict

exp/sp2pII-phase1.yaml

@@ -0,0 +1,49 @@
+common:
+  name: "sp2pII-phase1"
+  model: "style_based_pix2pixII"
+  gpu_ids: [0]
+  option_group:
+    - gpu_ids: [0]
+    - gpu_ids: [1]
+    - gpu_ids: [2]
+
+model:
+  ngf: 64
+
+dataset:
+  unpaired_trainA_folder: # source domain folder path(FFHQ)
+  unpaired_trainB_folder: # target domain folder path(AAHQ)
+  preprocess: ["resize"]
+  batch_size: 8
+  crop_size: 512
+  drop_last: true
+  load_size: 512
+
+training:
+  epoch_as_iter: true
+  n_epochs: 100000
+  n_epochs_decay: 10
+  print_freq: 1000
+  pretrained_model: "pretrained_models/ffhq_pretrain_res512_200000.pt"
+  save_epoch_freq: 5000
+  style_mixing_prob: 0.5
+  lambda_GAN: 1.0
+  lambda_ST: 1.0
+  lambda_L1: 1.0
+  option_group:
+    - lambda_Feat: 4.0
+    - lambda_Feat: 2.0
+    - lambda_Feat: 1.0
+  lr: 0.001
+  lr_policy: "linear"
+  beta1: 0.1
+
+testing:
+  num_test: 100000
+  preprocess: ["resize"]
+  load_size: 512
+  crop_size: 512
+  results_dir: "./results/sp2pII"
+  visual_names: ["fake_B"]
+  image_format: "png"
+  which_epoch: "latest"

exp/sp2pII-phase2.yaml

@@ -0,0 +1,49 @@
+common:
+  name: "sp2pII-phase2"
+  model: "style_based_pix2pixII"
+  gpu_ids: [0]
+  option_group:
+    - gpu_ids: [0]
+    - gpu_ids: [1]
+    - gpu_ids: [2]
+
+model:
+  ngf: 64
+
+dataset:
+  unpaired_trainA_folder: # source domain folder path(FFHQ)
+  unpaired_trainB_folder: # target domain folder path(AAHQ) 
+  preprocess: ["resize"]
+  batch_size: 8
+  crop_size: 512
+  drop_last: true
+  load_size: 512
+
+training:
+  epoch_as_iter: true
+  n_epochs: 300000
+  n_epochs_decay: 10
+  print_freq: 1000
+  phase: 2 
+  pretrained_model: "pretrained_models/ffhq_pretrain_res512_200000.pt"  # phase1 model
+  save_epoch_freq: 5000
+  style_mixing_prob: 0.5
+  lambda_GAN: 1.0
+  lambda_ST: 0.5 # 这个参数可以调整
+  option_group:
+    - data_aug_prob: 0.0
+    - data_aug_prob: 0.1
+    - data_aug_prob: 0.2
+  lr: 0.001
+  lr_policy: "linear"
+  beta1: 0.1
+
+testing:
+  num_test: 100000
+  preprocess: ["resize"]
+  load_size: 512
+  crop_size: 512
+  results_dir: "./results/sp2pII"
+  visual_names: ["fake_B"]
+  image_format: "png"
+  which_epoch: "latest"

exp/sp2pII-phase3.yaml

@@ -0,0 +1,50 @@
+common:
+  name: "sp2pII-phase3"
+  model: "style_based_pix2pixII"
+  gpu_ids: [0]
+  option_group:
+    - gpu_ids: [0]
+    - gpu_ids: [1]
+    - gpu_ids: [2]
+
+model:
+  ngf: 64
+
+dataset:
+  unpaired_trainA_folder: # source domain folder path(FFHQ) 
+  unpaired_trainB_folder: # target domain folder path(AAHQ) 
+  preprocess: ["resize"]
+  batch_size: 8
+  crop_size: 512
+  drop_last: true
+  load_size: 512
+
+training:
+  epoch_as_iter: true
+  n_epochs: 100000 # 这个收敛很快，10w iter就差不多了
+  n_epochs_decay: 10
+  print_freq: 1000
+  phase: 3
+  pretrained_model: "pretrained_models/phase2_pretrain_90000.pth"
+  save_epoch_freq: 5000
+  style_mixing_prob: 0.5
+  lambda_GAN: 1.0
+  lambda_ST: 1.0
+  lambda_L1: 1.0
+  option_group:
+    - lambda_Feat: 4.0
+    - lambda_Feat: 2.0
+    - lambda_Feat: 1.0
+  lr: 0.0002
+  lr_policy: "linear"
+  beta1: 0.9
+
+testing:
+  num_test: 100000
+  preprocess: ["resize"]
+  load_size: 512
+  crop_size: 512
+  results_dir: "./results/sp2pII"
+  visual_names: ["fake_B"]
+  image_format: "png"
+  which_epoch: "latest"

exp/sp2pII-phase4.yaml

@@ -0,0 +1,49 @@
+common:
+  name: sp2pII-phase4
+  model: style_based_pix2pixII
+  gpu_ids:
+  - 0
+
+dataset:
+  batch_size: 8
+  crop_size: 512
+  drop_last: true
+  load_size: 512
+  preprocess:
+  - resize
+  unpaired_trainA_folder: "/share/group_machongyang/project/hand_drawn/data/dataset/0909/trainPA/" # source domain folder path(FFHQ) 
+  unpaired_trainB_folder: "/share/group_machongyang/project/hand_drawn/data/dataset/0909/trainPA/" # source domain folder path(AAHQ) 
+model:
+  ngf: 64
+testing:
+  crop_size: 512
+  image_format: png
+  load_size: 512
+  num_test: 100000
+  preprocess:
+  - resize
+  results_dir: ./results/sp2pII
+  visual_names:
+  - fake_B
+  which_epoch: latest
+training:
+  beta1: 0.9
+  epoch_as_iter: true
+  lambda_Feat: 4.0
+  lambda_GAN: 1.0
+  lambda_L1: 1.0
+  lambda_ST: 0.5
+  lambda_CLIP: 1.0 # 这个参数需要调整
+  lambda_PROJ: 100.0 # 这个参数需要调整(仅使用image prompt情况)
+  ema: 0.99 # 1-1/n
+  text_prompt: "not existed"
+  image_prompt: "" # 如果这个文件存在就用image prompt, 否则用text prompt
+  lr: 0.0002 # 这个参数需要调整, 大概1e-5 ~ 2e-4之间
+  lr_policy: linear
+  n_epochs: 200 # 这个一般500 iter就够了
+  n_epochs_decay: 10
+  phase: 4
+  pretrained_model: pretrained_models/phase3_pretrain_10000.pth
+  print_freq: 50
+  save_epoch_freq: 200
+  style_mixing_prob: 0.5

logs/01_2023_09_07__18_32_26/events.out.tfevents.1694082748.aiplatform-wlf2-hi-12.idchb2az2.hb2.kwaidc.com.16044.0

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

logs/01_2023_09_12__14_54_32/events.out.tfevents.1694501684.aiplatform-wlf2-ge4-22.idchb2az2.hb2.kwaidc.com.76748.0

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

logs/01_2023_09_12__14_55_34/events.out.tfevents.1694501736.aiplatform-wlf2-ge4-22.idchb2az2.hb2.kwaidc.com.77369.0

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

logs/01_2023_09_12__15_03_47/events.out.tfevents.1694502229.aiplatform-wlf2-ge4-22.idchb2az2.hb2.kwaidc.com.77940.0

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

models/__init__.py

@@ -0,0 +1,68 @@
+"""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):         define networks used in our training.
+    -- self.visual_names (str list):        specify the images that you want to display and save.
+    -- 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 usage.
+
+Now you can use the model class by specifying flag '--model dummy'.
+See our template model class 'template_model.py' for more details.
+"""
+
+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)
+    #print(modellib)
+    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(config, DDP_device=None):
+    """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(config['common']['model'])
+    instance = model(config, DDP_device=DDP_device)
+    print("model [%s] was created" % type(instance).__name__)
+    return instance

models/base_model.py

@@ -0,0 +1,340 @@
+import os
+import torch
+from collections import OrderedDict
+from abc import ABC, abstractmethod
+from models.modules import networks
+from utils.util import check_path
+from  utils.net_size import calc_computation
+
+
+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, config, DDP_device=None):
+        """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 function, 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):         define networks used in our training.
+            -- self.visual_names (str list):        specify the images that you want to display and save.
+            -- 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.config = config
+        self.gpu_ids = config['common']['gpu_ids']
+        self.isTrain = config['common']['phase'] == 'train'
+        if DDP_device is None:
+            self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')  # get device name: CPU or GPU
+            self.DDP_device = None
+            self.on_cpu = (self.device.type == 'cpu')
+        else:
+            self.device = DDP_device
+            self.DDP_device = DDP_device
+            self.on_cpu = False
+        self.save_dir = os.path.join(config['training']['checkpoints_dir'], config['common']['name'])  # save all the checkpoints to save_dir
+        if config['dataset']['preprocess'] != 'scale_width':  # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
+            torch.backends.cudnn.benchmark = True
+        self.loss_names = []
+        self.model_names = []
+        self.visual_names = []
+        self.optimizers = []
+        self.image_paths = []
+        self.metric = 0  # used for learning rate policy 'plateau'
+        self.curr_epoch = 0
+        self.total_iters = 0
+        self.best_val_loss = 999999
+
+    @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 <configimize_parameters> and <test>."""
+        pass
+
+    @abstractmethod
+    def trace_jit(self, input):
+        """trace torchscript model for C++. Called by <trace_jit.py>"""
+        pass
+
+    @abstractmethod
+    def optimize_parameters(self):
+        """Calculate losses, gradients, and update network weights; called in every training iteration"""
+        pass
+
+    @abstractmethod
+    def eval_step(self):
+        """Forward and backward pass but without upgrading weights; called in every validation iteration"""
+        pass
+
+    def setup(self, config, DDP_device=None):
+        """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, config) for optimizer in self.optimizers]
+        if not self.isTrain:
+            load_suffix = '{}'.format(config['testing']['which_epoch'])
+            self.load_networks(load_suffix)
+        elif config['training']['continue_train']:
+            load_suffix = '{}'.format(config['training']['which_epoch'])
+            self.load_networks(load_suffix)
+        self.print_networks(config['common']['verbose'], DDP_device=DDP_device)
+
+    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 train(self):
+        """Make models train mode during train time"""
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                net.train()
+
+    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:
+            if self.config['training']['lr_policy'] == 'plateau':
+                scheduler.step(self.metric, epoch=self.curr_epoch)
+            else:
+                scheduler.step(epoch=self.curr_epoch)
+
+        # 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"""
+        if not self.isTrain and len(self.config['testing']['visual_names']) > 0:
+            visual_names = list(set(self.visual_names).intersection(set(self.config['testing']['visual_names'])))
+        else:
+            visual_names = self.visual_names
+        visual_ret = OrderedDict()
+        for name in visual_names:
+            if isinstance(name, str):
+                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):
+                errors_ret[name] = float(getattr(self, 'loss_' + name))  # float(...) works for both scalar tensor and float number
+        return errors_ret
+
+    def save_networks(self, epoch, val_loss=None):
+        """Save all the networks to the disk.
+
+        Parameters:
+            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
+        """
+        check_path(self.save_dir)
+        save_filename = 'epoch_%s.pth' % epoch if val_loss is None else 'best_val_epoch.pth'
+        checkpoint = {}
+        # save all the models
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                if len(self.gpu_ids) > 0 and torch.cuda.is_available():
+                    # if use DDP, save only on rank 0. If using dataparallel, second condition meets.
+                    if self.DDP_device == 0 or self.DDP_device is None:
+                        checkpoint[name+'_model'] = net.module.state_dict()
+                else:
+                    checkpoint[name+'_model'] = net.state_dict()
+
+        # save all the optimizers
+        optimizer_index = 0
+        for optimizer in self.optimizers:
+            checkpoint['optimizer_'+str(optimizer_index)] = optimizer.state_dict()
+            optimizer_index += 1
+
+        # save all the schedulers
+        scheduler_index = 0
+        for scheduler in self.schedulers:
+            checkpoint['scheduler_' + str(scheduler_index)] = scheduler.state_dict()
+            scheduler_index += 1
+
+        # save other information
+        checkpoint['epoch'] = self.curr_epoch
+        checkpoint['total_iters'] = self.total_iters
+        checkpoint['metric'] = self.metric
+        if val_loss is not None:
+            checkpoint['best_val_loss'] = val_loss
+
+        torch.save(checkpoint, os.path.join(self.save_dir, save_filename))
+
+    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, ckpt=None):
+        """Load all the networks from the disk.
+
+        Parameters:
+            epoch (str) -- current epoch; used in the file name 'epoch_%s.pth' % epoch. Models in the old format
+            with the names '%s_net_%s.pth' % (epoch, name) are also supported. Models in the new format takes priority.
+        """
+        load_filename = 'epoch_%s.pth' % epoch
+        if ckpt is None:
+            final_load_path = os.path.join(self.save_dir, load_filename)
+        else:
+            final_load_path = ckpt
+        if os.path.exists(final_load_path):
+            # new checkpoint format.
+            print('loading the model in new format from %s' % final_load_path)
+            if self.DDP_device is not None:
+                # unpack the tensors on GPU 0, then transfer to whatever device it needs to be on
+                map_location = {'cuda:%d' % 0: 'cuda:%d' % self.DDP_device}
+                checkpoint = torch.load(final_load_path, map_location=map_location)
+            else:
+                checkpoint = torch.load(final_load_path)
+            for k, v in checkpoint.items():
+                # load models
+                if 'model' in k:
+                    name = k.split('_model')[0]
+                    if not self.isTrain and 'D' in name: # does not load discriminator when not training
+                        continue
+                    if not hasattr(self, 'net' + name):
+                        continue
+                    net = getattr(self, 'net' + name)
+                    if isinstance(net, torch.nn.DataParallel) or isinstance(net, torch.nn.parallel.DistributedDataParallel):
+                        net = net.module
+
+                    # if you are using PyTorch newer than 0.4 (e.g., built from
+                    # GitHub source), you can remove str() on self.device
+                    if hasattr(v, '_metadata'):
+                        del v._metadata
+
+                    # patch InstanceNorm checkpoints prior to 0.4
+                    for key in list(v.keys()):  # need to copy keys here because we mutate in loop
+                        self.__patch_instance_norm_state_dict(v, net, key.split('.'))
+                    net.load_state_dict(v)
+                # load optimizers
+                elif 'optimizer' in k:
+                    if not self.isTrain:
+                        continue
+                    index = int(k.split('_')[-1])
+                    self.optimizers[index].load_state_dict(v)
+                # load schedulers
+                elif 'scheduler' in k:
+                    if not self.isTrain:
+                        continue
+                    index = int(k.split('_')[-1])
+                    self.schedulers[index].load_state_dict(v)
+                # load other stuffs
+                elif k == 'epoch':
+                    self.curr_epoch = int(v) + 1
+                elif k == 'total_iters':
+                    self.total_iters = int(v)
+                elif k == 'metric':
+                    self.metric = float(v)
+                elif k == 'best_val_loss':
+                    self.best_val_loss = float(v)
+                else:
+                    print('Checkpoint load error. Unrecognized parameter saved in checkpoint: ', k)
+            return
+
+        # old checkpoint format.
+        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) or isinstance(net, torch.nn.parallel.DistributedDataParallel):
+                    net = net.module
+                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=str(self.device))
+                if hasattr(state_dict, '_metadata'):
+                    del state_dict._metadata
+
+                # patch InstanceNorm checkpoints prior to 0.4
+                for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
+                    self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
+                net.load_state_dict(state_dict)
+
+    def print_networks(self, verbose, DDP_device=None):
+        """Print the total number of parameters in the network and (if verbose) network architecture
+
+        Parameters:
+            verbose (bool) -- if verbose: print the network architecture
+        """
+        if DDP_device is None or DDP_device == 0:
+            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))
+                    if 'G' in name:
+                        calc_computation(net, self.config['model']['input_nc'], self.config['dataset']['crop_size'],self.config['dataset']['crop_size'], DDP_device=DDP_device)
+            print('-----------------------------------------------')
+
+    def set_requires_grad(self, nets, requires_grad=False):
+        """Set requies_grad=Fasle 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

models/modules/networks.py

@@ -0,0 +1,1101 @@
+import re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.spectral_norm as spectral_norm
+from torch.nn import init
+from torch.autograd import Variable
+import functools
+from torch.optim import lr_scheduler
+from packaging import version
+import numpy as np
+
+###############################################################################
+# Helper Functions
+###############################################################################
+
+
+class Identity(nn.Module):
+    def forward(self, 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 == 'none':
+        def norm_layer(x): return Identity()
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+
+def get_scheduler(optimizer, config):
+    """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.n_epochs> epochs
+    and linearly decay the rate to zero over the next <opt.n_epochs_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 config['training']['lr_policy'] == 'linear':
+        def lambda_rule(epoch):
+            lr_l = 1.0 - max(0, epoch + 1 - config['training']['n_epochs']) / float(config['training']['n_epochs_decay'] + 1)
+            return lr_l
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif config['training']['lr_policy'] == 'step':
+        scheduler = lr_scheduler.StepLR(optimizer, step_size=config['training']['lr_decay_iters'], gamma=0.1)
+    elif config['training']['lr_policy'] == 'plateau':
+        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
+    elif config['training']['lr_policy'] == 'cosine':
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['training']['n_epochs'], eta_min=0)
+    else:
+        return NotImplementedError('learning rate policy [%s] is not implemented', config['training']['lr_policy'])
+    return scheduler
+
+
+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')
+            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.
+            if m.affine:
+                init.normal_(m.weight.data, 1.0, init_gain)
+                init.constant_(m.bias.data, 0.0)
+
+    if not init_type == 'none':
+        net.apply(init_func)  # apply the initialization function <init_func>
+
+
+def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[], DDP_device=None, find_unused_parameters=False):
+    """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.
+    """
+    init_weights(net, init_type, init_gain=init_gain)
+    if DDP_device is not None:
+        net.to(DDP_device)
+        net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[DDP_device], output_device=DDP_device,
+                                                         broadcast_buffers=False, find_unused_parameters=find_unused_parameters) # DDP multi-GPUs
+        if DDP_device == 0:
+            print("model initiated in DDP mode.")
+    elif gpu_ids is not None and len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        net.to(gpu_ids[0])
+        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
+        print("model initiated in dataparallel mode.")
+    return net
+
+
+def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02,
+             gpu_ids=[], DDP_device=None, find_unused_parameters=False):
+    """Create a generator
+
+    Parameters:
+        input_nc (int) -- the number of channels in input images
+        output_nc (int) -- the number of channels in output images
+        ngf (int) -- the number of filters in the last conv layer
+        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
+        norm (str) -- the name of normalization layers used in the network: batch | instance | none
+        use_dropout (bool) -- if use dropout layers.
+        init_type (str)    -- the name of our initialization method.
+        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Returns a generator
+
+    Our current implementation provides two types of generators:
+        U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
+        The original U-Net paper: https://arxiv.org/abs/1505.04597
+
+        Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
+        Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
+        We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).
+
+
+    The generator has been initialized by <init_net>. It uses RELU for non-linearity.
+    """
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netG == 'resnet_9blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
+    elif netG == 'resnet_6blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
+    elif netG == 'unet_128':
+        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_256':
+        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    else:
+        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
+    return init_net(net, init_type, init_gain, gpu_ids, DDP_device=DDP_device, find_unused_parameters=find_unused_parameters)
+
+def define_F(netF, netF_nc=256, channels=[], use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[], DDP_device=None, find_unused_parameters=False):
+    if netF == 'sample':
+        net = PatchSampleF(use_mlp=False, nc=netF_nc)
+    elif netF == 'mlp_sample':
+        net = PatchSampleF(use_mlp=True, nc=netF_nc)
+    else:
+        raise NotImplementedError('Projection model name [%s] is not recognized' % netF)
+    net.create_mlp(channels)
+    return init_net(net, init_type, init_gain, gpu_ids, DDP_device=DDP_device, find_unused_parameters=find_unused_parameters)
+
+def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02,
+             gpu_ids=[], DDP_device=None, find_unused_parameters=False):
+    """Create a discriminator
+
+    Parameters:
+        input_nc (int)     -- the number of channels in input images
+        ndf (int)          -- the number of filters in the first conv layer
+        netD (str)         -- the architecture's name: basic | n_layers | pixel
+        n_layers_D (int)   -- the number of conv layers in the discriminator; effective when netD=='n_layers'
+        norm (str)         -- the type of normalization layers used in the network.
+        init_type (str)    -- the name of the initialization method.
+        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Returns a discriminator
+
+    Our current implementation provides three types of discriminators:
+        [basic]: 'PatchGAN' classifier described in the original pix2pix paper.
+        It can classify whether 70×70 overlapping patches are real or fake.
+        Such a patch-level discriminator architecture has fewer parameters
+        than a full-image discriminator and can work on arbitrarily-sized images
+        in a fully convolutional fashion.
+
+        [n_layers]: With this mode, you can specify the number of conv layers in the discriminator
+        with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)
+
+        [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
+        It encourages greater color diversity but has no effect on spatial statistics.
+
+    The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
+    """
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netD == 'basic':  # default PatchGAN classifier
+        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
+    elif netD == 'n_layers':  # more options
+        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)
+    elif netD == 'pixel':     # classify if each pixel is real or fake
+        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
+    else:
+        raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
+    return init_net(net, init_type, init_gain, gpu_ids, DDP_device=DDP_device, find_unused_parameters=find_unused_parameters)
+
+def define_G_pix2pixHD(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1, 
+            n_blocks_local=3, norm='instance', gpu_ids=[], DDP_device=None, find_unused_parameters=False):
+    norm_layer = get_norm_layer(norm_type=norm)     
+    if netG == 'global':    
+        netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)       
+    elif netG == 'local':        
+        netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, 
+                                n_local_enhancers, n_blocks_local, norm_layer)
+    else:
+        raise('generator not implemented!')
+    return init_net(netG, 'normal', 0.02, gpu_ids, DDP_device=DDP_device, find_unused_parameters=find_unused_parameters)
+
+def define_D_pix2pixHD(input_nc, ndf, n_layers_D, norm='instance', use_sigmoid=False, num_D=1, getIntermFeat=False,
+                       gpu_ids=[], DDP_device=None, find_unused_parameters=False):
+    norm_layer = get_norm_layer(norm_type=norm)   
+    netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat)
+    return init_net(netD, 'normal', 0.02, gpu_ids, DDP_device=DDP_device, find_unused_parameters=find_unused_parameters)
+
+
+class Normalize(nn.Module):
+
+    def __init__(self, power=2):
+        super(Normalize, self).__init__()
+        self.power = power
+
+    def forward(self, x):
+        norm = x.pow(self.power).sum(1, keepdim=True).pow(1. / self.power)
+        out = x.div(norm + 1e-7)
+        return out
+
+
+class PatchSampleF(nn.Module):
+    def __init__(self, use_mlp=False, nc=256):
+        # potential issues: currently, we use the same patch_ids for multiple images in the batch
+        super(PatchSampleF, self).__init__()
+        self.l2norm = Normalize(2)
+        self.use_mlp = use_mlp
+        self.nc = nc
+
+    def create_mlp(self, channels):
+        if not self.use_mlp:
+            return
+        for mlp_id, ch in enumerate(channels):
+            mlp = nn.Sequential(*[nn.Linear(ch, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)])
+            setattr(self, 'mlp_%d' % mlp_id, mlp)
+
+    def forward(self, feats, num_patches=64, patch_ids=None):
+        return_ids = []
+        return_feats = []
+        for feat_id, feat in enumerate(feats):
+            B, H, W = feat.shape[0], feat.shape[2], feat.shape[3]
+            feat_reshape = feat.permute(0, 2, 3, 1).flatten(1, 2)
+            if num_patches > 0:
+                if patch_ids is not None:
+                    patch_id = patch_ids[feat_id]
+                else:
+                    patch_id = torch.randperm(feat_reshape.shape[1], device=feats[0].device)
+                    patch_id = patch_id[:int(min(num_patches, patch_id.shape[0]))]
+                x_sample = feat_reshape[:, patch_id, :].flatten(0, 1)
+            else:
+                x_sample = feat_reshape
+                patch_id = []
+            if self.use_mlp:
+                mlp = getattr(self, 'mlp_%d' % feat_id)
+                x_sample = mlp(x_sample)
+            return_ids.append(patch_id)
+            x_sample = self.l2norm(x_sample)
+
+            if num_patches == 0:
+                x_sample = x_sample.permute(0, 2, 1).reshape([B, x_sample.shape[-1], H, W])
+            return_feats.append(x_sample)
+        return return_feats, return_ids
+
+
+class LocalEnhancer(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=32, n_downsample_global=3, n_blocks_global=9, 
+                n_local_enhancers=1, n_blocks_local=3, norm_layer=nn.BatchNorm2d, padding_type='reflect'):        
+        super(LocalEnhancer, self).__init__()
+        self.n_local_enhancers = n_local_enhancers
+        
+        ###### global generator model #####           
+        ngf_global = ngf * (2**n_local_enhancers)
+        model_global = GlobalGenerator(input_nc, output_nc, ngf_global, n_downsample_global, n_blocks_global, norm_layer).model        
+        model_global = [model_global[i] for i in range(len(model_global)-3)] # get rid of final convolution layers        
+        self.model = nn.Sequential(*model_global)                
+
+        ###### local enhancer layers #####
+        for n in range(1, n_local_enhancers+1):
+            ### downsample            
+            ngf_global = ngf * (2**(n_local_enhancers-n))
+            model_downsample = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0), 
+                                norm_layer(ngf_global), nn.ReLU(True),
+                                nn.Conv2d(ngf_global, ngf_global * 2, kernel_size=3, stride=2, padding=1), 
+                                norm_layer(ngf_global * 2), nn.ReLU(True)]
+            ### residual blocks
+            model_upsample = []
+            for i in range(n_blocks_local):
+                model_upsample += [ResnetBlock(ngf_global * 2, padding_type=padding_type, norm_layer=norm_layer, use_dropout=False, use_bias=True)]
+
+            ### upsample
+            model_upsample += [nn.ConvTranspose2d(ngf_global * 2, ngf_global, kernel_size=3, stride=2, padding=1, output_padding=1), 
+                            norm_layer(ngf_global), nn.ReLU(True)]      
+
+            ### final convolution
+            if n == n_local_enhancers:                
+                model_upsample += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]                       
+            
+            setattr(self, 'model'+str(n)+'_1', nn.Sequential(*model_downsample))
+            setattr(self, 'model'+str(n)+'_2', nn.Sequential(*model_upsample))                  
+        
+        self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
+
+    def forward(self, input): 
+        ### create input pyramid
+        input_downsampled = [input]
+        for i in range(self.n_local_enhancers):
+            input_downsampled.append(self.downsample(input_downsampled[-1]))
+
+        ### output at coarest level
+        output_prev = self.model(input_downsampled[-1])        
+        ### build up one layer at a time
+        for n_local_enhancers in range(1, self.n_local_enhancers+1):
+            model_downsample = getattr(self, 'model'+str(n_local_enhancers)+'_1')
+            model_upsample = getattr(self, 'model'+str(n_local_enhancers)+'_2')            
+            input_i = input_downsampled[self.n_local_enhancers-n_local_enhancers]            
+            output_prev = model_upsample(model_downsample(input_i) + output_prev)
+        return output_prev
+
+class GlobalGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d, 
+                padding_type='reflect'):
+        assert(n_blocks >= 0)
+        super(GlobalGenerator, self).__init__()        
+        activation = nn.ReLU(True)        
+
+        model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
+        ### downsample
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+                      norm_layer(ngf * mult * 2), activation]
+
+        ### resnet blocks
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=False, use_bias=True)]
+
+        ### upsample         
+        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),
+                       norm_layer(int(ngf * mult / 2)), activation]
+        model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]        
+        self.model = nn.Sequential(*model)
+            
+    def forward(self, input):
+        return self.model(input)
+
+class MultiscaleDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, 
+                use_sigmoid=False, num_D=3, getIntermFeat=False):
+        super(MultiscaleDiscriminator, self).__init__()
+        self.num_D = num_D
+        self.n_layers = n_layers
+        self.getIntermFeat = getIntermFeat
+
+        for i in range(num_D):
+            netD = Pix2PixHDNLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat)
+            if getIntermFeat:                                
+                for j in range(n_layers+2):
+                    setattr(self, 'scale'+str(i)+'_layer'+str(j), getattr(netD, 'model'+str(j)))                                   
+            else:
+                setattr(self, 'layer'+str(i), netD.model)
+
+        self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
+
+    def singleD_forward(self, model, input):
+        if self.getIntermFeat:
+            result = [input]
+            for i in range(len(model)):
+                result.append(model[i](result[-1]))
+            return result[1:]
+        else:
+            return [model(input)]
+
+    def forward(self, input):        
+        num_D = self.num_D
+        result = []
+        input_downsampled = input
+        for i in range(num_D):
+            if self.getIntermFeat:
+                model = [getattr(self, 'scale'+str(num_D-1-i)+'_layer'+str(j)) for j in range(self.n_layers+2)]
+            else:
+                model = getattr(self, 'layer'+str(num_D-1-i))
+            result.append(self.singleD_forward(model, input_downsampled))
+            if i != (num_D-1):
+                input_downsampled = self.downsample(input_downsampled)
+        return result
+
+class Pix2PixHDNLayerDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, getIntermFeat=False):
+        super(Pix2PixHDNLayerDiscriminator, self).__init__()
+        self.getIntermFeat = getIntermFeat
+        self.n_layers = n_layers
+
+        kw = 4
+        padw = int(np.ceil((kw-1.0)/2))
+        sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]
+
+        nf = ndf
+        for n in range(1, n_layers):
+            nf_prev = nf
+            nf = min(nf * 2, 512)
+            sequence += [[
+                nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw),
+                norm_layer(nf), nn.LeakyReLU(0.2, True)
+            ]]
+
+        nf_prev = nf
+        nf = min(nf * 2, 512)
+        sequence += [[
+            nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
+            norm_layer(nf),
+            nn.LeakyReLU(0.2, True)
+        ]]
+
+        sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
+
+        if use_sigmoid:
+            sequence += [[nn.Sigmoid()]]
+
+        if getIntermFeat:
+            for n in range(len(sequence)):
+                setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
+        else:
+            sequence_stream = []
+            for n in range(len(sequence)):
+                sequence_stream += sequence[n]
+            self.model = nn.Sequential(*sequence_stream)
+
+    def forward(self, input):
+        if self.getIntermFeat:
+            res = [input]
+            for n in range(self.n_layers+2):
+                model = getattr(self, 'model'+str(n))
+                res.append(model(res[-1]))
+            return res[1:]
+        else:
+            return self.model(input)
+
+
+class MultiGANLoss(nn.Module):
+    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0,
+                 tensor=torch.cuda.FloatTensor):
+        super(MultiGANLoss, self).__init__()
+        self.real_label = target_real_label
+        self.fake_label = target_fake_label
+        self.real_label_var = None
+        self.fake_label_var = None
+        self.Tensor = tensor
+        if use_lsgan:
+            self.loss = nn.MSELoss()
+        else:
+            self.loss = nn.BCELoss()
+
+    def get_target_tensor(self, input, target_is_real):
+        target_tensor = None
+        if target_is_real:
+            create_label = ((self.real_label_var is None) or
+                            (self.real_label_var.numel() != input.numel()))
+            if create_label:
+                real_tensor = self.Tensor(input.size()).fill_(self.real_label)
+                self.real_label_var = Variable(real_tensor, requires_grad=False)
+            target_tensor = self.real_label_var
+        else:
+            create_label = ((self.fake_label_var is None) or
+                            (self.fake_label_var.numel() != input.numel()))
+            if create_label:
+                fake_tensor = self.Tensor(input.size()).fill_(self.fake_label)
+                self.fake_label_var = Variable(fake_tensor, requires_grad=False)
+            target_tensor = self.fake_label_var
+        return target_tensor
+
+    def __call__(self, input, target_is_real):
+        if isinstance(input[0], list):
+            loss = 0
+            for input_i in input:
+                pred = input_i[-1]
+                target_tensor = self.get_target_tensor(pred, target_is_real)
+                loss += self.loss(pred, target_tensor)
+            return loss
+        else:            
+            target_tensor = self.get_target_tensor(input[-1], target_is_real)
+            return self.loss(input[-1], target_tensor)
+
+##############################################################################
+# Classes
+##############################################################################
+class GANLoss(nn.Module):
+    """Define different GAN objectives.
+
+    The GANLoss class abstracts away the need to create the target label tensor
+    that has the same size as the input.
+    """
+
+    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
+        """ Initialize the GANLoss class.
+
+        Parameters:
+            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
+            target_real_label (bool) - - label for a real image
+            target_fake_label (bool) - - label of a fake image
+
+        Note: Do not use sigmoid as the last layer of Discriminator.
+        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
+        """
+        super(GANLoss, self).__init__()
+        self.register_buffer('real_label', torch.tensor(target_real_label))
+        self.register_buffer('fake_label', torch.tensor(target_fake_label))
+        self.gan_mode = gan_mode
+        if gan_mode == 'lsgan':
+            self.loss = nn.MSELoss()
+        elif gan_mode == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif gan_mode in ['wgangp']:
+            self.loss = None
+        else:
+            raise NotImplementedError('gan mode %s not implemented' % gan_mode)
+
+    def get_target_tensor(self, prediction, target_is_real):
+        """Create label tensors with the same size as the input.
+
+        Parameters:
+            prediction (tensor) - - tpyically the prediction from a discriminator
+            target_is_real (bool) - - if the ground truth label is for real images or fake images
+
+        Returns:
+            A label tensor filled with ground truth label, and with the size of the input
+        """
+
+        if target_is_real:
+            target_tensor = self.real_label
+        else:
+            target_tensor = self.fake_label
+        return target_tensor.expand_as(prediction)
+
+    def __call__(self, prediction, target_is_real):
+        """Calculate loss given Discriminator's output and grount truth labels.
+
+        Parameters:
+            prediction (tensor) - - tpyically the prediction output from a discriminator
+            target_is_real (bool) - - if the ground truth label is for real images or fake images
+
+        Returns:
+            the calculated loss.
+        """
+        if self.gan_mode in ['lsgan', 'vanilla']:
+            target_tensor = self.get_target_tensor(prediction, target_is_real)
+            loss = self.loss(prediction, target_tensor)
+        elif self.gan_mode == 'wgangp':
+            if target_is_real:
+                loss = nn.functional.softplus(-prediction).mean()
+            else:
+                loss = nn.functional.softplus(prediction).mean()
+        return loss
+
+
+class PatchNCELoss(nn.Module):
+    def __init__(self, batch_size, nce_T):
+        super().__init__()
+        self.batch_size = batch_size
+        self.nce_T = nce_T
+        self.cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction='none')
+        self.mask_dtype = torch.uint8 if version.parse(torch.__version__) < version.parse('1.2.0') else torch.bool
+
+    def forward(self, feat_q, feat_k):
+        batchSize = feat_q.shape[0]
+        dim = feat_q.shape[1]
+        feat_k = feat_k.detach()
+
+        # pos logit
+        l_pos = torch.bmm(feat_q.view(batchSize, 1, -1), feat_k.view(batchSize, -1, 1))
+        l_pos = l_pos.view(batchSize, 1)
+
+        # neg logit
+        batch_dim_for_bmm = self.batch_size
+
+        # reshape features to batch size
+        feat_q = feat_q.view(batch_dim_for_bmm, -1, dim)
+        feat_k = feat_k.view(batch_dim_for_bmm, -1, dim)
+        npatches = feat_q.size(1)
+        l_neg_curbatch = torch.bmm(feat_q, feat_k.transpose(2, 1))
+
+        # diagonal entries are similarity between same features, and hence meaningless.
+        # just fill the diagonal with very small number, which is exp(-10) and almost zero
+        diagonal = torch.eye(npatches, device=feat_q.device, dtype=self.mask_dtype)[None, :, :]
+        l_neg_curbatch.masked_fill_(diagonal, -10.0)
+        l_neg = l_neg_curbatch.view(-1, npatches)
+
+        out = torch.cat((l_pos, l_neg), dim=1) / self.nce_T
+
+        loss = self.cross_entropy_loss(out, torch.zeros(out.size(0), dtype=torch.long,
+                                                        device=feat_q.device))
+
+        return loss
+
+
+def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):
+    """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
+
+    Arguments:
+        netD (network)              -- discriminator network
+        real_data (tensor array)    -- real images
+        fake_data (tensor array)    -- generated images from the generator
+        device (str)                -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
+        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].
+        constant (float)            -- the constant used in formula ( | |gradient||_2 - constant)^2
+        lambda_gp (float)           -- weight for this loss
+
+    Returns the gradient penalty loss
+    """
+    if lambda_gp > 0.0:
+        if type == 'real':   # either use real images, fake images, or a linear interpolation of two.
+            interpolatesv = real_data
+        elif type == 'fake':
+            interpolatesv = fake_data
+        elif type == 'mixed':
+            alpha = torch.rand(real_data.shape[0], 1, device=device)
+            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
+            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
+        else:
+            raise NotImplementedError('{} not implemented'.format(type))
+        interpolatesv.requires_grad_(True)
+        disc_interpolates = netD(interpolatesv)
+        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,
+                                        grad_outputs=torch.ones(disc_interpolates.size()).to(device),
+                                        create_graph=True, retain_graph=True, only_inputs=True)
+        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data
+        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp        # added eps
+        return gradient_penalty, gradients
+    else:
+        return 0.0, None
+
+
+class ResnetGenerator(nn.Module):
+    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
+
+    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
+    """
+
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        """Construct a Resnet-based generator
+
+        Parameters:
+            input_nc (int)      -- the number of channels in input images
+            output_nc (int)     -- the number of channels in output images
+            ngf (int)           -- the number of filters in the last conv layer
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers
+            n_blocks (int)      -- the number of ResNet blocks
+            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
+        """
+        assert(n_blocks >= 0)
+        super(ResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+                norm_layer(ngf),
+                nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):  # add downsampling layers
+            mult = 2 ** i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
+                    norm_layer(ngf * mult * 2),
+                    nn.ReLU(True)]
+
+        mult = 2 ** n_downsampling
+        for i in range(n_blocks):       # add ResNet 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):  # add upsampling layers
+            mult = 2 ** (n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                        kernel_size=3, stride=2,
+                                        padding=1, output_padding=1,
+                                        bias=use_bias),
+                    norm_layer(int(ngf * mult / 2)),
+                    nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input, layers=[]):
+        if len(layers) > 0:
+            feat = input
+            feats = []
+            for layer_id, layer in enumerate(self.model):
+                feat = layer(feat)
+                if layer_id in layers:
+                    feats.append(feat)
+                if layer_id == layers[-1]:
+                    break
+            return feats
+        else:
+            """Standard forward"""
+            return self.model(input)
+
+
+class ResnetBlock(nn.Module):
+    """Define a Resnet block"""
+
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Initialize the Resnet block
+
+        A resnet block is a conv block with skip connections
+        We construct a conv block with build_conv_block function,
+        and implement skip connections in <forward> function.
+        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
+        """
+        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):
+        """Construct a convolutional block.
+
+        Parameters:
+            dim (int)           -- the number of channels in the conv layer.
+            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+            use_bias (bool)     -- if the conv layer uses bias or not
+
+        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
+        """
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        """Forward function (with skip connections)"""
+        out = x + self.conv_block(x)  # add skip connections
+        return out
+
+
+class UnetGenerator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(UnetGenerator, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer
+        for i in range(num_downs - 5):          # add intermediate layers with ngf * 8 filters
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        # gradually reduce the number of filters from ngf * 8 to ngf
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class UnetSkipConnectionBlock(nn.Module):
+    """Defines the Unet submodule with skip connection.
+        X -------------------identity----------------------
+        |-- downsampling -- |submodule| -- upsampling --|
+    """
+
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet submodule with skip connections.
+
+        Parameters:
+            outer_nc (int) -- the number of filters in the outer conv layer
+            inner_nc (int) -- the number of filters in the inner conv layer
+            input_nc (int) -- the number of channels in input images/features
+            submodule (UnetSkipConnectionBlock) -- previously defined submodules
+            outermost (bool)    -- if this module is the outermost module
+            innermost (bool)    -- if this module is the innermost module
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+        """
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                            stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:   # add skip connections
+            return torch.cat([x, self.model(x)], 1)
+
+
+# Creates SPADE normalization layer based on the given configuration
+# SPADE consists of two steps. First, it normalizes the activations using
+# your favorite normalization method, such as Batch Norm or Instance Norm.
+# Second, it applies scale and bias to the normalized output, conditioned on
+# the segmentation map.
+# The format of |config_text| is spade(norm)(ks), where
+# (norm) specifies the type of parameter-free normalization.
+#       (e.g. syncbatch, batch, instance)
+# (ks) specifies the size of kernel in the SPADE module (e.g. 3x3)
+# Example |config_text| will be spadesyncbatch3x3, or spadeinstance5x5.
+# Also, the other arguments are
+# |norm_nc|: the #channels of the normalized activations, hence the output dim of SPADE
+# |label_nc|: the #channels of the input semantic map, hence the input dim of SPADE
+class SPADE(nn.Module):
+    def __init__(self, config_text, norm_nc, label_nc):
+        super().__init__()
+
+        assert config_text.startswith('spade')
+        parsed = re.search(r'spade(\D+)(\d)x\d', config_text)
+        param_free_norm_type = str(parsed.group(1))
+        ks = int(parsed.group(2))
+
+        if param_free_norm_type == 'instance':
+            self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
+        elif param_free_norm_type == 'batch':
+            self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False)
+        elif param_free_norm_type == 'identity':
+            self.param_free_norm = nn.Identity()
+        else:
+            raise ValueError('%s is not a recognized param-free norm type in SPADE'
+                            % param_free_norm_type)
+
+        # The dimension of the intermediate embedding space. Yes, hardcoded.
+        nhidden = 128
+
+        pw = ks // 2
+        self.mlp_shared = nn.Sequential(
+            nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw),
+            nn.ReLU()
+        )
+        self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
+        self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
+
+    def forward(self, x, segmap):
+
+        # Part 1. generate parameter-free normalized activations
+        normalized = self.param_free_norm(x)
+
+        # Part 2. produce scaling and bias conditioned on semantic map
+        segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
+        actv = self.mlp_shared(segmap)
+        gamma = self.mlp_gamma(actv)
+        beta = self.mlp_beta(actv)
+
+        # apply scale and bias
+        out = normalized * (1 + gamma) + beta
+
+        return out
+
+
+# ResNet block that uses SPADE.
+# It differs from the ResNet block of pix2pixHD in that
+# it takes in the segmentation map as input, learns the skip connection if necessary,
+# and applies normalization first and then convolution.
+# This architecture seemed like a standard architecture for unconditional or
+# class-conditional GAN architecture using residual block.
+# The code was inspired from https://github.com/LMescheder/GAN_stability.
+class SPADEResnetBlock(nn.Module):
+    def __init__(self, fin, fout, config_str, semantic_nc):
+        super().__init__()
+        # Attributes
+        self.learned_shortcut = (fin != fout)
+        fmiddle = min(fin, fout)
+
+        # create conv layers
+        self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)
+        self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)
+        if self.learned_shortcut:
+            self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
+
+        # apply spectral norm if specified
+        if 'spectral' in config_str:
+            self.conv_0 = spectral_norm(self.conv_0)
+            self.conv_1 = spectral_norm(self.conv_1)
+            if self.learned_shortcut:
+                self.conv_s = spectral_norm(self.conv_s)
+
+        # define normalization layers
+        spade_config_str = config_str.replace('spectral', '')
+        self.norm_0 = SPADE(spade_config_str, fin, semantic_nc)
+        self.norm_1 = SPADE(spade_config_str, fmiddle, semantic_nc)
+        if self.learned_shortcut:
+            self.norm_s = SPADE(spade_config_str, fin, semantic_nc)
+
+    # note the resnet block with SPADE also takes in |seg|,
+    # the semantic segmentation map as input
+    def forward(self, x, seg):
+        x_s = self.shortcut(x, seg)
+
+        dx = self.conv_0(self.actvn(self.norm_0(x, seg)))
+        dx = self.conv_1(self.actvn(self.norm_1(dx, seg)))
+
+        out = x_s + dx
+
+        return out
+
+    def shortcut(self, x, seg):
+        if self.learned_shortcut:
+            x_s = self.conv_s(self.norm_s(x, seg))
+        else:
+            x_s = x
+        return x_s
+
+    def actvn(self, x):
+        return F.leaky_relu(x, 2e-1)
+
+
+class NLayerDiscriminator(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(NLayerDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # 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 PixelDiscriminator(nn.Module):
+    """Defines a 1x1 PatchGAN discriminator (pixelGAN)"""
+
+    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):
+        """Construct a 1x1 PatchGAN discriminator
+
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+        """
+        super(PixelDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        self.net = [
+            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
+            norm_layer(ndf * 2),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
+
+        self.net = nn.Sequential(*self.net)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.net(input)

models/modules/sr/light_model_270M.py

@@ -0,0 +1,347 @@
+import torch
+import torch.nn as nn
+from models.modules.pix2pixMini_module import *
+
+pack = Pack()
+unpack = UnPack()
+AVG = AvgQuant()
+mul = SliceMul()
+
+def get_int(channel_dict):
+    for i in range(len(channel_dict['down'])):
+        channel_dict['down'][i] = int(channel_dict['down'][i])
+    for i in range(len(channel_dict['backbone'])):
+        channel_dict['backbone'][i] = int(channel_dict['backbone'][i])
+    for i in range(len(channel_dict['up'])):
+        for j in range(len(channel_dict['up'][i])):
+            if channel_dict['up'][i][j] is not None:
+                channel_dict['up'][i][j] = int(channel_dict['up'][i][j])
+    return channel_dict
+
+def get_channel_dict(dict_name, ngf):
+    if dict_name is None:
+        raise('invalid channel_dict name')
+
+    if dict_name == '3G':
+        channel_dict = {
+            'n_blocks': 8,
+            'down': [ngf * 1, ngf * 2, ngf * 4, ngf * 8, ngf * 8],
+            'backbone': [ngf * 8, ngf * 8],
+            'hair_up': [
+                [None, None, None, ngf * 4, 2],
+                [None, None, None, ngf * 2, 2],
+                [None, None, None, ngf * 1, 2],
+                [None, None, None, ngf * 1, 1],
+            ],
+            'face_up': [
+                [None, None, None, ngf * 4, 2],
+                [None, None, None, ngf * 2, 2],
+                [None, None, None, ngf * 2, 2],
+                [None, None, None, ngf * 1, 1],
+            ],
+            'up': [
+                [None, None, None, ngf * 1, 1],
+            ],
+        }
+    elif dict_name == '270M':
+        channel_dict = {
+            'n_blocks': 4,
+            'down': [ngf * 2, ngf * 2, ngf * 3, ngf * 6, ngf * 4],
+            'backbone': [ngf * 4, ngf * 4],
+            'hair_up': [
+                [ngf * 4, None, None, ngf * 2, 1],
+                [ngf * 2, None, None, ngf * 2, 1],
+                [ngf * 2, None, ngf * 1, ngf * 1, 1],
+                [None, None, None, ngf * 1, 1],
+            ],
+            'face_up': [
+                [ngf * 4, ngf * 6, ngf * 3, ngf * 2, 1],
+                [ngf * 2, ngf * 3, ngf * 2, ngf * 2, 1],
+                [ngf * 2, ngf * 2, ngf * 1, ngf * 1, 1],
+                [None, None, None, ngf * 1, 1],
+            ],
+            'up': [
+                [ngf * 2, ngf * 1, ngf * 1, ngf * 2, 1],
+            ],
+        }
+    else:
+        raise('invalid_dict_name')
+    return get_int(channel_dict)
+
+chans = 8
+in_channels = 12
+out_channels = 16
+face_branch_out_channels = 12
+hair_branch_out_channels = 12
+
+class hair_face_model_old(nn.Module):
+    def __init__(self, ngf=chans, backbone_type='resnet', use_se=True, channel_dict_name = None, with_hair_branch=False, design=5):        
+        super().__init__()
+        self.design = design
+        self.with_hair_branch = with_hair_branch
+        channel_dict = get_channel_dict(channel_dict_name, ngf)
+        n_blocks = channel_dict['n_blocks']
+
+        self.inconv = ConvBlock(in_channels, channel_dict['down'][0], stride=1)
+
+        # Down-Sampling
+        self.DownBlock1 = ConvBlock(channel_dict['down'][0], channel_dict['down'][1], stride=2)
+        self.DownBlock2 = ConvBlock(channel_dict['down'][1], channel_dict['down'][2], stride=2)
+        self.DownBlock3 = ConvBlock(channel_dict['down'][2], channel_dict['down'][3], stride=2)
+        self.DownBlock4 = ConvBlock(channel_dict['down'][3], channel_dict['down'][4], stride=2)
+
+        # Down-Sampling Bottleneck
+        if backbone_type == 'resnet':
+            backbone_block = ResnetBlock
+        elif backbone_type == 'mobilenet':
+            backbone_block = InvertedBottleneck
+            n_blocks = n_blocks
+        else:
+            raise('invalid backbone type')
+        ResBlock = []
+        ResBlock += [backbone_block(channel_dict['down'][4], channel_dict['backbone'][0], use_bias=False, use_se=use_se)]
+        for i in range(1, n_blocks - 1):
+            ResBlock += [backbone_block(channel_dict['backbone'][0], channel_dict['backbone'][0], use_bias=False, use_se=use_se)]
+        ResBlock += [backbone_block(channel_dict['backbone'][0], channel_dict['backbone'][1], use_bias=False, use_se=use_se)]
+        self.ResBlock = nn.Sequential(*ResBlock)
+
+        self.HairUpBlock4 = UpBlock(channel_dict['backbone'][1],    None, channel_dict['hair_up'][0][0], None, channel_dict['hair_up'][0][2], channel_dict['hair_up'][0][3], num_conv=channel_dict['hair_up'][0][4])
+        self.HairUpBlock3 = UpBlock(channel_dict['hair_up'][0][3],  None, channel_dict['hair_up'][1][0], None, channel_dict['hair_up'][1][2], channel_dict['hair_up'][1][3], num_conv=channel_dict['hair_up'][1][4])
+        self.HairUpBlock2 = UpBlock(channel_dict['hair_up'][1][3],  None, channel_dict['hair_up'][2][0], None, channel_dict['hair_up'][2][2], channel_dict['hair_up'][2][3], num_conv=channel_dict['hair_up'][2][4])
+
+        self.FaceUpBlock4 = UpBlock(channel_dict['backbone'][1],    channel_dict['down'][3], channel_dict['face_up'][0][0], channel_dict['face_up'][0][1], channel_dict['face_up'][0][2], channel_dict['face_up'][0][3], num_conv=channel_dict['face_up'][0][4])
+        self.FaceUpBlock3 = UpBlock(channel_dict['face_up'][0][3],  channel_dict['down'][2], channel_dict['face_up'][1][0], channel_dict['face_up'][1][1], channel_dict['face_up'][1][2], channel_dict['face_up'][1][3], num_conv=channel_dict['face_up'][1][4])
+        self.FaceUpBlock2 = UpBlock(channel_dict['face_up'][1][3],  channel_dict['down'][1], channel_dict['face_up'][2][0], channel_dict['face_up'][2][1], channel_dict['face_up'][2][2], channel_dict['face_up'][2][3], num_conv=channel_dict['face_up'][2][4])
+
+        self.UpBlock1 = UpBlock(channel_dict['hair_up'][2][3] + channel_dict['face_up'][2][3],     channel_dict['down'][0], channel_dict['up'][0][0],     channel_dict['up'][0][1], channel_dict['up'][0][2], channel_dict['up'][0][3], num_conv=channel_dict['up'][0][4])
+        self.outconv = ConvOutBlock(channel_dict['up'][0][3], out_channels )
+
+        #self.shortcut_ratio = [1,1,1,1]
+
+        if self.with_hair_branch:
+            self.HairUpBlock1 = UpBlock(channel_dict['hair_up'][2][3], None, channel_dict['hair_up'][3][0], None, channel_dict['hair_up'][3][2],channel_dict['hair_up'][3][3])
+            self.Hairoutconv = ConvOutBlock(channel_dict['hair_up'][3][3], hair_branch_out_channels)
+
+            self.FaceUpBlock1 = UpBlock(channel_dict['face_up'][2][3], channel_dict['down'][0], channel_dict['face_up'][3][0], channel_dict['face_up'][3][1], channel_dict['face_up'][3][2], channel_dict['face_up'][3][3])
+            self.Faceoutconv = ConvOutBlock(channel_dict['face_up'][3][3], face_branch_out_channels)
+
+        self.up = UpsampleQuant(scale_factor=1.5, mode='bilinear')
+
+
+    def forward(self, x , with_hair=False):
+
+        design = self.design
+
+        x0 = self.inconv(x)
+        x1 = self.DownBlock1(x0)
+        x2 = self.DownBlock2(x1)
+        x3 = self.DownBlock3(x2)
+        x4 = self.DownBlock4(x3)
+        x5 = self.ResBlock(x4)
+
+        x6 = x5
+        hair = self.HairUpBlock4(x6)
+        hair = self.HairUpBlock3(hair)
+        hair = self.HairUpBlock2( hair)
+
+        face = self.FaceUpBlock4(x6, self.up(x3) if design == 0 else x3)
+        face = self.FaceUpBlock3(self.up(face) if design==1 else face,
+                                self.up(x2) if design <= 1 else x2)
+        face = self.FaceUpBlock2(self.up(face) if design == 2 else face,
+                                self.up(x1) if design <= 2 else x1)
+
+        hf_cat = torch.cat([hair,face], dim=1)
+
+        x7 = self.UpBlock1(hf_cat,  x0)
+
+        x7 = self.outconv( x7)
+        if not with_hair or not self.with_hair_branch:
+            return x7 
+        else:
+            hair = self.HairUpBlock1(hair)
+            hair = self.Hairoutconv(hair)
+            face = self.FaceUpBlock1(face, x0)
+            face = self.Faceoutconv(face)
+#             print(design == 5) #true
+            return x7 if design == 5 else x7, hair, face
+
+class hair_face_model(hair_face_model_old):
+    def __init__(self, **kwargs):
+        super(hair_face_model, self).__init__(**kwargs)
+
+        self.upconv = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode='bilinear'),
+            nn.Conv2d(4, 4, kernel_size=3, stride=1, padding=1, bias=True),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        x=pack(x)
+        x = super().forward(x)
+        x=unpack(x)
+        return self.upconv(x)
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, stride):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Sequential(
+            #nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=stride, padding=1, bias=False),
+            Conv2dQuant(in_ch, out_ch, kernel_size=3, stride=stride, padding=1, bias=True),
+            nn.BatchNorm2d(out_ch),
+            HardQuant(0, 4)
+            #nn.ReLU(False))
+        )
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class ConvOutBlock(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(ConvOutBlock, self).__init__()
+        self.conv = nn.Sequential(
+            #nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False),
+            Conv2dQuant(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False),
+            #nn.Tanh()
+            TanhOp(data_in_type='float', data_out_type='fixed')
+        )
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class UpBlock(nn.Module):
+    def __init__(self, in_ch1, in_ch2, mid_ch1, mid_ch2, mid_ch, out_ch, num_conv=1, use_bn=True):
+        super(UpBlock, self).__init__()
+
+        #self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.up = UpsampleQuant(scale_factor=2, mode='nearest')
+
+        ## branch_1
+        if mid_ch1 is None or in_ch1 == mid_ch1:
+            self.conv1 = None
+        else:
+            self.conv1 = nn.Sequential(
+                #nn.Conv2d(in_ch1, mid_ch1, 1, bias=False),
+                Conv2dQuant(in_ch1, mid_ch1, 1, bias=True),
+                #nn.ReLU(False),
+                HardQuant(0, 4)
+            )
+        if mid_ch1 is None:
+            mid_ch1 = in_ch1
+
+        if in_ch2 is None:
+            self.use_shortcut = False
+            self.conv2 = None
+        else:
+            self.use_shortcut = True
+            if mid_ch2 is None or in_ch2 == mid_ch2:
+                self.conv2 = None
+            else:
+                self.conv2 = nn.Sequential(
+                    #nn.Conv2d(in_ch2, mid_ch2, 1, bias=False),
+                    Conv2dQuant(in_ch2, mid_ch2, 1, bias=True),
+                    #nn.ReLU(False),
+                    HardQuant(0, 4)
+                )
+            if mid_ch2 is None:
+                mid_ch2 = in_ch2
+        #print(self.conv1 is None, self.conv2 is None)
+        combine_ch = mid_ch1
+        if self.use_shortcut:
+            combine_ch = combine_ch + mid_ch2
+        if mid_ch is None or combine_ch == mid_ch:
+            self.conv_combine = None
+            mid_ch = combine_ch
+        else:
+            self.conv_combine = nn.Sequential(
+                #nn.Conv2d(combine_ch, mid_ch, 1, bias=False),
+                Conv2dQuant(combine_ch, mid_ch, 1, bias=True),
+                #nn.ReLU(False),
+                HardQuant(0, 4)
+            )
+
+        conv_list = []
+        #conv_list.append(nn.Conv2d(mid_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False))
+        conv_list.append(Conv2dQuant(mid_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True))
+        if use_bn:
+            conv_list.append(nn.BatchNorm2d(out_ch))
+        #conv_list.append(nn.ReLU(False))
+        conv_list.append(HardQuant(0, 4))
+        for n in range(1, num_conv):
+            conv_list.append(ResnetBlock(out_ch, out_ch, use_bias=False, use_se = False, use_bn=use_bn))
+        self.conv = nn.Sequential(*conv_list)
+
+    def forward(self, x1, x2=None, ratio=None):
+
+        if self.conv1 is not None:
+            x1 = self.conv1(x1)
+        x1 = self.up(x1)
+
+        if self.use_shortcut:
+            if self.conv2 is not None:
+                x2 = self.conv2(x2)
+
+        if self.use_shortcut:
+            if ratio is None:
+                x = torch.cat([x1, x2], dim=1)
+            else:
+                x = torch.cat([x1, x2 * ratio], dim=1)
+        else:
+            x = x1
+
+        if self.conv_combine is not None:
+            x = self.conv_combine(x)
+
+        x = self.conv(x)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, dim_out, use_bias, use_se = False, use_bn=True):
+        super(ResnetBlock, self).__init__()
+        conv_block = []
+        #conv_block += [nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, bias=use_bias),]
+        conv_block += [Conv2dQuant(dim, dim, kernel_size=3, stride=1, padding=1, bias=True),]
+        if use_bn:
+            conv_block += [nn.BatchNorm2d(dim),]
+        #conv_block += [nn.ReLU(False)]
+        conv_block += [HardQuant(0, 4)]
+
+        #conv_block.append(nn.Conv2d(dim, dim_out, kernel_size=3, stride=1, padding=1, bias=use_bias))
+        conv_block.append(Conv2dQuant(dim, dim_out, kernel_size=3, stride=1, padding=1, bias=True))
+        
+        if use_bn:
+            conv_block.append(nn.BatchNorm2d(dim_out))
+            conv_block += [HardQuant(0, 4)]
+
+        if use_se:
+            conv_block.append(SqEx(dim_out, 4))
+
+        self.conv_block = nn.Sequential(*conv_block)
+
+        self.downsample = None
+        if dim != dim_out:
+            if use_bn:
+                self.downsample = nn.Sequential(
+                    #nn.Conv2d(dim, dim_out, kernel_size=1, stride=1, bias=use_bias),
+                    #nn.BatchNorm2d(dim_out),
+                    Conv2dQuant(dim, dim_out, kernel_size=1, stride=1, bias=True),
+                    nn.BatchNorm2d(dim_out),
+                )
+            else:
+                self.downsample = nn.Sequential(
+                    #nn.Conv2d(dim, dim_out, kernel_size=1, stride=1, bias=use_bias),
+                    Conv2dQuant(dim, dim_out, kernel_size=1, stride=1, bias=True),
+                )
+
+        #self.relu = nn.ReLU(False)
+        self.relu = HardQuant(0, 4)
+
+    def forward(self, x):
+        if self.downsample is None:
+            y = AVG(x, self.conv_block(x))
+        else:
+            y = AVG(self.downsample(x), self.conv_block(x))
+        #y = self.relu(y)
+        return y

models/modules/sr/light_model_470M.py

@@ -0,0 +1,442 @@
+import torch
+import torch.nn as nn
+from models.modules.pix2pixMini_module import *
+
+AVG = AvgQuant()
+mul = SliceMul()
+
+def get_int(channel_dict):
+    for i in range(len(channel_dict['down'])):
+        channel_dict['down'][i] = int(channel_dict['down'][i])
+    for i in range(len(channel_dict['backbone'])):
+        channel_dict['backbone'][i] = int(channel_dict['backbone'][i])
+    for i in range(len(channel_dict['up'])):
+        for j in range(len(channel_dict['up'][i])):
+            if channel_dict['up'][i][j] is not None:
+                channel_dict['up'][i][j] = int(channel_dict['up'][i][j])
+    return channel_dict
+
+def get_channel_dict(dict_name, ngf):
+    if dict_name is None:
+        raise('invalid channel_dict name')
+
+    if dict_name == '3G':
+        channel_dict = {
+            'n_blocks': 8,
+            'down': [ngf * 1, ngf * 2, ngf * 4, ngf * 8, ngf * 8],
+            'backbone': [ngf * 8, ngf * 8],
+            'hair_up': [
+                [None, None, None, ngf * 4, 2],
+                [None, None, None, ngf * 2, 2],
+                [None, None, None, ngf * 1, 2],
+                [None, None, None, ngf * 1, 1],
+            ],
+            'face_up': [
+                [None, None, None, ngf * 4, 2],
+                [None, None, None, ngf * 2, 2],
+                [None, None, None, ngf * 2, 2],
+                [None, None, None, ngf * 1, 1],
+            ],
+            'up': [
+                [None, None, None, ngf * 1, 1],
+            ],
+        }
+
+    elif dict_name == '470M':
+        channel_dict = {
+            'n_blocks': 4,
+            'down': [ngf * 1, ngf * 2, ngf * 3, ngf * 6, ngf * 4],
+            'backbone': [ngf * 4, ngf * 4],
+            'hair_up': [
+                [None, None, None, ngf * 2, 1],
+                [None, None, None, ngf * 2, 1],
+                [None, None, None, ngf * 1, 1],
+                [None, None, None, ngf * 1, 1],
+            ],
+            'face_up': [
+                [None, None, ngf * 4, ngf * 2, 1],
+                [None, None, ngf * 2, ngf * 2, 1],
+                [None, None, ngf * 2, ngf * 1, 1],
+                [None, None, None, ngf * 1, 1],
+            ],
+            'up': [
+                [None, None, ngf * 2, ngf * 1, 1],
+            ],
+        }
+    else:
+        raise('invalid_dict_name')
+    return get_int(channel_dict)
+
+chans = 8
+in_channels = 3
+out_channels = 4
+face_branch_out_channels = 3
+hair_branch_out_channels = 3
+
+def conv(numIn, numOut, k, s=1, p=0, relu=True, bn=False):
+    layers = []
+    layers.append(Conv2dQuant(numIn, numOut, k, s, p, bias=True))
+    if bn:
+        layers.append(nn.BatchNorm2d(numOut))
+        
+    if relu is True:
+        layers.append(HardQuant(0, 4))
+    return nn.Sequential(*layers)
+
+def mnconv(numIn, numOut, k, s=1, p=0, dilation=1, relu=True, bn = True):
+    if k < 2:
+        return conv(numIn, numOut, k, s, p, relu, bn)
+    layers = []
+    layers.append(Conv2dQuant(numIn, numIn, k, s, p, groups=numIn, dilation=dilation, bias=True))
+    layers.append(nn.BatchNorm2d(numIn))
+    layers.append(HardQuant(0, 4))
+    layers.append(conv(numIn, numOut, 1, 1, 0, relu, bn))
+    return nn.Sequential(*layers)
+
+class hair_face_model(nn.Module):
+    def __init__(self, ngf=chans, backbone_type='resnet', use_se=True, channel_dict_name = None, with_hair_branch=False, design=5):
+        
+        super().__init__()
+        self.design = design
+        self.with_hair_branch = with_hair_branch
+        channel_dict = get_channel_dict(channel_dict_name, ngf)
+        n_blocks = channel_dict['n_blocks']
+
+        self.inconv = ConvBlock(in_channels, channel_dict['down'][0], stride=1)
+        self.shortcut_ratio = [1,1,1,1]
+
+        # Down-Sampling
+        self.DownBlock1 = ConvBlock(channel_dict['down'][0], channel_dict['down'][1], stride=2)
+        self.DownBlock2 = ConvBlock(channel_dict['down'][1], channel_dict['down'][2], stride=2)
+        self.DownBlock3 = ConvBlock(channel_dict['down'][2], channel_dict['down'][3], stride=2)
+        self.DownBlock4 = ConvBlock(channel_dict['down'][3], channel_dict['down'][4], stride=2)
+
+        # Down-Sampling Bottleneck
+        if backbone_type == 'resnet':
+            backbone_block = ResnetBlock
+        elif backbone_type == 'mobilenet':
+            backbone_block = InvertedBottleneck
+            n_blocks = n_blocks
+        else:
+            raise('invalid backbone type')
+        ResBlock = []
+        ResBlock += [backbone_block(channel_dict['down'][4], channel_dict['backbone'][0], use_bias=False, use_se=use_se)]
+        for i in range(1, n_blocks - 1):
+            ResBlock += [backbone_block(channel_dict['backbone'][0], channel_dict['backbone'][0], use_bias=False, use_se=use_se)]
+        ResBlock += [backbone_block(channel_dict['backbone'][0], channel_dict['backbone'][1], use_bias=False, use_se=use_se)]
+        self.ResBlock = nn.Sequential(*ResBlock)
+
+        self.HairUpBlock4 = UpBlock(channel_dict['backbone'][1],    None, channel_dict['hair_up'][0][0], None, channel_dict['hair_up'][0][2], channel_dict['hair_up'][0][3], num_conv=channel_dict['hair_up'][0][4])
+        self.HairUpBlock3 = UpBlock(channel_dict['hair_up'][0][3],  None, channel_dict['hair_up'][1][0], None, channel_dict['hair_up'][1][2], channel_dict['hair_up'][1][3], num_conv=channel_dict['hair_up'][1][4])
+        self.HairUpBlock2 = UpBlock(channel_dict['hair_up'][1][3],  None, channel_dict['hair_up'][2][0], None, channel_dict['hair_up'][2][2], channel_dict['hair_up'][2][3], num_conv=channel_dict['hair_up'][2][4])
+
+        self.FaceUpBlock4 = UpBlock(channel_dict['backbone'][1],    channel_dict['down'][3], channel_dict['face_up'][0][0], channel_dict['face_up'][0][1], channel_dict['face_up'][0][2], channel_dict['face_up'][0][3], num_conv=channel_dict['face_up'][0][4])
+        self.FaceUpBlock3 = UpBlock(channel_dict['face_up'][0][3],  channel_dict['down'][2], channel_dict['face_up'][1][0], channel_dict['face_up'][1][1], channel_dict['face_up'][1][2], channel_dict['face_up'][1][3], num_conv=channel_dict['face_up'][1][4])
+        self.FaceUpBlock2 = UpBlock(channel_dict['face_up'][1][3],  channel_dict['down'][1], channel_dict['face_up'][2][0], channel_dict['face_up'][2][1], channel_dict['face_up'][2][2], channel_dict['face_up'][2][3], num_conv=channel_dict['face_up'][2][4])
+
+        self.UpBlock1 = mnUpBlock(channel_dict['hair_up'][2][3] + channel_dict['face_up'][2][3],     channel_dict['down'][0], channel_dict['up'][0][0],     channel_dict['up'][0][1], channel_dict['up'][0][2], channel_dict['up'][0][3], num_conv=channel_dict['up'][0][4])
+        self.outconv = mnConvOutBlock(channel_dict['up'][0][3], out_channels)
+
+        #self.shortcut_ratio = [1,1,1,1]
+
+        if self.with_hair_branch:
+            self.HairUpBlock1 = UpBlock(channel_dict['hair_up'][2][3], None, channel_dict['hair_up'][3][0], None, channel_dict['hair_up'][3][2],channel_dict['hair_up'][3][3])
+            self.Hairoutconv = ConvOutBlock(channel_dict['hair_up'][3][3], hair_branch_out_channels)
+
+            self.FaceUpBlock1 = UpBlock(channel_dict['face_up'][2][3], channel_dict['down'][0], channel_dict['face_up'][3][0], channel_dict['face_up'][3][1], channel_dict['face_up'][3][2], channel_dict['face_up'][3][3])
+            self.Faceoutconv = ConvOutBlock(channel_dict['face_up'][3][3], face_branch_out_channels)
+
+        self.up = UpsampleQuant(scale_factor=1.5, mode='bilinear')
+
+
+    def forward(self, x , with_hair=False):
+        x0 = self.inconv(x)
+        x1 = self.DownBlock1(x0)
+        x2 = self.DownBlock2(x1)
+        x3 = self.DownBlock3(x2)
+        x4 = self.DownBlock4(x3)
+        x = self.ResBlock(x4)
+
+        hair = self.HairUpBlock4(x)
+        hair = self.HairUpBlock3(hair)
+        hair = self.HairUpBlock2(hair)
+        face = self.FaceUpBlock4(x, x3, self.shortcut_ratio[0])
+        face = self.FaceUpBlock3(face, x2, self.shortcut_ratio[1])
+        face = self.FaceUpBlock2(face, x1, self.shortcut_ratio[2])
+
+        x = self.UpBlock1(torch.cat([hair,face], dim=1), x0, self.shortcut_ratio[3])
+#         print(self.outconv)
+        x = self.outconv(x)
+        if not with_hair or not self.with_hair_branch:
+            return x
+        else:
+            hair = self.HairUpBlock1(hair)
+            hair = self.Hairoutconv(hair)
+
+            face = self.FaceUpBlock1(face, x0, self.shortcut_ratio[3])
+            face = self.Faceoutconv(face)
+
+            return x, hair, face
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, stride):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Sequential(
+            #nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=stride, padding=1, bias=False),
+            Conv2dQuant(in_ch, out_ch, kernel_size=3, stride=stride, padding=1, bias=True),
+            nn.BatchNorm2d(out_ch),
+            HardQuant(0, 4)
+            #nn.ReLU(False))
+        )
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+
+class UpBlock(nn.Module):
+    def __init__(self, in_ch1, in_ch2, mid_ch1, mid_ch2, mid_ch, out_ch, num_conv=1, use_bn=True):
+        super(UpBlock, self).__init__()
+
+        #self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.up = UpsampleQuant(scale_factor=2, mode='nearest')
+
+        ## branch_1
+        if mid_ch1 is None or in_ch1 == mid_ch1:
+            self.conv1 = None
+        else:
+            self.conv1 = nn.Sequential(
+                #nn.Conv2d(in_ch1, mid_ch1, 1, bias=False),
+                Conv2dQuant(in_ch1, mid_ch1, 1, bias=True),
+                #nn.ReLU(False),
+                HardQuant(0, 4)
+            )
+        if mid_ch1 is None:
+            mid_ch1 = in_ch1
+
+        if in_ch2 is None:
+            self.use_shortcut = False
+            self.conv2 = None
+        else:
+            self.use_shortcut = True
+            if mid_ch2 is None or in_ch2 == mid_ch2:
+                self.conv2 = None
+            else:
+                self.conv2 = nn.Sequential(
+                    #nn.Conv2d(in_ch2, mid_ch2, 1, bias=False),
+                    Conv2dQuant(in_ch2, mid_ch2, 1, bias=True),
+                    #nn.ReLU(False),
+                    HardQuant(0, 4)
+                )
+            if mid_ch2 is None:
+                mid_ch2 = in_ch2
+        #print(self.conv1 is None, self.conv2 is None)
+        combine_ch = mid_ch1
+        if self.use_shortcut:
+            combine_ch = combine_ch + mid_ch2
+        if mid_ch is None or combine_ch == mid_ch:
+            self.conv_combine = None
+            mid_ch = combine_ch
+        else:
+            self.conv_combine = nn.Sequential(
+                #nn.Conv2d(combine_ch, mid_ch, 1, bias=False),
+                Conv2dQuant(combine_ch, mid_ch, 1, bias=True),
+                #nn.ReLU(False),
+                HardQuant(0, 4)
+            )
+
+        conv_list = []
+        #conv_list.append(nn.Conv2d(mid_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False))
+        conv_list.append(Conv2dQuant(mid_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True))
+#         conv_list.append(mnconv(mid_ch, out_ch, k=3, s=1, p=1))
+
+        if use_bn:
+            conv_list.append(nn.BatchNorm2d(out_ch))
+        #conv_list.append(nn.ReLU(False))
+        conv_list.append(HardQuant(0, 4))
+        for n in range(1, num_conv):
+            conv_list.append(ResnetBlock(out_ch, out_ch, use_bias=False, use_se = False, use_bn=use_bn))
+        self.conv = nn.Sequential(*conv_list)
+
+    def forward(self, x1, x2=None, ratio=None):
+
+        if self.conv1 is not None:
+            x1 = self.conv1(x1)
+        x1 = self.up(x1)
+
+        if self.use_shortcut:
+            if self.conv2 is not None:
+                x2 = self.conv2(x2)
+
+        if self.use_shortcut:
+            if ratio is None:
+                x = torch.cat([x1, x2], dim=1)
+            else:
+                x = torch.cat([x1, x2], dim=1)
+        else:
+            x = x1
+
+        if self.conv_combine is not None:
+            x = self.conv_combine(x)
+
+        x = self.conv(x)
+        return x
+
+class mnUpBlock(nn.Module):
+    def __init__(self, in_ch1, in_ch2, mid_ch1, mid_ch2, mid_ch, out_ch, num_conv=1, use_bn=True):
+        super(mnUpBlock, self).__init__()
+
+        #self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.up = UpsampleQuant(scale_factor=2, mode='nearest')
+
+        ## branch_1
+        if mid_ch1 is None or in_ch1 == mid_ch1:
+            self.conv1 = None
+        else:
+            self.conv1 = nn.Sequential(
+                #nn.Conv2d(in_ch1, mid_ch1, 1, bias=False),
+                Conv2dQuant(in_ch1, mid_ch1, 1, bias=True),
+                #nn.ReLU(False),
+                HardQuant(0, 4)
+            )
+        if mid_ch1 is None:
+            mid_ch1 = in_ch1
+
+        if in_ch2 is None:
+            self.use_shortcut = False
+            self.conv2 = None
+        else:
+            self.use_shortcut = True
+            if mid_ch2 is None or in_ch2 == mid_ch2:
+                self.conv2 = None
+            else:
+                self.conv2 = nn.Sequential(
+                    #nn.Conv2d(in_ch2, mid_ch2, 1, bias=False),
+                    Conv2dQuant(in_ch2, mid_ch2, 1, bias=True),
+                    #nn.ReLU(False),
+                    HardQuant(0, 4)
+                )
+            if mid_ch2 is None:
+                mid_ch2 = in_ch2
+        #print(self.conv1 is None, self.conv2 is None)
+        combine_ch = mid_ch1
+        if self.use_shortcut:
+            combine_ch = combine_ch + mid_ch2
+        if mid_ch is None or combine_ch == mid_ch:
+            self.conv_combine = None
+            mid_ch = combine_ch
+        else:
+            self.conv_combine = nn.Sequential(
+                #nn.Conv2d(combine_ch, mid_ch, 1, bias=False),
+                Conv2dQuant(combine_ch, mid_ch, 1, bias=True),
+                #nn.ReLU(False),
+                HardQuant(0, 4)
+            )
+
+        conv_list = []
+        #conv_list.append(nn.Conv2d(mid_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False))
+        # conv_list.append(Conv2dQuant(mid_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True))
+        conv_list.append(Conv2dQuant(mid_ch, mid_ch, kernel_size=3, stride=1, padding=1, groups=mid_ch,bias=True))
+        conv_list.append(nn.BatchNorm2d(mid_ch))
+        conv_list.append(HardQuant(0, 4))
+        conv_list.append(Conv2dQuant(mid_ch, out_ch, kernel_size=1,stride=1,padding=0))
+        # conv_list.append(mnconv(mid_ch, out_ch, k=3, s=1, p=1))
+
+        if use_bn:
+            conv_list.append(nn.BatchNorm2d(out_ch))
+        #conv_list.append(nn.ReLU(False))
+        conv_list.append(HardQuant(0, 4))
+        for n in range(1, num_conv):
+            conv_list.append(ResnetBlock(out_ch, out_ch, use_bias=False, use_se = False, use_bn=use_bn))
+        self.conv = nn.Sequential(*conv_list)
+
+    def forward(self, x1, x2=None, ratio=None):
+
+        if self.conv1 is not None:
+            x1 = self.conv1(x1)
+        x1 = self.up(x1)
+
+        if self.use_shortcut:
+            if self.conv2 is not None:
+                x2 = self.conv2(x2)
+
+        if self.use_shortcut:
+            if ratio is None:
+                x = torch.cat([x1, x2], dim=1)
+            else:
+                x = torch.cat([x1, x2], dim=1)
+        else:
+            x = x1
+
+        if self.conv_combine is not None:
+            x = self.conv_combine(x)
+
+        x = self.conv(x)
+        return x
+
+class mnConvOutBlock(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(mnConvOutBlock, self).__init__()
+        self.conv = nn.Sequential(
+            Conv2dQuant(in_ch, in_ch, kernel_size=3, stride=1, padding=1, groups=in_ch, bias=True),
+            nn.BatchNorm2d(in_ch),
+            HardQuant(0, 4),
+            Conv2dQuant(in_ch, out_ch, kernel_size=1, stride=1, padding=0, bias=False),
+            #nn.Tanh()
+            TanhOp(data_in_type='float', data_out_type='fixed'),
+            nn.Upsample(scale_factor=2, mode='bilinear'),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
+            nn.Tanh()
+        )
+    def forward(self, x0):
+        x0 = self.conv(x0)
+        return x0
+
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, dim_out, use_bias, use_se = False, use_bn=True):
+        super(ResnetBlock, self).__init__()
+        conv_block = []
+        #conv_block += [nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, bias=use_bias),]
+        conv_block += [Conv2dQuant(dim, dim, kernel_size=3, stride=1, padding=1, bias=True),]
+        if use_bn:
+            conv_block += [nn.BatchNorm2d(dim),]
+        #conv_block += [nn.ReLU(False)]
+        conv_block += [HardQuant(0, 4)]
+
+        #conv_block.append(nn.Conv2d(dim, dim_out, kernel_size=3, stride=1, padding=1, bias=use_bias))
+        conv_block.append(Conv2dQuant(dim, dim_out, kernel_size=3, stride=1, padding=1, bias=True))
+        if use_bn:
+            conv_block.append(nn.BatchNorm2d(dim_out))
+        conv_block += [HardQuant(0, 4)]
+
+        if use_se:
+            conv_block.append(SqEx(dim_out, 4))
+
+        self.conv_block = nn.Sequential(*conv_block)
+
+        self.downsample = None
+        if dim != dim_out:
+            if use_bn:
+                self.downsample = nn.Sequential(
+                    #nn.Conv2d(dim, dim_out, kernel_size=1, stride=1, bias=use_bias),
+                    #nn.BatchNorm2d(dim_out),
+                    Conv2dQuant(dim, dim_out, kernel_size=1, stride=1, bias=True),
+                    nn.BatchNorm2d(dim_out),
+                )
+            else:
+                self.downsample = nn.Sequential(
+                    #nn.Conv2d(dim, dim_out, kernel_size=1, stride=1, bias=use_bias),
+                    Conv2dQuant(dim, dim_out, kernel_size=1, stride=1, bias=True),
+                )
+
+        #self.relu = nn.ReLU(False)
+        #self.relu = HardQuant(0, 4)
+
+    def forward(self, x):
+        if self.downsample is None:
+            y = AVG(x, self.conv_block(x))
+        else:
+            y = AVG(self.downsample(x), self.conv_block(x))
+        #y = self.relu(y)
+        return y

models/modules/stylegan2/model.py

@@ -0,0 +1,716 @@
+import math
+import random
+import functools
+import operator
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.autograd import Function
+
+from .op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d, conv2d_gradfix
+
+
+class PixelNorm(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input):
+        return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)
+
+
+def make_kernel(k):
+    k = torch.tensor(k, dtype=torch.float32)
+
+    if k.ndim == 1:
+        k = k[None, :] * k[:, None]
+
+    k /= k.sum()
+
+    return k
+
+
+class Upsample(nn.Module):
+    def __init__(self, kernel, factor=2):
+        super().__init__()
+
+        self.factor = factor
+        kernel = make_kernel(kernel) * (factor ** 2)
+        self.register_buffer("kernel", kernel)
+
+        p = kernel.shape[0] - factor
+
+        pad0 = (p + 1) // 2 + factor - 1
+        pad1 = p // 2
+
+        self.pad = (pad0, pad1)
+
+    def forward(self, input):
+        out = upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad)
+
+        return out
+
+
+class Downsample(nn.Module):
+    def __init__(self, kernel, factor=2):
+        super().__init__()
+
+        self.factor = factor
+        kernel = make_kernel(kernel)
+        self.register_buffer("kernel", kernel)
+
+        p = kernel.shape[0] - factor
+
+        pad0 = (p + 1) // 2
+        pad1 = p // 2
+
+        self.pad = (pad0, pad1)
+
+    def forward(self, input):
+        out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)
+
+        return out
+
+
+class Blur(nn.Module):
+    def __init__(self, kernel, pad, upsample_factor=1):
+        super().__init__()
+
+        kernel = make_kernel(kernel)
+
+        if upsample_factor > 1:
+            kernel = kernel * (upsample_factor ** 2)
+
+        self.register_buffer("kernel", kernel)
+
+        self.pad = pad
+
+    def forward(self, input):
+        out = upfirdn2d(input, self.kernel, pad=self.pad)
+
+        return out
+
+
+class EqualConv2d(nn.Module):
+    def __init__(
+        self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True
+    ):
+        super().__init__()
+
+        self.weight = nn.Parameter(
+            torch.randn(out_channel, in_channel, kernel_size, kernel_size)
+        )
+        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
+
+        self.stride = stride
+        self.padding = padding
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_channel))
+
+        else:
+            self.bias = None
+
+    def forward(self, input):
+        out = conv2d_gradfix.conv2d(
+            input,
+            self.weight * self.scale,
+            bias=self.bias,
+            stride=self.stride,
+            padding=self.padding,
+        )
+
+        return out
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},"
+            f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})"
+        )
+
+
+class EqualLinear(nn.Module):
+    def __init__(
+        self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None
+    ):
+        super().__init__()
+
+        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
+
+        else:
+            self.bias = None
+
+        self.activation = activation
+
+        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+        self.lr_mul = lr_mul
+
+    def forward(self, input):
+        if self.activation:
+            out = F.linear(input, self.weight * self.scale)
+            out = fused_leaky_relu(out, self.bias * self.lr_mul)
+
+        else:
+            out = F.linear(
+                input, self.weight * self.scale, bias=self.bias * self.lr_mul
+            )
+
+        return out
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})"
+        )
+
+
+class ModulatedConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        style_dim,
+        demodulate=True,
+        upsample=False,
+        downsample=False,
+        blur_kernel=[1, 3, 3, 1],
+        fused=True,
+    ):
+        super().__init__()
+
+        self.eps = 1e-8
+        self.kernel_size = kernel_size
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.upsample = upsample
+        self.downsample = downsample
+
+        if upsample:
+            factor = 2
+            p = (len(blur_kernel) - factor) - (kernel_size - 1)
+            pad0 = (p + 1) // 2 + factor - 1
+            pad1 = p // 2 + 1
+
+            self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)
+
+        if downsample:
+            factor = 2
+            p = (len(blur_kernel) - factor) + (kernel_size - 1)
+            pad0 = (p + 1) // 2
+            pad1 = p // 2
+
+            self.blur = Blur(blur_kernel, pad=(pad0, pad1))
+
+        fan_in = in_channel * kernel_size ** 2
+        self.scale = 1 / math.sqrt(fan_in)
+        self.padding = kernel_size // 2
+
+        self.weight = nn.Parameter(
+            torch.randn(1, out_channel, in_channel, kernel_size, kernel_size)
+        )
+
+        self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)
+
+        self.demodulate = demodulate
+        self.fused = fused
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, "
+            f"upsample={self.upsample}, downsample={self.downsample})"
+        )
+
+    def forward(self, input, style):
+        batch, in_channel, height, width = input.shape
+
+        if not self.fused:
+            weight = self.scale * self.weight.squeeze(0)
+            style = self.modulation(style)
+
+            if self.demodulate:
+                w = weight.unsqueeze(0) * style.view(batch, 1, in_channel, 1, 1)
+                dcoefs = (w.square().sum((2, 3, 4)) + 1e-8).rsqrt()
+
+            input = input * style.reshape(batch, in_channel, 1, 1)
+
+            if self.upsample:
+                weight = weight.transpose(0, 1)
+                out = conv2d_gradfix.conv_transpose2d(
+                    input, weight, padding=0, stride=2
+                )
+                out = self.blur(out)
+
+            elif self.downsample:
+                input = self.blur(input)
+                out = conv2d_gradfix.conv2d(input, weight, padding=0, stride=2)
+
+            else:
+                out = conv2d_gradfix.conv2d(input, weight, padding=self.padding)
+
+            if self.demodulate:
+                out = out * dcoefs.view(batch, -1, 1, 1)
+
+            return out
+
+        style = self.modulation(style).view(batch, 1, in_channel, 1, 1)
+        weight = self.scale * self.weight * style
+
+        if self.demodulate:
+            demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
+            weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)
+
+        weight = weight.view(
+            batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size
+        )
+
+        if self.upsample:
+            input = input.view(1, batch * in_channel, height, width)
+            weight = weight.view(
+                batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size
+            )
+            weight = weight.transpose(1, 2).reshape(
+                batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size
+            )
+            out = conv2d_gradfix.conv_transpose2d(
+                input, weight, padding=0, stride=2, groups=batch
+            )
+            _, _, height, width = out.shape
+            out = out.view(batch, self.out_channel, height, width)
+            out = self.blur(out)
+
+        elif self.downsample:
+            input = self.blur(input)
+            _, _, height, width = input.shape
+            input = input.view(1, batch * in_channel, height, width)
+            out = conv2d_gradfix.conv2d(
+                input, weight, padding=0, stride=2, groups=batch
+            )
+            _, _, height, width = out.shape
+            out = out.view(batch, self.out_channel, height, width)
+
+        else:
+            input = input.view(1, batch * in_channel, height, width)
+            out = conv2d_gradfix.conv2d(
+                input, weight, padding=self.padding, groups=batch
+            )
+            _, _, height, width = out.shape
+            out = out.view(batch, self.out_channel, height, width)
+
+        return out
+
+
+class NoiseInjection(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.weight = nn.Parameter(torch.zeros(1))
+
+    def forward(self, image, noise=None):
+        if noise is None:
+            batch, _, height, width = image.shape
+            noise = image.new_empty(batch, 1, height, width).normal_()
+
+        return image + self.weight * noise
+
+
+class ConstantInput(nn.Module):
+    def __init__(self, channel, size=4):
+        super().__init__()
+
+        if type(size) is tuple:
+            self.input = nn.Parameter(torch.randn(1, channel, size[0], size[1]))
+        else:
+            self.input = nn.Parameter(torch.randn(1, channel, size, size))
+
+    def forward(self, input):
+        batch = input.shape[0]
+        out = self.input.repeat(batch, 1, 1, 1)
+
+        return out
+
+
+class StyledConv(nn.Module):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        style_dim,
+        upsample=False,
+        blur_kernel=[1, 3, 3, 1],
+        demodulate=True,
+    ):
+        super().__init__()
+
+        self.conv = ModulatedConv2d(
+            in_channel,
+            out_channel,
+            kernel_size,
+            style_dim,
+            upsample=upsample,
+            blur_kernel=blur_kernel,
+            demodulate=demodulate,
+        )
+
+        self.noise = NoiseInjection()
+        # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1))
+        # self.activate = ScaledLeakyReLU(0.2)
+        self.activate = FusedLeakyReLU(out_channel)
+
+    def forward(self, input, style, noise=None):
+        out = self.conv(input, style)
+        out = self.noise(out, noise=noise)
+        # out = out + self.bias
+        out = self.activate(out)
+
+        return out
+
+
+class ToRGB(nn.Module):
+    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
+        super().__init__()
+
+        if upsample:
+            self.upsample = Upsample(blur_kernel)
+
+        self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
+        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
+
+    def forward(self, input, style, skip=None):
+        out = self.conv(input, style)
+        out = out + self.bias
+
+        if skip is not None:
+            skip = self.upsample(skip)
+
+            out = out + skip
+
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+        self,
+        size,
+        style_dim,
+        n_mlp,
+        channel_multiplier=2,
+        blur_kernel=[1, 3, 3, 1],
+        lr_mlp=0.01,
+    ):
+        super().__init__()
+
+        self.size = size
+
+        self.style_dim = style_dim
+
+        layers = [PixelNorm()]
+
+        for i in range(n_mlp):
+            layers.append(
+                EqualLinear(
+                    style_dim, style_dim, lr_mul=lr_mlp, activation="fused_lrelu"
+                )
+            )
+
+        self.style = nn.Sequential(*layers)
+
+        self.channels = {
+            4: 512,
+            8: 512,
+            16: 512,
+            32: 512,
+            64: 256 * channel_multiplier,
+            128: 128 * channel_multiplier,
+            256: 64 * channel_multiplier,
+            512: 32 * channel_multiplier,
+            1024: 16 * channel_multiplier,
+        }
+
+        self.input = ConstantInput(self.channels[4])
+        self.conv1 = StyledConv(
+            self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel
+        )
+        self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)
+
+        self.log_size = int(math.log(size, 2))
+        self.num_layers = (self.log_size - 2) * 2 + 1
+
+        self.convs = nn.ModuleList()
+        self.upsamples = nn.ModuleList()
+        self.to_rgbs = nn.ModuleList()
+        self.noises = nn.Module()
+
+        in_channel = self.channels[4]
+
+        for layer_idx in range(self.num_layers):
+            res = (layer_idx + 5) // 2
+            shape = [1, 1, 2 ** res, 2 ** res]
+            self.noises.register_buffer(f"noise_{layer_idx}", torch.randn(*shape))
+
+        for i in range(3, self.log_size + 1):
+            out_channel = self.channels[2 ** i]
+
+            self.convs.append(
+                StyledConv(
+                    in_channel,
+                    out_channel,
+                    3,
+                    style_dim,
+                    upsample=True,
+                    blur_kernel=blur_kernel,
+                )
+            )
+
+            self.convs.append(
+                StyledConv(
+                    out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel
+                )
+            )
+
+            self.to_rgbs.append(ToRGB(out_channel, style_dim))
+
+            in_channel = out_channel
+
+        self.n_latent = self.log_size * 2 - 2
+
+    def make_noise(self):
+        device = self.input.input.device
+
+        noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]
+
+        for i in range(3, self.log_size + 1):
+            for _ in range(2):
+                noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device))
+
+        return noises
+
+    def mean_latent(self, n_latent):
+        latent_in = torch.randn(
+            n_latent, self.style_dim, device=self.input.input.device
+        )
+        latent = self.style(latent_in).mean(0, keepdim=True)
+
+        return latent
+
+    def get_latent(self, input):
+        return self.style(input)
+
+    def forward(
+        self,
+        styles,
+        return_latents=False,
+        inject_index=None,
+        truncation=1,
+        truncation_latent=None,
+        input_is_latent=False,
+        noise=None,
+        randomize_noise=True,
+    ):
+        if not input_is_latent:
+            styles = [self.style(s) for s in styles]
+
+        if noise is None:
+            if randomize_noise:
+                noise = [None] * self.num_layers
+            else:
+                noise = [
+                    getattr(self.noises, f"noise_{i}") for i in range(self.num_layers)
+                ]
+
+        if truncation < 1:
+            style_t = []
+
+            for style in styles:
+                style_t.append(
+                    truncation_latent + truncation * (style - truncation_latent)
+                )
+
+            styles = style_t
+
+        if len(styles) < 2:
+            inject_index = self.n_latent
+
+            if styles[0].ndim < 3:
+                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
+
+            else:
+                latent = styles[0]
+
+        else:
+            if inject_index is None:
+                inject_index = random.randint(1, self.n_latent - 1)
+
+            latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
+            latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)
+
+            latent = torch.cat([latent, latent2], 1)
+
+        out = self.input(latent)
+        out = self.conv1(out, latent[:, 0], noise=noise[0])
+
+        skip = self.to_rgb1(out, latent[:, 1])
+
+        i = 1
+        for conv1, conv2, noise1, noise2, to_rgb in zip(
+            self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
+        ):
+            out = conv1(out, latent[:, i], noise=noise1)
+            out = conv2(out, latent[:, i + 1], noise=noise2)
+            skip = to_rgb(out, latent[:, i + 2], skip)
+
+            i += 2
+
+        image = skip
+
+        if return_latents:
+            return image, latent
+
+        else:
+            return image, None
+
+
+class ConvLayer(nn.Sequential):
+    def __init__(
+        self,
+        in_channel,
+        out_channel,
+        kernel_size,
+        downsample=False,
+        blur_kernel=[1, 3, 3, 1],
+        bias=True,
+        activate=True,
+    ):
+        layers = []
+
+        if downsample:
+            factor = 2
+            p = (len(blur_kernel) - factor) + (kernel_size - 1)
+            pad0 = (p + 1) // 2
+            pad1 = p // 2
+
+            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
+
+            stride = 2
+            self.padding = 0
+
+        else:
+            stride = 1
+            self.padding = kernel_size // 2
+
+        layers.append(
+            EqualConv2d(
+                in_channel,
+                out_channel,
+                kernel_size,
+                padding=self.padding,
+                stride=stride,
+                bias=bias and not activate,
+            )
+        )
+
+        if activate:
+            layers.append(FusedLeakyReLU(out_channel, bias=bias))
+
+        super().__init__(*layers)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
+        super().__init__()
+
+        self.conv1 = ConvLayer(in_channel, in_channel, 3)
+        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
+
+        self.skip = ConvLayer(
+            in_channel, out_channel, 1, downsample=True, activate=False, bias=False
+        )
+
+    def forward(self, input):
+        out = self.conv1(input)
+        out = self.conv2(out)
+
+        skip = self.skip(input)
+        out = (out + skip) / math.sqrt(2)
+
+        return out
+
+
+class Discriminator(nn.Module):
+    def __init__(self, size, min_feats_size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
+        super().__init__()
+
+        channels = {
+            4: 512,
+            8: 512,
+            16: 512,
+            32: 512,
+            64: 256 * channel_multiplier,
+            128: 128 * channel_multiplier,
+            256: 64 * channel_multiplier,
+            512: 32 * channel_multiplier,
+            1024: 16 * channel_multiplier,
+        }
+
+        convs = [ConvLayer(3, channels[size], 1)]
+
+        log_size = int(math.log(size, 2))
+        if type(min_feats_size) is tuple:
+            fsize = min_feats_size[0] * min_feats_size[1]
+        else:
+            fsize = min_feats_size * min_feats_size
+
+        in_channel = channels[size]
+
+        for i in range(log_size, 2, -1):
+            out_channel = channels[2 ** (i - 1)]
+
+            convs.append(ResBlock(in_channel, out_channel, blur_kernel))
+
+            in_channel = out_channel
+
+        self.convs = nn.Sequential(*convs)
+
+        self.stddev_group = 4
+        self.stddev_feat = 1
+
+        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
+        self.final_linear = nn.Sequential(
+            EqualLinear(channels[4] * fsize, channels[4], activation="fused_lrelu"),
+            EqualLinear(channels[4], 1),
+        )
+
+    def forward(self, input, rtn_feats=False):
+        if rtn_feats:
+            feats = []
+            feat = input
+            for i, block in enumerate(self.convs):
+                feat = block(feat)
+                if i in [ 1, 3, 4, 5 ]:
+                    feats.append(feat)
+                if i == 5:
+                    break
+            return feats
+
+        out = self.convs(input)
+
+        batch, channel, height, width = out.shape
+        group = min(batch, self.stddev_group)
+        stddev = out.view(
+            group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
+        )
+        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
+        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
+        stddev = stddev.repeat(group, 1, height, width)
+        out = torch.cat([out, stddev], 1)
+
+        out = self.final_conv(out)
+
+        out = out.view(batch, -1)
+        out = self.final_linear(out)
+
+        return out
+

models/modules/stylegan2/non_leaking.py

@@ -0,0 +1,465 @@
+import math
+
+import torch
+from torch import autograd
+from torch.nn import functional as F
+import numpy as np
+
+# from distributed import reduce_sum
+from .op import upfirdn2d
+
+
+# class AdaptiveAugment:
+#     def __init__(self, ada_aug_target, ada_aug_len, update_every, device):
+#         self.ada_aug_target = ada_aug_target
+#         self.ada_aug_len = ada_aug_len
+#         self.update_every = update_every
+
+#         self.ada_update = 0
+#         self.ada_aug_buf = torch.tensor([0.0, 0.0], device=device)
+#         self.r_t_stat = 0
+#         self.ada_aug_p = 0
+
+#     @torch.no_grad()
+#     def tune(self, real_pred):
+#         self.ada_aug_buf += torch.tensor(
+#             (torch.sign(real_pred).sum().item(), real_pred.shape[0]),
+#             device=real_pred.device,
+#         )
+#         self.ada_update += 1
+
+#         if self.ada_update % self.update_every == 0:
+#             self.ada_aug_buf = reduce_sum(self.ada_aug_buf)
+#             pred_signs, n_pred = self.ada_aug_buf.tolist()
+
+#             self.r_t_stat = pred_signs / n_pred
+
+#             if self.r_t_stat > self.ada_aug_target:
+#                 sign = 1
+
+#             else:
+#                 sign = -1
+
+#             self.ada_aug_p += sign * n_pred / self.ada_aug_len
+#             self.ada_aug_p = min(1, max(0, self.ada_aug_p))
+#             self.ada_aug_buf.mul_(0)
+#             self.ada_update = 0
+
+#         return self.ada_aug_p
+
+
+SYM6 = (
+    0.015404109327027373,
+    0.0034907120842174702,
+    -0.11799011114819057,
+    -0.048311742585633,
+    0.4910559419267466,
+    0.787641141030194,
+    0.3379294217276218,
+    -0.07263752278646252,
+    -0.021060292512300564,
+    0.04472490177066578,
+    0.0017677118642428036,
+    -0.007800708325034148,
+)
+
+
+def translate_mat(t_x, t_y, device="cpu"):
+    batch = t_x.shape[0]
+
+    mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
+    translate = torch.stack((t_x, t_y), 1)
+    mat[:, :2, 2] = translate
+
+    return mat
+
+
+def rotate_mat(theta, device="cpu"):
+    batch = theta.shape[0]
+
+    mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
+    sin_t = torch.sin(theta)
+    cos_t = torch.cos(theta)
+    rot = torch.stack((cos_t, -sin_t, sin_t, cos_t), 1).view(batch, 2, 2)
+    mat[:, :2, :2] = rot
+
+    return mat
+
+
+def scale_mat(s_x, s_y, device="cpu"):
+    batch = s_x.shape[0]
+
+    mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
+    mat[:, 0, 0] = s_x
+    mat[:, 1, 1] = s_y
+
+    return mat
+
+
+def translate3d_mat(t_x, t_y, t_z):
+    batch = t_x.shape[0]
+
+    mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
+    translate = torch.stack((t_x, t_y, t_z), 1)
+    mat[:, :3, 3] = translate
+
+    return mat
+
+
+def rotate3d_mat(axis, theta):
+    batch = theta.shape[0]
+
+    u_x, u_y, u_z = axis
+
+    eye = torch.eye(3).unsqueeze(0)
+    cross = torch.tensor([(0, -u_z, u_y), (u_z, 0, -u_x), (-u_y, u_x, 0)]).unsqueeze(0)
+    outer = torch.tensor(axis)
+    outer = (outer.unsqueeze(1) * outer).unsqueeze(0)
+
+    sin_t = torch.sin(theta).view(-1, 1, 1)
+    cos_t = torch.cos(theta).view(-1, 1, 1)
+
+    rot = cos_t * eye + sin_t * cross + (1 - cos_t) * outer
+
+    eye_4 = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
+    eye_4[:, :3, :3] = rot
+
+    return eye_4
+
+
+def scale3d_mat(s_x, s_y, s_z):
+    batch = s_x.shape[0]
+
+    mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
+    mat[:, 0, 0] = s_x
+    mat[:, 1, 1] = s_y
+    mat[:, 2, 2] = s_z
+
+    return mat
+
+
+def luma_flip_mat(axis, i):
+    batch = i.shape[0]
+
+    eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
+    axis = torch.tensor(axis + (0,))
+    flip = 2 * torch.ger(axis, axis) * i.view(-1, 1, 1)
+
+    return eye - flip
+
+
+def saturation_mat(axis, i):
+    batch = i.shape[0]
+
+    eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
+    axis = torch.tensor(axis + (0,))
+    axis = torch.ger(axis, axis)
+    saturate = axis + (eye - axis) * i.view(-1, 1, 1)
+
+    return saturate
+
+
+def lognormal_sample(size, mean=0, std=1, device="cpu"):
+    return torch.empty(size, device=device).log_normal_(mean=mean, std=std)
+
+
+def category_sample(size, categories, device="cpu"):
+    category = torch.tensor(categories, device=device)
+    sample = torch.randint(high=len(categories), size=(size,), device=device)
+
+    return category[sample]
+
+
+def uniform_sample(size, low, high, device="cpu"):
+    return torch.empty(size, device=device).uniform_(low, high)
+
+
+def normal_sample(size, mean=0, std=1, device="cpu"):
+    return torch.empty(size, device=device).normal_(mean, std)
+
+
+def bernoulli_sample(size, p, device="cpu"):
+    return torch.empty(size, device=device).bernoulli_(p)
+
+
+def random_mat_apply(p, transform, prev, eye, device="cpu"):
+    size = transform.shape[0]
+    select = bernoulli_sample(size, p, device=device).view(size, 1, 1)
+    select_transform = select * transform + (1 - select) * eye
+
+    return select_transform @ prev
+
+
+def sample_affine(p, size, height, width, device="cpu"):
+    G = torch.eye(3, device=device).unsqueeze(0).repeat(size, 1, 1)
+    eye = G
+
+    # flip
+    param = category_sample(size, (0, 1))
+    Gc = scale_mat(1 - 2.0 * param, torch.ones(size), device=device)
+    G = random_mat_apply(p, Gc, G, eye, device=device)
+    # print('flip', G, scale_mat(1 - 2.0 * param, torch.ones(size)), sep='\n')
+
+    # 90 rotate
+    param = category_sample(size, (0, 3))
+    Gc = rotate_mat(-math.pi / 2 * param, device=device)
+    G = random_mat_apply(p, Gc, G, eye, device=device)
+    # print('90 rotate', G, rotate_mat(-math.pi / 2 * param), sep='\n')
+
+    # integer translate
+    param = uniform_sample((2, size), -0.125, 0.125)
+    param_height = torch.round(param[0] * height)
+    param_width = torch.round(param[1] * width)
+    Gc = translate_mat(param_width, param_height, device=device)
+    G = random_mat_apply(p, Gc, G, eye, device=device)
+    # print('integer translate', G, translate_mat(param_width, param_height), sep='\n')
+
+    # isotropic scale
+    param = lognormal_sample(size, std=0.2 * math.log(2))
+    Gc = scale_mat(param, param, device=device)
+    G = random_mat_apply(p, Gc, G, eye, device=device)
+    # print('isotropic scale', G, scale_mat(param, param), sep='\n')
+
+    p_rot = 1 - math.sqrt(1 - p)
+
+    # pre-rotate
+    param = uniform_sample(size, -math.pi, math.pi)
+    Gc = rotate_mat(-param, device=device)
+    G = random_mat_apply(p_rot, Gc, G, eye, device=device)
+    # print('pre-rotate', G, rotate_mat(-param), sep='\n')
+
+    # anisotropic scale
+    param = lognormal_sample(size, std=0.2 * math.log(2))
+    Gc = scale_mat(param, 1 / param, device=device)
+    G = random_mat_apply(p, Gc, G, eye, device=device)
+    # print('anisotropic scale', G, scale_mat(param, 1 / param), sep='\n')
+
+    # post-rotate
+    param = uniform_sample(size, -math.pi, math.pi)
+    Gc = rotate_mat(-param, device=device)
+    G = random_mat_apply(p_rot, Gc, G, eye, device=device)
+    # print('post-rotate', G, rotate_mat(-param), sep='\n')
+
+    # fractional translate
+    param = normal_sample((2, size), std=0.125)
+    Gc = translate_mat(param[1] * width, param[0] * height, device=device)
+    G = random_mat_apply(p, Gc, G, eye, device=device)
+    # print('fractional translate', G, translate_mat(param, param), sep='\n')
+
+    return G
+
+
+def sample_color(p, size):
+    C = torch.eye(4).unsqueeze(0).repeat(size, 1, 1)
+    eye = C
+    axis_val = 1 / math.sqrt(3)
+    axis = (axis_val, axis_val, axis_val)
+
+    # brightness
+    param = normal_sample(size, std=0.2)
+    Cc = translate3d_mat(param, param, param)
+    C = random_mat_apply(p, Cc, C, eye)
+
+    # contrast
+    param = lognormal_sample(size, std=0.5 * math.log(2))
+    Cc = scale3d_mat(param, param, param)
+    C = random_mat_apply(p, Cc, C, eye)
+
+    # luma flip
+    param = category_sample(size, (0, 1))
+    Cc = luma_flip_mat(axis, param)
+    C = random_mat_apply(p, Cc, C, eye)
+
+    # hue rotation
+    param = uniform_sample(size, -math.pi, math.pi)
+    Cc = rotate3d_mat(axis, param)
+    C = random_mat_apply(p, Cc, C, eye)
+
+    # saturation
+    param = lognormal_sample(size, std=1 * math.log(2))
+    Cc = saturation_mat(axis, param)
+    C = random_mat_apply(p, Cc, C, eye)
+
+    return C
+
+
+def make_grid(shape, x0, x1, y0, y1, device):
+    n, c, h, w = shape
+    grid = torch.empty(n, h, w, 3, device=device)
+    grid[:, :, :, 0] = torch.linspace(x0, x1, w, device=device)
+    grid[:, :, :, 1] = torch.linspace(y0, y1, h, device=device).unsqueeze(-1)
+    grid[:, :, :, 2] = 1
+
+    return grid
+
+
+def affine_grid(grid, mat):
+    n, h, w, _ = grid.shape
+    return (grid.view(n, h * w, 3) @ mat.transpose(1, 2)).view(n, h, w, 2)
+
+
+def get_padding(G, height, width, kernel_size):
+    device = G.device
+
+    cx = (width - 1) / 2
+    cy = (height - 1) / 2
+    cp = torch.tensor(
+        [(-cx, -cy, 1), (cx, -cy, 1), (cx, cy, 1), (-cx, cy, 1)], device=device
+    )
+    cp = G @ cp.T
+
+    pad_k = kernel_size // 4
+
+    pad = cp[:, :2, :].permute(1, 0, 2).flatten(1)
+    pad = torch.cat((-pad, pad)).max(1).values
+    pad = pad + torch.tensor([pad_k * 2 - cx, pad_k * 2 - cy] * 2, device=device)
+    pad = pad.max(torch.tensor([0.0, 0.0] * 2, device=device))
+    pad = pad.min(torch.tensor([width - 1.0, height - 1.0] * 2, device=device))
+
+    pad_x1, pad_y1, pad_x2, pad_y2 = pad.ceil().to(torch.int32)
+
+    return pad_x1, pad_x2, pad_y1, pad_y2
+
+
+def try_sample_affine_and_pad(img, p, kernel_size, G=None):
+    batch, _, height, width = img.shape
+
+    G_try = G
+
+    if G is None:
+        G_try = torch.inverse(sample_affine(p, batch, height, width))
+
+    pad_x1, pad_x2, pad_y1, pad_y2 = get_padding(G_try, height, width, kernel_size)
+
+    img_pad = F.pad(img, (pad_x1, pad_x2, pad_y1, pad_y2), mode="reflect")
+
+    return img_pad, G_try, (pad_x1, pad_x2, pad_y1, pad_y2)
+
+
+class GridSampleForward(autograd.Function):
+    @staticmethod
+    def forward(ctx, input, grid):
+        out = F.grid_sample(
+            input, grid, mode="bilinear", padding_mode="zeros", align_corners=False
+        )
+        ctx.save_for_backward(input, grid)
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, grid = ctx.saved_tensors
+        grad_input, grad_grid = GridSampleBackward.apply(grad_output, input, grid)
+
+        return grad_input, grad_grid
+
+
+class GridSampleBackward(autograd.Function):
+    @staticmethod
+    def forward(ctx, grad_output, input, grid):
+        op = torch._C._jit_get_operation("aten::grid_sampler_2d_backward")
+        grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False)
+        ctx.save_for_backward(grid)
+
+        return grad_input, grad_grid
+
+    @staticmethod
+    def backward(ctx, grad_grad_input, grad_grad_grid):
+        (grid,) = ctx.saved_tensors
+        grad_grad_output = None
+
+        if ctx.needs_input_grad[0]:
+            grad_grad_output = GridSampleForward.apply(grad_grad_input, grid)
+
+        return grad_grad_output, None, None
+
+
+grid_sample = GridSampleForward.apply
+
+
+def scale_mat_single(s_x, s_y):
+    return torch.tensor(((s_x, 0, 0), (0, s_y, 0), (0, 0, 1)), dtype=torch.float32)
+
+
+def translate_mat_single(t_x, t_y):
+    return torch.tensor(((1, 0, t_x), (0, 1, t_y), (0, 0, 1)), dtype=torch.float32)
+
+
+def random_apply_affine(img, p, G=None, antialiasing_kernel=SYM6):
+    kernel = antialiasing_kernel
+    len_k = len(kernel)
+
+    kernel = torch.as_tensor(kernel).to(img)
+    # kernel = torch.ger(kernel, kernel).to(img)
+    kernel_flip = torch.flip(kernel, (0,))
+
+    img_pad, G, (pad_x1, pad_x2, pad_y1, pad_y2) = try_sample_affine_and_pad(
+        img, p, len_k, G
+    )
+
+    G_inv = (
+        translate_mat_single((pad_x1 - pad_x2).item() / 2, (pad_y1 - pad_y2).item() / 2)
+        @ G
+    )
+    up_pad = (
+        (len_k + 2 - 1) // 2,
+        (len_k - 2) // 2,
+        (len_k + 2 - 1) // 2,
+        (len_k - 2) // 2,
+    )
+    img_2x = upfirdn2d(img_pad, kernel.unsqueeze(0), up=(2, 1), pad=(*up_pad[:2], 0, 0))
+    img_2x = upfirdn2d(img_2x, kernel.unsqueeze(1), up=(1, 2), pad=(0, 0, *up_pad[2:]))
+    G_inv = scale_mat_single(2, 2) @ G_inv @ scale_mat_single(1 / 2, 1 / 2)
+    G_inv = translate_mat_single(-0.5, -0.5) @ G_inv @ translate_mat_single(0.5, 0.5)
+    batch_size, channel, height, width = img.shape
+    pad_k = len_k // 4
+    shape = (batch_size, channel, (height + pad_k * 2) * 2, (width + pad_k * 2) * 2)
+    G_inv = (
+        scale_mat_single(2 / img_2x.shape[3], 2 / img_2x.shape[2])
+        @ G_inv
+        @ scale_mat_single(1 / (2 / shape[3]), 1 / (2 / shape[2]))
+    )
+    grid = F.affine_grid(G_inv[:, :2, :].to(img_2x), shape, align_corners=False)
+    img_affine = grid_sample(img_2x, grid)
+    d_p = -pad_k * 2
+    down_pad = (
+        d_p + (len_k - 2 + 1) // 2,
+        d_p + (len_k - 2) // 2,
+        d_p + (len_k - 2 + 1) // 2,
+        d_p + (len_k - 2) // 2,
+    )
+    img_down = upfirdn2d(
+        img_affine, kernel_flip.unsqueeze(0), down=(2, 1), pad=(*down_pad[:2], 0, 0)
+    )
+    img_down = upfirdn2d(
+        img_down, kernel_flip.unsqueeze(1), down=(1, 2), pad=(0, 0, *down_pad[2:])
+    )
+
+    return img_down, G
+
+
+def apply_color(img, mat):
+    batch = img.shape[0]
+    img = img.permute(0, 2, 3, 1)
+    mat_mul = mat[:, :3, :3].transpose(1, 2).view(batch, 1, 3, 3)
+    mat_add = mat[:, :3, 3].view(batch, 1, 1, 3)
+    img = img @ mat_mul + mat_add
+    img = img.permute(0, 3, 1, 2)
+
+    return img
+
+
+def random_apply_color(img, p, C=None):
+    if C is None:
+        C = sample_color(p, img.shape[0])
+
+    img = apply_color(img, C.to(img))
+
+    return img, C
+
+
+def augment(img, p, transform_matrix=(None, None)):
+    img, G = random_apply_affine(img, p, transform_matrix[0])
+    img, C = random_apply_color(img, p, transform_matrix[1])
+
+    return img, (G, C)

models/modules/stylegan2/op/__init__.py

@@ -0,0 +1,2 @@
+from .fused_act import FusedLeakyReLU, fused_leaky_relu
+from .upfirdn2d import upfirdn2d