Official PyTorch implementation of the PG 2023 paper: Multi-Modal Face Stylization with a Generative Prior
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/__init__.py
mmfs-paper/MMFS
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/base_config.py
mmfs-paper/MMFS
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)

configs/style_based_pix2pixII_config.py
mmfs-paper/MMFS
"""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/__init__.py
mmfs-paper/MMFS
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/static_data.py
mmfs-paper/MMFS
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/super_dataset.py
mmfs-paper/MMFS
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_data.py
mmfs-paper/MMFS
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

data/test_video_data.py
mmfs-paper/MMFS
"""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/__init__.py
mmfs-paper/MMFS
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/base_model.py
mmfs-paper/MMFS
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/networks.py
mmfs-paper/MMFS
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/model.py
mmfs-paper/MMFS
# code from https://github.com/omerbt/Splice/blob/master/models/extractor.py

import torch


def attn_cosine_sim(x, eps=1e-08):
    x = x[0]  # TEMP: getting rid of redundant dimension, TBF
    norm1 = x.norm(dim=2, keepdim=True)
    factor = torch.clamp(norm1 @ norm1.permute(0, 2, 1), min=eps)
    sim_matrix = (x @ x.permute(0, 2, 1)) / factor
    return sim_matrix


class VitExtractor:
    BLOCK_KEY = 'block'
    ATTN_KEY = 'attn'
    PATCH_IMD_KEY = 'patch_imd'
    QKV_KEY = 'qkv'
    KEY_LIST = [BLOCK_KEY, ATTN_KEY, PATCH_IMD_KEY, QKV_KEY]

    def __init__(self, model_name, device):
        self.model = torch.hub.load('facebookresearch/dino:main', model_name).to(device)
        self.model.eval()
        self.model_name = model_name
        self.hook_handlers = []
        self.layers_dict = {}
        self.outputs_dict = {}
        for key in VitExtractor.KEY_LIST:
            self.layers_dict[key] = []
            self.outputs_dict[key] = []
        self._init_hooks_data()

    def _init_hooks_data(self):
        self.layers_dict[VitExtractor.BLOCK_KEY] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
        self.layers_dict[VitExtractor.ATTN_KEY] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
        self.layers_dict[VitExtractor.QKV_KEY] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
        self.layers_dict[VitExtractor.PATCH_IMD_KEY] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
        for key in VitExtractor.KEY_LIST:
            # self.layers_dict[key] = kwargs[key] if key in kwargs.keys() else []
            self.outputs_dict[key] = []

    def _register_hooks(self, **kwargs):
        for block_idx, block in enumerate(self.model.blocks):
            if block_idx in self.layers_dict[VitExtractor.BLOCK_KEY]:
                self.hook_handlers.append(block.register_forward_hook(self._get_block_hook()))
            if block_idx in self.layers_dict[VitExtractor.ATTN_KEY]:
                self.hook_handlers.append(block.attn.attn_drop.register_forward_hook(self._get_attn_hook()))
            if block_idx in self.layers_dict[VitExtractor.QKV_KEY]:
                self.hook_handlers.append(block.attn.qkv.register_forward_hook(self._get_qkv_hook()))
            if block_idx in self.layers_dict[VitExtractor.PATCH_IMD_KEY]:
                self.hook_handlers.append(block.attn.register_forward_hook(self._get_patch_imd_hook()))

    def _clear_hooks(self):
        for handler in self.hook_handlers:
            handler.remove()
        self.hook_handlers = []

    def _get_block_hook(self):
        def _get_block_output(model, input, output):
            self.outputs_dict[VitExtractor.BLOCK_KEY].append(output)

        return _get_block_output

    def _get_attn_hook(self):
        def _get_attn_output(model, inp, output):
            self.outputs_dict[VitExtractor.ATTN_KEY].append(output)

        return _get_attn_output

    def _get_qkv_hook(self):
        def _get_qkv_output(model, inp, output):
            self.outputs_dict[VitExtractor.QKV_KEY].append(output)

        return _get_qkv_output

    # TODO: CHECK ATTN OUTPUT TUPLE
    def _get_patch_imd_hook(self):
        def _get_attn_output(model, inp, output):
            self.outputs_dict[VitExtractor.PATCH_IMD_KEY].append(output[0])

        return _get_attn_output

    def get_feature_from_input(self, input_img):  # List([B, N, D])
        self._register_hooks()
        self.model(input_img)
        feature = self.outputs_dict[VitExtractor.BLOCK_KEY]
        self._clear_hooks()
        self._init_hooks_data()
        return feature

    def get_qkv_feature_from_input(self, input_img):
        self._register_hooks()
        self.model(input_img)
        feature = self.outputs_dict[VitExtractor.QKV_KEY]
        self._clear_hooks()
        self._init_hooks_data()
        return feature

    def get_attn_feature_from_input(self, input_img):
        self._register_hooks()
        self.model(input_img)
        feature = self.outputs_dict[VitExtractor.ATTN_KEY]
        self._clear_hooks()
        self._init_hooks_data()
        return feature

    def get_patch_size(self):
        return 8 if "8" in self.model_name else 16

    def get_width_patch_num(self, input_img_shape):
        b, c, h, w = input_img_shape
        patch_size = self.get_patch_size()
        return w // patch_size

    def get_height_patch_num(self, input_img_shape):
        b, c, h, w = input_img_shape
        patch_size = self.get_patch_size()
        return h // patch_size

    def get_patch_num(self, input_img_shape):
        patch_num = 1 + (self.get_height_patch_num(input_img_shape) * self.get_width_patch_num(input_img_shape))
        return patch_num

    def get_head_num(self):
        if "dino" in self.model_name:
            return 6 if "s" in self.model_name else 12
        return 6 if "small" in self.model_name else 12

    def get_embedding_dim(self):
        if "dino" in self.model_name:
            return 384 if "s" in self.model_name else 768
        return 384 if "small" in self.model_name else 768

    def get_queries_from_qkv(self, qkv, input_img_shape):
        patch_num = self.get_patch_num(input_img_shape)
        head_num = self.get_head_num()
        embedding_dim = self.get_embedding_dim()
        q = qkv.reshape(patch_num, 3, head_num, embedding_dim // head_num).permute(1, 2, 0, 3)[0]
        return q

    def get_keys_from_qkv(self, qkv, input_img_shape):
        patch_num = self.get_patch_num(input_img_shape)
        head_num = self.get_head_num()
        embedding_dim = self.get_embedding_dim()
        k = qkv.reshape(patch_num, 3, head_num, embedding_dim // head_num).permute(1, 2, 0, 3)[1]
        return k

    def get_values_from_qkv(self, qkv, input_img_shape):
        patch_num = self.get_patch_num(input_img_shape)
        head_num = self.get_head_num()
        embedding_dim = self.get_embedding_dim()
        v = qkv.reshape(patch_num, 3, head_num, embedding_dim // head_num).permute(1, 2, 0, 3)[2]
        return v

    def get_keys_from_input(self, input_img, layer_num):
        qkv_features = self.get_qkv_feature_from_input(input_img)[layer_num]
        keys = self.get_keys_from_qkv(qkv_features, input_img.shape)
        return keys

    def get_keys_self_sim_from_input(self, input_img, layer_num):
        keys = self.get_keys_from_input(input_img, layer_num=layer_num)
        h, t, d = keys.shape
        concatenated_keys = keys.transpose(0, 1).reshape(t, h * d)
        ssim_map = attn_cosine_sim(concatenated_keys[None, None, ...])
        return ssim_map

models/modules/vit/extractor.py
mmfs-paper/MMFS
# code from https://github.com/omerbt/Splice/blob/master/util/losses.py

from torchvision.transforms import Resize
from torchvision import transforms
import torch
import torch.nn.functional as F

from .extractor import VitExtractor


class LossG:

    def __init__(self, device):
        super().__init__()

        self.extractor = VitExtractor(model_name='dino_vitb8', device=device)

        imagenet_norm = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        global_resize_transform = Resize(224, max_size=480)

        self.global_transform = transforms.Compose([ global_resize_transform, imagenet_norm ])

    def get_cls_token(self, inputs):
        tokens = []
        for x in inputs:
            x = self.global_transform(x).unsqueeze(0)
            cls_token = self.extractor.get_feature_from_input(x)[-1][0, 0, :]
            tokens.append(cls_token.unsqueeze(0))
        return torch.cat(tokens, dim=0) # shape: B x 768

    def calculate_global_ssim_loss(self, outputs, inputs):
        loss = 0.0
        for a, b in zip(inputs, outputs):  # avoid memory limitations
            a = self.global_transform(a)
            b = self.global_transform(b)
            with torch.no_grad():
                target_keys_self_sim = self.extractor.get_keys_self_sim_from_input(a.unsqueeze(0), layer_num=11)
            keys_ssim = self.extractor.get_keys_self_sim_from_input(b.unsqueeze(0), layer_num=11)
            loss += F.mse_loss(keys_ssim, target_keys_self_sim)
        return loss

    def calculate_crop_cls_loss(self, outputs, inputs):
        loss = 0.0
        for a, b in zip(outputs, inputs):  # avoid memory limitations
            a = self.global_transform(a).unsqueeze(0).to(device)
            b = self.global_transform(b).unsqueeze(0).to(device)
            cls_token = self.extractor.get_feature_from_input(a)[-1][0, 0, :]
            with torch.no_grad():
                target_cls_token = self.extractor.get_feature_from_input(b)[-1][0, 0, :]
            loss += F.mse_loss(cls_token, target_cls_token)
        return loss

    def calculate_global_id_loss(self, outputs, inputs):
        loss = 0.0
        for a, b in zip(inputs, outputs):
            a = self.global_transform(a)
            b = self.global_transform(b)
            with torch.no_grad():
                keys_a = self.extractor.get_keys_from_input(a.unsqueeze(0), 11)
            keys_b = self.extractor.get_keys_from_input(b.unsqueeze(0), 11)
            loss += F.mse_loss(keys_a, keys_b)
        return loss

models/modules/vit/losses.py
mmfs-paper/MMFS
import copy
import clip
import os
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from enum import Enum
from PIL import Image
from torch import autograd
from .base_model import BaseModel
from models.modules import networks
from models.modules.stylegan2.model import Generator, Discriminator, StyledConv, ToRGB, EqualLinear, ResBlock, ConvLayer, PixelNorm
from models.modules.stylegan2.op import conv2d_gradfix
from models.modules.stylegan2.non_leaking import augment
from models.modules.vit.losses import LossG


class TrainingPhase(Enum):
    ENCODER = 1
    BASE_MODEL = 2
    CLIP_MAPPING = 3
    FEW_SHOT = 4


class CLIPFeats2Wplus(nn.Module):

    def __init__(self, n_tokens=16, embedding_dim=512):
        super().__init__()

        self.position_embedding = nn.Parameter(embedding_dim ** -0.5 * torch.randn(n_tokens, embedding_dim))
        self.transformer = nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=embedding_dim, nhead=8, norm_first=True), num_layers=4)

    def forward(self, x):
        x_in = x.view(x.shape[0], 1, x.shape[1]) + self.position_embedding
        return F.leaky_relu(self.transformer(x_in.permute(1, 0, 2)), negative_slope=0.2)


class Stylizer(nn.Module):

    def __init__(self, ngf=64, phase=TrainingPhase.ENCODER, model_weights=None):
        super(Stylizer, self).__init__()

        # encoder
        self.encoder = nn.Sequential(
            ConvLayer(3, ngf, 3),          # 512
            ResBlock(ngf * 1, ngf * 1),    # 256
            ResBlock(ngf * 1, ngf * 2),    # 128
            ResBlock(ngf * 2, ngf * 4),    # 64
            ResBlock(ngf * 4, ngf * 8),    # 32
            ConvLayer(ngf * 8, ngf * 8, 3) # 32
        )

        # mapping network
        self.mapping_z = nn.Sequential(*([ PixelNorm() ] + [ EqualLinear(512, 512, activation='fused_lrelu', lr_mul=0.01) for _ in range(8) ]))

        # style-based decoder
        channels = {
            32 : ngf * 8,
            64 : ngf * 8,
            128: ngf * 4,
            256: ngf * 2,
            512: ngf * 1
        }
        self.decoder0 = StyledConv(channels[32], channels[32], 3, 512)
        self.to_rgb0 = ToRGB(channels[32], 512, upsample=False)
        for i in range(4):
            ichan = channels[2 ** (i + 5)]
            ochan = channels[2 ** (i + 6)]
            setattr(self, f'decoder{i + 1}a', StyledConv(ichan, ochan, 3, 512, upsample=True))
            setattr(self, f'decoder{i + 1}b', StyledConv(ochan, ochan, 3, 512))
            setattr(self, f'to_rgb{i + 1}', ToRGB(ochan, 512))
        self.n_latent = 10

        # random style for testing
        self.test_z = torch.randn(1, 512)

        # load pretrained model weights
        if phase == TrainingPhase.ENCODER:
            # load pretrained stylegan2 and freeze these params
            for param in self.mapping_z.parameters():
                param.requires_grad = False
            for i in range(4):
                for key in [f'decoder{i + 1}a', f'decoder{i + 1}b', f'to_rgb{i + 1}']:
                    for param in getattr(self, key).parameters():
                        param.requires_grad = False
            self.load_state_dict(self._convert_stylegan2_dict(model_weights), strict=False)
        elif phase == TrainingPhase.BASE_MODEL:
            # load pretrained encoder and stylegan2 decoder
            self.load_state_dict(model_weights)
        elif phase == TrainingPhase.CLIP_MAPPING:
            self.clip_mapper = CLIPFeats2Wplus(n_tokens=self.n_latent)
            # load pretraned base model and freeze all params except clip mapper
            self.load_state_dict(model_weights, strict=False)
            params = dict(self.named_parameters())
            for k in params.keys():
                if 'clip_mapper' in k:
                    print(f'{k} not freezed !')
                    continue
                params[k].requires_grad = False
        elif phase == TrainingPhase.FEW_SHOT:
            self.clip_mapper = CLIPFeats2Wplus(n_tokens=self.n_latent)
            # load pretrained base model and freeze encoder & mapping
            self.load_state_dict(model_weights)
            self.encoder.requires_grad_(False)
            self.mapping_z.requires_grad_(False)
            self.clip_mapper.requires_grad_(False)

    def _convert_stylegan2_dict(self, src):
        res = {}
        for k, v in src.items():
            if k.startswith('style.'):
                res[k.replace('style.', 'mapping_z.')] = v
            else:
                name, idx = k.split('.')[:2]
                if name == 'convs':
                    idx = int(idx)
                    if idx >= 6:
                        res[k.replace(f'{name}.{idx}.', f'decoder{idx // 2 - 2}{chr(97 + idx % 2)}.')] = v
                elif name == 'to_rgbs':
                    idx = int(idx)
                    if idx >= 3:
                        res[k.replace(f'{name}.{idx}.', f'to_rgb{idx - 2}.')] = v
        return res

    def get_styles(self, x, **kwargs):
        if len(kwargs) == 0:
            return self.mapping_z(self.test_z.to(x.device).repeat(x.shape[0], 1)).repeat(self.n_latent, 1, 1)
        elif 'mixing' in kwargs and kwargs['mixing']:
            w0 = self.mapping_z(torch.randn(x.shape[0], 512, device=x.device))
            w1 = self.mapping_z(torch.randn(x.shape[0], 512, device=x.device))
            inject_index = random.randint(1, self.n_latent - 1)
            return torch.cat([
                w0.repeat(inject_index, 1, 1),
                w1.repeat(self.n_latent - inject_index, 1, 1)
            ])
        elif 'z' in kwargs:
            return self.mapping_z(kwargs['z']).repeat(self.n_latent, 1, 1)
        elif 'clip_feats' in kwargs:
            return self.clip_mapper(kwargs['clip_feats'])
        else:
            z = torch.randn(x.shape[0], 512, device=x.device)
            return self.mapping_z(z).repeat(self.n_latent, 1, 1)

    def forward(self, x, **kwargs):
        # encode
        feat = self.encoder(x)

        # get style code
        styles = self.get_styles(x, **kwargs)

        # style-based generate
        feat = self.decoder0(feat, styles[0])
        out = self.to_rgb0(feat, styles[1])
        for i in range(4):
            feat = getattr(self, f'decoder{i + 1}a')(feat, styles[i * 2 + 1])
            feat = getattr(self, f'decoder{i + 1}b')(feat, styles[i * 2 + 2])
            out = getattr(self, f'to_rgb{i + 1}')(feat, styles[i * 2 + 3], out)

        return F.hardtanh(out)


class StyleBasedPix2PixIIModel(BaseModel):
    """
    This class implements the Style-Based Pix2Pix model version II.
    """

    def __init__(self, config, DDP_device=None):
        BaseModel.__init__(self, config, DDP_device=DDP_device)

        self.d_reg_freq = 16
        self.lambda_r1 = 10
        self.step = 0
        self.phase = TrainingPhase(config['training']['phase'])

        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        if self.phase == TrainingPhase.ENCODER:
            self.loss_names = ['G', 'G_L1', 'G_Feat']
        elif self.phase == TrainingPhase.BASE_MODEL:
            self.loss_names = ['G', 'G_ST', 'G_GAN', 'D']
        elif self.phase == TrainingPhase.CLIP_MAPPING:
            self.loss_names = ['G', 'G_L1', 'G_Feat']
        elif self.phase == TrainingPhase.FEW_SHOT:
            self.loss_names = ['G', 'G_ST', 'G_CLIP', 'G_PROJ']

        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
        self.visual_names = ['real_A', 'real_B', 'fake_B']

        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
        if self.isTrain:
            self.model_names = ['G', 'G_ema', 'D']
        else:  # during test time, only load Gs
            self.model_names = ['G_ema']

        self.data_aug_prob = config['training']['data_aug_prob']

        min_feats_size = tuple(config['model']['min_feats_size'])

        def __init_net(model):
            return networks.init_net(model, init_type='none', init_gain=0.0, gpu_ids=self.gpu_ids,
                DDP_device=self.DDP_device, find_unused_parameters=config['training']['find_unused_parameters'])

        if self.phase == TrainingPhase.ENCODER: # train a encoder for stylegan2
            # load and init pretrained stylegan2
            model_dict = torch.load(config['training']['pretrained_model'], map_location='cpu')
            self.stylegan2 = Generator(512, 512, 8)
            self.stylegan2.load_state_dict(model_dict['g'])
            self.stylegan2 = __init_net(self.stylegan2)
            self.stylegan2.eval()
            self.stylegan2.requires_grad_(False)

            # init netG
            self.netG = Stylizer(ngf=config['model']['ngf'], phase=self.phase, model_weights=model_dict['g'])
            self.netG = __init_net(self.netG)

            # init netD
            self.netD = Discriminator(min(min_feats_size) * 128, min_feats_size)
            self.netD.load_state_dict(model_dict['d'])
            self.netD = __init_net(self.netD)
            self.netD.eval()
            self.netD.requires_grad_(False)

        elif self.phase == TrainingPhase.BASE_MODEL: # finetune the whole model
            model_dict = torch.load(config['training']['pretrained_model'], map_location='cpu')

            # init netG
            self.netG = Stylizer(ngf=config['model']['ngf'], phase=self.phase, model_weights=model_dict['G_ema_model'])
            self.netG = __init_net(self.netG)

            # init netD
            self.netD = Discriminator(min(min_feats_size) * 128, min_feats_size)
            self.netD.load_state_dict(model_dict['D_model'])
            self.netD = __init_net(self.netD)

        elif self.phase == TrainingPhase.CLIP_MAPPING or self.phase == TrainingPhase.FEW_SHOT: # train the clip mapper or zero/one shot finetune
            # init CLIP
            self.clip_model, self.pil_to_tensor = clip.load('ViT-B/32', device=self.device)
            self.clip_model.eval()
            self.clip_model.requires_grad_(False)

            model_dict = torch.load(config['training']['pretrained_model'], map_location='cpu')

            # init netG
            self.netG = Stylizer(ngf=config['model']['ngf'], phase=self.phase, model_weights=model_dict['G_ema_model'])
            self.netG = __init_net(self.netG)

            # init netD
            self.netD = Discriminator(min(min_feats_size) * 128, min_feats_size)
            self.netD.load_state_dict(model_dict['D_model'])
            self.netD = __init_net(self.netD)
            self.netD.eval()
            self.netD.requires_grad_(False)

        if self.phase == TrainingPhase.FEW_SHOT: # set hook to get clip vit tokens
            def clip_vit_hook(model, feat_in, feat_out):
                self.clip_vit_tokens = feat_out[1:].permute(1, 0, 2).float() # remove cls token
            self.clip_model.visual.transformer.resblocks[3].register_forward_hook(clip_vit_hook)

        # create netG ema
        self.netG_ema = copy.deepcopy(self.netG)
        self.netG_ema.eval()
        self.netG_ema.requires_grad_(False)
        self.ema(self.netG_ema, self.netG, 0.0)

        # CLIP mean & std
        self.clip_mean = torch.tensor((0.48145466, 0.4578275, 0.40821073), device=self.device).view(1, 3, 1, 1)
        self.clip_std = torch.tensor((0.26862954, 0.26130258, 0.27577711), device=self.device).view(1, 3, 1, 1)

        if self.isTrain:
            # define loss functions
            if self.phase == TrainingPhase.ENCODER:
                self.criterionL1 = nn.L1Loss()
            elif self.phase == TrainingPhase.BASE_MODEL:
                self.criterionStyleGAN = networks.GANLoss('wgangp').to(self.device)
                self.vitLoss = LossG(self.device)
            elif self.phase == TrainingPhase.CLIP_MAPPING:
                self.criterionL1 = nn.L1Loss()
            elif self.phase == TrainingPhase.FEW_SHOT:
                self.criterionL1 = nn.L1Loss()
                self.vitLoss = LossG(self.device)
                self.cosineSim = nn.CosineSimilarity(dim=1)

            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=config['training']['lr'], betas=(config['training']['beta1'], 0.999))
            d_reg_ratio = self.d_reg_freq / (self.d_reg_freq + 1)
            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=config['training']['lr'] * d_reg_ratio, betas=(config['training']['beta1'] ** d_reg_ratio, 0.999 ** d_reg_ratio))
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)

    def ema(self, tgt, src, decay=0.999):
        param_tgt = dict(tgt.named_parameters())
        param_src = dict(src.named_parameters())
        for key in param_tgt.keys():
            param_tgt[key].data.mul_(decay).add_(param_src[key].data, alpha=1.0 - decay)

    def preprocess_clip_image(self, x, size):
        x = x * 0.5 + 0.5
        x = F.interpolate(x, (size, size), mode='bilinear', antialias=True, align_corners=False)
        return (x - self.clip_mean) / self.clip_std

    def set_input(self, input):
        if self.phase == TrainingPhase.ENCODER:
            # sample via stylegan2
            self.z = torch.randn(self.config['dataset']['batch_size'], 512, device=self.device)
            with torch.no_grad():
                self.real_A = F.hardtanh(self.stylegan2.forward([self.z])[0])
                self.real_B = self.real_A.clone()
        elif self.phase == TrainingPhase.BASE_MODEL:
            if self.config['common']['phase'] == 'test':
                self.real_A = input['test_A'].to(self.device)
                self.real_B = input['test_B'].to(self.device)
                self.image_paths = input['test_A_path']
            else:
                self.real_A = input['unpaired_A'].to(self.device)
                self.real_B = input['unpaired_B'].to(self.device)
                self.image_paths = input['unpaired_A_path']
        elif self.phase == TrainingPhase.CLIP_MAPPING:
            self.real_A = input['unpaired_A'].to(self.device)
            with torch.no_grad():
                self.real_B = self.netG_ema(self.real_A, mixing=random.random() < self.config['training']['style_mixing_prob'])
                self.clip_feats = self.clip_model.encode_image(self.preprocess_clip_image(self.real_B, self.clip_model.visual.input_resolution))
                self.clip_feats /= self.clip_feats.norm(dim=1, keepdim=True)
        elif self.phase == TrainingPhase.FEW_SHOT:
            self.real_A = input['unpaired_A'].to(self.device)
            self.real_B = self.real_A
            if not hasattr(self, 'clip_feats'):
                with torch.no_grad():
                    if os.path.isfile(self.config['training']['image_prompt']):
                        image = self.pil_to_tensor(Image.open(self.config['training']['image_prompt'])).unsqueeze(0).to(self.device)
                        self.clip_feats = self.clip_model.encode_image(image)
                        ref_tokens = self.clip_vit_tokens
                        ref_tokens /= ref_tokens.norm(dim=2, keepdim=True)
                        D = ref_tokens.shape[2]
                        ref_tokens = ref_tokens.reshape(-1, D).permute(1, 0)
                        U, _, _ = torch.linalg.svd(ref_tokens, full_matrices=False)
                        self.UUT = U @ U.permute(1, 0)
                        self.use_image_prompt = True
                    else:
                        text = clip.tokenize(self.config['training']['text_prompt']).to(self.device)
                        self.clip_feats = self.clip_model.encode_text(text)
                        self.use_image_prompt = False
                        # get source text prompt feature
                        text = clip.tokenize(self.config['training']['src_text_prompt']).to(self.device)
                        self.src_clip_feats = self.clip_model.encode_text(text)
                        self.src_clip_feats /= self.src_clip_feats.norm(dim=1, keepdim=True)
                        self.src_clip_feats = self.src_clip_feats.repeat(self.config['dataset']['batch_size'], 1)
                    self.clip_feats /= self.clip_feats.norm(dim=1, keepdim=True)
                    self.clip_feats = self.clip_feats.repeat(self.config['dataset']['batch_size'], 1)
            # get direction in clip space
            with torch.no_grad():
                self.real_A_clip_feats = self.clip_model.encode_image(self.preprocess_clip_image(self.real_A, self.clip_model.visual.input_resolution))
                self.real_A_clip_feats /= self.real_A_clip_feats.norm(dim=1, keepdim=True)
                if self.use_image_prompt:
                    self.src_clip_feats = self.real_A_clip_feats.mean(dim=0, keepdim=True).repeat(self.config['dataset']['batch_size'], 1)
                self.clip_feats_dir = self.clip_feats - self.src_clip_feats

    def forward(self, use_ema=False):
        if self.phase == TrainingPhase.ENCODER:
            if use_ema:
                self.fake_B = self.netG_ema(self.real_A, z=self.z)
            else:
                self.fake_B = self.netG(self.real_A, z=self.z)
        elif self.phase == TrainingPhase.BASE_MODEL:
            if not self.isTrain:
                self.fake_B = self.netG_ema(self.real_A, mixing=False)
            elif use_ema:
                self.fake_B = self.netG_ema(self.real_A, mixing=random.random() < self.config['training']['style_mixing_prob'])
            else:
                self.fake_B = self.netG(self.real_A, mixing=random.random() < self.config['training']['style_mixing_prob'])
        elif self.phase == TrainingPhase.CLIP_MAPPING or self.phase == TrainingPhase.FEW_SHOT:
            if use_ema:
                self.fake_B = self.netG_ema(self.real_A, clip_feats=self.clip_feats)
            else:
                self.fake_B = self.netG(self.real_A, clip_feats=self.clip_feats)

    def backward_D_r1(self):
        self.real_B.requires_grad = True
        if self.data_aug_prob == 0.0:
            real_aug = self.real_B
        else:
            real_aug, _ = augment(self.real_B, self.data_aug_prob)
        real_pred = self.netD(real_aug)
        with conv2d_gradfix.no_weight_gradients():
            grad, = autograd.grad(outputs=real_pred.sum(), inputs=real_aug, create_graph=True)
        r1_loss = grad.pow(2).reshape(grad.shape[0], -1).sum(1).mean()
        (r1_loss * self.lambda_r1 / 2 * self.d_reg_freq + 0 * real_pred[0]).backward()

    def backward_D(self, backward=True):
        if self.data_aug_prob == 0.0:
            loss_fake = self.criterionStyleGAN(self.netD(self.fake_B.detach()), False)
            loss_real = self.criterionStyleGAN(self.netD(self.real_B), True)
        else:
            fake_aug, _ = augment(self.fake_B.detach(), self.data_aug_prob)
            real_aug, _ = augment(self.real_B, self.data_aug_prob)
            loss_fake = self.criterionStyleGAN(self.netD(fake_aug), False)
            loss_real = self.criterionStyleGAN(self.netD(real_aug), True)
        self.loss_D = (loss_fake + loss_real) * 0.5

        if backward:
            self.loss_D.backward()

    def backward_G(self, backward=True):
        self.loss_G = 0

        if self.phase == TrainingPhase.ENCODER or self.phase == TrainingPhase.CLIP_MAPPING:
            self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B)
            with torch.no_grad():
                real_feats = self.netD(self.real_B, rtn_feats=True)
            fake_feats = self.netD(self.fake_B, rtn_feats=True)
            self.loss_G_Feat = sum([ self.criterionL1(fake, real) for fake, real in zip(fake_feats, real_feats) ])

            self.loss_G += self.loss_G_L1 * self.config['training']['lambda_L1']
            self.loss_G += self.loss_G_Feat * self.config['training']['lambda_Feat']

        elif self.phase == TrainingPhase.BASE_MODEL:
            self.loss_G_ST = self.vitLoss.calculate_global_ssim_loss(self.fake_B * 0.5 + 0.5, self.real_A * 0.5 + 0.5)
            if self.data_aug_prob == 0.0:
                self.loss_G_GAN = self.criterionStyleGAN(self.netD(self.fake_B), True)
            else:
                fake_aug, _ = augment(self.fake_B, self.data_aug_prob)
                self.loss_G_GAN = self.criterionStyleGAN(self.netD(fake_aug), True)

            self.loss_G += self.loss_G_ST * self.config['training']['lambda_ST']
            self.loss_G += self.loss_G_GAN * self.config['training']['lambda_GAN']

        elif self.phase == TrainingPhase.FEW_SHOT:
            self.loss_G_ST = self.vitLoss.calculate_global_ssim_loss(self.fake_B * 0.5 + 0.5, self.real_A * 0.5 + 0.5)
            fake_clip_feats = self.clip_model.encode_image(self.preprocess_clip_image(self.fake_B, self.clip_model.visual.input_resolution))
            fake_clip_feats = fake_clip_feats / fake_clip_feats.norm(dim=1, keepdim=True)
            fake_clip_feats_dir = fake_clip_feats - self.real_A_clip_feats
            self.loss_G_CLIP = (1.0 - self.cosineSim(fake_clip_feats_dir, self.clip_feats_dir)).mean()
            if self.use_image_prompt:
                fake_tokens = self.clip_vit_tokens
                fake_tokens = fake_tokens / fake_tokens.norm(dim=2, keepdim=True)
                D = fake_tokens.shape[2]
                fake_tokens = fake_tokens.reshape(-1, D).permute(1, 0)
                self.loss_G_PROJ = self.criterionL1(self.UUT @ fake_tokens, fake_tokens)
            else:
                self.loss_G_PROJ = 0.0

            self.loss_G += self.loss_G_ST * self.config['training']['lambda_ST']
            self.loss_G += self.loss_G_CLIP * self.config['training']['lambda_CLIP']
            self.loss_G += self.loss_G_PROJ * self.config['training']['lambda_PROJ']

        if backward:
            self.loss_G.backward()

    def optimize_parameters(self):
        # forward
        self.forward()

        if not self.phase == TrainingPhase.BASE_MODEL:
            # only G
            self.optimizer_G.zero_grad()
            self.backward_G()
            self.optimizer_G.step()

            # update G_ema
            self.ema(self.netG_ema, self.netG, decay=self.config['training']['ema'])

        else:
            # G
            self.set_requires_grad([self.netD], False)
            self.optimizer_G.zero_grad()
            self.backward_G()
            self.optimizer_G.step()
            # D
            self.set_requires_grad([self.netD], True)
            self.optimizer_D.zero_grad()
            self.backward_D()
            self.optimizer_D.step()

            # update G_ema
            self.ema(self.netG_ema, self.netG, decay=self.config['training']['ema'])

            # r1 reg
            self.step += 1
            if self.step % self.d_reg_freq == 0:
                self.optimizer_D.zero_grad()
                self.backward_D_r1()
                self.optimizer_D.step()

    def eval_step(self):
        self.forward(use_ema=True)
        self.backward_G(False)
        if self.phase == TrainingPhase.BASE_MODEL:
            self.backward_D(False)
        self.step += 1

    def trace_jit(self, input):
        self.netG = self.netG.module.cpu()
        traced_script_module = torch.jit.trace(self.netG, input)
        dummy_output = self.netG_ema(input)
        dummy_output_traced = traced_script_module(input)
        return traced_script_module, dummy_output, dummy_output_traced

