Upload 194 files

models/stylegan2_ffhq1024.pth

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

networks/biggan/__init__.py

@@ -0,0 +1,6 @@
+from .config import BigGANConfig
+from .model import BigGAN
+from .file_utils import PYTORCH_PRETRAINED_BIGGAN_CACHE, cached_path
+from .utils import (truncated_noise_sample, save_as_images,
+                    convert_to_images, display_in_terminal,
+                    one_hot_from_int, one_hot_from_names)

networks/biggan/config.py

@@ -0,0 +1,70 @@
+# coding: utf-8
+"""
+BigGAN config.
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+import copy
+import json
+
+class BigGANConfig(object):
+    """ Configuration class to store the configuration of a `BigGAN`. 
+        Defaults are for the 128x128 model.
+        layers tuple are (up-sample in the layer ?, input channels, output channels)
+    """
+    def __init__(self,
+                 output_dim=128,
+                 z_dim=128,
+                 class_embed_dim=128,
+                 channel_width=128,
+                 num_classes=1000,
+                 layers=[(False, 16, 16),
+                         (True, 16, 16),
+                         (False, 16, 16),
+                         (True, 16, 8),
+                         (False, 8, 8),
+                         (True, 8, 4),
+                         (False, 4, 4),
+                         (True, 4, 2),
+                         (False, 2, 2),
+                         (True, 2, 1)],
+                 attention_layer_position=8,
+                 eps=1e-4,
+                 n_stats=51):
+        """Constructs BigGANConfig. """
+        self.output_dim = output_dim
+        self.z_dim = z_dim
+        self.class_embed_dim = class_embed_dim
+        self.channel_width = channel_width
+        self.num_classes = num_classes
+        self.layers = layers
+        self.attention_layer_position = attention_layer_position
+        self.eps = eps
+        self.n_stats = n_stats
+
+    @classmethod
+    def from_dict(cls, json_object):
+        """Constructs a `BigGANConfig` from a Python dictionary of parameters."""
+        config = BigGANConfig()
+        for key, value in json_object.items():
+            config.__dict__[key] = value
+        return config
+
+    @classmethod
+    def from_json_file(cls, json_file):
+        """Constructs a `BigGANConfig` from a json file of parameters."""
+        with open(json_file, "r", encoding='utf-8') as reader:
+            text = reader.read()
+        return cls.from_dict(json.loads(text))
+
+    def __repr__(self):
+        return str(self.to_json_string())
+
+    def to_dict(self):
+        """Serializes this instance to a Python dictionary."""
+        output = copy.deepcopy(self.__dict__)
+        return output
+
+    def to_json_string(self):
+        """Serializes this instance to a JSON string."""
+        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

networks/biggan/convert.sh

@@ -0,0 +1,21 @@
+# Copyright (c) 2019-present, Thomas Wolf, Huggingface Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+set -e
+set -x
+
+models="128 256 512"
+
+mkdir -p models/model_128
+mkdir -p models/model_256
+mkdir -p models/model_512
+
+# Convert TF Hub models.
+for model in $models
+do
+    pytorch_pretrained_biggan --model_type $model --tf_model_path models/model_$model --pt_save_path models/model_$model
+done

networks/biggan/convert_tf_to_pytorch.py

@@ -0,0 +1,312 @@
+# coding: utf-8
+"""
+Convert a TF Hub model for BigGAN in a PT one.
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+from itertools import chain
+
+import os
+import argparse
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.functional import normalize
+
+from .model import BigGAN, WEIGHTS_NAME, CONFIG_NAME
+from .config import BigGANConfig
+
+logger = logging.getLogger(__name__)
+
+
+def extract_batch_norm_stats(tf_model_path, batch_norm_stats_path=None):
+    try:
+        import numpy as np
+        import tensorflow as tf
+        import tensorflow_hub as hub
+    except ImportError:
+        raise ImportError("Loading a TensorFlow models in PyTorch, requires TensorFlow and TF Hub to be installed. "
+                          "Please see https://www.tensorflow.org/install/ for installation instructions for TensorFlow. "
+                          "And see https://github.com/tensorflow/hub for installing Hub. "
+                          "Probably pip install tensorflow tensorflow-hub")
+    tf.reset_default_graph()
+    logger.info('Loading BigGAN module from: {}'.format(tf_model_path))
+    module = hub.Module(tf_model_path)
+    inputs = {k: tf.placeholder(v.dtype, v.get_shape().as_list(), k)
+              for k, v in module.get_input_info_dict().items()}
+    output = module(inputs)
+
+    initializer = tf.global_variables_initializer()
+    sess = tf.Session()
+    stacks = sum(((i*10 + 1, i*10 + 3, i*10 + 6, i*10 + 8) for i in range(50)), ())
+    numpy_stacks = []
+    for i in stacks:
+        logger.info("Retrieving module_apply_default/stack_{}".format(i))
+        try:
+            stack_var = tf.get_default_graph().get_tensor_by_name("module_apply_default/stack_%d:0" % i)
+        except KeyError:
+            break  # We have all the stats
+        numpy_stacks.append(sess.run(stack_var))
+
+    if batch_norm_stats_path is not None:
+        torch.save(numpy_stacks, batch_norm_stats_path)
+    else:
+        return numpy_stacks
+
+
+def build_tf_to_pytorch_map(model, config):
+    """ Build a map from TF variables to PyTorch modules. """
+    tf_to_pt_map = {}
+
+    # Embeddings and GenZ
+    tf_to_pt_map.update({'linear/w/ema_0.9999': model.embeddings.weight,
+                         'Generator/GenZ/G_linear/b/ema_0.9999': model.generator.gen_z.bias,
+                         'Generator/GenZ/G_linear/w/ema_0.9999': model.generator.gen_z.weight_orig,
+                         'Generator/GenZ/G_linear/u0': model.generator.gen_z.weight_u})
+
+    # GBlock blocks
+    model_layer_idx = 0
+    for i, (up, in_channels, out_channels) in enumerate(config.layers):
+        if i == config.attention_layer_position:
+            model_layer_idx += 1
+        layer_str = "Generator/GBlock_%d/" % i if i > 0 else "Generator/GBlock/"
+        layer_pnt = model.generator.layers[model_layer_idx]
+        for i in range(4):  #  Batchnorms
+            batch_str = layer_str + ("BatchNorm_%d/" % i if i > 0 else "BatchNorm/")
+            batch_pnt = getattr(layer_pnt, 'bn_%d' % i)
+            for name in ('offset', 'scale'):
+                sub_module_str = batch_str + name + "/"
+                sub_module_pnt = getattr(batch_pnt, name)
+                tf_to_pt_map.update({sub_module_str + "w/ema_0.9999": sub_module_pnt.weight_orig,
+                                     sub_module_str + "u0": sub_module_pnt.weight_u})
+        for i in range(4):  # Convolutions
+            conv_str = layer_str + "conv%d/" % i
+            conv_pnt = getattr(layer_pnt, 'conv_%d' % i)
+            tf_to_pt_map.update({conv_str + "b/ema_0.9999": conv_pnt.bias,
+                                 conv_str + "w/ema_0.9999": conv_pnt.weight_orig,
+                                 conv_str + "u0": conv_pnt.weight_u})
+        model_layer_idx += 1
+
+    # Attention block
+    layer_str = "Generator/attention/"
+    layer_pnt = model.generator.layers[config.attention_layer_position]
+    tf_to_pt_map.update({layer_str + "gamma/ema_0.9999": layer_pnt.gamma})
+    for pt_name, tf_name in zip(['snconv1x1_g', 'snconv1x1_o_conv', 'snconv1x1_phi', 'snconv1x1_theta'],
+                                ['g/', 'o_conv/', 'phi/', 'theta/']):
+        sub_module_str = layer_str + tf_name
+        sub_module_pnt = getattr(layer_pnt, pt_name)
+        tf_to_pt_map.update({sub_module_str + "w/ema_0.9999": sub_module_pnt.weight_orig,
+                             sub_module_str + "u0": sub_module_pnt.weight_u})
+
+    # final batch norm and conv to rgb
+    layer_str = "Generator/BatchNorm/"
+    layer_pnt = model.generator.bn
+    tf_to_pt_map.update({layer_str + "offset/ema_0.9999": layer_pnt.bias,
+                         layer_str + "scale/ema_0.9999": layer_pnt.weight})
+    layer_str = "Generator/conv_to_rgb/"
+    layer_pnt = model.generator.conv_to_rgb
+    tf_to_pt_map.update({layer_str + "b/ema_0.9999": layer_pnt.bias,
+                         layer_str + "w/ema_0.9999": layer_pnt.weight_orig,
+                         layer_str + "u0": layer_pnt.weight_u})
+    return tf_to_pt_map
+
+
+def load_tf_weights_in_biggan(model, config, tf_model_path, batch_norm_stats_path=None):
+    """ Load tf checkpoints and standing statistics in a pytorch model
+    """
+    try:
+        import numpy as np
+        import tensorflow as tf
+    except ImportError:
+        raise ImportError("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
+            "https://www.tensorflow.org/install/ for installation instructions.")
+    # Load weights from TF model
+    checkpoint_path = tf_model_path + "/variables/variables"
+    init_vars = tf.train.list_variables(checkpoint_path)
+    from pprint import pprint
+    pprint(init_vars)
+
+    # Extract batch norm statistics from model if needed
+    if batch_norm_stats_path:
+        stats = torch.load(batch_norm_stats_path)
+    else:
+        logger.info("Extracting batch norm stats")
+        stats = extract_batch_norm_stats(tf_model_path)
+
+    # Build TF to PyTorch weights loading map
+    tf_to_pt_map = build_tf_to_pytorch_map(model, config)
+
+    tf_weights = {}
+    for name in tf_to_pt_map.keys():
+        array = tf.train.load_variable(checkpoint_path, name)
+        tf_weights[name] = array
+        # logger.info("Loading TF weight {} with shape {}".format(name, array.shape))
+
+    # Load parameters
+    with torch.no_grad():
+        pt_params_pnt = set()
+        for name, pointer in tf_to_pt_map.items():
+            array = tf_weights[name]
+            if pointer.dim() == 1:
+                if pointer.dim() < array.ndim:
+                    array = np.squeeze(array)
+            elif pointer.dim() == 2:  # Weights
+                array = np.transpose(array)
+            elif pointer.dim() == 4:  # Convolutions
+                array = np.transpose(array, (3, 2, 0, 1))
+            else:
+                raise "Wrong dimensions to adjust: " + str((pointer.shape, array.shape))
+            if pointer.shape != array.shape:
+                raise ValueError("Wrong dimensions: " + str((pointer.shape, array.shape)))
+            logger.info("Initialize PyTorch weight {} with shape {}".format(name, pointer.shape))
+            pointer.data = torch.from_numpy(array) if isinstance(array, np.ndarray) else torch.tensor(array)
+            tf_weights.pop(name, None)
+            pt_params_pnt.add(pointer.data_ptr())
+
+        # Prepare SpectralNorm buffers by running one step of Spectral Norm (no need to train the model):
+        for module in model.modules():
+            for n, buffer in module.named_buffers():
+                if n == 'weight_v':
+                    weight_mat = module.weight_orig
+                    weight_mat = weight_mat.reshape(weight_mat.size(0), -1)
+                    u = module.weight_u
+
+                    v = normalize(torch.mv(weight_mat.t(), u), dim=0, eps=config.eps)
+                    buffer.data = v
+                    pt_params_pnt.add(buffer.data_ptr())
+
+                    u = normalize(torch.mv(weight_mat, v), dim=0, eps=config.eps)
+                    module.weight_u.data = u
+                    pt_params_pnt.add(module.weight_u.data_ptr())
+
+        # Load batch norm statistics
+        index = 0
+        for layer in model.generator.layers:
+            if not hasattr(layer, 'bn_0'):
+                continue
+            for i in range(4):  #  Batchnorms
+                bn_pointer = getattr(layer, 'bn_%d' % i)
+                pointer = bn_pointer.running_means
+                if pointer.shape != stats[index].shape:
+                    raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
+                pointer.data = torch.from_numpy(stats[index])
+                pt_params_pnt.add(pointer.data_ptr())
+
+                pointer = bn_pointer.running_vars
+                if pointer.shape != stats[index+1].shape:
+                    raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
+                pointer.data = torch.from_numpy(stats[index+1])
+                pt_params_pnt.add(pointer.data_ptr())
+
+                index += 2
+
+        bn_pointer = model.generator.bn
+        pointer = bn_pointer.running_means
+        if pointer.shape != stats[index].shape:
+            raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
+        pointer.data = torch.from_numpy(stats[index])
+        pt_params_pnt.add(pointer.data_ptr())
+
+        pointer = bn_pointer.running_vars
+        if pointer.shape != stats[index+1].shape:
+            raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
+        pointer.data = torch.from_numpy(stats[index+1])
+        pt_params_pnt.add(pointer.data_ptr())
+
+    remaining_params = list(n for n, t in chain(model.named_parameters(), model.named_buffers()) \
+                            if t.data_ptr() not in pt_params_pnt)
+
+    logger.info("TF Weights not copied to PyTorch model: {} -".format(', '.join(tf_weights.keys())))
+    logger.info("Remanining parameters/buffers from PyTorch model: {} -".format(', '.join(remaining_params)))
+
+    return model
+
+
+BigGAN128 = BigGANConfig(output_dim=128, z_dim=128, class_embed_dim=128, channel_width=128, num_classes=1000,
+                         layers=[(False, 16, 16),
+                                 (True, 16, 16),
+                                 (False, 16, 16),
+                                 (True, 16, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 4),
+                                 (False, 4, 4),
+                                 (True, 4, 2),
+                                 (False, 2, 2),
+                                 (True, 2, 1)],
+                         attention_layer_position=8, eps=1e-4, n_stats=51)
+
+BigGAN256 = BigGANConfig(output_dim=256, z_dim=128, class_embed_dim=128, channel_width=128, num_classes=1000,
+                         layers=[(False, 16, 16),
+                                 (True, 16, 16),
+                                 (False, 16, 16),
+                                 (True, 16, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 4),
+                                 (False, 4, 4),
+                                 (True, 4, 2),
+                                 (False, 2, 2),
+                                 (True, 2, 1)],
+                         attention_layer_position=8, eps=1e-4, n_stats=51)
+
+BigGAN512 = BigGANConfig(output_dim=512, z_dim=128, class_embed_dim=128, channel_width=128, num_classes=1000,
+                         layers=[(False, 16, 16),
+                                 (True, 16, 16),
+                                 (False, 16, 16),
+                                 (True, 16, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 8),
+                                 (False, 8, 8),
+                                 (True, 8, 4),
+                                 (False, 4, 4),
+                                 (True, 4, 2),
+                                 (False, 2, 2),
+                                 (True, 2, 1),
+                                 (False, 1, 1),
+                                 (True, 1, 1)],
+                         attention_layer_position=8, eps=1e-4, n_stats=51)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert a BigGAN TF Hub model in a PyTorch model")
+    parser.add_argument("--model_type", type=str, default="", required=True,
+                        help="BigGAN model type (128, 256, 512)")
+    parser.add_argument("--tf_model_path", type=str, default="", required=True,
+                        help="Path of the downloaded TF Hub model")
+    parser.add_argument("--pt_save_path", type=str, default="",
+                        help="Folder to save the PyTorch model (default: Folder of the TF Hub model)")
+    parser.add_argument("--batch_norm_stats_path", type=str, default="",
+                        help="Path of previously extracted batch norm statistics")
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+
+    if not args.pt_save_path:
+        args.pt_save_path = args.tf_model_path
+
+    if args.model_type == "128":
+        config = BigGAN128
+    elif args.model_type == "256":
+        config = BigGAN256
+    elif args.model_type == "512":
+        config = BigGAN512
+    else:
+        raise ValueError("model_type should be one of 128, 256 or 512")
+
+    model = BigGAN(config)
+    model = load_tf_weights_in_biggan(model, config, args.tf_model_path, args.batch_norm_stats_path)
+
+    model_save_path = os.path.join(args.pt_save_path, WEIGHTS_NAME)
+    config_save_path = os.path.join(args.pt_save_path, CONFIG_NAME)
+
+    logger.info("Save model dump to {}".format(model_save_path))
+    torch.save(model.state_dict(), model_save_path)
+    logger.info("Save configuration file to {}".format(config_save_path))
+    with open(config_save_path, "w", encoding="utf-8") as f:
+        f.write(config.to_json_string())
+
+if __name__ == "__main__":
+    main()

networks/biggan/download_tf.sh

@@ -0,0 +1,21 @@
+# Copyright (c) 2019-present, Thomas Wolf, Huggingface Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+set -e
+set -x
+
+models="128 256 512"
+
+mkdir -p models/model_128
+mkdir -p models/model_256
+mkdir -p models/model_512
+
+# Download TF Hub models.
+for model in $models
+do
+    curl -L "https://tfhub.dev/deepmind/biggan-deep-$model/1?tf-hub-format=compressed" | tar -zxvC models/model_$model
+done

networks/biggan/file_utils.py

@@ -0,0 +1,232 @@
+"""
+Utilities for working with the local dataset cache.
+This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp
+Copyright by the AllenNLP authors.
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+import json
+import logging
+import os
+import shutil
+import tempfile
+from functools import wraps
+from hashlib import sha256
+import sys
+from io import open
+
+import requests
+from tqdm import tqdm
+
+try:
+    from urllib.parse import urlparse
+except ImportError:
+    from urlparse import urlparse
+
+try:
+    from pathlib import Path
+    PYTORCH_PRETRAINED_BIGGAN_CACHE = Path(os.getenv('PYTORCH_PRETRAINED_BIGGAN_CACHE',
+                                                   Path.home() / '.pytorch_pretrained_biggan'))
+except (AttributeError, ImportError):
+    PYTORCH_PRETRAINED_BIGGAN_CACHE = os.getenv('PYTORCH_PRETRAINED_BIGGAN_CACHE',
+                                              os.path.join(os.path.expanduser("~"), '.pytorch_pretrained_biggan'))
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+
+def url_to_filename(url, etag=None):
+    """
+    Convert `url` into a hashed filename in a repeatable way.
+    If `etag` is specified, append its hash to the url's, delimited
+    by a period.
+    """
+    url_bytes = url.encode('utf-8')
+    url_hash = sha256(url_bytes)
+    filename = url_hash.hexdigest()
+
+    if etag:
+        etag_bytes = etag.encode('utf-8')
+        etag_hash = sha256(etag_bytes)
+        filename += '.' + etag_hash.hexdigest()
+
+    return filename
+
+
+def filename_to_url(filename, cache_dir=None):
+    """
+    Return the url and etag (which may be ``None``) stored for `filename`.
+    Raise ``EnvironmentError`` if `filename` or its stored metadata do not exist.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BIGGAN_CACHE
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    cache_path = os.path.join(cache_dir, filename)
+    if not os.path.exists(cache_path):
+        raise EnvironmentError("file {} not found".format(cache_path))
+
+    meta_path = cache_path + '.json'
+    if not os.path.exists(meta_path):
+        raise EnvironmentError("file {} not found".format(meta_path))
+
+    with open(meta_path, encoding="utf-8") as meta_file:
+        metadata = json.load(meta_file)
+    url = metadata['url']
+    etag = metadata['etag']
+
+    return url, etag
+
+
+def cached_path(url_or_filename, cache_dir=None):
+    """
+    Given something that might be a URL (or might be a local path),
+    determine which. If it's a URL, download the file and cache it, and
+    return the path to the cached file. If it's already a local path,
+    make sure the file exists and then return the path.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BIGGAN_CACHE
+    if sys.version_info[0] == 3 and isinstance(url_or_filename, Path):
+        url_or_filename = str(url_or_filename)
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    parsed = urlparse(url_or_filename)
+
+    if parsed.scheme in ('http', 'https', 's3'):
+        # URL, so get it from the cache (downloading if necessary)
+        return get_from_cache(url_or_filename, cache_dir)
+    elif os.path.exists(url_or_filename):
+        # File, and it exists.
+        return url_or_filename
+    elif parsed.scheme == '':
+        # File, but it doesn't exist.
+        raise EnvironmentError("file {} not found".format(url_or_filename))
+    else:
+        # Something unknown
+        raise ValueError("unable to parse {} as a URL or as a local path".format(url_or_filename))
+
+
+def split_s3_path(url):
+    """Split a full s3 path into the bucket name and path."""
+    parsed = urlparse(url)
+    if not parsed.netloc or not parsed.path:
+        raise ValueError("bad s3 path {}".format(url))
+    bucket_name = parsed.netloc
+    s3_path = parsed.path
+    # Remove '/' at beginning of path.
+    if s3_path.startswith("/"):
+        s3_path = s3_path[1:]
+    return bucket_name, s3_path
+
+
+def s3_request(func):
+    """
+    Wrapper function for s3 requests in order to create more helpful error
+    messages.
+    """
+
+    @wraps(func)
+    def wrapper(url, *args, **kwargs):
+        try:
+            return func(url, *args, **kwargs)
+        except ClientError as exc:
+            if int(exc.response["Error"]["Code"]) == 404:
+                raise EnvironmentError("file {} not found".format(url))
+            else:
+                raise
+
+    return wrapper
+
+
+def http_get(url, temp_file):
+    req = requests.get(url, stream=True)
+    content_length = req.headers.get('Content-Length')
+    total = int(content_length) if content_length is not None else None
+    progress = tqdm(unit="B", total=total)
+    for chunk in req.iter_content(chunk_size=1024):
+        if chunk: # filter out keep-alive new chunks
+            progress.update(len(chunk))
+            temp_file.write(chunk)
+    progress.close()
+
+
+def get_from_cache(url, cache_dir=None):
+    """
+    Given a URL, look for the corresponding dataset in the local cache.
+    If it's not there, download it. Then return the path to the cached file.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BIGGAN_CACHE
+    if sys.version_info[0] == 3 and isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    if not os.path.exists(cache_dir):
+        os.makedirs(cache_dir)
+
+    # Get eTag to add to filename, if it exists.
+    if url.startswith("s3://"):
+        print('Not supported due to colab demo. Sorry!')
+        raise
+    else:
+        response = requests.head(url, allow_redirects=True)
+        if response.status_code != 200:
+            raise IOError("HEAD request failed for url {} with status code {}"
+                          .format(url, response.status_code))
+        etag = response.headers.get("ETag")
+
+    filename = url_to_filename(url, etag)
+
+    # get cache path to put the file
+    cache_path = os.path.join(cache_dir, filename)
+
+    if not os.path.exists(cache_path):
+        # Download to temporary file, then copy to cache dir once finished.
+        # Otherwise you get corrupt cache entries if the download gets interrupted.
+        with tempfile.NamedTemporaryFile() as temp_file:
+            logger.info("%s not found in cache, downloading to %s", url, temp_file.name)
+
+            # GET file object
+            if url.startswith("s3://"):
+                print('Not supported due to colab demo. Sorry!')
+                raise
+            else:
+                http_get(url, temp_file)
+
+            # we are copying the file before closing it, so flush to avoid truncation
+            temp_file.flush()
+            # shutil.copyfileobj() starts at the current position, so go to the start
+            temp_file.seek(0)
+
+            logger.info("copying %s to cache at %s", temp_file.name, cache_path)
+            with open(cache_path, 'wb') as cache_file:
+                shutil.copyfileobj(temp_file, cache_file)
+
+            logger.info("creating metadata file for %s", cache_path)
+            meta = {'url': url, 'etag': etag}
+            meta_path = cache_path + '.json'
+            with open(meta_path, 'w', encoding="utf-8") as meta_file:
+                json.dump(meta, meta_file)
+
+            logger.info("removing temp file %s", temp_file.name)
+
+    return cache_path
+
+
+def read_set_from_file(filename):
+    '''
+    Extract a de-duped collection (set) of text from a file.
+    Expected file format is one item per line.
+    '''
+    collection = set()
+    with open(filename, 'r', encoding='utf-8') as file_:
+        for line in file_:
+            collection.add(line.rstrip())
+    return collection
+
+
+def get_file_extension(path, dot=True, lower=True):
+    ext = os.path.splitext(path)[1]
+    ext = ext if dot else ext[1:]
+    return ext.lower() if lower else ext

networks/biggan/model.py