models/style_based_pix2pixII_model.py
mmfs-paper/MMFS
import argparse
import os
import cv2
from tqdm import tqdm

from utils.util import *
from data import CustomDataLoader
from data.super_dataset import SuperDataset
from models import create_model
from configs import parse_config

if __name__ == '__main__':
    # parse arguments
    parser = argparse.ArgumentParser(description='Style Master')
    parser.add_argument('--cfg_file', type=str, default='./exp/sp2pII-phase2.yaml')
    parser.add_argument('--test_img', type=str, default='', help='path to your test img')
    parser.add_argument('--test_video', type=str, default='')
    parser.add_argument('--test_folder', type=str, default='') # ./example/source
    parser.add_argument('--ckpt', type=str, default='') # ./pretrained_models/phase2_pretrain_90000.pth
    parser.add_argument('--overwrite_output_dir', type=str, default='') # ./example/outputs/multi-model
    parser.add_argument('--gpus', type=str, default='0')
    args = parser.parse_args()

    # parse config
    config = parse_config(args.cfg_file)

    # fix gpu ordering
    gpu_string = ','.join(map(str, config['common']['gpu_ids']))
    gpu_ids_fix = list(range(len(config['common']['gpu_ids'])))  # wants GPU ids match nvidia-smi output order

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
    #os.environ['CUDA_VISIBLE_DEVICES'] = gpu_string
    config['common']['gpu_ids'] = gpu_ids_fix

    # hard-code some parameters for test
    config['common']['phase'] = 'test'
    config['dataset']['n_threads'] = 0   # test code only supports num_threads = 0
    config['dataset']['batch_size'] = 1    # test code only supports batch_size = 1
    config['dataset']['serial_batches'] = True  # disable data shuffling; comment this line if results on randomly chosen images are needed.
    config['dataset']['no_flip'] = True    # no flip; comment this line if results on flipped images are needed.

    # override data augmentation
    config['dataset']['load_size'] = config['testing']['load_size']
    config['dataset']['crop_size'] = config['testing']['crop_size']
    config['dataset']['preprocess'] = config['testing']['preprocess']

    # add testing path
    config['testing']['test_img'] = None if args.test_img == '' else args.test_img
    config['testing']['test_video'] = None if args.test_video == '' else args.test_video
    config['testing']['test_folder'] = args.test_folder

    config['training']['pretrained_model'] = args.ckpt
    dataset = SuperDataset(config)
    dataloader = CustomDataLoader(config, dataset)

    model = create_model(config)      # create a model given opt.model and other options
    model.load_networks(0, ckpt=args.ckpt)

    model.eval()

    if args.overwrite_output_dir != '':
        save_path = args.overwrite_output_dir
    else:
        save_path = os.path.join(config['testing']['results_dir'], os.path.splitext(os.path.split(args.cfg_file)[1])[0],
                                config['common']['name'])

    if not os.path.exists(save_path):
        os.makedirs(save_path)

    def reduce(x):
        return reduce(x[0]) if not type(x) is str else x

    ext_name = config['testing']['image_format']
    use_input_format = (ext_name == 'input')
    output_video = (not config['testing']['test_video'] is None)
    vw_dict = {}
    video_paths = []

    for i, data in enumerate(tqdm(dataloader)):
        if i >= config['testing']['num_test']:  # only apply our model to opt.num_test images.
            break
        model.set_input(data)  # unpack data from data loader
        model.test()           # run inference
        visuals = model.get_current_visuals()  # get image results
        img_path = model.get_image_paths()     # get image paths

        # save result
        items = os.path.splitext(os.path.split(reduce(img_path))[1])
        img_fn = items[0]
        if use_input_format:
            ext_name = items[1][1:]

        for k, v in visuals.items():
            if not output_video:
                tensor2file(v, os.path.join(save_path, img_fn + '_' + k), ext_name)
            else:
                img = tensor2im(v)
                if not k in vw_dict:
                    h, w = img.shape[:2]
                    video_path = os.path.join(save_path, k + '_.mp4')
                    video_paths.append(video_path)
                    vw_dict[k] = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'mp4v'), 30.0, (w, h))
                vw_dict[k].write(img[:,:,::-1])

    for _, v in vw_dict.items():
        v.release()

    # convert to libx264
    for video_path in video_paths:
        os.system('ffmpeg -i {} -c:v libx264 {}'.format(video_path, video_path[:-5] + '.mp4'))
        os.system('rm {}'.format(video_path))

test.py
mmfs-paper/MMFS
import argparse
import os
import numpy as np
from PIL import Image
from tqdm import tqdm

import clip
import torch
from torchvision.transforms import Compose, Resize, ToTensor, Normalize, InterpolationMode
from models.style_based_pix2pixII_model import Stylizer, TrainingPhase

if __name__ == '__main__':
    # define & parse args
    parser = argparse.ArgumentParser(description='sp2pII test')
    parser.add_argument('--ckpt', type=str, default='') # ./checkpoints/watercolor_painting.pth
    parser.add_argument('--in_folder', type=str, default='') # ./example/source
    parser.add_argument('--out_folder', type=str, default='') # ./example/outputs/zero-shot/watercolor_painting
    parser.add_argument('--phase', type=int, default=3)
    parser.add_argument('--txt_prompt', type=str, default='') # watercolor painting
    parser.add_argument('--img_prompt', type=str, default='') # ./example/reference/04.png 
    parser.add_argument('--device', type=str, default='cuda:0')
    args = parser.parse_args()
    args.phase = TrainingPhase(args.phase)

    os.makedirs(args.out_folder, exist_ok=True)

    # init model
    state_dict = torch.load(args.ckpt, map_location='cpu')
    model = Stylizer(ngf=64, phase=args.phase, model_weights=state_dict['G_ema_model'])
    model.to(args.device)
    model.eval()
    model.requires_grad_(False)

    clip_model, img_preprocess = clip.load('ViT-B/32', device=args.device)
    clip_model.eval()
    clip_model.requires_grad_(False)

    # image transform for stylizer
    img_transform = Compose([
        Resize((512, 512), interpolation=InterpolationMode.LANCZOS),
        ToTensor(),
        Normalize([0.5], [0.5])
    ])

    # get clip features
    with torch.no_grad():
        if os.path.isfile(args.img_prompt):
            img = img_preprocess(Image.open(args.img_prompt)).unsqueeze(0).to(args.device)
            clip_feats = clip_model.encode_image(img)
        else:
            text = clip.tokenize(args.txt_prompt).to(args.device)
            clip_feats = clip_model.encode_text(text)
        clip_feats /= clip_feats.norm(dim=1, keepdim=True)

    # enum image files
    files = os.listdir(args.in_folder)
    for fn in tqdm(files):
        prefix, ext = os.path.splitext(fn)
        if not ext.lower() in ['.png', '.jpg', '.jpeg']:
            continue

        # load image & to tensor
        img = Image.open(os.path.join(args.in_folder, fn))
        if not img.mode == 'RGB':
            img = img.convert('RGB')
        img = img_transform(img).unsqueeze(0).to(args.device)

        # stylize it !
        with torch.no_grad():
            if args.phase == TrainingPhase.CLIP_MAPPING:
                res = model(img, clip_feats=clip_feats)

        # save image
        res = res.cpu().numpy()[0]
        res = np.transpose(res, (1, 2, 0)) * 0.5 + 0.5
        Image.fromarray((res * 255).astype(np.uint8)).save(os.path.join(args.out_folder, prefix + '.png'))

test_sp2pII.py
mmfs-paper/MMFS
import argparse
import time
import datetime
from data import CustomDataLoader
from data.super_dataset import SuperDataset

from models import create_model
from configs import parse_config
from utils.util import print_losses, check_path, make_grid, AverageMeter
from utils.data_utils import check_old_config_val_possible
import os
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.utils.tensorboard import SummaryWriter
import copy
import sys