@@ -0,0 +1,352 @@
+# coding: utf-8
+""" BigGAN PyTorch model.
+    From "Large Scale GAN Training for High Fidelity Natural Image Synthesis"
+    By Andrew Brocky, Jeff Donahuey and Karen Simonyan.
+    https://openreview.net/forum?id=B1xsqj09Fm
+
+    PyTorch version implemented from the computational graph of the TF Hub module for BigGAN.
+    Some part of the code are adapted from https://github.com/brain-research/self-attention-gan
+
+    This version only comprises the generator (since the discriminator's weights are not released).
+    This version only comprises the "deep" version of BigGAN (see publication).
+"""
+from __future__ import (absolute_import, division, print_function, unicode_literals)
+
+import os
+import logging
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import BigGANConfig
+from .file_utils import cached_path
+
+logger = logging.getLogger(__name__)
+
+PRETRAINED_MODEL_ARCHIVE_MAP = {
+    'biggan-deep-128': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-128-pytorch_model.bin",
+    'biggan-deep-256': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-256-pytorch_model.bin",
+    'biggan-deep-512': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-512-pytorch_model.bin",
+}
+
+PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    'biggan-deep-128': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-128-config.json",
+    'biggan-deep-256': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-256-config.json",
+    'biggan-deep-512': "https://s3.amazonaws.com/models.huggingface.co/biggan/biggan-deep-512-config.json",
+}
+
+WEIGHTS_NAME = 'pytorch_model.bin'
+CONFIG_NAME = 'config.json'
+
+
+def snconv2d(eps=1e-12, **kwargs):
+    return nn.utils.spectral_norm(nn.Conv2d(**kwargs), eps=eps)
+
+def snlinear(eps=1e-12, **kwargs):
+    return nn.utils.spectral_norm(nn.Linear(**kwargs), eps=eps)
+
+def sn_embedding(eps=1e-12, **kwargs):
+    return nn.utils.spectral_norm(nn.Embedding(**kwargs), eps=eps)
+
+class SelfAttn(nn.Module):
+    """ Self attention Layer"""
+    def __init__(self, in_channels, eps=1e-12):
+        super(SelfAttn, self).__init__()
+        self.in_channels = in_channels
+        self.snconv1x1_theta = snconv2d(in_channels=in_channels, out_channels=in_channels//8,
+                                        kernel_size=1, bias=False, eps=eps)
+        self.snconv1x1_phi = snconv2d(in_channels=in_channels, out_channels=in_channels//8,
+                                      kernel_size=1, bias=False, eps=eps)
+        self.snconv1x1_g = snconv2d(in_channels=in_channels, out_channels=in_channels//2,
+                                    kernel_size=1, bias=False, eps=eps)
+        self.snconv1x1_o_conv = snconv2d(in_channels=in_channels//2, out_channels=in_channels,
+                                         kernel_size=1, bias=False, eps=eps)
+        self.maxpool = nn.MaxPool2d(2, stride=2, padding=0)
+        self.softmax  = nn.Softmax(dim=-1)
+        self.gamma = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        _, ch, h, w = x.size()
+        # Theta path
+        theta = self.snconv1x1_theta(x)
+        theta = theta.view(-1, ch//8, h*w)
+        # Phi path
+        phi = self.snconv1x1_phi(x)
+        phi = self.maxpool(phi)
+        phi = phi.view(-1, ch//8, h*w//4)
+        # Attn map
+        attn = torch.bmm(theta.permute(0, 2, 1), phi)
+        attn = self.softmax(attn)
+        # g path
+        g = self.snconv1x1_g(x)
+        g = self.maxpool(g)
+        g = g.view(-1, ch//2, h*w//4)
+        # Attn_g - o_conv
+        attn_g = torch.bmm(g, attn.permute(0, 2, 1))
+        attn_g = attn_g.view(-1, ch//2, h, w)
+        attn_g = self.snconv1x1_o_conv(attn_g)
+        # Out
+        out = x + self.gamma*attn_g
+        return out
+
+
+class BigGANBatchNorm(nn.Module):
+    """ This is a batch norm module that can handle conditional input and can be provided with pre-computed
+        activation means and variances for various truncation parameters.
+
+        We cannot just rely on torch.batch_norm since it cannot handle
+        batched weights (pytorch 1.0.1). We computate batch_norm our-self without updating running means and variances.
+        If you want to train this model you should add running means and variance computation logic.
+    """
+    def __init__(self, num_features, condition_vector_dim=None, n_stats=51, eps=1e-4, conditional=True):
+        super(BigGANBatchNorm, self).__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.conditional = conditional
+
+        # We use pre-computed statistics for n_stats values of truncation between 0 and 1
+        self.register_buffer('running_means', torch.zeros(n_stats, num_features))
+        self.register_buffer('running_vars', torch.ones(n_stats, num_features))
+        self.step_size = 1.0 / (n_stats - 1)
+
+        if conditional:
+            assert condition_vector_dim is not None
+            self.scale = snlinear(in_features=condition_vector_dim, out_features=num_features, bias=False, eps=eps)
+            self.offset = snlinear(in_features=condition_vector_dim, out_features=num_features, bias=False, eps=eps)
+        else:
+            self.weight = torch.nn.Parameter(torch.Tensor(num_features))
+            self.bias = torch.nn.Parameter(torch.Tensor(num_features))
+
+    def forward(self, x, truncation, condition_vector=None):
+        # Retreive pre-computed statistics associated to this truncation
+        coef, start_idx = math.modf(truncation / self.step_size)
+        start_idx = int(start_idx)
+        if coef != 0.0:  # Interpolate
+            running_mean = self.running_means[start_idx] * coef + self.running_means[start_idx + 1] * (1 - coef)
+            running_var = self.running_vars[start_idx] * coef + self.running_vars[start_idx + 1] * (1 - coef)
+        else:
+            running_mean = self.running_means[start_idx]
+            running_var = self.running_vars[start_idx]
+
+        if self.conditional:
+            running_mean = running_mean.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+            running_var = running_var.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+
+            weight = 1 + self.scale(condition_vector).unsqueeze(-1).unsqueeze(-1)
+            bias = self.offset(condition_vector).unsqueeze(-1).unsqueeze(-1)
+
+            out = (x - running_mean) / torch.sqrt(running_var + self.eps) * weight + bias
+        else:
+            out = F.batch_norm(x, running_mean, running_var, self.weight, self.bias,
+                               training=False, momentum=0.0, eps=self.eps)
+
+        return out
+
+
+class GenBlock(nn.Module):
+    def __init__(self, in_size, out_size, condition_vector_dim, reduction_factor=4, up_sample=False,
+                 n_stats=51, eps=1e-12):
+        super(GenBlock, self).__init__()
+        self.up_sample = up_sample
+        self.drop_channels = (in_size != out_size)
+        middle_size = in_size // reduction_factor
+
+        self.bn_0 = BigGANBatchNorm(in_size, condition_vector_dim, n_stats=n_stats, eps=eps, conditional=True)
+        self.conv_0 = snconv2d(in_channels=in_size, out_channels=middle_size, kernel_size=1, eps=eps)
+
+        self.bn_1 = BigGANBatchNorm(middle_size, condition_vector_dim, n_stats=n_stats, eps=eps, conditional=True)
+        self.conv_1 = snconv2d(in_channels=middle_size, out_channels=middle_size, kernel_size=3, padding=1, eps=eps)
+
+        self.bn_2 = BigGANBatchNorm(middle_size, condition_vector_dim, n_stats=n_stats, eps=eps, conditional=True)
+        self.conv_2 = snconv2d(in_channels=middle_size, out_channels=middle_size, kernel_size=3, padding=1, eps=eps)
+
+        self.bn_3 = BigGANBatchNorm(middle_size, condition_vector_dim, n_stats=n_stats, eps=eps, conditional=True)
+        self.conv_3 = snconv2d(in_channels=middle_size, out_channels=out_size, kernel_size=1, eps=eps)
+
+        self.relu = nn.ReLU()
+
+    def forward(self, x, cond_vector, truncation):
+        x0 = x
+
+        x = self.bn_0(x, truncation, cond_vector)
+        x = self.relu(x)
+        x = self.conv_0(x)
+
+        x = self.bn_1(x, truncation, cond_vector)
+        x = self.relu(x)
+        if self.up_sample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv_1(x)
+
+        x = self.bn_2(x, truncation, cond_vector)
+        x = self.relu(x)
+        x = self.conv_2(x)
+
+        x = self.bn_3(x, truncation, cond_vector)
+        x = self.relu(x)
+        x = self.conv_3(x)
+
+        if self.drop_channels:
+            new_channels = x0.shape[1] // 2
+            x0 = x0[:, :new_channels, ...]
+        if self.up_sample:
+            x0 = F.interpolate(x0, scale_factor=2, mode='nearest')
+
+        out = x + x0
+        return out
+
+class Generator(nn.Module):
+    def __init__(self, config):
+        super(Generator, self).__init__()
+        self.config = config
+        ch = config.channel_width
+        condition_vector_dim = config.z_dim * 2
+
+        self.gen_z = snlinear(in_features=condition_vector_dim,
+                              out_features=4 * 4 * 16 * ch, eps=config.eps)
+
+        layers = []
+        for i, layer in enumerate(config.layers):
+            if i == config.attention_layer_position:
+                layers.append(SelfAttn(ch*layer[1], eps=config.eps))
+            layers.append(GenBlock(ch*layer[1],
+                                   ch*layer[2],
+                                   condition_vector_dim,
+                                   up_sample=layer[0],
+                                   n_stats=config.n_stats,
+                                   eps=config.eps))
+        self.layers = nn.ModuleList(layers)
+
+        self.bn = BigGANBatchNorm(ch, n_stats=config.n_stats, eps=config.eps, conditional=False)
+        self.relu = nn.ReLU()
+        self.conv_to_rgb = snconv2d(in_channels=ch, out_channels=ch, kernel_size=3, padding=1, eps=config.eps)
+        self.tanh = nn.Tanh()
+
+    def forward(self, cond_vector, truncation, z=None, start=0, stop=None):
+        # We use this conversion step to be able to use TF weights:
+        # TF convention on shape is [batch, height, width, channels]
+        # PT convention on shape is [batch, channels, height, width]
+        if start == 0 and z is None:
+            z = self.gen_z(cond_vector)
+            z = z.view(-1, 4, 4, 16 * self.config.channel_width)
+            z = z.permute(0, 3, 1, 2).contiguous()
+
+        if stop is None: stop = len(self.layers)
+
+        # for i, layer in enumerate(self.layers):
+        for i in range(start, stop):
+            if isinstance(self.layers[i], GenBlock):
+                z = self.layers[i](z, cond_vector, truncation)
+            else:
+                z = self.layers[i](z)
+
+        if stop == len(self.layers):
+            z = self.bn(z, truncation)
+            z = self.relu(z)
+            z = self.conv_to_rgb(z)
+            z = z[:, :3, ...]
+            z = self.tanh(z)
+
+        # for i, layer in enumerate(self.layers):
+        #     if isinstance(layer, GenBlock):
+        #         z = layer(z, cond_vector, truncation)
+        #     else:
+        #         z = layer(z)
+
+        # z = self.bn(z, truncation)
+        # z = self.relu(z)
+        # z = self.conv_to_rgb(z)
+        # z = z[:, :3, ...]
+        # z = self.tanh(z)
+        return z
+
+class BigGAN(nn.Module):
+    """BigGAN Generator."""
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None, *inputs, **kwargs):
+        if pretrained_model_name_or_path in PRETRAINED_MODEL_ARCHIVE_MAP:
+            model_file = PRETRAINED_MODEL_ARCHIVE_MAP[pretrained_model_name_or_path]
+            config_file = PRETRAINED_CONFIG_ARCHIVE_MAP[pretrained_model_name_or_path]
+        else:
+            model_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+            config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
+
+        try:
+            resolved_model_file = cached_path(model_file, cache_dir=cache_dir)
+            resolved_config_file = cached_path(config_file, cache_dir=cache_dir)
+        except EnvironmentError:
+            logger.error("Wrong model name, should be a valid path to a folder containing "
+                         "a {} file and a {} file or a model name in {}".format(
+                         WEIGHTS_NAME, CONFIG_NAME, PRETRAINED_MODEL_ARCHIVE_MAP.keys()))
+            raise
+
+        logger.info("loading model {} from cache at {}".format(pretrained_model_name_or_path, resolved_model_file))
+
+        # Load config
+        config = BigGANConfig.from_json_file(resolved_config_file)
+        logger.info("Model config {}".format(config))
+
+        # Instantiate model.
+        model = cls(config, *inputs, **kwargs)
+        state_dict = torch.load(resolved_model_file, map_location='cpu' if not torch.cuda.is_available() else None)
+        model.load_state_dict(state_dict, strict=False)
+        return model
+
+    def __init__(self, config):
+        super(BigGAN, self).__init__()
+        self.config = config
+        self.embeddings = nn.Linear(config.num_classes, config.z_dim, bias=False)
+        self.generator = Generator(config)
+        # self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        # print(f'device: {self.device}')
+        # self.generator.to(self.device)
+
+    def forward(self, z, class_label, truncation, cond_vector=None, start=0, stop=None):
+        assert 0 < truncation <= 1
+
+        results = {}
+        if start == 0 and cond_vector is None:
+            embed = self.embeddings(class_label)
+            cond_vector = torch.cat((z, embed), dim=1)
+            results['cond_vector'] = cond_vector
+
+        results['z'] = self.generator(cond_vector, truncation, z=None if start == 0 else z, start=start, stop=stop)
+        return results
+
+
+if __name__ == "__main__":
+    import PIL
+    from .utils import truncated_noise_sample, save_as_images, one_hot_from_names
+    from .convert_tf_to_pytorch import load_tf_weights_in_biggan
+
+    load_cache = False
+    cache_path = './saved_model.pt'
+    config = BigGANConfig()
+    model = BigGAN(config)
+    if not load_cache:
+        model = load_tf_weights_in_biggan(model, config, './models/model_128/', './models/model_128/batchnorms_stats.bin')
+        torch.save(model.state_dict(), cache_path)
+    else:
+        model.load_state_dict(torch.load(cache_path))
+
+    model.eval()
+
+    truncation = 0.4
+    noise = truncated_noise_sample(batch_size=2, truncation=truncation)
+    label = one_hot_from_names('diver', batch_size=2)
+
+    # Tests
+    # noise = np.zeros((1, 128))
+    # label = [983]
+
+    noise = torch.tensor(noise, dtype=torch.float)
+    label = torch.tensor(label, dtype=torch.float)
+    with torch.no_grad():
+        outputs = model(noise, label, truncation)
+    print(outputs.shape)
+
+    save_as_images(outputs)

networks/biggan/utils.py

@@ -0,0 +1,216 @@
+# coding: utf-8
+""" BigGAN utilities to prepare truncated noise samples and convert/save/display output images.
+    Also comprise ImageNet utilities to prepare one hot input vectors for ImageNet classes.
+    We use Wordnet so you can just input a name in a string and automatically get a corresponding
+    imagenet class if it exists (or a hypo/hypernym exists in imagenet).
+"""
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import json
+import logging
+from io import BytesIO
+
+import numpy as np
+from scipy.stats import truncnorm
+
+logger = logging.getLogger(__name__)
+
+NUM_CLASSES = 1000
+
+
+def truncated_noise_sample(batch_size=1, dim_z=128, truncation=1., seed=None):
+    """ Create a truncated noise vector.
+        Params:
+            batch_size: batch size.
+            dim_z: dimension of z
+            truncation: truncation value to use
+            seed: seed for the random generator
+        Output:
+            array of shape (batch_size, dim_z)
+    """
+    state = None if seed is None else np.random.RandomState(seed)
+    values = truncnorm.rvs(-2, 2, size=(batch_size, dim_z), random_state=state).astype(np.float32)
+    return truncation * values
+
+
+def convert_to_images(obj):
+    """ Convert an output tensor from BigGAN in a list of images.
+        Params:
+            obj: tensor or numpy array of shape (batch_size, channels, height, width)
+        Output:
+            list of Pillow Images of size (height, width)
+    """
+    try:
+        import PIL
+    except ImportError:
+        raise ImportError("Please install Pillow to use images: pip install Pillow")
+
+    if not isinstance(obj, np.ndarray):
+        obj = obj.detach().numpy()
+
+    obj = obj.transpose((0, 2, 3, 1))
+    obj = np.clip(((obj + 1) / 2.0) * 256, 0, 255)
+
+    img = []
+    for i, out in enumerate(obj):
+        out_array = np.asarray(np.uint8(out), dtype=np.uint8)
+        img.append(PIL.Image.fromarray(out_array))
+    return img
+
+
+def save_as_images(obj, file_name='output'):
+    """ Convert and save an output tensor from BigGAN in a list of saved images.
+        Params:
+            obj: tensor or numpy array of shape (batch_size, channels, height, width)
+            file_name: path and beggingin of filename to save.
+                Images will be saved as `file_name_{image_number}.png`
+    """
+    img = convert_to_images(obj)
+
+    for i, out in enumerate(img):
+        current_file_name = file_name + '_%d.png' % i
+        logger.info("Saving image to {}".format(current_file_name))
+        out.save(current_file_name, 'png')
+
+
+def display_in_terminal(obj):
+    """ Convert and display an output tensor from BigGAN in the terminal.
+        This function use `libsixel` and will only work in a libsixel-compatible terminal.
+        Please refer to https://github.com/saitoha/libsixel for more details.
+
+        Params:
+            obj: tensor or numpy array of shape (batch_size, channels, height, width)
+            file_name: path and beggingin of filename to save.
+                Images will be saved as `file_name_{image_number}.png`
+    """
+    try:
+        import PIL
+        from libsixel import (sixel_output_new, sixel_dither_new, sixel_dither_initialize,
+                              sixel_dither_set_palette, sixel_dither_set_pixelformat,
+                              sixel_dither_get, sixel_encode, sixel_dither_unref,
+                              sixel_output_unref, SIXEL_PIXELFORMAT_RGBA8888,
+                              SIXEL_PIXELFORMAT_RGB888, SIXEL_PIXELFORMAT_PAL8,
+                              SIXEL_PIXELFORMAT_G8, SIXEL_PIXELFORMAT_G1)
+    except ImportError:
+        raise ImportError("Display in Terminal requires Pillow, libsixel "
+                          "and a libsixel compatible terminal. "
+                          "Please read info at https://github.com/saitoha/libsixel "
+                          "and install with pip install Pillow libsixel-python")
+
+    s = BytesIO()
+
+    images = convert_to_images(obj)
+    widths, heights = zip(*(i.size for i in images))
+
+    output_width = sum(widths)
+    output_height = max(heights)
+
+    output_image = PIL.Image.new('RGB', (output_width, output_height))
+
+    x_offset = 0
+    for im in images:
+        output_image.paste(im, (x_offset,0))
+        x_offset += im.size[0]
+
+    try:
+        data = output_image.tobytes()
+    except NotImplementedError:
+        data = output_image.tostring()
+    output = sixel_output_new(lambda data, s: s.write(data), s)
+
+    try:
+        if output_image.mode == 'RGBA':
+            dither = sixel_dither_new(256)
+            sixel_dither_initialize(dither, data, output_width, output_height, SIXEL_PIXELFORMAT_RGBA8888)
+        elif output_image.mode == 'RGB':
+            dither = sixel_dither_new(256)
+            sixel_dither_initialize(dither, data, output_width, output_height, SIXEL_PIXELFORMAT_RGB888)
+        elif output_image.mode == 'P':
+            palette = output_image.getpalette()
+            dither = sixel_dither_new(256)
+            sixel_dither_set_palette(dither, palette)
+            sixel_dither_set_pixelformat(dither, SIXEL_PIXELFORMAT_PAL8)
+        elif output_image.mode == 'L':
+            dither = sixel_dither_get(SIXEL_BUILTIN_G8)
+            sixel_dither_set_pixelformat(dither, SIXEL_PIXELFORMAT_G8)
+        elif output_image.mode == '1':
+            dither = sixel_dither_get(SIXEL_BUILTIN_G1)
+            sixel_dither_set_pixelformat(dither, SIXEL_PIXELFORMAT_G1)
+        else:
+            raise RuntimeError('unexpected output_image mode')
+        try:
+            sixel_encode(data, output_width, output_height, 1, dither, output)
+            print(s.getvalue().decode('ascii'))
+        finally:
+            sixel_dither_unref(dither)
+    finally:
+        sixel_output_unref(output)
+
+
+def one_hot_from_int(int_or_list, batch_size=1):
+    """ Create a one-hot vector from a class index or a list of class indices.
+        Params:
+            int_or_list: int, or list of int, of the imagenet classes (between 0 and 999)
+            batch_size: batch size.
+                If int_or_list is an int create a batch of identical classes.
+                If int_or_list is a list, we should have `len(int_or_list) == batch_size`
+        Output:
+            array of shape (batch_size, 1000)
+    """
+    if isinstance(int_or_list, int):
+        int_or_list = [int_or_list]
+
+    if len(int_or_list) == 1 and batch_size > 1:
+        int_or_list = [int_or_list[0]] * batch_size
+
+    assert batch_size == len(int_or_list)
+
+    array = np.zeros((batch_size, NUM_CLASSES), dtype=np.float32)
+    for i, j in enumerate(int_or_list):
+        array[i, j] = 1.0
+    return array
+
+
+def one_hot_from_names(class_name_or_list, batch_size=1):
+    """ Create a one-hot vector from the name of an imagenet class ('tennis ball', 'daisy', ...).
+        We use NLTK's wordnet search to try to find the relevant synset of ImageNet and take the first one.
+        If we can't find it direcly, we look at the hyponyms and hypernyms of the class name.
+
+        Params:
+            class_name_or_list: string containing the name of an imagenet object or a list of such strings (for a batch).
+        Output:
+            array of shape (batch_size, 1000)
+    """
+    try:
+        from nltk.corpus import wordnet as wn
+    except ImportError:
+        raise ImportError("You need to install nltk to use this function")
+
+    if not isinstance(class_name_or_list, (list, tuple)):
+        class_name_or_list = [class_name_or_list]
+    else:
+        batch_size = max(batch_size, len(class_name_or_list))
+
+    classes = []
+    for class_name in class_name_or_list:
+        class_name = class_name.replace(" ", "_")
+
+        original_synsets = wn.synsets(class_name)
+        original_synsets = list(filter(lambda s: s.pos() == 'n', original_synsets))  # keep only names
+        if not original_synsets:
+            return None
+
+        possible_synsets = list(filter(lambda s: s.offset() in IMAGENET, original_synsets))
+        if possible_synsets:
+            classes.append(IMAGENET[possible_synsets[0].offset()])
+        else:
+            # try hypernyms and hyponyms
+            possible_synsets = sum([s.hypernyms() + s.hyponyms() for s in original_synsets], [])
+            possible_synsets = list(filter(lambda s: s.offset() in IMAGENET, possible_synsets))
+            if possible_synsets:
+                classes.append(IMAGENET[possible_synsets[0].offset()])
+
+    return one_hot_from_int(classes, batch_size=batch_size)
+
+
+IMAGENET = {1440764: 0, 1443537: 1, 1484850: 2, 1491361: 3, 1494475: 4, 1496331: 5, 1498041: 6, 1514668: 7, 1514859: 8, 1518878: 9, 1530575: 10, 1531178: 11, 1532829: 12, 1534433: 13, 1537544: 14, 1558993: 15, 1560419: 16, 1580077: 17, 1582220: 18, 1592084: 19, 1601694: 20, 1608432: 21, 1614925: 22, 1616318: 23, 1622779: 24, 1629819: 25, 1630670: 26, 1631663: 27, 1632458: 28, 1632777: 29, 1641577: 30, 1644373: 31, 1644900: 32, 1664065: 33, 1665541: 34, 1667114: 35, 1667778: 36, 1669191: 37, 1675722: 38, 1677366: 39, 1682714: 40, 1685808: 41, 1687978: 42, 1688243: 43, 1689811: 44, 1692333: 45, 1693334: 46, 1694178: 47, 1695060: 48, 1697457: 49, 1698640: 50, 1704323: 51, 1728572: 52, 1728920: 53, 1729322: 54, 1729977: 55, 1734418: 56, 1735189: 57, 1737021: 58, 1739381: 59, 1740131: 60, 1742172: 61, 1744401: 62, 1748264: 63, 1749939: 64, 1751748: 65, 1753488: 66, 1755581: 67, 1756291: 68, 1768244: 69, 1770081: 70, 1770393: 71, 1773157: 72, 1773549: 73, 1773797: 74, 1774384: 75, 1774750: 76, 1775062: 77, 1776313: 78, 1784675: 79, 1795545: 80, 1796340: 81, 1797886: 82, 1798484: 83, 1806143: 84, 1806567: 85, 1807496: 86, 1817953: 87, 1818515: 88, 1819313: 89, 1820546: 90, 1824575: 91, 1828970: 92, 1829413: 93, 1833805: 94, 1843065: 95, 1843383: 96, 1847000: 97, 1855032: 98, 1855672: 99, 1860187: 100, 1871265: 101, 1872401: 102, 1873310: 103, 1877812: 104, 1882714: 105, 1883070: 106, 1910747: 107, 1914609: 108, 1917289: 109, 1924916: 110, 1930112: 111, 1943899: 112, 1944390: 113, 1945685: 114, 1950731: 115, 1955084: 116, 1968897: 117, 1978287: 118, 1978455: 119, 1980166: 120, 1981276: 121, 1983481: 122, 1984695: 123, 1985128: 124, 1986214: 125, 1990800: 126, 2002556: 127, 2002724: 128, 2006656: 129, 2007558: 130, 2009229: 131, 2009912: 132, 2011460: 133, 2012849: 134, 2013706: 135, 2017213: 136, 2018207: 137, 2018795: 138, 2025239: 139, 2027492: 140, 2028035: 141, 2033041: 142, 2037110: 143, 2051845: 144, 2056570: 145, 2058221: 146, 2066245: 147, 2071294: 148, 2074367: 149, 2077923: 150, 2085620: 151, 2085782: 152, 2085936: 153, 2086079: 154, 2086240: 155, 2086646: 156, 2086910: 157, 2087046: 158, 2087394: 159, 2088094: 160, 2088238: 161, 2088364: 162, 2088466: 163, 2088632: 164, 2089078: 165, 2089867: 166, 2089973: 167, 2090379: 168, 2090622: 169, 2090721: 170, 2091032: 171, 2091134: 172, 2091244: 173, 2091467: 174, 2091635: 175, 2091831: 176, 2092002: 177, 2092339: 178, 2093256: 179, 2093428: 180, 2093647: 181, 2093754: 182, 2093859: 183, 2093991: 184, 2094114: 185, 2094258: 186, 2094433: 187, 2095314: 188, 2095570: 189, 2095889: 190, 2096051: 191, 2096177: 192, 2096294: 193, 2096437: 194, 2096585: 195, 2097047: 196, 2097130: 197, 2097209: 198, 2097298: 199, 2097474: 200, 2097658: 201, 2098105: 202, 2098286: 203, 2098413: 204, 2099267: 205, 2099429: 206, 2099601: 207, 2099712: 208, 2099849: 209, 2100236: 210, 2100583: 211, 2100735: 212, 2100877: 213, 2101006: 214, 2101388: 215, 2101556: 216, 2102040: 217, 2102177: 218, 2102318: 219, 2102480: 220, 2102973: 221, 2104029: 222, 2104365: 223, 2105056: 224, 2105162: 225, 2105251: 226, 2105412: 227, 2105505: 228, 2105641: 229, 2105855: 230, 2106030: 231, 2106166: 232, 2106382: 233, 2106550: 234, 2106662: 235, 2107142: 236, 2107312: 237, 2107574: 238, 2107683: 239, 2107908: 240, 2108000: 241, 2108089: 242, 2108422: 243, 2108551: 244, 2108915: 245, 2109047: 246, 2109525: 247, 2109961: 248, 2110063: 249, 2110185: 250, 2110341: 251, 2110627: 252, 2110806: 253, 2110958: 254, 2111129: 255, 2111277: 256, 2111500: 257, 2111889: 258, 2112018: 259, 2112137: 260, 2112350: 261, 2112706: 262, 2113023: 263, 2113186: 264, 2113624: 265, 2113712: 266, 2113799: 267, 2113978: 268, 2114367: 269, 2114548: 270, 2114712: 271, 2114855: 272, 2115641: 273, 2115913: 274, 2116738: 275, 2117135: 276, 2119022: 277, 2119789: 278, 2120079: 279, 2120505: 280, 2123045: 281, 2123159: 282, 2123394: 283, 2123597: 284, 2124075: 285, 2125311: 286, 2127052: 287, 2128385: 288, 2128757: 289, 2128925: 290, 2129165: 291, 2129604: 292, 2130308: 293, 2132136: 294, 2133161: 295, 2134084: 296, 2134418: 297, 2137549: 298, 2138441: 299, 2165105: 300, 2165456: 301, 2167151: 302, 2168699: 303, 2169497: 304, 2172182: 305, 2174001: 306, 2177972: 307, 2190166: 308, 2206856: 309, 2219486: 310, 2226429: 311, 2229544: 312, 2231487: 313, 2233338: 314, 2236044: 315, 2256656: 316, 2259212: 317, 2264363: 318, 2268443: 319, 2268853: 320, 2276258: 321, 2277742: 322, 2279972: 323, 2280649: 324, 2281406: 325, 2281787: 326, 2317335: 327, 2319095: 328, 2321529: 329, 2325366: 330, 2326432: 331, 2328150: 332, 2342885: 333, 2346627: 334, 2356798: 335, 2361337: 336, 2363005: 337, 2364673: 338, 2389026: 339, 2391049: 340, 2395406: 341, 2396427: 342, 2397096: 343, 2398521: 344, 2403003: 345, 2408429: 346, 2410509: 347, 2412080: 348, 2415577: 349, 2417914: 350, 2422106: 351, 2422699: 352, 2423022: 353, 2437312: 354, 2437616: 355, 2441942: 356, 2442845: 357, 2443114: 358, 2443484: 359, 2444819: 360, 2445715: 361, 2447366: 362, 2454379: 363, 2457408: 364, 2480495: 365, 2480855: 366, 2481823: 367, 2483362: 368, 2483708: 369, 2484975: 370, 2486261: 371, 2486410: 372, 2487347: 373, 2488291: 374, 2488702: 375, 2489166: 376, 2490219: 377, 2492035: 378, 2492660: 379, 2493509: 380, 2493793: 381, 2494079: 382, 2497673: 383, 2500267: 384, 2504013: 385, 2504458: 386, 2509815: 387, 2510455: 388, 2514041: 389, 2526121: 390, 2536864: 391, 2606052: 392, 2607072: 393, 2640242: 394, 2641379: 395, 2643566: 396, 2655020: 397, 2666196: 398, 2667093: 399, 2669723: 400, 2672831: 401, 2676566: 402, 2687172: 403, 2690373: 404, 2692877: 405, 2699494: 406, 2701002: 407, 2704792: 408, 2708093: 409, 2727426: 410, 2730930: 411, 2747177: 412, 2749479: 413, 2769748: 414, 2776631: 415, 2777292: 416, 2782093: 417, 2783161: 418, 2786058: 419, 2787622: 420, 2788148: 421, 2790996: 422, 2791124: 423, 2791270: 424, 2793495: 425, 2794156: 426, 2795169: 427, 2797295: 428, 2799071: 429, 2802426: 430, 2804414: 431, 2804610: 432, 2807133: 433, 2808304: 434, 2808440: 435, 2814533: 436, 2814860: 437, 2815834: 438, 2817516: 439, 2823428: 440, 2823750: 441, 2825657: 442, 2834397: 443, 2835271: 444, 2837789: 445, 2840245: 446, 2841315: 447, 2843684: 448, 2859443: 449, 2860847: 450, 2865351: 451, 2869837: 452, 2870880: 453, 2871525: 454, 2877765: 455, 2879718: 456, 2883205: 457, 2892201: 458, 2892767: 459, 2894605: 460, 2895154: 461, 2906734: 462, 2909870: 463, 2910353: 464, 2916936: 465, 2917067: 466, 2927161: 467, 2930766: 468, 2939185: 469, 2948072: 470, 2950826: 471, 2951358: 472, 2951585: 473, 2963159: 474, 2965783: 475, 2966193: 476, 2966687: 477, 2971356: 478, 2974003: 479, 2977058: 480, 2978881: 481, 2979186: 482, 2980441: 483, 2981792: 484, 2988304: 485, 2992211: 486, 2992529: 487, 2999410: 488, 3000134: 489, 3000247: 490, 3000684: 491, 3014705: 492, 3016953: 493, 3017168: 494, 3018349: 495, 3026506: 496, 3028079: 497, 3032252: 498, 3041632: 499, 3042490: 500, 3045698: 501, 3047690: 502, 3062245: 503, 3063599: 504, 3063689: 505, 3065424: 506, 3075370: 507, 3085013: 508, 3089624: 509, 3095699: 510, 3100240: 511, 3109150: 512, 3110669: 513, 3124043: 514, 3124170: 515, 3125729: 516, 3126707: 517, 3127747: 518, 3127925: 519, 3131574: 520, 3133878: 521, 3134739: 522, 3141823: 523, 3146219: 524, 3160309: 525, 3179701: 526, 3180011: 527, 3187595: 528, 3188531: 529, 3196217: 530, 3197337: 531, 3201208: 532, 3207743: 533, 3207941: 534, 3208938: 535, 3216828: 536, 3218198: 537, 3220513: 538, 3223299: 539, 3240683: 540, 3249569: 541, 3250847: 542, 3255030: 543, 3259280: 544, 3271574: 545, 3272010: 546, 3272562: 547, 3290653: 548, 3291819: 549, 3297495: 550, 3314780: 551, 3325584: 552, 3337140: 553, 3344393: 554, 3345487: 555, 3347037: 556, 3355925: 557, 3372029: 558, 3376595: 559, 3379051: 560, 3384352: 561, 3388043: 562, 3388183: 563, 3388549: 564, 3393912: 565, 3394916: 566, 3400231: 567, 3404251: 568, 3417042: 569, 3424325: 570, 3425413: 571, 3443371: 572, 3444034: 573, 3445777: 574, 3445924: 575, 3447447: 576, 3447721: 577, 3450230: 578, 3452741: 579, 3457902: 580, 3459775: 581, 3461385: 582, 3467068: 583, 3476684: 584, 3476991: 585, 3478589: 586, 3481172: 587, 3482405: 588, 3483316: 589, 3485407: 590, 3485794: 591, 3492542: 592, 3494278: 593, 3495258: 594, 3496892: 595, 3498962: 596, 3527444: 597, 3529860: 598, 3530642: 599, 3532672: 600, 3534580: 601, 3535780: 602, 3538406: 603, 3544143: 604, 3584254: 605, 3584829: 606, 3590841: 607, 3594734: 608, 3594945: 609, 3595614: 610, 3598930: 611, 3599486: 612, 3602883: 613, 3617480: 614, 3623198: 615, 3627232: 616, 3630383: 617, 3633091: 618, 3637318: 619, 3642806: 620, 3649909: 621, 3657121: 622, 3658185: 623, 3661043: 624, 3662601: 625, 3666591: 626, 3670208: 627, 3673027: 628, 3676483: 629, 3680355: 630, 3690938: 631, 3691459: 632, 3692522: 633, 3697007: 634, 3706229: 635, 3709823: 636, 3710193: 637, 3710637: 638, 3710721: 639, 3717622: 640, 3720891: 641, 3721384: 642, 3724870: 643, 3729826: 644, 3733131: 645, 3733281: 646, 3733805: 647, 3742115: 648, 3743016: 649, 3759954: 650, 3761084: 651, 3763968: 652, 3764736: 653, 3769881: 654, 3770439: 655, 3770679: 656, 3773504: 657, 3775071: 658, 3775546: 659, 3776460: 660, 3777568: 661, 3777754: 662, 3781244: 663, 3782006: 664, 3785016: 665, 3786901: 666, 3787032: 667, 3788195: 668, 3788365: 669, 3791053: 670, 3792782: 671, 3792972: 672, 3793489: 673, 3794056: 674, 3796401: 675, 3803284: 676, 3804744: 677, 3814639: 678, 3814906: 679, 3825788: 680, 3832673: 681, 3837869: 682, 3838899: 683, 3840681: 684, 3841143: 685, 3843555: 686, 3854065: 687, 3857828: 688, 3866082: 689, 3868242: 690, 3868863: 691, 3871628: 692, 3873416: 693, 3874293: 694, 3874599: 695, 3876231: 696, 3877472: 697, 3877845: 698, 3884397: 699, 3887697: 700, 3888257: 701, 3888605: 702, 3891251: 703, 3891332: 704, 3895866: 705, 3899768: 706, 3902125: 707, 3903868: 708, 3908618: 709, 3908714: 710, 3916031: 711, 3920288: 712, 3924679: 713, 3929660: 714, 3929855: 715, 3930313: 716, 3930630: 717, 3933933: 718, 3935335: 719, 3937543: 720, 3938244: 721, 3942813: 722, 3944341: 723, 3947888: 724, 3950228: 725, 3954731: 726, 3956157: 727, 3958227: 728, 3961711: 729, 3967562: 730, 3970156: 731, 3976467: 732, 3976657: 733, 3977966: 734, 3980874: 735, 3982430: 736, 3983396: 737, 3991062: 738, 3992509: 739, 3995372: 740, 3998194: 741, 4004767: 742, 4005630: 743, 4008634: 744, 4009552: 745, 4019541: 746, 4023962: 747, 4026417: 748, 4033901: 749, 4033995: 750, 4037443: 751, 4039381: 752, 4040759: 753, 4041544: 754, 4044716: 755, 4049303: 756, 4065272: 757, 4067472: 758, 4069434: 759, 4070727: 760, 4074963: 761, 4081281: 762, 4086273: 763, 4090263: 764, 4099969: 765, 4111531: 766, 4116512: 767, 4118538: 768, 4118776: 769, 4120489: 770, 4125021: 771, 4127249: 772, 4131690: 773, 4133789: 774, 4136333: 775, 4141076: 776, 4141327: 777, 4141975: 778, 4146614: 779, 4147183: 780, 4149813: 781, 4152593: 782, 4153751: 783, 4154565: 784, 4162706: 785, 4179913: 786, 4192698: 787, 4200800: 788, 4201297: 789, 4204238: 790, 4204347: 791, 4208210: 792, 4209133: 793, 4209239: 794, 4228054: 795, 4229816: 796, 4235860: 797, 4238763: 798, 4239074: 799, 4243546: 800, 4251144: 801, 4252077: 802, 4252225: 803, 4254120: 804, 4254680: 805, 4254777: 806, 4258138: 807, 4259630: 808, 4263257: 809, 4264628: 810, 4265275: 811, 4266014: 812, 4270147: 813, 4273569: 814, 4275548: 815, 4277352: 816, 4285008: 817, 4286575: 818, 4296562: 819, 4310018: 820, 4311004: 821, 4311174: 822, 4317175: 823, 4325704: 824, 4326547: 825, 4328186: 826, 4330267: 827, 4332243: 828, 4335435: 829, 4336792: 830, 4344873: 831, 4346328: 832, 4347754: 833, 4350905: 834, 4355338: 835, 4355933: 836, 4356056: 837, 4357314: 838, 4366367: 839, 4367480: 840, 4370456: 841, 4371430: 842, 4371774: 843, 4372370: 844, 4376876: 845, 4380533: 846, 4389033: 847, 4392985: 848, 4398044: 849, 4399382: 850, 4404412: 851, 4409515: 852, 4417672: 853, 4418357: 854, 4423845: 855, 4428191: 856, 4429376: 857, 4435653: 858, 4442312: 859, 4443257: 860, 4447861: 861, 4456115: 862, 4458633: 863, 4461696: 864, 4462240: 865, 4465501: 866, 4467665: 867, 4476259: 868, 4479046: 869, 4482393: 870, 4483307: 871, 4485082: 872, 4486054: 873, 4487081: 874, 4487394: 875, 4493381: 876, 4501370: 877, 4505470: 878, 4507155: 879, 4509417: 880, 4515003: 881, 4517823: 882, 4522168: 883, 4523525: 884, 4525038: 885, 4525305: 886, 4532106: 887, 4532670: 888, 4536866: 889, 4540053: 890, 4542943: 891, 4548280: 892, 4548362: 893, 4550184: 894, 4552348: 895, 4553703: 896, 4554684: 897, 4557648: 898, 4560804: 899, 4562935: 900, 4579145: 901, 4579432: 902, 4584207: 903, 4589890: 904, 4590129: 905, 4591157: 906, 4591713: 907, 4592741: 908, 4596742: 909, 4597913: 910, 4599235: 911, 4604644: 912, 4606251: 913, 4612504: 914, 4613696: 915, 6359193: 916, 6596364: 917, 6785654: 918, 6794110: 919, 6874185: 920, 7248320: 921, 7565083: 922, 7579787: 923, 7583066: 924, 7584110: 925, 7590611: 926, 7613480: 927, 7614500: 928, 7615774: 929, 7684084: 930, 7693725: 931, 7695742: 932, 7697313: 933, 7697537: 934, 7711569: 935, 7714571: 936, 7714990: 937, 7715103: 938, 7716358: 939, 7716906: 940, 7717410: 941, 7717556: 942, 7718472: 943, 7718747: 944, 7720875: 945, 7730033: 946, 7734744: 947, 7742313: 948, 7745940: 949, 7747607: 950, 7749582: 951, 7753113: 952, 7753275: 953, 7753592: 954, 7754684: 955, 7760859: 956, 7768694: 957, 7802026: 958, 7831146: 959, 7836838: 960, 7860988: 961, 7871810: 962, 7873807: 963, 7875152: 964, 7880968: 965, 7892512: 966, 7920052: 967, 7930864: 968, 7932039: 969, 9193705: 970, 9229709: 971, 9246464: 972, 9256479: 973, 9288635: 974, 9332890: 975, 9399592: 976, 9421951: 977, 9428293: 978, 9468604: 979, 9472597: 980, 9835506: 981, 10148035: 982, 10565667: 983, 11879895: 984, 11939491: 985, 12057211: 986, 12144580: 987, 12267677: 988, 12620546: 989, 12768682: 990, 12985857: 991, 12998815: 992, 13037406: 993, 13040303: 994, 13044778: 995, 13052670: 996, 13054560: 997, 13133613: 998, 15075141: 999}

networks/genforce/.gitignore

@@ -0,0 +1,29 @@
+__pycache__/
+*.py[cod]
+
+/.vscode/
+/.idea/
+*.sw[pon]
+
+/data/
+/work_dirs/
+*.jpg
+*.png
+*.jpeg
+*.gif
+*.avi
+*.mp4
+
+*.npy
+*.txt
+*.json
+*.log
+*.html
+*.tar
+*.zip
+events.*
+
+*.pth
+*.pkl
+*.h5
+*.dat

networks/genforce/LICENSE

@@ -0,0 +1,18 @@
+Copyright (c) 2020 GenForce
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

networks/genforce/MODEL_ZOO.md

@@ -0,0 +1,131 @@
+# Model Zoo
+
+## Pre-trained Models
+
+First of all, we thank the following repositories for their work on high-quality image synthesis
+
+- [PGGAN](https://github.com/tkarras/progressive_growing_of_gans)
+- [StyleGAN](https://github.com/NVlabs/stylegan)
+- [StyleGAN2](https://github.com/NVlabs/stylegan2)
+
+Please download the models you need and save them to `checkpoints/`.
+
+| PGGAN Official | | | |
+| :-- | :-- | :-- | :-- |
+| *Face*
+| [celebahq-1024x1024](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EW_3jQ6E7xlKvCSHYrbmkQQBAB8tgIv5W5evdT6-GuXiWw?e=gRifVa&download=1)
+| *Indoor Scene*
+| [bedroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EUZQWGz2GT5Bh_GJLalP63IBvCsXDTOxDFIC_ZBsmoEacA?e=VNXiDb&download=1)       | [livingroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Efzh6qQv6QtCm0YN1lulH-YByqdE3AqlI-E6US_hXMuiig?e=ppdyB2&download=1) | [diningroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EcLb3_hGUkdClompZo27xk0BNmotgbFqdIeu-ZOGJsBMRg?e=xjYpN3&download=1)  | [kitchen-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ESCyg6hpNn1LlHVX_un1wLsBZAORUNkW9MO2kU1X5kafAQ?e=09TbGC&download=1)
+| *Outdoor Scene*
+| [churchoutdoor-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EQ8cKujs2TVGjCL_j6bsnk8BqD9REF2ME2lBnpbTPsqIvA?e=zH55fT&download=1) | [tower-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EeyBJvgRVGJClKr1KKYDF_cBT1FDepRU1-GLqYNh8W9-fQ?e=nrpa5N&download=1)      | [bridge-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EZ2QScfPy19PiDERLJQ3gPMBP4WmvZHwhNFLzfaP2YD8hQ?e=bef1U9&download=1)
+| *Other Scene*
+| [restaurant-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ERvJ4pz8jgtMrcuJXUfcOQEBDugZ099_TetCQs-9-ILCVg?e=qYsVdQ&download=1)    | [classroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EUU9SCOPUxhMoUS4Ceo9kl0BQkVK7d69lA-JeOP-zOWvXw?e=YIB4no&download=1)  | [conferenceroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EX8AF0_6NoJAl5vKFewHWnsBk0r4PK4WsqsMrJyj84TrqQ?e=oNQIZS&download=1)
+| *Animal*
+| [person-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EWu4SqR42YpCoqsVJOcM2cMBcdfXA0j5wZ2hno9X0R9ydQ?e=KuDRns&download=1)        | [cat-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EQdveyUNOMtAue52n6BxoHoB6Yup5-PTvBDmyfUn7Un4Hw?e=7acGbT&download=1)        | [dog-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ESaKyXA5fGlOvXJYDDFbT2kB9c0HlXh9n_wnyhiP05nhow?e=d4aKDV&download=1)         | [bird-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Ef2p4Pd3AKVCmSm00YikCIABhylh2dLPaFjPfPVn3RiTXA?e=9bRitp&download=1)
+| [horse-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EXwCPdv6XqJFtuvFFoswRScBmLJbhKzaC5D_iovl1GFOTw?e=WDdD77&download=1)         | [sheep-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ER6J5EKjAUNFtm9VwLf-uUsBZ5dnqxeKsPxY9ijiPtMhcQ?e=OKtfva&download=1)      | [cow-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ERZLxw7N7xJPm72FyePTbpcByzrr0pH-Fg7qyLt5tYGXwQ?e=ovIPCl&download=1)
+| *Transportation*
+| [car-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EfGc2we47aFDtAY1548pRvsByIju-uXRbkZEFpJotuPKZw?e=DQqVj8&download=1)           | [bicycle-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Ed1dN_FgwmdBgeNWhaRUry8BgwT88-n2ppicSDPx-f7f_Q?e=bxTxnf&download=1)    | [motorbike-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EV3yQdeJXIdPjZbMO0mp2-MBJbKuuBdypzBL4gnedO57Dw?e=tXdvtD&download=1)   | [bus-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Ed7-OYLnq0RCqRlM8qK8wZ8B87dz_NUxIKBrvyFUwRCEbg?e=VP5bmX&download=1)
+| [train-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EedE2cozKOVAkhvbdLd4SfwBknFW8vWZnKiqgeIBbAvCCA?e=BrLpTl&download=1)         | [boat-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Eb39waqQFr9Bp4wO0rC5NHwB0Vz2NGCuqbRPucguBIkDrg?e=lddSyL&download=1)       | [airplane-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Ee6FzIx3KjNDhxrS5mDvpCEB3iQ7TgErmKhbwbV-eF07iw?e=xflPXa&download=1)
+| *Furniture*
+| [bottle-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EWhoy2AFCTZGtEG1UoayWjcB9Kdc_wreJ8p4RlBB93nbNg?e=DMZceU&download=1)        | [chair-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EbQRTfwdostBhXG30Uacn7ABsEUFa-tEW3oxiM5zDYQbRw?e=FkB7T0&download=1)      | [pottedplant-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EWg7hnoGATBOuJvXWr4m7CQBJL9o7nqnD6nOMRhtH2SKXg?e=Zi3hjD&download=1) | [tvmonitor-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EVXwttoJVtBMuhHNDdK3cMwBdMiZARJV38PMTsL6whnFlA?e=RbG0ru&download=1)
+| [diningtable-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EXVzBkbmTCVImMtuHLCTBeMBXZmv0RWyx5KXQQAe7-7D5w?e=6RYSnm&download=1)   | [sofa-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EaADQYDXwY9NrzbiUFcRYRgBOu1GdJMG8YgNZZmbNjbn-Q?e=DqKrXG&download=1)
+
+| StyleGAN Official | | | |
+| :-- | :--: | :--: | :--: |
+| Model (Dataset) | Training Samples | Training Duration (K Images) | FID
+| [ffhq-1024x1024](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EdfMxgb0hU9BoXwiR3dqYDEBowCSEF1IcsW3n4kwfoZ9OQ?e=VwIV58&download=1)     |    70,000 | 25,000 | 4.40 |
+| [celebahq-1024x1024](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EcCdXHddE7FOvyfmqeOyc9ABqVuWh8PQYFnV6JM1CXvFig?e=1nUYZ5&download=1) |    30,000 | 25,000 | 5.06 |
+| [bedroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Ea6RBPddjcRNoFMXm8AyEBcBUHdlRNtjtclNKFe89amjBw?e=Og8Vff&download=1)    | 3,033,042 | 70,000 | 2.65 |
+| [cat-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EVjX8u9HuehLip3z0hRfIHcB7QtoFkTB7NiRDb8nrKOl2w?e=lHcp1B&download=1)        | 1,657,266 | 70,000 | 8.53 |
+| [car-512x384](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EcRJNNzzUzJGjI2X53S9HjkBhXkKT5JRd6Q3IIhCY1AyRw?e=FvMRNj&download=1)        | 5,520,756 | 46,000 | 3.27 |
+
+| StyleGAN Ours | | | |
+| :-- | :--: | :--: | :--: |
+| Model (Dataset) | Training Samples | Training Duration (K Images) | FID
+| *Face ("partial" means faces are not fully aligned to center)*
+| [celeba_partial-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ET2etKNzMS9JmHj5j60fqMcBRJfQfYNvqUrujaIXxCvKDQ?e=QReLE6&download=1) |     103,706 | 50,000 |  7.03 |
+| [ffhq-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ES-NAUCC2qdHg87BftvlBiQBVpbJ8-005Q4TNr5KrOxQEw?e=00AnWt&download=1)           |      70,000 | 25,000 |  5.70 |
+| [ffhq-512x512](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EZYrrwOiEgVOg-PfGv7QTegBzFQ9yq2v7o1WxNq5JJ9KNA?e=SZU8PI&download=1)           |      70,000 | 25,000 |  5.15 |
+| *LSUN Indoor Scene*
+| [livingroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EfFCYLHjqbFDmjOvCCFJgDcBZ1QYgETfZJxp4ZTHjLxZBg?e=InVd0n&download=1)     |   1,315,802 | 30,000 |  5.16 |
+| [diningroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ERsUza_hSFRIm4iZCag7P0kBQ9EIdfQKByw4QYt_ay97lg?e=Cimh7S&download=1)     |     657,571 | 25,000 |  4.13 |
+| [kitchen-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ERcYvoingQNKix35lUs0vUkBQQkAZMp1rtDxjwNlOJAoaA?e=a1Tcwr&download=1)        |   1,000,000 | 30,000 |  5.06 |
+| *LSUN Indoor Scene Mixture*
+| [apartment-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EfurPNSB2BRFtXdqGkmDD6YBwyKN8YK2v7nKwnJQdsbf6A?e=w3oYa4&download=1)      | 4 * 200,000 | 60,000 |  4.18 |
+| *LSUN Outdoor Scene*
+| [church-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ETMgG1_d06tAlbUkJD1qA9IBaLZ9zJKPkG2kO-4jxhVV5w?e=Dbkb7o&download=1)         |     126,227 | 30,000 |  4.82 |
+| [tower-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Ebm9QMgqB2VDqyIE5rFhreEBgZ_RyKcRf8bQ333K453u3w?e=if8sDj&download=1)          |     708,264 | 30,000 |  5.99 |
+| [bridge-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Ed9QM6OP9sVHnazSp4cqPSEBb-ALfBPXRxP1hD7FsTYh8w?e=3vv06p&download=1)         |     818,687 | 25,000 |  6.42 |
+| *LSUN Other Scene*
+| [restaurant-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/ESDhYr01WtlEvBNFrVpFezcB2l9lF1rBYuHFoeNpBr5B7A?e=uFWFNh&download=1)     |     626,331 | 50,000 |  4.03 |
+| [classroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EbWnI3oto9NPk-lxwZlWqPQB2atWpGiTWMIT59MzF9ij9Q?e=KvcNBg&download=1)      |     168,103 | 50,000 | 10.10 |
+| [conferenceroom-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Eb1gVi3pGa9PgJ4XYYu_6yABQZ0ZcGDak4FEHaTHaeYFzw?e=0BeE8t&download=1) |     229,069 | 50,000 |  6.20 |
+
+| StyleGAN Third-Party | |
+| :-- | :--: |
+| Model (Dataset) | Source |
+| [animeface-512x512](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EWDWflY6lBpGgX0CGQpd2Z4B5wTEVamTOA9JRYne7zdCvA?e=tOzgYA&download=1)     | [link](https://www.gwern.net/Faces#portrait-results)
+| [animeportrait-512x512](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EXBvhTBi-v5NsnQtrxhFEKsBin4xg-Dud9Jr62AEwFTIxg?e=bMGK7r&download=1) | [link](https://www.gwern.net/Faces#portrait-results)
+| [artface-512x512](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/Eca0OiGqhyZMmoPbKahSBWQBWvcAH4q2CE3zdZJflp2jkQ?e=h4rWAm&download=1)       | [link](https://github.com/ak9250/stylegan-art)
+
+| StyleGAN2 Official | | | |
+| :-- | :--: | :--: | :--: |
+| Model (Dataset) | Training Samples | Training Duration (K Images) | FID
+| [ffhq-1024x1024](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EX0DNWiBvl5FuOQTF4oMPBYBNSalcxTK0AbLwBn9Y3vfgg?e=Q0sZit&download=1) |    70,000 |  25,000 | 2.84 |
+| [church-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EQzDtJUdQ4ROunMGn2sZouEBmNeFX4QWvxjermVE5cZvNA?e=tQ7r9r&download=1) |   126,227 |  48,000 | 3.86 |
+| [cat-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EUKXeBwUUbZJr6kup7PW4ekBx2-vmTp8FjcGb10v8bgJxQ?e=nkerMF&download=1)    | 1,657,266 |  88,000 | 6.93 |
+| [horse-256x256](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EconoT6tb69OuAIqfXRtGlsBZz4vBx01UmmFO-JAS356Jg?e=bcSCC4&download=1)  | 2,000,340 | 100,000 | 3.43 |
+| [car-512x384](https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EYSnUsxU8KJFuMHhZm-JLWoB0nHxdlbrLHNZ_Qkoe3b9LA?e=Ycjp5A&download=1)    | 5,520,756 |  57,000 | 2.32 |
+
+## Training Datasets
+
+- [MNIST](http://yann.lecun.com/exdb/mnist/) (60,000 training samples and 10,000 test samples on 10 digital numbers)
+- [SVHN](http://ufldl.stanford.edu/housenumbers/) (73,257 training samples, 26,032 testing samples, and 531,131 additional samples on 10 digital numbers)
+- [CIFAR10](https://www.cs.toronto.edu/~kriz/cifar.html) (50,000 training samples and 10,000 test samples on 10 classes)
+- [CIFAR100](https://www.cs.toronto.edu/~kriz/cifar.html) (50,000 training samples and 10,000 test samples on 100 classes)
+- [ImageNet](http://www.image-net.org/) (1,281,167 training samples, 50,000 validation samples, and 100,100 testing samples on 1000 classes)
+- [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) (202,599 samples from 10,177 identities, with 5 landmarks and 40 binary facial attributes)
+- [CelebA-HQ](https://github.com/tkarras/progressive_growing_of_gans) (30,000 samples)
+- [FF-HQ](https://github.com/NVlabs/ffhq-dataset) (70,000 samples)
+- [LSUN](https://github.com/fyu/lsun) (see statistical information below)
+- [Places](http://places2.csail.mit.edu/) (around 1.8M training samples covering 365 classes)
+- [Cityscapes](https://www.cityscapes-dataset.com/) (2,975 training samples, 19998 extra training samples (one broken), 500 validation samples, and 1,525 test samples)
+- [Streetscapes](http://streetscore.media.mit.edu/data.html)
+
+Statistical information of [LSUN](https://github.com/fyu/lsun) dataset is summarized as follows:
+
+| LSUN Datasets Stats    |                   |      |
+| :--                    | :--:              | :--: |
+| Name                   | Number of Samples | Size |
+| *Scenes*
+| bedroom (train)        |         3,033,042 |  43G |
+| bridge (train)         |           818,687 |  15G |
+| churchoutdoor (train)  |           126,227 |   2G |
+| classroom (train)      |           168,103 |   3G |
+| conferenceroom (train) |           229,069 |   4G |
+| diningroom (train)     |           657,571 |  11G |
+| kitchen (train)        |         2,212,277 |  33G |
+| livingroom (train)     |         1,315,802 |  21G |
+| restaurant (train)     |           626,331 |  13G |
+| tower (train)          |           708,264 |  11G |
+| *Objects*
+| airplane               |         1,530,696 |  34G |
+| bicycle                |         3,347,211 | 129G |
+| bird                   |         2,310,362 |  65G |
+| boat                   |         2,651,165 |  86G |
+| bottle                 |         3,202,760 |  64G |
+| bus                    |           695,891 |  24G |
+| car                    |         5,520,756 | 173G |
+| cat                    |         1,657,266 |  42G |
+| chair                  |         5,037,807 | 116G |
+| cow                    |           377,379 |  15G |
+| diningtable            |         1,537,123 |  48G |
+| dog                    |         5,054,817 | 145G |
+| horse                  |         2,000,340 |  69G |
+| motorbike              |         1,194,101 |  42G |
+| person                 |        18,890,816 | 477G |
+| pottedplant            |         1,104,859 |  43G |
+| sheep                  |           418,983 |  18G |
+| sofa                   |         2,365,870 |  56G |
+| train                  |         1,148,020 |  43G |
+| tvmonitor              |         2,463,284 |  46G |

networks/genforce/README.md

@@ -0,0 +1,169 @@
+# GenForce Lib for Generative Modeling
+
+An efficient PyTorch library for deep generative modeling. May the Generative Force (GenForce) be with You.
+
+![image](./teaser.gif)
+
+## Updates
+
+- **Encoder Training:** We support training encoders on top of pre-trained GANs for GAN inversion.
+- **Model Converters:** You can easily migrate your already started projects to this repository. Please check [here](./converters/README.md) for more details.
+
+## Highlights
+
+- **Distributed** training framework.
+- **Fast** training speed.
+- **Modular** design for prototyping new models.
+- **Model zoo** containing a rich set of pretrained GAN models, with [Colab live demo](https://colab.research.google.com/github/genforce/genforce/blob/master/docs/synthesize_demo.ipynb) to play.
+
+## Installation
+
+1. Create a virtual environment via `conda`.
+
+   ```shell
+   conda create -n genforce python=3.7
+   conda activate genforce
+   ```
+
+2. Install `cuda` and `cudnn`. (We use `CUDA 10.0` in case you would like to use `TensorFlow 1.15` for model conversion.)
+
+   ```shell
+   conda install cudatoolkit=10.0 cudnn=7.6.5
+   ```
+
+3. Install `torch` and `torchvision`.
+
+   ```shell
+   pip install torch==1.7 torchvision==0.8
+   ```
+
+4. Install requirements
+
+   ```shell
+   pip install -r requirements.txt
+   ```
+
+## Quick Demo
+
+We provide a quick training demo, `scripts/stylegan_training_demo.py`, which allows to train StyleGAN on a toy dataset (500 animeface images with 64 x 64 resolution). Try it via
+
+```shell
+./scripts/stylegan_training_demo.sh
+```
+
+We also provide an inference demo, `synthesize.py`, which allows to synthesize images with pre-trained models. Generated images can be found at `work_dirs/synthesis_results/`. Try it via
+
+```shell
+python synthesize.py stylegan_ffhq1024
+```
+
+You can also play the demo at [Colab](https://colab.research.google.com/github/genforce/genforce/blob/master/docs/synthesize_demo.ipynb).
+
+## Play with GANs
+
+### Test
+
+Pre-trained models can be found at [model zoo](MODEL_ZOO.md).
+
+- On local machine:
+
+  ```shell
+  GPUS=8
+  CONFIG=configs/stylegan_ffhq256_val.py
+  WORK_DIR=work_dirs/stylegan_ffhq256_val
+  CHECKPOINT=checkpoints/stylegan_ffhq256.pth
+  ./scripts/dist_test.sh ${GPUS} ${CONFIG} ${WORK_DIR} ${CHECKPOINT}
+  ```
+
+- Using `slurm`:
+
+  ```shell
+  CONFIG=configs/stylegan_ffhq256_val.py
+  WORK_DIR=work_dirs/stylegan_ffhq256_val
+  CHECKPOINT=checkpoints/stylegan_ffhq256.pth
+  GPUS=8 ./scripts/slurm_test.sh ${PARTITION} ${JOB_NAME} \
+      ${CONFIG} ${WORK_DIR} ${CHECKPOINT}
+  ```
+
+### Train
+
+All log files in the training process, such as log message, checkpoints, synthesis snapshots, etc, will be saved to the work directory.
+
+- On local machine:
+
+  ```shell
+  GPUS=8
+  CONFIG=configs/stylegan_ffhq256.py
+  WORK_DIR=work_dirs/stylegan_ffhq256_train
+  ./scripts/dist_train.sh ${GPUS} ${CONFIG} ${WORK_DIR} \
+      [--options additional_arguments]
+  ```
+
+- Using `slurm`:
+
+  ```shell
+  CONFIG=configs/stylegan_ffhq256.py
+  WORK_DIR=work_dirs/stylegan_ffhq256_train
+  GPUS=8 ./scripts/slurm_train.sh ${PARTITION} ${JOB_NAME} \
+      ${CONFIG} ${WORK_DIR} \
+      [--options additional_arguments]
+  ```
+
+## Play with Encoders for GAN Inversion
+
+### Train
+
+- On local machine:
+
+  ```shell
+  GPUS=8
+  CONFIG=configs/stylegan_ffhq256_encoder_y.py
+  WORK_DIR=work_dirs/stylegan_ffhq256_encoder_y
+  ./scripts/dist_train.sh ${GPUS} ${CONFIG} ${WORK_DIR} \
+      [--options additional_arguments]
+  ```
+
+
+- Using `slurm`:
+
+  ```shell
+  CONFIG=configs/stylegan_ffhq256_encoder_y.py
+  WORK_DIR=work_dirs/stylegan_ffhq256_encoder_y
+  GPUS=8 ./scripts/slurm_train.sh ${PARTITION} ${JOB_NAME} \
+      ${CONFIG} ${WORK_DIR} \
+      [--options additional_arguments]
+  ```
+## Contributors
+
+| Member                                      | Module |
+| :--                                         | :--    |
+|[Yujun Shen](http://shenyujun.github.io/)    | models and running controllers
+|[Yinghao Xu](https://justimyhxu.github.io/)  | runner and loss functions
+|[Ceyuan Yang](http://ceyuan.me/)             | data loader
+|[Jiapeng Zhu](https://zhujiapeng.github.io/) | evaluation metrics
+|[Bolei Zhou](http://bzhou.ie.cuhk.edu.hk/)   | cheerleader
+
+**NOTE:** The above form only lists the person in charge for each module. We help each other a lot and develop as a **TEAM**.
+
+*We welcome external contributors to join us for improving this library.*
+
+## License
+
+The project is under the [MIT License](./LICENSE).
+
+## Acknowledgement
+
+We thank [PGGAN](https://github.com/tkarras/progressive_growing_of_gans), [StyleGAN](https://github.com/NVlabs/stylegan), [StyleGAN2](https://github.com/NVlabs/stylegan2), [StyleGAN2-ADA](https://github.com/NVlabs/stylegan2-ada) for their work on high-quality image synthesis. We thank [IDInvert](https://github.com/genforce/idinvert) and [GHFeat](https://github.com/genforce/ghfeat) for their contribution to GAN inversion. We also thank [MMCV](https://github.com/open-mmlab/mmcv) for the inspiration on the design of controllers.
+
+## BibTex
+
+We open source this library to the community to facilitate the research of generative modeling. If you do like our work and use the codebase or models for your research, please cite our work as follows.
+
+```bibtex
+@misc{genforce2020,
+  title =        {GenForce},
+  author =       {Shen, Yujun and Xu, Yinghao and Yang, Ceyuan and Zhu, Jiapeng and Zhou, Bolei},
+  howpublished = {\url{https://github.com/genforce/genforce}},
+  year =         {2020}
+}
+```

networks/genforce/configs/stylegan_demo.py

@@ -0,0 +1,61 @@
+# python3.7
+"""Configuration for StyleGAN training demo.
+
+All settings are particularly used for one replica (GPU), such as `batch_size`
+and `num_workers`.
+"""
+
+runner_type = 'StyleGANRunner'
+gan_type = 'stylegan'
+resolution = 64
+batch_size = 4
+val_batch_size = 32
+total_img = 100_000
+
+# Training dataset is repeated at the beginning to avoid loading dataset
+# repeatedly at the end of each epoch. This can save some I/O time.
+data = dict(
+    num_workers=4,
+    repeat=500,
+    train=dict(root_dir='data/demo.zip', data_format='zip',
+               resolution=resolution, mirror=0.5),
+    val=dict(root_dir='data/demo.zip', data_format='zip',
+             resolution=resolution),
+)
+
+controllers = dict(
+    RunningLogger=dict(every_n_iters=10),
+    ProgressScheduler=dict(
+        every_n_iters=1, init_res=8, minibatch_repeats=4,
+        lod_training_img=5_000, lod_transition_img=5_000,
+        batch_size_schedule=dict(res4=64, res8=32, res16=16, res32=8),
+    ),
+    Snapshoter=dict(every_n_iters=500, first_iter=True, num=200),
+    FIDEvaluator=dict(every_n_iters=5000, first_iter=True, num=50000),
+    Checkpointer=dict(every_n_iters=5000, first_iter=True),
+)
+
+modules = dict(
+    discriminator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        lr=dict(lr_type='FIXED'),
+        opt=dict(opt_type='Adam', base_lr=1e-3, betas=(0.0, 0.99)),
+        kwargs_train=dict(),
+        kwargs_val=dict(),
+    ),
+    generator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        lr=dict(lr_type='FIXED'),
+        opt=dict(opt_type='Adam', base_lr=1e-3, betas=(0.0, 0.99)),
+        kwargs_train=dict(w_moving_decay=0.995, style_mixing_prob=0.9,
+                          trunc_psi=1.0, trunc_layers=0, randomize_noise=True),
+        kwargs_val=dict(trunc_psi=1.0, trunc_layers=0, randomize_noise=False),
+        g_smooth_img=10000,
+    )
+)
+
+loss = dict(
+    type='LogisticGANLoss',
+    d_loss_kwargs=dict(r1_gamma=10.0),
+    g_loss_kwargs=dict(),
+)

networks/genforce/configs/stylegan_ffhq1024.py

@@ -0,0 +1,63 @@
+# python3.7
+"""Configuration for training StyleGAN on FF-HQ (1024) dataset.
+
+All settings are particularly used for one replica (GPU), such as `batch_size`
+and `num_workers`.
+"""
+
+runner_type = 'StyleGANRunner'
+gan_type = 'stylegan'
+resolution = 1024
+batch_size = 4
+val_batch_size = 16
+total_img = 25000_000
+
+# Training dataset is repeated at the beginning to avoid loading dataset
+# repeatedly at the end of each epoch. This can save some I/O time.
+data = dict(
+    num_workers=4,
+    repeat=500,
+    # train=dict(root_dir='data/ffhq', resolution=resolution, mirror=0.5),
+    # val=dict(root_dir='data/ffhq', resolution=resolution),
+    train=dict(root_dir='data/ffhq.zip', data_format='zip',
+               resolution=resolution, mirror=0.5),
+    val=dict(root_dir='data/ffhq.zip', data_format='zip',
+             resolution=resolution),
+)
+
+controllers = dict(
+    RunningLogger=dict(every_n_iters=10),
+    ProgressScheduler=dict(
+        every_n_iters=1, init_res=8, minibatch_repeats=4,
+        lod_training_img=600_000, lod_transition_img=600_000,
+        batch_size_schedule=dict(res4=64, res8=32, res16=16, res32=8),
+    ),
+    Snapshoter=dict(every_n_iters=500, first_iter=True, num=200),
+    FIDEvaluator=dict(every_n_iters=5000, first_iter=True, num=50000),
+    Checkpointer=dict(every_n_iters=5000, first_iter=True),
+)
+
+modules = dict(
+    discriminator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        lr=dict(lr_type='FIXED'),
+        opt=dict(opt_type='Adam', base_lr=1e-3, betas=(0.0, 0.99)),
+        kwargs_train=dict(),
+        kwargs_val=dict(),
+    ),
+    generator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        lr=dict(lr_type='FIXED'),
+        opt=dict(opt_type='Adam', base_lr=1e-3, betas=(0.0, 0.99)),
+        kwargs_train=dict(w_moving_decay=0.995, style_mixing_prob=0.9,
+                          trunc_psi=1.0, trunc_layers=0, randomize_noise=True),
+        kwargs_val=dict(trunc_psi=1.0, trunc_layers=0, randomize_noise=False),
+        g_smooth_img=10_000,
+    )
+)
+
+loss = dict(
+    type='LogisticGANLoss',
+    d_loss_kwargs=dict(r1_gamma=10.0),
+    g_loss_kwargs=dict(),
+)

networks/genforce/configs/stylegan_ffhq1024_val.py

@@ -0,0 +1,29 @@
+# python3.7
+"""Configuration for testing StyleGAN on FF-HQ (1024) dataset.
+
+All settings are particularly used for one replica (GPU), such as `batch_size`
+and `num_workers`.
+"""
+
+runner_type = 'StyleGANRunner'
+gan_type = 'stylegan'
+resolution = 1024
+batch_size = 16
+
+data = dict(
+    num_workers=4,
+    # val=dict(root_dir='data/ffhq', resolution=resolution),
+    val=dict(root_dir='data/ffhq.zip', data_format='zip',
+             resolution=resolution),
+)
+
+modules = dict(
+    discriminator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        kwargs_val=dict(),
+    ),
+    generator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        kwargs_val=dict(trunc_psi=0.7, trunc_layers=8, randomize_noise=False),
+    )
+)

networks/genforce/configs/stylegan_ffhq256.py

@@ -0,0 +1,63 @@
+# python3.7
+"""Configuration for training StyleGAN on FF-HQ (256) dataset.
+
+All settings are particularly used for one replica (GPU), such as `batch_size`
+and `num_workers`.
+"""
+
+runner_type = 'StyleGANRunner'
+gan_type = 'stylegan'
+resolution = 256
+batch_size = 4
+val_batch_size = 64
+total_img = 25000_000
+
+# Training dataset is repeated at the beginning to avoid loading dataset
+# repeatedly at the end of each epoch. This can save some I/O time.
+data = dict(
+    num_workers=4,
+    repeat=500,
+    # train=dict(root_dir='data/ffhq', resolution=resolution, mirror=0.5),
+    # val=dict(root_dir='data/ffhq', resolution=resolution),
+    train=dict(root_dir='data/ffhq.zip', data_format='zip',
+               resolution=resolution, mirror=0.5),
+    val=dict(root_dir='data/ffhq.zip', data_format='zip',
+             resolution=resolution),
+)
+
+controllers = dict(
+    RunningLogger=dict(every_n_iters=10),
+    ProgressScheduler=dict(
+        every_n_iters=1, init_res=8, minibatch_repeats=4,
+        lod_training_img=600_000, lod_transition_img=600_000,
+        batch_size_schedule=dict(res4=64, res8=32, res16=16, res32=8),
+    ),
+    Snapshoter=dict(every_n_iters=500, first_iter=True, num=200),
+    FIDEvaluator=dict(every_n_iters=5000, first_iter=True, num=50000),
+    Checkpointer=dict(every_n_iters=5000, first_iter=True),
+)
+
+modules = dict(
+    discriminator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        lr=dict(lr_type='FIXED'),
+        opt=dict(opt_type='Adam', base_lr=1e-3, betas=(0.0, 0.99)),
+        kwargs_train=dict(),
+        kwargs_val=dict(),
+    ),
+    generator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        lr=dict(lr_type='FIXED'),
+        opt=dict(opt_type='Adam', base_lr=1e-3, betas=(0.0, 0.99)),
+        kwargs_train=dict(w_moving_decay=0.995, style_mixing_prob=0.9,
+                          trunc_psi=1.0, trunc_layers=0, randomize_noise=True),
+        kwargs_val=dict(trunc_psi=1.0, trunc_layers=0, randomize_noise=False),
+        g_smooth_img=10_000,
+    )
+)
+
+loss = dict(
+    type='LogisticGANLoss',
+    d_loss_kwargs=dict(r1_gamma=10.0),
+    g_loss_kwargs=dict(),
+)

networks/genforce/configs/stylegan_ffhq256_encoder_y.py

@@ -0,0 +1,73 @@
+# python3.7
+"""Configuration for training StyleGAN Encoder on FF-HQ (256) dataset.
+
+All settings are particularly used for one replica (GPU), such as `batch_size`
+and `num_workers`.
+"""
+
+gan_model_path = 'checkpoints/stylegan_ffhq256.pth'
+perceptual_model_path = 'checkpoints/vgg16.pth'
+
+runner_type = 'EncoderRunner'
+gan_type = 'stylegan'
+resolution = 256
+batch_size = 12
+val_batch_size = 25
+total_img = 14000_000
+space_of_latent = 'y'
+
+# Training dataset is repeated at the beginning to avoid loading dataset
+# repeatedly at the end of each epoch. This can save some I/O time.
+data = dict(
+    num_workers=4,
+    repeat=500,
+    # train=dict(root_dir='data/ffhq', resolution=resolution, mirror=0.5),
+    # val=dict(root_dir='data/ffhq', resolution=resolution),
+    train=dict(root_dir='data/', data_format='list',
+               image_list_path='data/ffhq/ffhq_train_list.txt',
+               resolution=resolution, mirror=0.5),
+    val=dict(root_dir='data/', data_format='list',
+             image_list_path='./data/ffhq/ffhq_val_list.txt',
+             resolution=resolution),
+)
+
+controllers = dict(
+    RunningLogger=dict(every_n_iters=50),
+    Snapshoter=dict(every_n_iters=10000, first_iter=True, num=200),
+    Checkpointer=dict(every_n_iters=10000, first_iter=False),
+)
+
+modules = dict(
+    discriminator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        lr=dict(lr_type='ExpSTEP', decay_factor=0.8, decay_step=36458 // 2),
+        opt=dict(opt_type='Adam', base_lr=1e-4, betas=(0.9, 0.99)),
+        kwargs_train=dict(),
+        kwargs_val=dict(),
+    ),
+    generator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution, repeat_w=True),
+        kwargs_val=dict(randomize_noise=False),
+    ),
+    encoder=dict(
+        model=dict(gan_type=gan_type, resolution=resolution, network_depth=18,
+                   latent_dim = [1024] * 8 + [512, 512, 256, 256, 128, 128],
+                   num_latents_per_head=[4, 4, 6],
+                   use_fpn=True,
+                   fpn_channels=512,
+                   use_sam=True,
+                   sam_channels=512),
+        lr=dict(lr_type='ExpSTEP', decay_factor=0.8, decay_step=36458 // 2),
+        opt=dict(opt_type='Adam', base_lr=1e-4, betas=(0.9, 0.99)),
+        kwargs_train=dict(),
+        kwargs_val=dict(),
+    ),
+)
+
+loss = dict(
+    type='EncoderLoss',
+    d_loss_kwargs=dict(r1_gamma=10.0),
+    e_loss_kwargs=dict(adv_lw=0.08, perceptual_lw=5e-5),
+    perceptual_kwargs=dict(output_layer_idx=23,
+                           pretrained_weight_path=perceptual_model_path),
+)

networks/genforce/configs/stylegan_ffhq256_val.py

@@ -0,0 +1,29 @@
+# python3.7
+"""Configuration for testing StyleGAN on FF-HQ (256) dataset.
+
+All settings are particularly used for one replica (GPU), such as `batch_size`
+and `num_workers`.
+"""
+
+runner_type = 'StyleGANRunner'
+gan_type = 'stylegan'
+resolution = 256
+batch_size = 64
+
+data = dict(
+    num_workers=4,
+    # val=dict(root_dir='data/ffhq', resolution=resolution),
+    val=dict(root_dir='data/ffhq.zip', data_format='zip',
+             resolution=resolution),
+)
+
+modules = dict(
+    discriminator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        kwargs_val=dict(),
+    ),
+    generator=dict(
+        model=dict(gan_type=gan_type, resolution=resolution),
+        kwargs_val=dict(trunc_psi=0.7, trunc_layers=8, randomize_noise=False),
+    )
+)

networks/genforce/convert_model.py

@@ -0,0 +1,77 @@
+"""Script to convert officially released models to match this repository."""
+
+import argparse
+
+from converters import convert_pggan_weight
+from converters import convert_stylegan_weight
+from converters import convert_stylegan2_weight
+from converters import convert_stylegan2ada_tf_weight
+from converters import convert_stylegan2ada_pth_weight
+
+
+def parse_args():
+    """Parses arguments."""
+    parser = argparse.ArgumentParser(description='Convert pre-trained models.')
+    parser.add_argument('model_type', type=str,
+                        choices=['pggan', 'stylegan', 'stylegan2',
+                                 'stylegan2ada_tf', 'stylegan2ada_pth'],
+                        help='Type of the model to convert')
+    parser.add_argument('--source_model_path', type=str, required=True,
+                        help='Path to load the model for conversion.')
+    parser.add_argument('--target_model_path', type=str, default=None,
+                        help='Path to save the converted model. If not '
+                             'specified, the model will be saved to the same '
+                             'directory of the source model.')
+    parser.add_argument('--test_num', type=int, default=10,
+                        help='Number of test samples used to check the '
+                             'precision of the converted model. (default: 10)')
+    parser.add_argument('--save_test_image', action='store_true',
+                        help='Whether to save the test image. (default: False)')
+    parser.add_argument('--verbose_log', action='store_true',
+                        help='Whether to print verbose log. (default: False)')
+    return parser.parse_args()
+
+
+def main():
+    """Main function."""
+    args = parse_args()
+    if args.target_model_path is None:
+        args.target_model_path = args.source_model_path.replace('.pkl', '.pth')
+
+    if args.model_type == 'pggan':
+        convert_pggan_weight(tf_weight_path=args.source_model_path,
+                             pth_weight_path=args.target_model_path,
+                             test_num=args.test_num,
+                             save_test_image=args.save_test_image,
+                             verbose=args.verbose_log)
+    elif args.model_type == 'stylegan':
+        convert_stylegan_weight(tf_weight_path=args.source_model_path,
+                                pth_weight_path=args.target_model_path,
+                                test_num=args.test_num,
+                                save_test_image=args.save_test_image,
+                                verbose=args.verbose_log)
+    elif args.model_type == 'stylegan2':
+        convert_stylegan2_weight(tf_weight_path=args.source_model_path,
+                                 pth_weight_path=args.target_model_path,
+                                 test_num=args.test_num,
+                                 save_test_image=args.save_test_image,
+                                 verbose=args.verbose_log)
+    elif args.model_type == 'stylegan2ada_tf':
+        convert_stylegan2ada_tf_weight(tf_weight_path=args.source_model_path,
+                                       pth_weight_path=args.target_model_path,
+                                       test_num=args.test_num,
+                                       save_test_image=args.save_test_image,
+                                       verbose=args.verbose_log)
+    elif args.model_type == 'stylegan2ada_pth':
+        convert_stylegan2ada_pth_weight(src_weight_path=args.source_model_path,
+                                        dst_weight_path=args.target_model_path,
+                                        test_num=args.test_num,
+                                        save_test_image=args.save_test_image,
+                                        verbose=args.verbose_log)
+    else:
+        raise NotImplementedError(f'Model type `{args.model_type}` is not '
+                                  f'supported!')
+
+
+if __name__ == '__main__':
+    main()

networks/genforce/datasets/README.md

@@ -0,0 +1,24 @@
+# Data Preparation
+
+## Data Format
+
+Currently, our dataloader is able to load data from
+
+- a directory that is full of images (support using [`turbojpeg`](https://pypi.org/project/PyTurboJPEG/) to speed up decoding images.)
+- a `lmdb` file
+- an image list
+- a compressed file (i.e., `zip` package)
+
+by modifying `data_format` in the configuration.
+
+**NOTE:** For some computing clusters whose I/O speed may be slow, we recommend the `zip` format for two reasons. First, `zip` file is easy to create. Second, this can load a large file at one time instead of loading small files repeatedly.
+
+## Data Sampling
+
+Considering that most generative models are trained in the unit of iterations instead of epochs, we change the default data loader to an *iter-based* one. Besides, the original distributed data sampler is also modified to make the shuffling correspond to iteration instead of epoch.
+
+**NOTE:** In order to reduce the data re-loading cost between epochs, we manually extend the length of sampled indices to make it much more efficient.
+
+## Data Augmentation
+
+To better align with the original implementation of PGGAN and StyleGAN (i.e., models that require progressive training), we support progressive resize in `transforms.py`, which downsamples images with the maximum resize factor of 2 at each time.

networks/genforce/datasets/__init__.py

@@ -0,0 +1,7 @@
+# python3.7
+"""Collects datasets and data loaders."""
+
+from .datasets import BaseDataset
+from .dataloaders import IterDataLoader
+
+__all__ = ['BaseDataset', 'IterDataLoader']

networks/genforce/datasets/dataloaders.py

@@ -0,0 +1,128 @@
+# python3.7
+"""Contains the class of data loader."""
+
+import argparse
+
+from torch.utils.data import DataLoader
+from .distributed_sampler import DistributedSampler
+from .datasets import BaseDataset
+
+
+__all__ = ['IterDataLoader']
+
+
+class IterDataLoader(object):
+    """Iteration-based data loader."""
+
+    def __init__(self,
+                 dataset,
+                 batch_size,
+                 shuffle=True,
+                 num_workers=1,
+                 current_iter=0,
+                 repeat=1):
+        """Initializes the data loader.
+
+        Args:
+            dataset: The dataset to load data from.
+            batch_size: The batch size on each GPU.
+            shuffle: Whether to shuffle the data. (default: True)
+            num_workers: Number of data workers for each GPU. (default: 1)
+            current_iter: The current number of iterations. (default: 0)
+            repeat: The repeating number of the whole dataloader. (default: 1)
+        """
+        self._dataset = dataset
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.num_workers = num_workers
+        self._dataloader = None
+        self.iter_loader = None
+        self._iter = current_iter
+        self.repeat = repeat
+        self.build_dataloader()
+
+    def build_dataloader(self):
+        """Builds data loader."""
+        dist_sampler = DistributedSampler(self._dataset,
+                                          shuffle=self.shuffle,
+                                          current_iter=self._iter,
+                                          repeat=self.repeat)
+
+        self._dataloader = DataLoader(self._dataset,
+                                      batch_size=self.batch_size,
+                                      shuffle=(dist_sampler is None),
+                                      num_workers=self.num_workers,
+                                      drop_last=self.shuffle,
+                                      pin_memory=True,
+                                      sampler=dist_sampler)
+        self.iter_loader = iter(self._dataloader)
+
+
+    def overwrite_param(self, batch_size=None, resolution=None):
+        """Overwrites some parameters for progressive training."""
+        if (not batch_size) and (not resolution):
+            return
+        if (batch_size == self.batch_size) and (
+                resolution == self.dataset.resolution):
+            return
+        if batch_size:
+            self.batch_size = batch_size
+        if resolution:
+            self._dataset.resolution = resolution
+        self.build_dataloader()
+
+    @property
+    def iter(self):
+        """Returns the current iteration."""
+        return self._iter
+
+    @property
+    def dataset(self):
+        """Returns the dataset."""
+        return self._dataset
+
+    @property
+    def dataloader(self):
+        """Returns the data loader."""
+        return self._dataloader
+
+    def __next__(self):
+        try:
+            data = next(self.iter_loader)
+            self._iter += 1
+        except StopIteration:
+            self._dataloader.sampler.__reset__(self._iter)
+            self.iter_loader = iter(self._dataloader)
+            data = next(self.iter_loader)
+            self._iter += 1
+        return data
+
+    def __len__(self):
+        return len(self._dataloader)
+
+
+def dataloader_test(root_dir, test_num=10):
+    """Tests data loader."""
+    res = 2
+    bs = 2
+    dataset = BaseDataset(root_dir=root_dir, resolution=res)
+    dataloader = IterDataLoader(dataset=dataset,
+                                batch_size=bs,
+                                shuffle=False)
+    for _ in range(test_num):
+        data_batch = next(dataloader)
+        image = data_batch['image']
+        assert image.shape == (bs, 3, res, res)
+        res *= 2
+        bs += 1
+        dataloader.overwrite_param(batch_size=bs, resolution=res)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Test Data Loader.')
+    parser.add_argument('root_dir', type=str,
+                        help='Root directory of the dataset.')
+    parser.add_argument('--test_num', type=int, default=10,
+                        help='Number of tests. (default: %(default)s)')
+    args = parser.parse_args()
+    dataloader_test(args.root_dir, args.test_num)

networks/genforce/datasets/datasets.py

@@ -0,0 +1,239 @@
+# python3.7
+"""Contains the class of dataset."""
+
+import os
+import pickle
+import string
+import zipfile
+import numpy as np
+import cv2
+import lmdb
+
+import torch
+from torch.utils.data import Dataset
+
+from .transforms import progressive_resize_image
+from .transforms import crop_resize_image
+from .transforms import resize_image
+from .transforms import normalize_image
+
+try:
+    import turbojpeg
+    BASE_DIR = os.path.dirname(os.path.relpath(__file__))
+    LIBRARY_NAME = 'libturbojpeg.so.0'
+    LIBRARY_PATH = os.path.join(BASE_DIR, LIBRARY_NAME)
+    jpeg = turbojpeg.TurboJPEG(LIBRARY_PATH)
+except ImportError:
+    jpeg = None
+
+__all__ = ['BaseDataset']
+
+_FORMATS_ALLOWED = ['dir', 'lmdb', 'list', 'zip']
+
+
+class ZipLoader(object):
+    """Defines a class to load zip file.
+
+    This is a static class, which is used to solve the problem that different
+    data workers can not share the same memory.
+    """
+    files = dict()
+
+    @staticmethod
+    def get_zipfile(file_path):
+        """Fetches a zip file."""
+        zip_files = ZipLoader.files
+        if file_path not in zip_files:
+            zip_files[file_path] = zipfile.ZipFile(file_path, 'r')
+        return zip_files[file_path]
+
+    @staticmethod
+    def get_image(file_path, image_path):
+        """Decodes an image from a particular zip file."""
+        zip_file = ZipLoader.get_zipfile(file_path)
+        image_str = zip_file.read(image_path)
+        image_np = np.frombuffer(image_str, np.uint8)
+        image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
+        return image
+
+
+class LmdbLoader(object):
+    """Defines a class to load lmdb file.
+
+    This is a static class, which is used to solve lmdb loading error
+    when num_workers > 0
+    """
+    files = dict()
+
+    @staticmethod
+    def get_lmdbfile(file_path):
+        """Fetches a lmdb file"""
+        lmdb_files = LmdbLoader.files
+        if 'env' not in lmdb_files:
+            env = lmdb.open(file_path,
+                            max_readers=1,
+                            readonly=True,
+                            lock=False,
+                            readahead=False,
+                            meminit=False)
+            with env.begin(write=False) as txn:
+                num_samples = txn.stat()['entries']
+            cache_file = '_cache_' + ''.join(
+                c for c in file_path if c in string.ascii_letters)
+            if os.path.isfile(cache_file):
+                keys = pickle.load(open(cache_file, "rb"))
+            else:
+                with env.begin(write=False) as txn:
+                    keys = [key for key, _ in txn.cursor()]
+                pickle.dump(keys, open(cache_file, "wb"))
+            lmdb_files['env'] = env
+            lmdb_files['num_samples'] = num_samples
+            lmdb_files['keys'] = keys
+        return lmdb_files
+
+    @staticmethod
+    def get_image(file_path, idx):
+        """Decodes an image from a particular lmdb file"""
+        lmdb_files = LmdbLoader.get_lmdbfile(file_path)
+        env = lmdb_files['env']
+        keys = lmdb_files['keys']
+        with env.begin(write=False) as txn:
+            imagebuf = txn.get(keys[idx])
+        image_np = np.frombuffer(imagebuf, np.uint8)
+        image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
+        return image
+
+
+class BaseDataset(Dataset):
+    """Defines the base dataset class.
+
+    This class supports loading data from a full-of-image folder, a lmdb
+    database, or an image list. Images will be pre-processed based on the given
+    `transform` function before fed into the data loader.
+
+    NOTE: The loaded data will be returned as a directory, where there must be
+    a key `image`.
+    """
+    def __init__(self,
+                 root_dir,
+                 resolution,
+                 data_format='dir',
+                 image_list_path=None,
+                 mirror=0.0,
+                 progressive_resize=True,
+                 crop_resize_resolution=-1,
+                 transform=normalize_image,
+                 transform_kwargs=None,
+                 **_unused_kwargs):
+        """Initializes the dataset.
+
+        Args:
+            root_dir: Root directory containing the dataset.
+            resolution: The resolution of the returned image.
+            data_format: Format the dataset is stored. Supports `dir`, `lmdb`,
+                and `list`. (default: `dir`)
+            image_list_path: Path to the image list. This field is required if
+                `data_format` is `list`. (default: None)
+            mirror: The probability to do mirror augmentation. (default: 0.0)
+            progressive_resize: Whether to resize images progressively.
+                (default: True)
+            crop_resize_resolution: The resolution of the output after crop
+                and resize. (default: -1)
+            transform: The transform function for pre-processing.
+                (default: `datasets.transforms.normalize_image()`)
+            transform_kwargs: The additional arguments for the `transform`
+                function. (default: None)
+
+        Raises:
+            ValueError: If the input `data_format` is not supported.
+            NotImplementedError: If the input `data_format` is not implemented.
+        """
+        if data_format.lower() not in _FORMATS_ALLOWED:
+            raise ValueError(f'Invalid data format `{data_format}`!\n'
+                             f'Supported formats: {_FORMATS_ALLOWED}.')
+
+        self.root_dir = root_dir
+        self.resolution = resolution
+        self.data_format = data_format.lower()
+        self.image_list_path = image_list_path
+        self.mirror = np.clip(mirror, 0.0, 1.0)
+        self.progressive_resize = progressive_resize
+        self.crop_resize_resolution = crop_resize_resolution
+        self.transform = transform
+        self.transform_kwargs = transform_kwargs or dict()
+
+        if self.data_format == 'dir':
+            self.image_paths = sorted(os.listdir(self.root_dir))
+            self.num_samples = len(self.image_paths)
+        elif self.data_format == 'lmdb':
+            lmdb_file = LmdbLoader.get_lmdbfile(self.root_dir)
+            self.num_samples = lmdb_file['num_samples']
+        elif self.data_format == 'list':
+            self.metas = []
+            assert os.path.isfile(self.image_list_path)
+            with open(self.image_list_path) as f:
+                for line in f:
+                    fields = line.rstrip().split(' ')
+                    if len(fields) == 1:
+                        self.metas.append((fields[0], None))
+                    else:
+                        assert len(fields) == 2
+                        self.metas.append((fields[0], int(fields[1])))
+            self.num_samples = len(self.metas)
+        elif self.data_format == 'zip':
+            zip_file = ZipLoader.get_zipfile(self.root_dir)
+            image_paths = [f for f in zip_file.namelist()
+                           if ('.jpg' in f or '.jpeg' in f or '.png' in f)]
+            self.image_paths = sorted(image_paths)
+            self.num_samples = len(self.image_paths)
+        else:
+            raise NotImplementedError(f'Not implemented data format '
+                                      f'`{self.data_format}`!')
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, idx):
+        data = dict()
+
+        # Load data.
+        if self.data_format == 'dir':
+            image_path = self.image_paths[idx]
+            try:
+                in_file = open(os.path.join(self.root_dir, image_path), 'rb')
+                image = jpeg.decode(in_file.read())
+            except:  # pylint: disable=bare-except
+                image = cv2.imread(os.path.join(self.root_dir, image_path))
+        elif self.data_format == 'lmdb':
+            image = LmdbLoader.get_image(self.root_dir, idx)
+        elif self.data_format == 'list':
+            image_path, label = self.metas[idx]
+            image = cv2.imread(os.path.join(self.root_dir, image_path))
+            label = None if label is None else torch.LongTensor(label)
+            # data.update({'label': label})
+        elif self.data_format == 'zip':
+            image_path = self.image_paths[idx]
+            image = ZipLoader.get_image(self.root_dir, image_path)
+        else:
+            raise NotImplementedError(f'Not implemented data format '
+                                      f'`{self.data_format}`!')
+
+        image = image[:, :, ::-1]  # Converts BGR (cv2) to RGB.
+
+        # Transform image.
+        if self.crop_resize_resolution > 0:
+            image = crop_resize_image(image, self.crop_resize_resolution)
+        if self.progressive_resize:
+            image = progressive_resize_image(image, self.resolution)
+        image = image.transpose(2, 0, 1).astype(np.float32)
+        if np.random.uniform() < self.mirror:
+            image = image[:, :, ::-1]  # CHW
+        image = torch.FloatTensor(image.copy())
+        if not self.progressive_resize:
+            image = resize_image(image, self.resolution)
+
+        if self.transform is not None:
+            image = self.transform(image, **self.transform_kwargs)
+        data.update({'image': image})
+
+        return data

networks/genforce/datasets/distributed_sampler.py

@@ -0,0 +1,144 @@
+# python3.7
+"""Contains the distributed data sampler.
+
+This file is mostly borrowed from `torch/utils/data/distributed.py`.
+
+However, sometimes, initialize the data loader and data sampler can be time
+consuming (since it will load a large amount of data at one time). To avoid
+re-initializing the data loader again and again, we modified the sampler to
+support loading the data for only one time and then repeating the data loader.
+Please use the class member `repeat` to control how many times you want the
+data load to repeat. After `repeat` times, the data will be re-loaded.
+
+NOTE: The number of repeat times should not be very large, especially when there
+are too many samples in the dataset. We recommend to set `repeat = 500` for
+datasets with ~50K samples.
+"""
+
+# pylint: disable=line-too-long
+
+import math
+from typing import TypeVar, Optional, Iterator
+
+import torch
+from torch.utils.data import Sampler, Dataset
+import torch.distributed as dist
+
+
+T_co = TypeVar('T_co', covariant=True)
+
+
+class DistributedSampler(Sampler):
+    r"""Sampler that restricts data loading to a subset of the dataset.
+
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such a case, each
+    process can pass a :class:`~torch.utils.data.DistributedSampler` instance as a
+    :class:`~torch.utils.data.DataLoader` sampler, and load a subset of the
+    original dataset that is exclusive to it.
+
+    .. note::
+        Dataset is assumed to be of constant size.
+
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (int, optional): Number of processes participating in
+            distributed training. By default, :attr:`rank` is retrieved from the
+            current distributed group.
+        rank (int, optional): Rank of the current process within :attr:`num_replicas`.
+            By default, :attr:`rank` is retrieved from the current distributed
+            group.
+        shuffle (bool, optional): If ``True`` (default), sampler will shuffle the
+            indices.
+        seed (int, optional): random seed used to shuffle the sampler if
+            :attr:`shuffle=True`. This number should be identical across all
+            processes in the distributed group. Default: ``0``.
+        drop_last (bool, optional): if ``True``, then the sampler will drop the
+            tail of the data to make it evenly divisible across the number of
+            replicas. If ``False``, the sampler will add extra indices to make
+            the data evenly divisible across the replicas. Default: ``False``.
+        current_iter (int, optional): Number of current iteration. Default: ``0``.
+        repeat (int, optional): Repeating number of the whole dataloader. Default: ``1000``.
+
+    .. warning::
+        In distributed mode, calling the :meth:`set_epoch` method at
+        the beginning of each epoch **before** creating the :class:`DataLoader` iterator
+        is necessary to make shuffling work properly across multiple epochs. Otherwise,
+        the same ordering will be always used.
+
+    """
+
+    def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None,
+                 rank: Optional[int] = None, shuffle: bool = True,
+                 seed: int = 0, drop_last: bool = False, current_iter: int = 0,
+                 repeat: int = 1000) -> None:
+        super().__init__(None)
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.iter = current_iter
+        self.drop_last = drop_last
+
+        # NOTE: self.dataset_length is `repeat X len(self.dataset)`
+        self.repeat = repeat
+        self.dataset_length = len(self.dataset) * self.repeat
+
+        if self.drop_last and self.dataset_length % self.num_replicas != 0:
+            # Split to nearest available length that is evenly divisible.
+            # This is to ensure each rank receives the same amount of data when
+            # using this Sampler.
+            self.num_samples = math.ceil(
+                (self.dataset_length - self.num_replicas) / self.num_replicas
+            )
+        else:
+            self.num_samples = math.ceil(self.dataset_length / self.num_replicas)
+
+
+        self.total_size = self.num_samples * self.num_replicas
+        self.shuffle = shuffle
+        self.seed = seed
+        self.__generate_indices__()
+
+    def __generate_indices__(self) -> None:
+        g = torch.Generator()
+        indices_bank = []
+        for iter_ in range(self.iter, self.iter + self.repeat):
+            g.manual_seed(self.seed + iter_)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            indices_bank.extend(indices)
+        self.indices = indices_bank
+
+    def __iter__(self) -> Iterator[T_co]:
+        if self.shuffle:
+            # deterministically shuffle based on iter and seed
+            indices = self.indices
+        else:
+            indices = list(range(self.dataset_length))
+
+        if not self.drop_last:
+            # add extra samples to make it evenly divisible
+            indices += indices[:(self.total_size - len(indices))]
+        else:
+            # remove tail of data to make it evenly divisible.
+            indices = indices[:self.total_size]
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        return iter(indices)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __reset__(self, iteration: int) -> None:
+        self.iter = iteration
+        self.__generate_indices__()
+
+# pylint: enable=line-too-long

networks/genforce/datasets/transforms.py

@@ -0,0 +1,201 @@
+"""Contains transform functions."""
+
+import cv2
+import numpy as np
+import PIL.Image
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+__all__ = [
+    'crop_resize_image', 'progressive_resize_image', 'resize_image',
+    'normalize_image', 'normalize_latent_code', 'ImageResizing',
+    'ImageNormalization', 'LatentCodeNormalization',
+]
+
+
+def crop_resize_image(image, size):
+    """Crops a square patch and then resizes it to the given size.
+
+    Args:
+        image: The input image to crop and resize.
+        size: An integer, indicating the target size.
+
+    Returns:
+        An image with target size.
+
+    Raises:
+        TypeError: If the input `image` is not with type `numpy.ndarray`.
+        ValueError: If the input `image` is not with shape [H, W, C].
+    """
+    if not isinstance(image, np.ndarray):
+        raise TypeError(f'Input image should be with type `numpy.ndarray`, '
+                        f'but `{type(image)}` is received!')
+    if image.ndim != 3:
+        raise ValueError(f'Input image should be with shape [H, W, C], '
+                         f'but `{image.shape}` is received!')
+
+    height, width, channel = image.shape
+    short_side = min(height, width)
+    image = image[(height - short_side) // 2:(height + short_side) // 2,
+                  (width - short_side) // 2:(width + short_side) // 2]
+    pil_image = PIL.Image.fromarray(image)
+    pil_image = pil_image.resize((size, size), PIL.Image.ANTIALIAS)
+    image = np.asarray(pil_image)
+    assert image.shape == (size, size, channel)
+    return image
+
+
+def progressive_resize_image(image, size):
+    """Resizes image to target size progressively.
+
+    Different from normal resize, this function will reduce the image size
+    progressively. In each step, the maximum reduce factor is 2.
+
+    NOTE: This function can only handle square images, and can only be used for
+    downsampling.
+
+    Args:
+        image: The input (square) image to resize.
+        size: An integer, indicating the target size.
+
+    Returns:
+        An image with target size.
+
+    Raises:
+        TypeError: If the input `image` is not with type `numpy.ndarray`.
+        ValueError: If the input `image` is not with shape [H, W, C].
+    """
+    if not isinstance(image, np.ndarray):
+        raise TypeError(f'Input image should be with type `numpy.ndarray`, '
+                        f'but `{type(image)}` is received!')
+    if image.ndim != 3:
+        raise ValueError(f'Input image should be with shape [H, W, C], '
+                         f'but `{image.shape}` is received!')
+
+    height, width, channel = image.shape
+    assert height == width
+    assert height >= size
+    num_iters = int(np.log2(height) - np.log2(size))
+    for _ in range(num_iters):
+        height = max(height // 2, size)
+        image = cv2.resize(image, (height, height),
+                           interpolation=cv2.INTER_LINEAR)
+    assert image.shape == (size, size, channel)
+    return image
+
+
+def resize_image(image, size):
+    """Resizes image to target size.
+
+    NOTE: We use adaptive average pooing for image resizing. Instead of bilinear
+    interpolation, average pooling is able to acquire information from more
+    pixels, such that the resized results can be with higher quality.
+
+    Args:
+        image: The input image tensor, with shape [C, H, W], to resize.
+        size: An integer or a tuple of integer, indicating the target size.
+
+    Returns:
+        An image tensor with target size.
+
+    Raises:
+        TypeError: If the input `image` is not with type `torch.Tensor`.
+        ValueError: If the input `image` is not with shape [C, H, W].
+    """
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f'Input image should be with type `torch.Tensor`, '
+                        f'but `{type(image)}` is received!')
+    if image.ndim != 3:
+        raise ValueError(f'Input image should be with shape [C, H, W], '
+                         f'but `{image.shape}` is received!')
+
+    image = F.adaptive_avg_pool2d(image.unsqueeze(0), size).squeeze(0)
+    return image
+
+
+def normalize_image(image, mean=127.5, std=127.5):
+    """Normalizes image by subtracting mean and dividing std.
+
+    Args:
+        image: The input image tensor to normalize.
+        mean: The mean value to subtract from the input tensor. (default: 127.5)
+        std: The standard deviation to normalize the input tensor. (default:
+            127.5)
+
+    Returns:
+        A normalized image tensor.
+
+    Raises:
+        TypeError: If the input `image` is not with type `torch.Tensor`.
+    """
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f'Input image should be with type `torch.Tensor`, '
+                        f'but `{type(image)}` is received!')
+    out = (image - mean) / std
+    return out
+
+
+def normalize_latent_code(latent_code, adjust_norm=True):
+    """Normalizes latent code.
+
+    NOTE: The latent code will always be normalized along the last axis.
+    Meanwhile, if `adjust_norm` is set as `True`, the norm of the result will be
+    adjusted to `sqrt(latent_code.shape[-1])` in order to avoid too small value.
+
+    Args:
+        latent_code: The input latent code tensor to normalize.
+        adjust_norm: Whether to adjust the norm of the output. (default: True)
+
+    Returns:
+        A normalized latent code tensor.
+
+    Raises:
+        TypeError: If the input `latent_code` is not with type `torch.Tensor`.
+    """
+    if not isinstance(latent_code, torch.Tensor):
+        raise TypeError(f'Input latent code should be with type '
+                        f'`torch.Tensor`, but `{type(latent_code)}` is '
+                        f'received!')
+    dim = latent_code.shape[-1]
+    norm = latent_code.pow(2).sum(-1, keepdim=True).pow(0.5)
+    out = latent_code / norm
+    if adjust_norm:
+        out = out * (dim ** 0.5)
+    return out
+
+
+class ImageResizing(nn.Module):
+    """Implements the image resizing layer."""
+
+    def __init__(self, size):
+        super().__init__()
+        self.size = size
+
+    def forward(self, image):
+        return resize_image(image, self.size)
+
+
+class ImageNormalization(nn.Module):
+    """Implements the image normalization layer."""
+
+    def __init__(self, mean=127.5, std=127.5):
+        super().__init__()
+        self.mean = mean
+        self.std = std
+
+    def forward(self, image):
+        return normalize_image(image, self.mean, self.std)
+
+
+class LatentCodeNormalization(nn.Module):
+    """Implements the latent code normalization layer."""
+
+    def __init__(self, adjust_norm=True):
+        super().__init__()
+        self.adjust_norm = adjust_norm
+
+    def forward(self, latent_code):
+        return normalize_latent_code(latent_code, self.adjust_norm)

networks/genforce/metrics/README.md

@@ -0,0 +1,18 @@
+# Evaluation Metrics
+
+Frechet Inception Distance (FID) is commonly used to evaluate generative model. It employs an [Inception Model](https://arxiv.org/abs/1512.00567) (pretrained on ImageNet) to extract features from both real and synthesized images.
+
+## Inception Model
+
+For [PGGAN](https://github.com/tkarras/progressive_growing_of_gans), [StyleGAN](https://github.com/NVlabs/stylegan), etc, they use inception model from the [TensorFlow Models](https://github.com/tensorflow/models) repository, whose implementation is slightly different from that of `torchvision`. Hence, to make the evaluation metric comparable between different training frameworks (i.e., PyTorch and TensorFlow), we modify `torchvision/models/inception.py` as `inception.py`. The ported pre-trained weight is borrowed from [this repo](https://github.com/mseitzer/pytorch-fid).
+
+**NOTE:** We also support using the model from `torchvision` to compute the FID. However, please be aware that the FID value from `torchvision` is usually ~1.5 smaller than that from the TensorFlow model.
+
+Please use the following code to choose which model to use.
+
+```python
+from metrics.inception import build_inception_model
+
+inception_model_tf = build_inception_model(align_tf=True)
+inception_model_pth = build_inception_model(align_tf=False)
+```

networks/genforce/metrics/fid.py

@@ -0,0 +1,59 @@
+# python3.7
+"""Contains the functions to compute Frechet Inception Distance (FID).
+
+FID metric is introduced in paper
+
+GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash
+Equilibrium. Heusel et al. NeurIPS 2017.
+
+See details at https://arxiv.org/pdf/1706.08500.pdf
+"""
+
+import numpy as np
+import scipy.linalg
+
+__all__ = ['extract_feature', 'compute_fid']
+
+
+def extract_feature(inception_model, images):
+    """Extracts feature from input images with given model.
+
+    NOTE: The input images are assumed to be with pixel range [-1, 1].
+
+    Args:
+        inception_model: The model used to extract features.
+        images: The input image tensor to extract features from.
+
+    Returns:
+        A `numpy.ndarray`, containing the extracted features.
+    """
+    features = inception_model(images, output_logits=False)
+    features = features.detach().cpu().numpy()
+    assert features.ndim == 2 and features.shape[1] == 2048
+    return features
+
+
+def compute_fid(fake_features, real_features):
+    """Computes FID based on the features extracted from fake and real data.
+
+    Given the mean and covariance (m_f, C_f) of fake data and (m_r, C_r) of real
+    data, the FID metric can be computed by
+
+    d^2 = ||m_f - m_r||_2^2 + Tr(C_f + C_r - 2(C_f C_r)^0.5)
+
+    Args:
+        fake_features: The features extracted from fake data.
+        real_features: The features extracted from real data.
+
+    Returns:
+        A real number, suggesting the FID value.
+    """
+
+    m_f = np.mean(fake_features, axis=0)
+    C_f = np.cov(fake_features, rowvar=False)
+    m_r = np.mean(real_features, axis=0)
+    C_r = np.cov(real_features, rowvar=False)
+
+    fid = np.sum((m_f - m_r) ** 2) + np.trace(
+        C_f + C_r - 2 * scipy.linalg.sqrtm(np.dot(C_f, C_r)))
+    return np.real(fid)

networks/genforce/metrics/inception.py

@@ -0,0 +1,520 @@
+# python3.7
+"""Contains the Inception V3 model.
+
+This file is mostly borrowed from `torchvision/models/inception.py`.
+
+Inception model is widely used to compute FID or IS metric for evaluating
+generative models. However, the pre-trained models from torchvision is slightly
+different from the TensorFlow version.
+
+http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+
+In particular:
+
+(1) The number of classes in TensorFlow model is 1008 instead of 1000.
+(2) The avg_pool() layers in TensorFlow model does not include the padded zero.
+(3) The last Inception E Block in TensorFlow model use max_pool() instead of
+    avg_pool().
+
+Hence, to algin the evaluation results with those from TensorFlow
+implementation, we modified the inception model to support both versions. Please
+use `align_tf` argument to control the version.
+"""
+
+# pylint: disable=line-too-long
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+# pylint: disable=super-with-arguments
+# pylint: disable=consider-merging-isinstance
+# pylint: disable=import-outside-toplevel
+# pylint: disable=no-else-return
+
+from collections import namedtuple
+import warnings
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.jit.annotations import Optional
+from torch import Tensor
+from torchvision.models.utils import load_state_dict_from_url
+
+
+__all__ = ['build_inception_model', 'Inception3', 'inception_v3', 'InceptionOutputs', '_InceptionOutputs']
+
+model_urls = {
+    # Inception v3 ported from TensorFlow
+    'inception_v3_google': 'https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth',
+
+    # Inception v3 ported from http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+    # This model is provided by https://github.com/mseitzer/pytorch-fid
+    'tf_inception_v3': 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
+}
+
+InceptionOutputs = namedtuple('InceptionOutputs', ['logits', 'aux_logits'])
+InceptionOutputs.__annotations__ = {'logits': torch.Tensor, 'aux_logits': Optional[torch.Tensor]}
+
+# Script annotations failed with _GoogleNetOutputs = namedtuple ...
+# _InceptionOutputs set here for backwards compat
+_InceptionOutputs = InceptionOutputs
+
+
+def build_inception_model(align_tf=True):
+    """Builds Inception V3 model.
+
+    This model is particular used for inference, such that `requires_grad` and
+    `mode` will both be set as `False`.
+
+    Args:
+        align_tf: Whether to align the implementation with TensorFlow version. (default: True)
+
+    Returns:
+        A `torch.nn.Module` with pre-trained weight.
+    """
+    if align_tf:
+        num_classes = 1008
+        model_url = model_urls['tf_inception_v3']
+    else:
+        num_classes = 1000
+        model_url = model_urls['inception_v3_google']
+    model = Inception3(num_classes=num_classes,
+                       aux_logits=False,
+                       transform_input=False,
+                       align_tf=align_tf)
+    state_dict = load_state_dict_from_url(model_url)
+    model.load_state_dict(state_dict, strict=False)
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+    return model
+
+
+def inception_v3(pretrained=False, progress=True, **kwargs):
+    r"""Inception v3 model architecture from
+    `"Rethinking the Inception Architecture for Computer Vision" <http://arxiv.org/abs/1512.00567>`_.
+
+    .. note::
+        **Important**: In contrast to the other models the inception_v3 expects tensors with a size of
+        N x 3 x 299 x 299, so ensure your images are sized accordingly.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+        aux_logits (bool): If True, add an auxiliary branch that can improve training.
+            Default: *True*
+        transform_input (bool): If True, preprocesses the input according to the method with which it
+            was trained on ImageNet. Default: *False*
+    """
+    if pretrained:
+        if 'transform_input' not in kwargs:
+            kwargs['transform_input'] = True
+        if 'aux_logits' in kwargs:
+            original_aux_logits = kwargs['aux_logits']
+            kwargs['aux_logits'] = True
+        else:
+            original_aux_logits = True
+        model = Inception3(**kwargs)
+        state_dict = load_state_dict_from_url(model_urls['inception_v3_google'],
+                                              progress=progress)
+        model.load_state_dict(state_dict)
+        if not original_aux_logits:
+            model.aux_logits = False
+            del model.AuxLogits
+        return model
+
+    return Inception3(**kwargs)
+
+
+class Inception3(nn.Module):
+
+    def __init__(self, num_classes=1000, aux_logits=True, transform_input=False,
+                 inception_blocks=None, init_weights=True, align_tf=True):
+        super(Inception3, self).__init__()
+        if inception_blocks is None:
+            inception_blocks = [
+                BasicConv2d, InceptionA, InceptionB, InceptionC,
+                InceptionD, InceptionE, InceptionAux
+            ]
+        assert len(inception_blocks) == 7
+        conv_block = inception_blocks[0]
+        inception_a = inception_blocks[1]
+        inception_b = inception_blocks[2]
+        inception_c = inception_blocks[3]
+        inception_d = inception_blocks[4]
+        inception_e = inception_blocks[5]
+        inception_aux = inception_blocks[6]
+
+        self.aux_logits = aux_logits
+        self.transform_input = transform_input
+        self.align_tf = align_tf
+        self.Conv2d_1a_3x3 = conv_block(3, 32, kernel_size=3, stride=2)
+        self.Conv2d_2a_3x3 = conv_block(32, 32, kernel_size=3)
+        self.Conv2d_2b_3x3 = conv_block(32, 64, kernel_size=3, padding=1)
+        self.Conv2d_3b_1x1 = conv_block(64, 80, kernel_size=1)
+        self.Conv2d_4a_3x3 = conv_block(80, 192, kernel_size=3)
+        self.Mixed_5b = inception_a(192, pool_features=32, align_tf=self.align_tf)
+        self.Mixed_5c = inception_a(256, pool_features=64, align_tf=self.align_tf)
+        self.Mixed_5d = inception_a(288, pool_features=64, align_tf=self.align_tf)
+        self.Mixed_6a = inception_b(288)
+        self.Mixed_6b = inception_c(768, channels_7x7=128, align_tf=self.align_tf)
+        self.Mixed_6c = inception_c(768, channels_7x7=160, align_tf=self.align_tf)
+        self.Mixed_6d = inception_c(768, channels_7x7=160, align_tf=self.align_tf)
+        self.Mixed_6e = inception_c(768, channels_7x7=192, align_tf=self.align_tf)
+        if aux_logits:
+            self.AuxLogits = inception_aux(768, num_classes)
+        self.Mixed_7a = inception_d(768)
+        self.Mixed_7b = inception_e(1280, align_tf=self.align_tf)
+        self.Mixed_7c = inception_e(2048, use_max_pool=self.align_tf)
+        self.fc = nn.Linear(2048, num_classes)
+        if init_weights:
+            for m in self.modules():
+                if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                    import scipy.stats as stats
+                    stddev = m.stddev if hasattr(m, 'stddev') else 0.1
+                    X = stats.truncnorm(-2, 2, scale=stddev)
+                    values = torch.as_tensor(X.rvs(m.weight.numel()), dtype=m.weight.dtype)
+                    values = values.view(m.weight.size())
+                    with torch.no_grad():
+                        m.weight.copy_(values)
+                elif isinstance(m, nn.BatchNorm2d):
+                    nn.init.constant_(m.weight, 1)
+                    nn.init.constant_(m.bias, 0)
+
+    def _transform_input(self, x):
+        if self.transform_input:
+            x_ch0 = torch.unsqueeze(x[:, 0], 1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
+            x_ch1 = torch.unsqueeze(x[:, 1], 1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
+            x_ch2 = torch.unsqueeze(x[:, 2], 1) * (0.225 / 0.5) + (0.406 - 0.5) / 0.5
+            x = torch.cat((x_ch0, x_ch1, x_ch2), 1)
+        return x
+
+    def _forward(self, x, output_logits=False):
+        # Upsample if necessary
+        if x.shape[2] != 299 or x.shape[3] != 299:
+            x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)
+
+        # N x 3 x 299 x 299
+        x = self.Conv2d_1a_3x3(x)
+        # N x 32 x 149 x 149
+        x = self.Conv2d_2a_3x3(x)
+        # N x 32 x 147 x 147
+        x = self.Conv2d_2b_3x3(x)
+        # N x 64 x 147 x 147
+        x = F.max_pool2d(x, kernel_size=3, stride=2)
+        # N x 64 x 73 x 73
+        x = self.Conv2d_3b_1x1(x)
+        # N x 80 x 73 x 73
+        x = self.Conv2d_4a_3x3(x)
+        # N x 192 x 71 x 71
+        x = F.max_pool2d(x, kernel_size=3, stride=2)
+        # N x 192 x 35 x 35
+        x = self.Mixed_5b(x)
+        # N x 256 x 35 x 35
+        x = self.Mixed_5c(x)
+        # N x 288 x 35 x 35
+        x = self.Mixed_5d(x)
+        # N x 288 x 35 x 35
+        x = self.Mixed_6a(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6b(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6c(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6d(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6e(x)
+        # N x 768 x 17 x 17
+        aux_defined = self.training and self.aux_logits
+        if aux_defined:
+            aux = self.AuxLogits(x)
+        else:
+            aux = None
+        # N x 768 x 17 x 17
+        x = self.Mixed_7a(x)
+        # N x 1280 x 8 x 8
+        x = self.Mixed_7b(x)
+        # N x 2048 x 8 x 8
+        x = self.Mixed_7c(x)
+        # N x 2048 x 8 x 8
+        # Adaptive average pooling
+        x = F.adaptive_avg_pool2d(x, (1, 1))
+        # N x 2048 x 1 x 1
+        x = F.dropout(x, training=self.training)
+        # N x 2048 x 1 x 1
+        x = torch.flatten(x, 1)
+        # N x 2048
+        if output_logits:
+            x = self.fc(x)
+            # N x 1000 (num_classes)
+        return x, aux
+
+    @torch.jit.unused
+    def eager_outputs(self, x, aux):
+        # type: (Tensor, Optional[Tensor]) -> InceptionOutputs
+        if self.training and self.aux_logits:
+            return InceptionOutputs(x, aux)
+        else:
+            return x
+
+    def forward(self, x, output_logits=False):
+        x = self._transform_input(x)
+        x, aux = self._forward(x, output_logits)
+        aux_defined = self.training and self.aux_logits
+        if torch.jit.is_scripting():
+            if not aux_defined:
+                warnings.warn("Scripted Inception3 always returns Inception3 Tuple")
+            return InceptionOutputs(x, aux)
+        else:
+            return self.eager_outputs(x, aux)
+
+
+class InceptionA(nn.Module):
+
+    def __init__(self, in_channels, pool_features, conv_block=None, align_tf=False):
+        super(InceptionA, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch1x1 = conv_block(in_channels, 64, kernel_size=1)
+
+        self.branch5x5_1 = conv_block(in_channels, 48, kernel_size=1)
+        self.branch5x5_2 = conv_block(48, 64, kernel_size=5, padding=2)
+
+        self.branch3x3dbl_1 = conv_block(in_channels, 64, kernel_size=1)
+        self.branch3x3dbl_2 = conv_block(64, 96, kernel_size=3, padding=1)
+        self.branch3x3dbl_3 = conv_block(96, 96, kernel_size=3, padding=1)
+
+        self.branch_pool = conv_block(in_channels, pool_features, kernel_size=1)
+        self.pool_include_padding = not align_tf
+
+    def _forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch5x5 = self.branch5x5_1(x)
+        branch5x5 = self.branch5x5_2(branch5x5)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=self.pool_include_padding)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionB(nn.Module):
+
+    def __init__(self, in_channels, conv_block=None):
+        super(InceptionB, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch3x3 = conv_block(in_channels, 384, kernel_size=3, stride=2)
+
+        self.branch3x3dbl_1 = conv_block(in_channels, 64, kernel_size=1)
+        self.branch3x3dbl_2 = conv_block(64, 96, kernel_size=3, padding=1)
+        self.branch3x3dbl_3 = conv_block(96, 96, kernel_size=3, stride=2)
+
+    def _forward(self, x):
+        branch3x3 = self.branch3x3(x)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=2)
+
+        outputs = [branch3x3, branch3x3dbl, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionC(nn.Module):
+
+    def __init__(self, in_channels, channels_7x7, conv_block=None, align_tf=False):
+        super(InceptionC, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch1x1 = conv_block(in_channels, 192, kernel_size=1)
+
+        c7 = channels_7x7
+        self.branch7x7_1 = conv_block(in_channels, c7, kernel_size=1)
+        self.branch7x7_2 = conv_block(c7, c7, kernel_size=(1, 7), padding=(0, 3))
+        self.branch7x7_3 = conv_block(c7, 192, kernel_size=(7, 1), padding=(3, 0))
+
+        self.branch7x7dbl_1 = conv_block(in_channels, c7, kernel_size=1)
+        self.branch7x7dbl_2 = conv_block(c7, c7, kernel_size=(7, 1), padding=(3, 0))
+        self.branch7x7dbl_3 = conv_block(c7, c7, kernel_size=(1, 7), padding=(0, 3))
+        self.branch7x7dbl_4 = conv_block(c7, c7, kernel_size=(7, 1), padding=(3, 0))
+        self.branch7x7dbl_5 = conv_block(c7, 192, kernel_size=(1, 7), padding=(0, 3))
+
+        self.branch_pool = conv_block(in_channels, 192, kernel_size=1)
+        self.pool_include_padding = not align_tf
+
+    def _forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch7x7 = self.branch7x7_1(x)
+        branch7x7 = self.branch7x7_2(branch7x7)
+        branch7x7 = self.branch7x7_3(branch7x7)
+
+        branch7x7dbl = self.branch7x7dbl_1(x)
+        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
+
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=self.pool_include_padding)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionD(nn.Module):
+
+    def __init__(self, in_channels, conv_block=None):
+        super(InceptionD, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch3x3_1 = conv_block(in_channels, 192, kernel_size=1)
+        self.branch3x3_2 = conv_block(192, 320, kernel_size=3, stride=2)
+
+        self.branch7x7x3_1 = conv_block(in_channels, 192, kernel_size=1)
+        self.branch7x7x3_2 = conv_block(192, 192, kernel_size=(1, 7), padding=(0, 3))
+        self.branch7x7x3_3 = conv_block(192, 192, kernel_size=(7, 1), padding=(3, 0))
+        self.branch7x7x3_4 = conv_block(192, 192, kernel_size=3, stride=2)
+
+    def _forward(self, x):
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = self.branch3x3_2(branch3x3)
+
+        branch7x7x3 = self.branch7x7x3_1(x)
+        branch7x7x3 = self.branch7x7x3_2(branch7x7x3)
+        branch7x7x3 = self.branch7x7x3_3(branch7x7x3)
+        branch7x7x3 = self.branch7x7x3_4(branch7x7x3)
+
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=2)
+        outputs = [branch3x3, branch7x7x3, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionE(nn.Module):
+
+    def __init__(self, in_channels, conv_block=None, align_tf=False, use_max_pool=False):
+        super(InceptionE, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch1x1 = conv_block(in_channels, 320, kernel_size=1)
+
+        self.branch3x3_1 = conv_block(in_channels, 384, kernel_size=1)
+        self.branch3x3_2a = conv_block(384, 384, kernel_size=(1, 3), padding=(0, 1))
+        self.branch3x3_2b = conv_block(384, 384, kernel_size=(3, 1), padding=(1, 0))
+
+        self.branch3x3dbl_1 = conv_block(in_channels, 448, kernel_size=1)
+        self.branch3x3dbl_2 = conv_block(448, 384, kernel_size=3, padding=1)
+        self.branch3x3dbl_3a = conv_block(384, 384, kernel_size=(1, 3), padding=(0, 1))
+        self.branch3x3dbl_3b = conv_block(384, 384, kernel_size=(3, 1), padding=(1, 0))
+
+        self.branch_pool = conv_block(in_channels, 192, kernel_size=1)
+        self.pool_include_padding = not align_tf
+        self.use_max_pool = use_max_pool
+
+    def _forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        if self.use_max_pool:
+            branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+        else:
+            branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                       count_include_pad=self.pool_include_padding)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionAux(nn.Module):
+
+    def __init__(self, in_channels, num_classes, conv_block=None):
+        super(InceptionAux, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.conv0 = conv_block(in_channels, 128, kernel_size=1)
+        self.conv1 = conv_block(128, 768, kernel_size=5)
+        self.conv1.stddev = 0.01
+        self.fc = nn.Linear(768, num_classes)
+        self.fc.stddev = 0.001
+
+    def forward(self, x):
+        # N x 768 x 17 x 17
+        x = F.avg_pool2d(x, kernel_size=5, stride=3)
+        # N x 768 x 5 x 5
+        x = self.conv0(x)
+        # N x 128 x 5 x 5
+        x = self.conv1(x)
+        # N x 768 x 1 x 1
+        # Adaptive average pooling
+        x = F.adaptive_avg_pool2d(x, (1, 1))
+        # N x 768 x 1 x 1
+        x = torch.flatten(x, 1)
+        # N x 768
+        x = self.fc(x)
+        # N x 1000
+        return x
+
+
+class BasicConv2d(nn.Module):
+
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
+        self.bn = nn.BatchNorm2d(out_channels, eps=0.001)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return F.relu(x, inplace=True)
+
+# pylint: enable=line-too-long
+# pylint: enable=missing-function-docstring
+# pylint: enable=missing-class-docstring
+# pylint: enable=super-with-arguments
+# pylint: enable=consider-merging-isinstance
+# pylint: enable=import-outside-toplevel
+# pylint: enable=no-else-return

networks/genforce/models/__init__.py

@@ -0,0 +1,131 @@
+# python3.7
+"""Collects all available models together."""
+
+from .model_zoo import MODEL_ZOO
+from .pggan_generator import PGGANGenerator
+from .pggan_discriminator import PGGANDiscriminator
+from .stylegan_generator import StyleGANGenerator
+from .stylegan_discriminator import StyleGANDiscriminator
+from .stylegan2_generator import StyleGAN2Generator
+from .stylegan2_discriminator import StyleGAN2Discriminator
+from .encoder import EncoderNet
+from .perceptual_model import PerceptualModel
+
+__all__ = [
+    'MODEL_ZOO', 'PGGANGenerator', 'PGGANDiscriminator', 'StyleGANGenerator',
+    'StyleGANDiscriminator', 'StyleGAN2Generator', 'StyleGAN2Discriminator',
+    'EncoderNet', 'PerceptualModel', 'build_generator', 'build_discriminator',
+    'build_encoder', 'build_perceptual', 'build_model'
+]
+
+_GAN_TYPES_ALLOWED = ['pggan', 'stylegan', 'stylegan2']
+_MODULES_ALLOWED = ['generator', 'discriminator', 'encoder', 'perceptual']
+
+
+def build_generator(gan_type, resolution, **kwargs):
+    """Builds generator by GAN type.
+
+    Args:
+        gan_type: GAN type to which the generator belong.
+        resolution: Synthesis resolution.
+        **kwargs: Additional arguments to build the generator.
+
+    Raises:
+        ValueError: If the `gan_type` is not supported.
+        NotImplementedError: If the `gan_type` is not implemented.
+    """
+    if gan_type not in _GAN_TYPES_ALLOWED:
+        raise ValueError(f'Invalid GAN type: `{gan_type}`!\n'
+                         f'Types allowed: {_GAN_TYPES_ALLOWED}.')
+
+    if gan_type == 'pggan':
+        return PGGANGenerator(resolution, **kwargs)
+    if gan_type == 'stylegan':
+        return StyleGANGenerator(resolution, **kwargs)
+    if gan_type == 'stylegan2':
+        return StyleGAN2Generator(resolution, **kwargs)
+    raise NotImplementedError(f'Unsupported GAN type `{gan_type}`!')
+
+
+def build_discriminator(gan_type, resolution, **kwargs):
+    """Builds discriminator by GAN type.
+
+    Args:
+        gan_type: GAN type to which the discriminator belong.
+        resolution: Synthesis resolution.
+        **kwargs: Additional arguments to build the discriminator.
+
+    Raises:
+        ValueError: If the `gan_type` is not supported.
+        NotImplementedError: If the `gan_type` is not implemented.
+    """
+    if gan_type not in _GAN_TYPES_ALLOWED:
+        raise ValueError(f'Invalid GAN type: `{gan_type}`!\n'
+                         f'Types allowed: {_GAN_TYPES_ALLOWED}.')
+
+    if gan_type == 'pggan':
+        return PGGANDiscriminator(resolution, **kwargs)
+    if gan_type == 'stylegan':
+        return StyleGANDiscriminator(resolution, **kwargs)
+    if gan_type == 'stylegan2':
+        return StyleGAN2Discriminator(resolution, **kwargs)
+    raise NotImplementedError(f'Unsupported GAN type `{gan_type}`!')
+
+
+def build_encoder(gan_type, resolution, **kwargs):
+    """Builds encoder by GAN type.
+
+    Args:
+        gan_type: GAN type to which the encoder belong.
+        resolution: Input resolution for encoder.
+        **kwargs: Additional arguments to build the encoder.
+
+    Raises:
+        ValueError: If the `gan_type` is not supported.
+        NotImplementedError: If the `gan_type` is not implemented.
+    """
+    if gan_type not in _GAN_TYPES_ALLOWED:
+        raise ValueError(f'Invalid GAN type: `{gan_type}`!\n'
+                         f'Types allowed: {_GAN_TYPES_ALLOWED}.')
+
+    if gan_type in ['stylegan', 'stylegan2']:
+        return EncoderNet(resolution, **kwargs)
+
+    raise NotImplementedError(f'Unsupported GAN type `{gan_type}` for encoder!')
+
+
+def build_perceptual(**kwargs):
+    """Builds perceptual model.
+
+    Args:
+        **kwargs: Additional arguments to build the encoder.
+    """
+    return PerceptualModel(**kwargs)
+
+
+def build_model(gan_type, module, resolution, **kwargs):
+    """Builds a GAN module (generator/discriminator/etc).
+
+    Args:
+        gan_type: GAN type to which the model belong.
+        module: GAN module to build, such as generator or discrimiantor.
+        resolution: Synthesis resolution.
+        **kwargs: Additional arguments to build the discriminator.
+
+    Raises:
+        ValueError: If the `module` is not supported.
+        NotImplementedError: If the `module` is not implemented.
+    """
+    if module not in _MODULES_ALLOWED:
+        raise ValueError(f'Invalid module: `{module}`!\n'
+                         f'Modules allowed: {_MODULES_ALLOWED}.')
+
+    if module == 'generator':
+        return build_generator(gan_type, resolution, **kwargs)
+    if module == 'discriminator':
+        return build_discriminator(gan_type, resolution, **kwargs)
+    if module == 'encoder':
+        return build_encoder(gan_type, resolution, **kwargs)
+    if module == 'perceptual':
+        return build_perceptual(**kwargs)
+    raise NotImplementedError(f'Unsupported module `{module}`!')