def main():
    # parse arguments
    parser = argparse.ArgumentParser(description='Style Master')
    parser.add_argument('--cfg_file', type=str, default='') # ./exp/sp2pII-phase4.yaml
    
    args = parser.parse_args()

    # parse config
    config = parse_config(args.cfg_file)

    for group in config:
        print(group + ':')
        for k, v in config[group].items():
            print('  {}: {}'.format(k, v))

    # we want GPU ids match nvidia-smi output order, so do some manipulations here.
    # GPU ids need to always start from 0, but the system variable CUDA_VISIBLE_DEVICES can be set to e.g. GPU 2 and 3.
    gpu_string = ','.join(map(str, config['common']['gpu_ids']))
    gpu_ids_fix = list(range(len(config['common']['gpu_ids'])))
    os.environ['CUDA_VISIBLE_DEVICES'] = gpu_string
    config['common']['gpu_ids'] = gpu_ids_fix

    if config['training']['DDP']:
        num_nodes = config['training']['num_nodes']
        num_gpus = len(config['common']['gpu_ids'])
        config['training']['world_size'] = num_gpus * num_nodes
        os.environ['MASTER_ADDR'] = config['training']['DDP_address']
        os.environ['MASTER_PORT'] = config['training']['DDP_port']
        mp.spawn(train_val, nprocs=num_gpus, args=(config,))
    else:
        # under DP mode, shall set batch size to (actual effective batch size * num_gpu)
        config['dataset']['batch_size'] *= len(config['common']['gpu_ids'])
        if 'patch_batch_size' in config['dataset']:
            config['dataset']['patch_batch_size'] *= len(config['common']['gpu_ids'])
        train_val(None, config)

# GPU parameter is automatically filled when using DDP. It is an irrelevant placeholder if not using DDP.
def train_val(gpu, config):
    import torch
    if config['training']['val']:
        config_val = copy.deepcopy(config)
        config_val['common']['phase'] = 'val'

    if config['training']['DDP']:
        dist.init_process_group(
            backend='nccl',
            init_method='env://',
            world_size=config['training']['world_size'],
            rank=gpu
        )
        torch.cuda.set_device(gpu)  # DDP requirement

    # Dataset and dataloader construction.
    # prepare train data
    train_dataset = SuperDataset(config, shuffle=True, check_all_data=config['dataset']['check_all_data'],
                                 DDP_device=gpu)
    if config['dataset']['train_data'] == {}:
        print("old data config format detected. Converting to new format.")
        train_dataset.config = train_dataset.convert_old_config_to_new()

    train_dataset.static_data.load_static_data()
    train_dataset.static_data.create_transforms()
    if train_dataset.check_all_data:
        train_dataset.check_data()

    if len(train_dataset) == 0:
        if gpu == 0 or gpu is None:
            print("Train set has 0 data samples. Exiting.")
        sys.exit(0)

    _, train_video_dataset = train_dataset.split_data('count', 5, mode='copy_partial')

    # prepare val data
    force_use_train_data = False
    if config['training']['val']:
        val_dataset = SuperDataset(config_val)
        if len(config_val['dataset']['train_data']) == 0:
            if check_old_config_val_possible(config_val):
                val_dataset.convert_old_config_to_new()
                val_dataset.static_data.load_static_data()
            else:
                force_use_train_data = True
        else:
            val_dataset.static_data.load_static_data()


        if len(val_dataset) == 0 or force_use_train_data:
            if gpu == 0 or gpu is None:
                print("Validation set has 0 data samples. Using part of training data for validation.")
            validation_ratio = config['training']['val_percent']/100

            train_dataset, val_dataset = train_dataset.split_data('ratio', validation_ratio, mode='split')

            if len(val_dataset) == 0:
                print("There are too few training data to establish a validation set. "
                      "Use the training set as validation set.")
                val_dataset = copy.deepcopy(train_dataset)

        val_dataset.static_data.create_transforms()
        if val_dataset.check_all_data:
            val_dataset.check_data()

    # print dataset info
    if gpu == 0 or gpu is None:
        print("--------train dataset static data content-----------")
        for i, cnt in enumerate(train_dataset.static_data.len_of_groups):
            print("%s:   %d" % (train_dataset.static_data.group_names[i], cnt))
        print("----------------------------------------")

        if config['training']['val']:
            print("--------val dataset static data content-----------")
            for i, cnt in enumerate(val_dataset.static_data.len_of_groups):
                print("group %s:   %d" % (val_dataset.static_data.group_names[i], cnt))
            print("----------------------------------------")


    # prepare dataloaders
    train_dataloader = CustomDataLoader(config, train_dataset, DDP_gpu=gpu, drop_last=config['dataset']['drop_last'])
    if config['training']['val']:
        val_dataloader = CustomDataLoader(config_val, val_dataset, DDP_gpu=gpu, drop_last=config['dataset']['drop_last'])
    config_train_video = copy.deepcopy(config)
    config_train_video['dataset']['serial_batches'] = True
    config_train_video['batch_size'] = 1
    train_video_dataloader = CustomDataLoader(config_train_video, train_video_dataset, DDP_gpu=None, drop_last=False)

    if config['training']['DDP']:
        model = create_model(config, DDP_device=gpu)      # create a DDP model given opt.model and other options
        model.setup(config, DDP_device=gpu)  # regular setup: load and print networks; create schedulers
    else:
        model = create_model(config)  # create a model (singleGPU or dataparallel) given opt.model and other options
        model.setup(config)               # regular setup: load and print networks; create schedulers
    total_iters = model.total_iters                # the total number of training iterations

    # visualization setups
    if gpu == 0 or gpu is None:
        log_dir = os.path.join(config['training']['log_dir'], config['common']['name'] + '_' +
                               datetime.datetime.now().strftime("%Y_%m_%d__%H_%M_%S"))
        check_path(log_dir)
        writer = SummaryWriter(log_dir)
    train_loss_averages = {}
    val_loss_averages = {}

    if config['training']['epoch_as_iter']:
        iter_ub = config['training']['n_epochs'] + config['training']['n_epochs_decay']
        while total_iters < iter_ub:
            # Training loop
            if gpu == 0 or gpu is None:
                print('-------------------------------Training----------------------------------')
            model.train()
            quit = False
            for i, data in enumerate(train_dataloader):
                total_iters += 1
                if total_iters >= iter_ub:
                    quit = True
                    break
                model.curr_epoch = total_iters
                model.total_iters = total_iters

                model.set_input(data)
                model.optimize_parameters()

                losses = model.get_current_losses()
                for k, v in losses.items():
                    if k not in train_loss_averages:
                        train_loss_averages[k] = AverageMeter()
                    train_loss_averages[k].update(v)

                if (total_iters % config['training']['print_freq'] == 0) and ((gpu == 0) or (gpu is None)):
                    print_losses(total_iters, 1, 1, losses, train_loss_averages)
                    tmp = make_grid(model)
                    for j in range(0, len(tmp[0])):
                        img_grid, name = tmp[0][j], tmp[1][j]
                        if config['training']['use_new_log']:
                            writer.add_image('Training/' + name, img_grid, total_iters)
                        else:
                            writer.add_image('iteration ' + str(total_iters) + ' training ' + name, img_grid)
                    for k, v in losses.items():
                        writer.add_scalar('Training/' + k, train_loss_averages[k].avg(), total_iters)

                if (total_iters % config['training']['save_latest_freq'] == 0) and ((gpu == 0) or (gpu is None)):
                    print('saving the latest model (total_iters %d)' % total_iters)
                    model.save_networks('latest')

                if (total_iters % config['training']['save_epoch_freq'] == 0) and ((gpu == 0) or (gpu is None)):
                    print('saving the model at the end of iters %d' % total_iters)
                    model.save_networks('latest')
                    #model.save_networks(total_iters)

                model.update_learning_rate()
                for k, v in losses.items():
                    train_loss_averages[k].clear()

            if quit:
                exit(0)

            if config['training']['val']:
                # Validation loop
                if gpu == 0 or gpu is None:
                    print('-------------------------------Validating----------------------------------')
                model.eval()
                for i, data in enumerate(val_dataloader):
                    with torch.no_grad():
                        model.set_input(data)
                        model.eval_step()

                        losses = model.get_current_losses()
                        for k, v in losses.items():
                            if k not in val_loss_averages:
                                val_loss_averages[k] = AverageMeter()
                            val_loss_averages[k].update(v)

                if gpu == 0 or gpu is None:
                    print_losses(total_iters, 1, 1, losses, val_loss_averages, mode='validating')
                    tmp = make_grid(model)
                    for j in range(0, len(tmp[0])):
                        img_grid, name = tmp[0][j], tmp[1][j]
                        if config['training']['use_new_log']:
                            writer.add_image('Validation/' + name, img_grid, total_iters)
                        else:
                            writer.add_image('iteration ' + str(total_iters) + ' validating ' + name, img_grid)
                    for k, v in losses.items():
                        writer.add_scalar('Validation/' + k, val_loss_averages[k].avg(), total_iters)

                main_loss = 'G' if 'G' in losses else model.loss_names[0]
                if val_loss_averages[main_loss].avg() < model.best_val_loss and (gpu == 0 or gpu is None):
                    model.best_val_loss = val_loss_averages[main_loss].avg()
                    print('New validation best loss. saving the model.')
                    model.save_networks('', val_loss=model.best_val_loss)

                for k, v in losses.items():
                    val_loss_averages[k].clear()

            if config['training']['save_training_progress']:
                # produce images on the same images every epoch to visualize how training is progressing.
                if gpu == 0 or gpu is None:
                    for i, data in enumerate(train_video_dataloader):
                        with torch.no_grad():
                            model.set_input(data)
                            model.forward()
                            tmp = make_grid(model)
                            for j in range(0, len(tmp[0])):
                                img_grid, name = tmp[0][j], tmp[1][j]
                                if config['training']['use_new_log']:
                                    writer.add_image('Training Video/' + name + ' ' + str(i), img_grid, total_iters)
                                else:
                                    writer.add_image('epoch 0 iteration ' + str(total_iters) + ' training_video ' + name + ' ' + str(i), img_grid)


    for epoch in range(model.curr_epoch, config['training']['n_epochs'] + config['training']['n_epochs_decay'] + 1):    # outer loop for different epochs; we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>
        epoch_start_time = time.time()  # timer for entire epoch
        iter_data_time = time.time()    # timer for data loading per iteration
        epoch_iter = 0                  # the number of training iterations in current epoch, reset to 0 every epoch
        model.curr_epoch = epoch

        # Training loop
        if gpu == 0 or gpu is None:
            print('-------------------------------Training----------------------------------')
        model.train()
        for i, data in enumerate(train_dataloader):  # inner loop within one epoch
            iter_start_time = time.time()  # timer for computation per iteration

            total_iters += 1
            epoch_iter += 1
            model.set_input(data)         # unpack data from dataset and apply preprocessing
            model.optimize_parameters()   # calculate loss functions, get gradients, update network weights

            losses = model.get_current_losses()
            for k, v in losses.items():
                if k not in train_loss_averages:
                    train_loss_averages[k] = AverageMeter()
                train_loss_averages[k].update(v)

            if epoch_iter % config['training']['print_freq'] == 0 and ((gpu == 0) or (gpu is None)):
                print_losses(epoch, epoch_iter, (len(train_dataset) + config['dataset']['batch_size'] - 1) // config['dataset']['batch_size'], losses, train_loss_averages)
                tmp = make_grid(model)
                for j in range(0, len(tmp[0])):
                    img_grid, name = tmp[0][j], tmp[1][j]
                    if config['training']['use_new_log']:
                        writer.add_image('Training/' + name, img_grid, total_iters)
                    else:
                        writer.add_image('epoch ' + str(epoch) + ' iteration ' + str(total_iters) + ' training ' + name, img_grid)
                for k, v in losses.items():
                    writer.add_scalar('Training/' + k, train_loss_averages[k].avg(), total_iters)

            if total_iters % config['training']['save_latest_freq'] == 0 and ((gpu == 0) or (gpu is None)):   # cache our latest model every <save_latest_freq> iterations
                print('saving the latest model (epoch %d, total_iters %d)' % (epoch, total_iters))
                model.save_networks('latest')

            iter_data_time = time.time()

        if gpu == 0 or gpu is None:
            print_losses(epoch, epoch_iter, (len(train_dataset) + config['dataset']['batch_size'] - 1) // config['dataset']['batch_size'], losses, train_loss_averages)
            print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, config['training']['n_epochs'] + config['training']['n_epochs_decay'], time.time() - epoch_start_time))
            tmp = make_grid(model)
            for j in range(0, len(tmp[0])):
                img_grid, name = tmp[0][j], tmp[1][j]
                if config['training']['use_new_log']:
                    writer.add_image('Training/' + name, img_grid, total_iters)
                else:
                    writer.add_image('epoch ' + str(epoch) + ' iteration ' + str(total_iters) + ' training ' + name, img_grid)
            for k, v in losses.items():
                writer.add_scalar('Training/' + k, train_loss_averages[k].avg(), total_iters)

        if epoch % config['training']['save_epoch_freq'] == 0 and ((gpu == 0) or (gpu is None)):              # cache our model every <save_epoch_freq> epochs
            print('saving the model at the end of epoch %d, iters %d' % (epoch, total_iters))
            model.save_networks('latest')
            model.save_networks(epoch)

        model.update_learning_rate()  # update learning rates at the end of every epoch.
        for k, v in losses.items():
            train_loss_averages[k].clear()

        if config['training']['val']:
            # Validation loop
            if gpu == 0 or gpu is None:
                print('-------------------------------Validating----------------------------------')
            model.eval()
            epoch_iter = 0
            for i, data in enumerate(val_dataloader):  # inner loop within one epoch
                epoch_iter += 1

                with torch.no_grad():
                    model.set_input(data)
                    model.eval_step()

                    losses = model.get_current_losses()
                    for k, v in losses.items():
                        if k not in val_loss_averages:
                            val_loss_averages[k] = AverageMeter()
                        val_loss_averages[k].update(v)

            if gpu == 0 or gpu is None:
                print_losses(epoch, epoch_iter, (len(val_dataset) + config['dataset']['batch_size'] - 1)
                             // config['dataset']['batch_size'], losses, val_loss_averages, mode='validating')
                tmp = make_grid(model)
                for j in range(0, len(tmp[0])):
                    img_grid, name = tmp[0][j], tmp[1][j]
                    if config['training']['use_new_log']:
                        writer.add_image('Validation/' + name, img_grid, total_iters)
                    else:
                        writer.add_image('epoch ' + str(epoch) + ' iteration ' + str(total_iters) + ' validating ' + name, img_grid)
                for k, v in losses.items():
                    writer.add_scalar('Validation/' + k, val_loss_averages[k].avg(), total_iters)

            main_loss = 'G' if 'G' in losses else model.loss_names[0]
            if val_loss_averages[main_loss].avg() < model.best_val_loss and gpu == 0 or gpu is None:
                model.best_val_loss = val_loss_averages[main_loss].avg()
                print('New validation best loss. saving the model.')
                model.save_networks('', val_loss=model.best_val_loss)

            for k, v in losses.items():
                val_loss_averages[k].clear()

        if config['training']['save_training_progress']:
            # produce images on the same images every epoch to visualize how training is progressing.
            if gpu == 0 or gpu is None:
                for i, data in enumerate(train_video_dataloader):
                    with torch.no_grad():
                        model.set_input(data)
                        model.forward()
                        tmp = make_grid(model)
                        for j in range(0, len(tmp[0])):
                            img_grid, name = tmp[0][j], tmp[1][j]
                            if config['training']['use_new_log']:
                                writer.add_image('Training Video/' + name + ' ' + str(i), img_grid, total_iters)
                            else:
                                writer.add_image('epoch ' + str(epoch) + ' iteration ' + str(total_iters) + ' training_video ' + name + ' ' + str(i), img_grid)

    # If we shut down process now, writer could save incomplete data. Wait a bit to let it finish.
    time.sleep(5)

if __name__ == '__main__':
    main()

train.py
mmfs-paper/MMFS
from packaging import version
import random
import numpy as np
from PIL import Image, ImageFilter, ImageOps
from torchvision.transforms.transforms import Lambda, Compose
from torchvision.transforms import functional as F
from collections.abc import Iterable
import torch, torchvision
import numbers
import copy

if version.parse(torchvision.__version__) <= version.parse('0.7.0'):
    from torchvision.transforms.transforms import _get_image_size

def check_input_type_perform_action(input, type, action, *args, **kwargs):
    output = input
    if isinstance(input, list):
        for i in range(0, len(input)):
            if type is None:
                if input[i] is not None:  # do not combine with last line, to avoid calling isinstance on None.
                    output[i] = action(input[i], *args, **kwargs)
            elif isinstance(input[i], type):
                output[i] = action(input[i], *args, **kwargs)
    elif type is None:
        if input is not None:
            output = action(input, *args, **kwargs)
    elif isinstance(input, type):
        output = action(input, *args, **kwargs)
    return output


"""
Most of these functions are imported from torchvision.transforms.transforms and edited to support 2 or more inputs.
"""

class JointCompose(object):
    """
    Composes several transforms together.
    """

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, input1, input2):
        for t in self.transforms:
            input1, input2 = t(input1, input2)
        return input1, input2


class Grayscale(object):

    def __init__(self, input1_output_channels=1, input2_output_channels=1):
        self.input1_output_channels = input1_output_channels
        self.input2_output_channels = input2_output_channels

    def __call__(self, input1, input2):
        output1 = F.to_grayscale(input1, num_output_channels=self.input1_output_channels) if self.input1_output_channels == 1 else input1
        output2 = check_input_type_perform_action(input2, Image.Image, F.to_grayscale, num_output_channels=self.input2_output_channels) \
            if self.input2_output_channels == 1 else input2
        return output1, output2


class Resize(object):

    def __init__(self, size, interpolation=Image.BILINEAR):
        assert isinstance(size, int) or (isinstance(size, Iterable) and len(size) == 2)
        self.size = size
        self.interpolation = interpolation

    def __call__(self, input1, input2):
        output1 = F.resize(input1, self.size, self.interpolation)
        output2 = check_input_type_perform_action(input2, Image.Image, F.resize, self.size, self.interpolation)
        return output1, output2


class ScaleWidth:

    def __init__(self, target_size, method=Image.BICUBIC):
        self.target_size = target_size
        self.method = method

    def scalewidth(self, img):
        ow, oh = img.size
        w = self.target_size
        h = int(self.target_size * oh / ow)
        img_resized = img.resize((w, h), self.method)

        if h > w:
            # if resized image's height is larger than its width, crop the center
            left = 0
            top = h // 2 - self.target_size // 2
            right = self.target_size
            bottom = top + self.target_size
            img_resized = img_resized.crop((left, top, right, bottom))
        elif h < w:
            # pad the heights
            delta_w = self.target_size - w
            delta_h = self.target_size - h
            padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
            img_resized = ImageOps.expand(img_resized, padding)

        return img_resized

    def __call__(self, input1, input2):
        output1 = self.scalewidth(input1)
        output2 = check_input_type_perform_action(input2, Image.Image, self.scalewidth)
        return output1, output2


class RandomCrop(object):

    def __init__(self, size, padding=None, pad_if_needed=False, fill=0, padding_mode='constant'):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
        self.padding = padding
        self.pad_if_needed = pad_if_needed
        self.fill = fill
        self.padding_mode = padding_mode

    @staticmethod
    def get_params(img, output_size):
        if version.parse(torchvision.__version__) <= version.parse('0.7.0'):
            w, h = _get_image_size(img)
        else:
            w, h = F._get_image_size(img)
        th, tw = output_size
        if w == tw and h == th:
            return 0, 0, h, w

        i = random.randint(0, h - th)
        j = random.randint(0, w - tw)
        return i, j, th, tw

    def pad(self, img):
        if self.padding is not None:
            img = F.pad(img, self.padding, self.fill, self.padding_mode)

        # pad the width if needed
        if self.pad_if_needed and img.size[0] < self.size[1]:
            img = F.pad(img, (self.size[1] - img.size[0], 0), self.fill, self.padding_mode)
        # pad the height if needed
        if self.pad_if_needed and img.size[1] < self.size[0]:
            img = F.pad(img, (0, self.size[0] - img.size[1]), self.fill, self.padding_mode)

        return img

    def get_crop_range(self, img):
        return self.get_params(img, self.size)

    def pad_and_crop(self, input, i, j, h, w):
        return F.crop(self.pad(input), i, j, h, w)

    def __call__(self, input1, input2):
        output1 = self.pad(input1)
        i, j, h, w = self.get_crop_range(output1)
        output1 = F.crop(output1, i, j, h, w)
        output2 = check_input_type_perform_action(input2, Image.Image, self.pad_and_crop, i, j, h, w)
        return output1, output2


class Crop:

    def __init__(self, pos, size):
        self.pos = pos
        self.size = size

    def crop(self, img):
        ow, oh = img.size
        x1, y1 = self.pos
        tw = th = self.size
        if (ow > tw or oh > th):
            return img.crop((x1, y1, x1 + tw, y1 + th))
        return img

    def __call__(self, input1, input2):
        output1 = self.crop(input1)
        output2 = check_input_type_perform_action(input2, Image.Image, self.crop)
        return output1, output2


class ColorJitter(object):

    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        self.brightness = self._check_input(brightness, 'brightness')
        self.contrast = self._check_input(contrast, 'contrast')
        self.saturation = self._check_input(saturation, 'saturation')
        self.hue = self._check_input(hue, 'hue', center=0, bound=(-0.5, 0.5),
                                     clip_first_on_zero=False)

    def _check_input(self, value, name, center=1, bound=(0, float('inf')), clip_first_on_zero=True):
        if isinstance(value, numbers.Number):
            if value < 0:
                raise ValueError("If {} is a single number, it must be non negative.".format(name))
            value = [center - value, center + value]
            if clip_first_on_zero:
                value[0] = max(value[0], 0)
        elif isinstance(value, (tuple, list)) and len(value) == 2:
            if not bound[0] <= value[0] <= value[1] <= bound[1]:
                raise ValueError("{} values should be between {}".format(name, bound))
        else:
            raise TypeError("{} should be a single number or a list/tuple with lenght 2.".format(name))

        # if value is 0 or (1., 1.) for brightness/contrast/saturation
        # or (0., 0.) for hue, do nothing
        if value[0] == value[1] == center:
            value = None
        return value

    @staticmethod
    def get_params(brightness, contrast, saturation, hue):
        transforms = []

        if brightness is not None:
            brightness_factor = random.uniform(brightness[0], brightness[1])
            transforms.append(Lambda(lambda img: F.adjust_brightness(img, brightness_factor)))

        if contrast is not None:
            contrast_factor = random.uniform(contrast[0], contrast[1])
            transforms.append(Lambda(lambda img: F.adjust_contrast(img, contrast_factor)))

        if saturation is not None:
            saturation_factor = random.uniform(saturation[0], saturation[1])
            transforms.append(Lambda(lambda img: F.adjust_saturation(img, saturation_factor)))

        if hue is not None:
            hue_factor = random.uniform(hue[0], hue[1])
            transforms.append(Lambda(lambda img: F.adjust_hue(img, hue_factor)))

        random.shuffle(transforms)
        transform = Compose(transforms)

        return transform

    def __call__(self, input1, input2):
        transform = self.get_params(self.brightness, self.contrast,
                                    self.saturation, self.hue)
        output1 = transform(input1)
        output2 = check_input_type_perform_action(input2, Image.Image, transform)
        return output1, output2


class RandomAffine(object):

    def __init__(self, degrees, translate=None, scale=None, shear=None, resample=False, fillcolor=0):
        if isinstance(degrees, numbers.Number):
            if degrees < 0:
                raise ValueError("If degrees is a single number, it must be positive.")
            self.degrees = (-degrees, degrees)
        else:
            assert isinstance(degrees, (tuple, list)) and len(degrees) == 2, \
                "degrees should be a list or tuple and it must be of length 2."
            self.degrees = degrees

        if translate is not None:
            assert isinstance(translate, (tuple, list)) and len(translate) == 2, \
                "translate should be a list or tuple and it must be of length 2."
            for t in translate:
                if not (0.0 <= t <= 1.0):
                    raise ValueError("translation values should be between 0 and 1")
        self.translate = translate

        if scale is not None:
            assert isinstance(scale, (tuple, list)) and len(scale) == 2, \
                "scale should be a list or tuple and it must be of length 2."
            for s in scale:
                if s <= 0:
                    raise ValueError("scale values should be positive")
        self.scale = scale

        if shear is not None:
            if isinstance(shear, numbers.Number):
                if shear < 0:
                    raise ValueError("If shear is a single number, it must be positive.")
                self.shear = (-shear, shear)
            else:
                assert isinstance(shear, (tuple, list)) and \
                    (len(shear) == 2 or len(shear) == 4), \
                    "shear should be a list or tuple and it must be of length 2 or 4."
                # X-Axis shear with [min, max]
                if len(shear) == 2:
                    self.shear = [shear[0], shear[1], 0., 0.]
                elif len(shear) == 4:
                    self.shear = [s for s in shear]
        else:
            self.shear = shear

        self.resample = resample
        self.fillcolor = fillcolor

    @staticmethod
    def get_params(degrees, translate, scale_ranges, shears, img_size):
        angle = random.uniform(degrees[0], degrees[1])
        if translate is not None:
            max_dx = translate[0] * img_size[0]
            max_dy = translate[1] * img_size[1]
            translations = (np.round(random.uniform(-max_dx, max_dx)),
                            np.round(random.uniform(-max_dy, max_dy)))
        else:
            translations = (0, 0)

        if scale_ranges is not None:
            scale = random.uniform(scale_ranges[0], scale_ranges[1])
        else:
            scale = 1.0

        if shears is not None:
            if len(shears) == 2:
                shear = [random.uniform(shears[0], shears[1]), 0.]
            elif len(shears) == 4:
                shear = [random.uniform(shears[0], shears[1]),
                         random.uniform(shears[2], shears[3])]
        else:
            shear = 0.0

        return angle, translations, scale, shear

    def __call__(self, input1, input2):
        params = self.get_params(self.degrees, self.translate, self.scale, self.shear, input1.size)
        output1 = F.affine(input1, *params, resample=self.resample, fillcolor=self.fillcolor)
        output2 = check_input_type_perform_action(input2, Image.Image, F.affine, *params, resample=self.resample, fillcolor=self.fillcolor)
        return output1, output2


class RandomRotation(object):
    def __init__(self, degrees, resample=False, expand=False, center=None, fill=None):
        if isinstance(degrees, numbers.Number):
            if degrees < 0:
                raise ValueError("If degrees is a single number, it must be positive.")
            self.degrees = (-degrees, degrees)
        else:
            if len(degrees) != 2:
                raise ValueError("If degrees is a sequence, it must be of len 2.")
            self.degrees = degrees

        self.resample = resample
        self.expand = expand
        self.center = center
        self.fill = fill

    @staticmethod
    def get_params(degrees):
        angle = random.uniform(degrees[0], degrees[1])
        return angle

    def __call__(self, input1, input2):
        angle = self.get_params(self.degrees)
        output1 = F.rotate(input1, angle, self.resample, self.expand, self.center, self.fill)
        output2 = check_input_type_perform_action(input2, Image.Image, F.rotate, angle, self.resample, self.expand, self.center, self.fill)
        return output1, output2


class RandomBlur:
    def __init__(self, blur_chance):
        self.blur_chance = blur_chance

    def get_params(self):
        if self.blur_chance > random.random():
            kernel = random.randint(3, 12)
            while kernel % 2 != 1:
                kernel = random.randint(3, 12)
        else:
            kernel = None
        return kernel

    def blur(self, image, kernel):
        image = image.filter(ImageFilter.GaussianBlur(radius=kernel))
        return image

    def __call__(self, input1, input2):
        kernel = self.get_params()
        if kernel is None:
            return input1, input2
        else:
            output1 = self.blur(input1, kernel)
            output2 = check_input_type_perform_action(input2, Image.Image, self.blur, kernel)
            return output1, output2


class NoiseTransform:
    """code is partly from http://www.xiaoliangbai.com/2016/09/09/more-on-image-noise-generation and edited by Oliver."""

    def __init__(self, noise_type):
        self.noise_type = noise_type

    def get_params(self, image):
        params = []
        image_np = np.array(image)
        row, col, ch = image_np.shape
        if random.random() < 0.5:
            return None
        if self.noise_type == "gauss":
            mean = 0.0
            std = random.uniform(0.001, 0.3)
            gauss = np.random.normal(mean, std, (row, col, ch))
            gauss = gauss.reshape(row, col, ch)
            params.append(gauss)
            return params
        elif self.noise_type == "s&p":
            s_vs_p = 0.5
            amount = random.uniform(0.001, 0.01)

            # Generate Salt '1' noise
            num_salt = np.ceil(amount * image_np.size * s_vs_p)
            coords = [np.random.randint(0, i - 1, int(num_salt))
                      for i in image_np.shape]
            coords[2] = np.random.randint(0, 3, int(num_salt))
            params.append(copy.deepcopy(coords))

            # Generate Pepper '0' noise
            num_pepper = np.ceil(amount * image_np.size * (1. - s_vs_p))
            coords = [np.random.randint(0, i - 1, int(num_pepper))
                      for i in image_np.shape]
            params.append(copy.deepcopy(coords))
            return params
        elif self.noise_type == "poisson":
            noisy = np.random.poisson(image_np)
            params.append(noisy)
            return params
        elif self.noise_type == "speckle":
            factor = random.uniform(0.01, 0.4)
            gauss = np.random.randn(row, col, ch)
            gauss = gauss.reshape(row, col, ch) * factor
            params.append(gauss)
            return params
        elif self.noise_type == "band":
            smaller_dim = min(col, row)
            num_bands = random.randrange(smaller_dim // 2, smaller_dim)
            scale = random.uniform(1.0, 10.0)

            offset = np.zeros(image_np.shape).astype(np.float64)

            # horizontal branding
            num_list = list(range(image.width))  # list of integers from 0 to image width-1
            # adjust this boundaries to fit your needs
            random.shuffle(num_list)
            horizontal_bands = num_list[:num_bands]
            for w in horizontal_bands:
                offset[w, :, :] += random.uniform(-1, 1) * scale

            # vertical branding
            num_list = list(range(image.height))  # list of integers from 0 to image height-1
            # adjust this boundaries to fit your needs
            random.shuffle(num_list)
            vertical_bands = num_list[:num_bands]
            for h in vertical_bands:
                offset[:, h, :] += random.uniform(-1, 1) * scale

            params.append(offset)
            return params
        else:
            return params

    def apply(self, image, params):
        """
        image: ndarray (input image data. It will be converted to float)
        """
        if params is None:
            return image
        image_np = np.array(image)
        if self.noise_type == "gauss":
            gauss = params[0]
            noisy = image_np + image_np * gauss
            noisy = np.clip(noisy, 0, 255)
            return Image.fromarray(noisy.astype('uint8'))
        elif self.noise_type == "s&p":
            out = image_np
            # Generate Salt '1' noise
            coords = params[0]
            out[tuple(coords)] = 255
            # Generate Pepper '0' noise
            coords = params[1]
            out[tuple(coords)] = 0
            out = np.clip(out, 0, 255)
            return Image.fromarray(out.astype('uint8'))
        elif self.noise_type == "poisson":
            noisy = params[0]
            noisy = np.clip(noisy, 0, 255)
            return Image.fromarray(noisy.astype('uint8'))
        elif self.noise_type == "speckle":
            gauss = params[0]
            noisy = image_np + image_np * gauss
            noisy = np.clip(noisy, 0, 255)
            return Image.fromarray(noisy.astype('uint8'))
        elif self.noise_type == "band":
            offset = params[0]
            noisy = image_np + offset
            noisy = np.clip(noisy, 0, 255)
            return Image.fromarray(noisy.astype('uint8'))
        else:
            return image

    def __call__(self, input1, input2):
        params = self.get_params(input1)
        output1 = self.apply(input1, params)
        output2 = check_input_type_perform_action(input2, Image.Image, self.apply, params)
        return output1, output2


class MakePower2:
    def __init__(self, base, method=Image.BICUBIC):
        self.base = base
        self.method = method
        self.print_size_warning = PrintSizeWarning()

    def apply(self, img):
        ow, oh = img.size
        h = int(round(oh / self.base) * self.base)
        w = int(round(ow / self.base) * self.base)
        if h == oh and w == ow:
            return img

        self.print_size_warning(ow, oh, w, h)
        return img.resize((w, h), self.method)

    def __call__(self, input1, input2):
        output1 = self.apply(input1)
        output2 = check_input_type_perform_action(input2, Image.Image, self.apply)
        return output1, output2


class RandomHorizontalFlip(object):
    """Horizontally flip the given PIL Image randomly with a given probability.

    Args:
        p (float): probability of the image being flipped. Default value is 0.5
    """

    def __init__(self, p=0.5):
        self.p = p

    def get_params(self):
        if random.random() < self.p:
            return True
        else:
            return False

    def __call__(self, input1, input2):
        flip = self.get_params()
        if flip:
            output1 = F.hflip(input1)
            output2 = check_input_type_perform_action(input2, Image.Image, F.hflip)
        else:
            output1, output2 = input1, input2
        return output1, output2


class Flip:
    def __init__(self, flip):
        self.flip = flip

    def transpose(self, input):
        return input.transpose(Image.FLIP_LEFT_RIGHT)

    def __call__(self, input1, input2):
        if self.flip:
            output1 = input1.transpose(Image.FLIP_LEFT_RIGHT)
            output2 = check_input_type_perform_action(input2, Image.Image, self.transpose)
        else:
            output1, output2 = input1, input2
        return output1, output2


class ToTensor(object):
    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.

    Converts a PIL Image or numpy.ndarray (H x W x C) in the range
    [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]
    if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1)
    or if the numpy.ndarray has dtype = np.uint8

    In the other cases, tensors are returned without scaling.
    """

    def __call__(self, input1, input2):
        output1 = F.to_tensor(input1)
        output2 = check_input_type_perform_action(input2, None, F.to_tensor)
        return output1, output2


class Normalize(object):
    """Normalize a tensor image with mean and standard deviation.
    Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform
    will normalize each channel of the input ``torch.*Tensor`` i.e.
    ``output[channel] = (input[channel] - mean[channel]) / std[channel]``

    .. note::
        This transform acts out of place, i.e., it does not mutate the input tensor.

    Args:
        mean (sequence): Sequence of means for each channel.
        std (sequence): Sequence of standard deviations for each channel.
        inplace(bool,optional): Bool to make this operation in-place.

    """

    def __init__(self, first_input_mean, first_input_std, second_input_mean=None, second_input_std=None, inplace=False):
        self.first_input_mean = first_input_mean
        self.first_input_std = first_input_std
        self.second_input_mean = second_input_mean if second_input_mean is not None else first_input_mean
        self.second_input_std = second_input_std if second_input_std is not None else first_input_std
        self.inplace = inplace

    def __call__(self, tensor1, tensor2):
        """
        Args:
            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.

        Returns:
            Tensor: Normalized Tensor image.
        """
        output1 = F.normalize(tensor1, self.first_input_mean, self.first_input_std, self.inplace)
        output2 = check_input_type_perform_action(tensor2, None, F.normalize, self.second_input_mean, self.second_input_std, self.inplace)
        return output1, output2


class PrintSizeWarning:
    def __init__(self):
        self.has_printed = False

    def __call__(self, ow, oh, w, h):
        if not self.has_printed:
            print("The image size needs to be a multiple of 4. "
                "The loaded image size was (%d, %d), so it was adjusted to "
                "(%d, %d). This adjustment will be done to all images "
                "whose sizes are not multiples of 4" % (ow, oh, w, h))
            self.has_printed = True


class ImagePathToImage:
    """Convert an image path to an image.

    Parameters:
        filename  --  the input file path.
    """

    def load_img(self, path):
        return Image.open(path).convert('RGB')

    def __call__(self, filename1, filename2):
        img1 = self.load_img(filename1)
        img2 = check_input_type_perform_action(filename2, None, self.load_img)
        return img1, img2


class NumpyToTensor:
    """Convert a numpy array to a tensor.

    Parameters:
        filename  --  the input file path.
    """

    def load_numpy(self, filename):
        npy = np.load(filename)
        if isinstance(npy, np.lib.npyio.NpzFile):
            npy = npy['data']
        if len(npy.shape) == 2:
            npy = np.tile(npy, (1, 1, 1))
        else:
            npy = np.transpose(npy, (2, 0, 1))
        return torch.from_numpy(npy).float()

    def __call__(self, filename1, filename2):
        tensor1 = self.load_numpy(filename1)
        tensor2 = check_input_type_perform_action(filename2, None, self.load_numpy)
        return tensor1, tensor2

utils/augmentation.py
mmfs-paper/MMFS
import utils.augmentation as transforms
import sys
from PIL import Image
import numpy as np
import random
import cv2
import os


class Transforms():
    def __init__(self, config, input_grayscale_flag=False, output_grayscale_flag=False, method=Image.BICUBIC, convert=True):
        self.config = config
        self.input_grayscale_flag = input_grayscale_flag
        self.output_grayscale_flag = output_grayscale_flag
        self.method = method
        self.convert = convert
        self.transform_list = []

    def create_transforms_from_list(self, preprocess_list):
        if self.input_grayscale_flag:
            if self.output_grayscale_flag:
                self.transform_list.append(transforms.Grayscale())
            else:
                self.transform_list.append(transforms.Grayscale(1, 3))
        elif self.output_grayscale_flag:
            self.transform_list.append(transforms.Grayscale(3, 1))

        if 'resize' in preprocess_list:
            if self.config['dataset']['load_size'] < 10000:
                osize = [self.config['dataset']['load_size'], self.config['dataset']['load_size']]
            else:
                osize = [self.config['dataset']['load_size'] // 10000, self.config['dataset']['load_size'] % 10000]
            self.transform_list.append(transforms.Resize(osize, self.method))
        elif 'scale_width' in preprocess_list:
            self.transform_list.append(transforms.ScaleWidth(self.config['dataset']['load_size'], self.method))

        if 'crop' in preprocess_list:
            if 'crop_pos' in self.config['dataset']:
                self.transform_list.append(transforms.Crop(self.config['dataset']['crop_pos'], self.config['dataset']['crop_size']))
            else:
                self.transform_list.append(transforms.RandomCrop(self.config['dataset']['crop_size']))

        if 'add_lighting' in preprocess_list:
            self.transform_list.append(transforms.ColorJitter())

        if 'random_affine' in preprocess_list:
            self.transform_list.append(transforms.RandomAffine(20, translate=(0.2, 0.2), scale=(0.2, 0.2)))

        if 'random_rotate' in preprocess_list:
            self.transform_list.append(transforms.RandomRotation(20))

        if 'random_blur' in preprocess_list:
            self.transform_list.append(transforms.RandomBlur(0.2))

        if 'add_gauss_noise' in preprocess_list:
            self.transform_list.append(transforms.NoiseTransform("gauss"))
        if 'add_s&p_noise' in preprocess_list:
            self.transform_list.append(transforms.NoiseTransform("s&p"))
        if 'add_poisson_noise' in preprocess_list:
            self.transform_list.append(transforms.NoiseTransform("poisson"))
        if 'add_speckle_noise' in preprocess_list:
            self.transform_list.append(transforms.NoiseTransform("speckle"))
        if 'add_band_noise' in preprocess_list:
            self.transform_list.append(transforms.NoiseTransform("band"))

        if preprocess_list == 'none':
            self.transform_list.append(transforms.MakePower2(base=4, method=self.method))

        if not self.config['dataset']['no_flip']:
            if 'flip' in self.config['dataset']:
                self.transform_list.append(transforms.Flip(self.config['dataset']['flip']))
            else:
                self.transform_list.append(transforms.RandomHorizontalFlip())

        if self.convert:
            self.transform_list += [transforms.ToTensor()]
            if self.input_grayscale_flag:
                if self.output_grayscale_flag:
                    self.transform_list += [transforms.Normalize((0.5,), (0.5,))]
                else:
                    self.transform_list += [transforms.Normalize((0.5,), (0.5,), (0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
            elif self.output_grayscale_flag:
                self.transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), (0.5,), (0.5,))]
            else:
                self.transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]

    def get_transforms(self):
        return self.transform_list

    def compose_transforms(self):
        return transforms.JointCompose(self.transform_list)


def check_create_shuffled_order(data_list, order):
    # returns the order used to shuffle all paired data
    if order is None:  # Does not perform shuffling. Return normal order.
        order = np.arange(0, len(data_list)).tolist()
    else:
        if not isinstance(order, list):  # order is -1, which means has not been created.
            order = np.arange(0, len(data_list)).tolist()  # create the shuffle order.
            random.shuffle(order)
        # otherwise shuffle order already exists and we do nothing.
    return order


def check_equal_length(list1, list2, data):
    if len(list1) != len(list2):
        print("different length in paired data types. Please double check your data.")
        print("length of current data type: ", len(list1))
        print("----------------current lengths for all data types-------------------")
        for k, v in data.items():
            print("%s:   %d" % (k, len(v)))
        sys.exit()

def check_img_loaded(path):
    img = cv2.imread(path)
    if img is None or img.size == 0:
        print("image loading failed for " + path + '. Please double check.')
        return False
    return True

def check_numpy_loaded(path):
    try:
        arr = np.load(path)
    except Exception as e:
        print("numpy loading failed for " + path + '. Please double check.')
        return False
    return True

# custom, paired, numpy_paired, unpaired, numpy_unpaired, landmark
def check_old_config_val_possible(old_style_config):
    for data_type in old_style_config['dataset']['data_type']:
        if data_type == 'custom':
            if old_style_config['dataset']['custom_val_data'] == {}:
                return False
        elif data_type == 'paired' or data_type == 'numpy_paired':
            keyword = ''.join(data_type.split('_'))
            filelist_not_exist = old_style_config['dataset']['paired_val_filelist'] == ''
            filefolders_not_exist = old_style_config['dataset']['paired_valA_folder'] == '' or \
                old_style_config['dataset']['paired_valB_folder'] == ''
            dataroot_contains_no_val_folders = not os.path.exists(
                os.path.join(old_style_config['dataset']['dataroot'], 'val' + keyword + 'A')) \
                                              or not os.path.exists(
                os.path.join(old_style_config['dataset']['dataroot'], 'val' + keyword + 'B'))
            if filelist_not_exist and filefolders_not_exist and dataroot_contains_no_val_folders:
                return False
        elif data_type == 'unpaired' or data_type == 'numpy_unpaired':
            keyword = ''.join(data_type.split('_'))
            filelist_not_exist = old_style_config['dataset']['unpaired_valA_filelist'] == '' or \
                                 old_style_config['dataset']['unpaired_valB_filelist'] == ''
            filefolders_not_exist = old_style_config['dataset']['unpaired_valA_folder'] == '' or \
                                    old_style_config['dataset']['unpaired_valB_folder'] == ''
            dataroot_contains_no_val_folders = not os.path.exists(
                os.path.join(old_style_config['dataset']['dataroot'], 'val' + keyword + 'A')) \
                                               or not os.path.exists(
                os.path.join(old_style_config['dataset']['dataroot'], 'val' + keyword + 'B'))
            if filelist_not_exist and filefolders_not_exist and dataroot_contains_no_val_folders:
                return False
        elif data_type == 'landmark':
            filelist_not_exist = old_style_config['dataset']['paired_val_filelist'] == ''
            filefolders_not_exist = old_style_config['dataset']['paired_valA_folder'] == '' or \
                old_style_config['dataset']['paired_valB_folder'] == '' or \
                not os.path.exists(old_style_config['dataset']['paired_valA_lmk_folder']) or \
                not os.path.exists(old_style_config['dataset']['paired_valB_lmk_folder'])
            dataroot_contains_no_val_folders = not os.path.exists(
                os.path.join(old_style_config['dataset']['dataroot'], 'valpairedA_lmk')) \
                                               or not os.path.exists(
                os.path.join(old_style_config['dataset']['dataroot'], 'valpairedB_lmk'))
            if filelist_not_exist and filefolders_not_exist and dataroot_contains_no_val_folders:
                return False

    return True

utils/data_utils.py
mmfs-paper/MMFS
#!/usr/bin/python
# -*- encoding: utf-8 -*-

import os
import torch
import torch.nn as nn
import torch.nn.functional as F

from .resnet import Resnet18
from models.modules.networks import init_net


class ConvBNReLU(nn.Module):
    def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
        super(ConvBNReLU, self).__init__()
        self.conv = nn.Conv2d(in_chan,
                out_chan,
                kernel_size = ks,
                stride = stride,
                padding = padding,
                bias = False)
        self.bn = nn.BatchNorm2d(out_chan)
        self.init_weight()

    def forward(self, x):
        x = self.conv(x)
        x = F.relu(self.bn(x))
        return x

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

class BiSeNetOutput(nn.Module):
    def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):
        super(BiSeNetOutput, self).__init__()
        self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
        self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)
        self.init_weight()

    def forward(self, x):
        x = self.conv(x)
        x = self.conv_out(x)
        return x

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

    def get_params(self):
        wd_params, nowd_params = [], []
        for _, module in self.named_modules():
            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
                wd_params.append(module.weight)
                if not module.bias is None:
                    nowd_params.append(module.bias)
            elif isinstance(module, nn.BatchNorm2d):
                nowd_params += list(module.parameters())
        return wd_params, nowd_params


class AttentionRefinementModule(nn.Module):
    def __init__(self, in_chan, out_chan, *args, **kwargs):
        super(AttentionRefinementModule, self).__init__()
        self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
        self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)
        self.bn_atten = nn.BatchNorm2d(out_chan)
        self.sigmoid_atten = nn.Sigmoid()
        self.init_weight()

    def forward(self, x):
        feat = self.conv(x)
        atten = F.avg_pool2d(feat, feat.size()[2:])
        atten = self.conv_atten(atten)
        atten = self.bn_atten(atten)
        atten = self.sigmoid_atten(atten)
        out = torch.mul(feat, atten)
        return out

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)


class ContextPath(nn.Module):
    def __init__(self, *args, **kwargs):
        super(ContextPath, self).__init__()
        self.resnet = Resnet18()
        self.arm16 = AttentionRefinementModule(256, 128)
        self.arm32 = AttentionRefinementModule(512, 128)
        self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
        self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
        self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)

        self.init_weight()

    def forward(self, x):
        H0, W0 = x.size()[2:]
        feat8, feat16, feat32 = self.resnet(x)
        H8, W8 = feat8.size()[2:]
        H16, W16 = feat16.size()[2:]
        H32, W32 = feat32.size()[2:]

        avg = F.avg_pool2d(feat32, feat32.size()[2:])
        avg = self.conv_avg(avg)
        avg_up = F.interpolate(avg, (H32, W32), mode='nearest')

        feat32_arm = self.arm32(feat32)
        feat32_sum = feat32_arm + avg_up
        feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
        feat32_up = self.conv_head32(feat32_up)

        feat16_arm = self.arm16(feat16)
        feat16_sum = feat16_arm + feat32_up
        feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
        feat16_up = self.conv_head16(feat16_up)

        return feat8, feat16_up, feat32_up  # x8, x8, x16

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

    def get_params(self):
        wd_params, nowd_params = [], []
        for _, module in self.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                wd_params.append(module.weight)
                if not module.bias is None:
                    nowd_params.append(module.bias)
            elif isinstance(module, nn.BatchNorm2d):
                nowd_params += list(module.parameters())
        return wd_params, nowd_params


### This is not used, since I replace this with the resnet feature with the same size
class SpatialPath(nn.Module):
    def __init__(self, *args, **kwargs):
        super(SpatialPath, self).__init__()
        self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)
        self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
        self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
        self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)
        self.init_weight()

    def forward(self, x):
        feat = self.conv1(x)
        feat = self.conv2(feat)
        feat = self.conv3(feat)
        feat = self.conv_out(feat)
        return feat

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

    def get_params(self):
        wd_params, nowd_params = [], []
        for _, module in self.named_modules():
            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
                wd_params.append(module.weight)
                if not module.bias is None:
                    nowd_params.append(module.bias)
            elif isinstance(module, nn.BatchNorm2d):
                nowd_params += list(module.parameters())
        return wd_params, nowd_params


class FeatureFusionModule(nn.Module):
    def __init__(self, in_chan, out_chan, *args, **kwargs):
        super(FeatureFusionModule, self).__init__()
        self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
        self.conv1 = nn.Conv2d(out_chan,
                out_chan//4,
                kernel_size = 1,
                stride = 1,
                padding = 0,
                bias = False)
        self.conv2 = nn.Conv2d(out_chan//4,
                out_chan,
                kernel_size = 1,
                stride = 1,
                padding = 0,
                bias = False)
        self.relu = nn.ReLU(inplace=True)
        self.sigmoid = nn.Sigmoid()
        self.init_weight()

    def forward(self, fsp, fcp):
        fcat = torch.cat([fsp, fcp], dim=1)
        feat = self.convblk(fcat)
        atten = F.avg_pool2d(feat, feat.size()[2:])
        atten = self.conv1(atten)
        atten = self.relu(atten)
        atten = self.conv2(atten)
        atten = self.sigmoid(atten)
        feat_atten = torch.mul(feat, atten)
        feat_out = feat_atten + feat
        return feat_out

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

    def get_params(self):
        wd_params, nowd_params = [], []
        for _, module in self.named_modules():
            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
                wd_params.append(module.weight)
                if not module.bias is None:
                    nowd_params.append(module.bias)
            elif isinstance(module, nn.BatchNorm2d):
                nowd_params += list(module.parameters())
        return wd_params, nowd_params


class BiSeNet(nn.Module):
    def __init__(self, n_classes, *args, **kwargs):
        super(BiSeNet, self).__init__()
        self.cp = ContextPath()
        ## here self.sp is deleted
        self.ffm = FeatureFusionModule(256, 256)
        self.conv_out = BiSeNetOutput(256, 256, n_classes)
        self.conv_out16 = BiSeNetOutput(128, 64, n_classes)
        self.conv_out32 = BiSeNetOutput(128, 64, n_classes)
        self.init_weight()

    def forward(self, x):
        H, W = x.size()[2:]
        feat_res8, feat_cp8, feat_cp16 = self.cp(x)  # here return res3b1 feature
        feat_sp = feat_res8  # use res3b1 feature to replace spatial path feature
        feat_fuse = self.ffm(feat_sp, feat_cp8)

        feat_out = self.conv_out(feat_fuse)
        feat_out16 = self.conv_out16(feat_cp8)
        feat_out32 = self.conv_out32(feat_cp16)

        feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
        feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
        feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
        return feat_out, feat_out16, feat_out32

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

    def get_params(self):
        wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
        for _, child in self.named_children():
            child_wd_params, child_nowd_params = child.get_params()
            if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput):
                lr_mul_wd_params += child_wd_params
                lr_mul_nowd_params += child_nowd_params
            else:
                wd_params += child_wd_params
                nowd_params += child_nowd_params
        return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params


class PartWeightsGenerator():
    def __init__(self, gpu_ids, DDP_device):
        super(PartWeightsGenerator, self).__init__()

        # init face parsing network
        self.net = BiSeNet(n_classes=19)
        self.net = init_net(self.net, gpu_ids=gpu_ids, DDP_device=DDP_device)
        if isinstance(self.net, torch.nn.DataParallel) or isinstance(self.net, torch.nn.parallel.DistributedDataParallel):
            self.net = self.net.module
        cur_folder = os.path.split(__file__)[0]
        self.net.load_state_dict(torch.load(os.path.join(cur_folder, 'face_parsing.pth'), map_location=lambda storage, loc: storage))
        self.net.eval()

        # init some tensors
        self.mu = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
        self.sigma = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
        self.gauss_kernel = torch.tensor([1,4,7,4,1,4,16,26,16,4,7,26,41,26,7,4,16,26,16,4,1,4,7,4,1],).view(1, 1, 5, 5) / 273.0

        if len(gpu_ids) > 0 or not DDP_device is None:
            device = DDP_device if not DDP_device is None else gpu_ids[0]
            self.mu = self.mu.to(device)
            self.sigma = self.sigma.to(device)
            self.gauss_kernel = self.gauss_kernel.to(device)

        # init attributes list
        self.atts_list = [
            'background', #0
            'skin', #1
            'left_brow', #2
            'right_brow', #3
            'left_eye', #4
            'right_eye', #5
            'eye_glasses', #6
            'left_ear', #7
            'right_ear', #8
            'ear_rings', #9
            'nose', #10
            'teeth', #11
            'upper_lip', #12
            'lower_lip', #13
            'neck', #14
            'necklace', #15
            'cloth', #16
            'hair', #17
            'hat', #18
        ]

    def generate_masks(self, img):
        with torch.no_grad():
            _, _, h, w = img.size()
            img_512 = F.interpolate(img, size=512, mode='bilinear')
            pred_512 = self.net((img_512 * 0.5 + 0.5 - self.mu) / self.sigma)[0]
            pred = F.interpolate(pred_512, size=(h, w), mode='bilinear')
            pred = pred.argmax(1, keepdim=True)

            skin_mask = torch.zeros_like(pred).float()
            skin_mask[pred == 1] = 1.0
            skin_mask[pred == 7] = 1.0
            skin_mask[pred == 8] = 1.0
            skin_mask[pred == 10] = 1.0
            skin_mask[pred == 14] = 1.0
            skin_mask = F.conv2d(skin_mask, self.gauss_kernel, padding=2)

            eye_mask = torch.zeros_like(pred).float()
            eye_mask[pred == 4] = 1.0
            eye_mask[pred == 5] = 1.0
            eye_mask = F.conv2d(eye_mask, self.gauss_kernel, padding=2)

            mouth_mask = torch.zeros_like(pred).float()
            mouth_mask[pred == 11] = 1.0
            mouth_mask[pred == 12] = 1.0
            mouth_mask[pred == 13] = 1.0
            mouth_mask = F.conv2d(mouth_mask, self.gauss_kernel, padding=2)

            hair_mask = torch.zeros_like(pred).float()
            hair_mask[pred == 17] = 1.0
            hair_mask = F.conv2d(hair_mask, self.gauss_kernel, padding=2)

            return skin_mask, eye_mask, mouth_mask, hair_mask


    def generate_weights(self, img, weights_dict, blur=True):
        with torch.no_grad():
            _, _, h, w = img.size()
            img_512 = F.interpolate(img, size=512, mode='bilinear')
            pred_512 = self.net((img_512 * 0.5 + 0.5 - self.mu) / self.sigma)[0]
            pred = F.interpolate(pred_512, size=(h, w), mode='bilinear')
            pred = pred.argmax(1, keepdim=True)
            weights = torch.ones_like(pred).float()
            for idx, att in enumerate(self.atts_list):
                if att in weights_dict:
                    weights[pred == idx] = weights_dict[att]
            if blur:
                return F.conv2d(weights, self.gauss_kernel, padding=2)
            else:
                return weights


class GradWeightFunc(torch.autograd.Function):

    @staticmethod
    def forward(ctx, img, weights):
        ctx.param = (weights, )
        return img

    @staticmethod
    def backward(ctx, grad):
        weights = ctx.param[0]
        return weights * grad, None


class GradWeightLayer(nn.Module):

    def __init__(self):
        super(GradWeightLayer, self).__init__()

    def forward(self, img, weights):
        return GradWeightFunc.apply(img, weights)

utils/face_parsing/face_parsing.py
mmfs-paper/MMFS
#!/usr/bin/python
# -*- encoding: utf-8 -*-

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.model_zoo as modelzoo

resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'


def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)


class BasicBlock(nn.Module):
    def __init__(self, in_chan, out_chan, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(in_chan, out_chan, stride)
        self.bn1 = nn.BatchNorm2d(out_chan)
        self.conv2 = conv3x3(out_chan, out_chan)
        self.bn2 = nn.BatchNorm2d(out_chan)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = None
        if in_chan != out_chan or stride != 1:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_chan, out_chan,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_chan),
                )

    def forward(self, x):
        residual = self.conv1(x)
        residual = F.relu(self.bn1(residual))
        residual = self.conv2(residual)
        residual = self.bn2(residual)

        shortcut = x
        if self.downsample is not None:
            shortcut = self.downsample(x)

        out = shortcut + residual
        out = self.relu(out)
        return out


def create_layer_basic(in_chan, out_chan, bnum, stride=1):
    layers = [BasicBlock(in_chan, out_chan, stride=stride)]
    for i in range(bnum-1):
        layers.append(BasicBlock(out_chan, out_chan, stride=1))
    return nn.Sequential(*layers)


class Resnet18(nn.Module):
    def __init__(self):
        super(Resnet18, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
        self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
        self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
        self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
        self.init_weight()

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(self.bn1(x))
        x = self.maxpool(x)

        x = self.layer1(x)
        feat8 = self.layer2(x) # 1/8
        feat16 = self.layer3(feat8) # 1/16
        feat32 = self.layer4(feat16) # 1/32
        return feat8, feat16, feat32

    def init_weight(self):
        state_dict = modelzoo.load_url(resnet18_url)
        self_state_dict = self.state_dict()
        for k, v in state_dict.items():
            if 'fc' in k: continue
            self_state_dict.update({k: v})
        self.load_state_dict(self_state_dict)

    def get_params(self):
        wd_params, nowd_params = [], []
        for _, module in self.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                wd_params.append(module.weight)
                if not module.bias is None:
                    nowd_params.append(module.bias)
            elif isinstance(module,  nn.BatchNorm2d):
                nowd_params += list(module.parameters())
        return wd_params, nowd_params

utils/face_parsing/resnet.py
mmfs-paper/MMFS
from __future__ import absolute_import
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division

import torch
import torch.nn as nn
from torch.autograd import Variable
from functools import reduce
import operator
import pdb
#from new_modules import Cat

count_ops = 0
count_params = 0


def get_num_gen(gen):
    return sum(1 for x in gen)


def is_pruned(layer):
    try:
        layer.mask
        return True
    except AttributeError:
        return False


def is_leaf(model):
    return get_num_gen(model.children()) == 0


def get_layer_info(layer):
    layer_str = str(layer)
    type_name = layer_str[:layer_str.find('(')].strip()
    return type_name


def get_layer_param(model):
    return sum([reduce(operator.mul, i.size(), 1) for i in model.parameters()])


### The input batch size should be 1 to call this function
def measure_layer(layer, x):
    global count_ops, count_params
    delta_ops = 0
    delta_params = 0
    multi_add = 1
    type_name = get_layer_info(layer)
    ### ops_conv
    if type_name in ['Conv2d', 'Conv2dQuant', 'ConvTranspose2d']:
        out_h = int((x.size()[2] + 2 * layer.padding[0] - layer.kernel_size[0]) /
                    layer.stride[0] + 1)
        out_w = int((x.size()[3] + 2 * layer.padding[1] - layer.kernel_size[1]) /
                    layer.stride[1] + 1)
        delta_ops = layer.in_channels * layer.out_channels * layer.kernel_size[0] *  \
                layer.kernel_size[1] * out_h * out_w / layer.groups * multi_add
        print('Layer:', layer)
        print('Number of parameters: %.3f M' % (delta_ops / 1e6))
        delta_params = get_layer_param(layer)

    ### ops_nonlinearity
    elif type_name in ['ReLU']:
        #delta_ops = x.numel()
        delta_ops = 0.0
        delta_params = get_layer_param(layer)

    ### ops_pooling
    elif type_name in ['AvgPool2d']:
        in_w = x.size()[2]
        kernel_ops = layer.kernel_size * layer.kernel_size
        out_w = int((in_w + 2 * layer.padding - layer.kernel_size) / layer.stride + 1)
        out_h = int((in_w + 2 * layer.padding - layer.kernel_size) / layer.stride + 1)
        #delta_ops = x.size()[0] * x.size()[1] * out_w * out_h * kernel_ops
        delta_ops = 0.0
        delta_params = get_layer_param(layer)

    elif type_name in ['AdaptiveAvgPool2d']:
        #delta_ops = x.size()[0] * x.size()[1] * x.size()[2] * x.size()[3]
        delta_ops = 0.0
        delta_params = get_layer_param(layer)

    ### ops_linear
    elif type_name in ['Linear']:
        weight_ops = layer.weight.numel() * multi_add
        bias_ops = layer.bias.numel()
        #delta_ops = x.size()[0] * (weight_ops + bias_ops)
        delta_ops = 0.0
        delta_params = get_layer_param(layer)

    ### ops_nothing
    elif type_name in ['BatchNorm2d', 'Dropout2d', 'DropChannel', 'Dropout', 'Upsample', 'Softmax']:
        delta_params = get_layer_param(layer)
        delta_ops = 0.0

    elif type_name in ['Upsample', 'Hardtanh', 'QuantizedHardtanh', 'MaxPool2d', 'Cat', 'AvgQuant']:
        delta_params = 0
        delta_ops = 0.0

    elif type_name in ['InstanceNorm2d', 'LeakyReLU', 'Tanh', 'ReflectionPad2d']:
        delta_params = get_layer_param(layer)
        delta_ops = 0.0
    ### unknown layer type

    else:
        raise TypeError('unknown layer type: %s' % type_name)

    count_ops += delta_ops
    count_params += delta_params
    return


def measure_model(model, H, W):
    global count_ops, count_params
    count_ops = 0
    count_params = 0
    data = Variable(torch.zeros(1, 3, H, W))

    def should_measure(x):
        return is_leaf(x) or is_pruned(x)

    def modify_forward(model):
        for child in model.children():
            if should_measure(child):
                def new_forward(m):
                    def lambda_forward(x):
                        measure_layer(m, x)
                        return m.old_forward(x)
                    return lambda_forward
                child.old_forward = child.forward
                child.forward = new_forward(child)
            else:
                modify_forward(child)

    def restore_forward(model):
        for child in model.children():
            # leaf node
            if is_leaf(child) and hasattr(child, 'old_forward'):
                child.forward = child.old_forward
                child.old_forward = None
            else:
                restore_forward(child)

    modify_forward(model)
    model.forward(data)
    restore_forward(model)

    return count_ops, count_params

def calc_computation(pose_model, ch, H=224, W=224, DDP_device=None):
    list_conv = []
    list_hook = []
    def conv_hook(self, input, output):
        batch_size, input_channels, input_height, input_width = input[0].size()
        output_channels, output_height, output_width = output[0].size()
        kernel_ops = self.kernel_size[0] * self.kernel_size[1] * (self.in_channels / self.groups)
        bias_ops = 1 if self.bias is not None else 0
        params = output_channels * (kernel_ops + bias_ops)
        flops = batch_size * params * output_height * output_width
        list_conv.append(flops)
    def calc(net):
        childrens = list(net.children())
        if not childrens:
            if isinstance(net, nn.Conv2d):# | isinstance(net, nn.ConvTranspose2d):
                list_hook.append(net.register_forward_hook(conv_hook))
            return
        for c in childrens:
            calc(c)
    calc(pose_model)
    if next(pose_model.parameters()).is_cuda:
        if DDP_device is None:
            y = pose_model((torch.ones(2, ch, H, W).cuda()))
        else:
            y = pose_model((torch.ones(2, ch, H, W).to(DDP_device)))
    else:
        y = pose_model((torch.ones(2, ch, H, W)))
    print('total_compuation:', sum(list_conv) / 2e6, 'M')
    [item.remove() for item in list_hook]

utils/net_size.py
mmfs-paper/MMFS
"""This module contains simple helper functions """
from __future__ import print_function
import torch
import numpy as np
from PIL import Image
import os
import torchvision
import sys
import shutil
import datetime


def check_path_is_static_data(path):
    last_extension = path.split(".")[-1]
    acceptable_extensions = ['png', 'jpg', 'jpeg', 'npy', 'npz']
    if last_extension.lower() in acceptable_extensions:
        return True
    return False

def check_path(path):
    try:
        if not os.path.exists(path):
            os.makedirs(path)
    except OSError as e:
        print(e)


def tensor2file(input_image, file_path, ext_name):
    """Convert a tensor into a file.

    Parameters:
        input_image  --  the input image tensor
        file_path    --  the file path without extension name
    """
    if not isinstance(input_image, np.ndarray):
        if isinstance(input_image, torch.Tensor):  # get the data from a variable
            image_tensor = input_image.data
        else:
            return input_image
        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
        if image_numpy.shape[0] == 1:  # grayscale to RGB
            image_numpy = np.tile(image_numpy, (3, 1, 1))
        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
    else:  # if it is a numpy array, do nothing
        image_numpy = input_image

    if image_numpy.shape[2] <= 3:
        image_numpy = image_numpy.astype(np.uint8)
        # save as image
        image_pil = Image.fromarray(image_numpy)
        image_pil.save(file_path + '.' + ext_name)
    else:
        # save as numpy
        np.save(file_path + '.npy', image_numpy)


def tensor2im(input_image, imtype=np.uint8):
    """"Converts a Tensor array into a numpy image array.

    Parameters:
        input_image (tensor) --  the input image tensor array
        imtype (type)        --  the desired type of the converted numpy array
    """
    if not isinstance(input_image, np.ndarray):
        if isinstance(input_image, torch.Tensor):  # get the data from a variable
            image_tensor = input_image.data
        else:
            return input_image
        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
        if image_numpy.shape[0] == 1:  # grayscale to RGB
            image_numpy = np.tile(image_numpy, (3, 1, 1))
        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
        if image_numpy.shape[2] > 3: # clip to 3 channel
            print('Warning: the channel count of output image exceeds 3.')
            image_numpy = image_numpy[:,:,:3]
    else:  # if it is a numpy array, do nothing
        image_numpy = input_image
    return image_numpy.astype(imtype)


def diagnose_network(net, name='network'):
    """Calculate and print the mean of average absolute(gradients)

    Parameters:
        net (torch network) -- Torch network
        name (str) -- the name of the network
    """
    mean = 0.0
    count = 0
    for param in net.parameters():
        if param.grad is not None:
            mean += torch.mean(torch.abs(param.grad.data))
            count += 1
    if count > 0:
        mean = mean / count
    print(name)
    print(mean)


def save_image(image_numpy, image_path, aspect_ratio=1.0):
    """Save a numpy image to the disk

    Parameters:
        image_numpy (numpy array) -- input numpy array
        image_path (str)          -- the path of the image
    """

    image_pil = Image.fromarray(image_numpy)
    h, w, _ = image_numpy.shape

    if aspect_ratio > 1.0:
        image_pil = image_pil.resize((h, int(w * aspect_ratio)), Image.BICUBIC)
    if aspect_ratio < 1.0:
        image_pil = image_pil.resize((int(h / aspect_ratio), w), Image.BICUBIC)
    image_pil.save(image_path)


def print_numpy(x, val=True, shp=False):
    """Print the mean, min, max, median, std, and size of a numpy array

    Parameters:
        val (bool) -- if print the values of the numpy array
        shp (bool) -- if print the shape of the numpy array
    """
    x = x.astype(np.float64)
    if shp:
        print('shape,', x.shape)
    if val:
        x = x.flatten()
        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))


def mkdirs(paths):
    """create empty directories if they don't exist

    Parameters:
        paths (str list) -- a list of directory paths
    """
    if isinstance(paths, list) and not isinstance(paths, str):
        for path in paths:
            mkdir(path)
    else:
        mkdir(paths)


def mkdir(path):
    """create a single empty directory if it didn't exist

    Parameters:
        path (str) -- a single directory path
    """
    if not os.path.exists(path):
        os.makedirs(path)

def print_losses(epoch, iters, len_dataset, current_losses, average_losses, mode='training'):
    """print current losses on console"""
    if mode=='training':
        message = mode + ': (epoch: %d, iters: %d / %d) ' % (epoch, iters, len_dataset)
    else:
        message = mode + ': (epoch: %d) ' % epoch
    for k, v in current_losses.items():
        message += '%s: %.3f ' % (k, v)
    for k, v in average_losses.items():
        message += 'average %s: %.3f ' % (k, v.avg())
    message += datetime.datetime.now().strftime("%Y_%m_%d %H:%M:%S")
    print(message)  # print the message

def make_grid(model):
    """
    create an image grid to be visualized by tensorboard.
    """
    visuals = model.get_current_visuals()
    names, grids = [], []
    for name, img in visuals.items():
        names.append(name)
        grid = torchvision.utils.make_grid(img[:,:3,:,:], nrow=img.size()[0], normalize=True)
        grids.append(grid)
    return grids, names

class AverageMeter:
    def __init__(self):
        self.sum = 0
        self.num_item = 0

    def update(self, value):
        self.sum += value
        self.num_item += 1

    def avg(self):
        return self.sum / self.num_item

    def clear(self):
        self.sum = 0
        self.num_item = 0


utils/util.py
mmfs-paper/MMFS