# python 3.7
"""Contains basic configurations for models used in this project.

Please download the public released models from the following repositories
OR train your own models, and then put them into the folder
`pretrain/tensorflow`.

PGGAN: https://github.com/tkarras/progressive_growing_of_gans
StyleGAN: https://github.com/NVlabs/stylegan
StyleGAN2: https://github.com/NVlabs/stylegan2

NOTE: Any new model should be registered in `MODEL_POOL` before used.
"""

import os
import copy
import json

BASE_DIR = os.path.dirname(os.path.relpath(__file__))

MODEL_DIR = os.path.join(BASE_DIR, 'pretrain')
PTH_MODEL_DIR = 'pytorch'
TF_MODEL_DIR = 'tensorflow'



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"

if not os.path.exists(os.path.join(MODEL_DIR, PTH_MODEL_DIR)):
  os.makedirs(os.path.join(MODEL_DIR, PTH_MODEL_DIR))


# pylint: disable=line-too-long
MODEL_POOL = {
    # PGGAN Official.
    'pggan_celebahq': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_celebahq1024_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-celebahq-1024x1024.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'celebahq',
        'z_space_dim': 512,
        'resolution': 1024,
        'fused_scale': False,
    },
    'pggan_celebahq_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_celebahq1024_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-celebahq-1024x1024.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan_disc',
        'dataset_name': 'celebahq',
        'resolution': 1024,
        'fused_scale': False,
    },
    'pggan_bedroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bedroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bedroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bedroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bedroom_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bedroom256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bedroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan_disc',
        'dataset_name': 'lsun-bedroom',
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_livingroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_livingroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-livingroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-livingroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_diningroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_diningroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-dining_room-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-diningroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_kitchen': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_kitchen256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-kitchen-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-kitchen',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_churchoutdoor': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_churchoutdoor256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-churchoutdoor-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-churchoutdoor',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_churchoutdoor_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_churchoutdoor256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-churchoutdoor-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan_disc',
        'dataset_name': 'lsun-churchoutdoor',
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_tower': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_tower256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-tower-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-tower',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bridge': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bridge256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bridge-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bridge',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_restaurant': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_restaurant256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-restaurant-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-restaurant',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_classroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_classroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-classroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-classroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_conferenceroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_conferenceroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-conferenceroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-conferenceroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_person': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_person256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-person-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-person',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_cat': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_cat256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-cat-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-cat',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_dog': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_dog256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-dog-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-dog',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bird': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bird256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bird-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bird',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_horse': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_horse256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-horse-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-horse',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_horse_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_horse256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-horse-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan_disc',
        'dataset_name': 'lsun-horse',
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_sheep': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_sheep256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-sheep-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-sheep',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_cow': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_cow256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-cow-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-cow',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_car': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_car256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-car-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-car',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bicycle': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bicycle256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bicycle-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bicycle',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_motorbike': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_motorbike256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-motorbike-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-motorbike',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bus': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bus256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bus-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bus',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_train': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_train256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-train-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-train',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_boat': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_boat256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-boat-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-boat',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_airplane': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_airplane256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-airplane-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-airplane',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bottle': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bottle256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bottle-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bottle',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_chair': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_chair256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-chair-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-chair',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_pottedplant': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_pottedplant256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-pottedplant-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-pottedplant',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_tvmonitor': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_tvmonitor256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-tvmonitor-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-tvmonitor',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_diningtable': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_diningtable256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-diningtable-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-diningtable',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_sofa': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_sofa256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-sofa-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-sofa',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },

    # StyleGAN Official.
    'stylegan_ffhq': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_ffhq1024_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-ffhq-1024x1024.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'ffhq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 1024,
        'fused_scale': 'auto',
    },
    'stylegan_celebahq': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_celebahq1024_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-celebahq-1024x1024.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'celebahq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 1024,
        'fused_scale': 'auto',
    },
    'stylegan_bedroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_bedroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-bedrooms-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-bedroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_cat': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_cat256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-cats-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-cat',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_car': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_car512_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-cars-512x384.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-car',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 512,
        'fused_scale': 'auto',
    },

    # StyleGAN Self-Training.
    'stylegan_ffhq256': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_ffhq256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-ffhq-256x256-025000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'ffhq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_ffhq512': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_ffhq512_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-ffhq-512x512-025000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'ffhq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 512,
        'fused_scale': 'auto',
    },
    'stylegan_livingroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_livingroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-livingroom-256x256-030000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-livingroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_diningroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_diningroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-diningroom-256x256-025000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-diningroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_kitchen': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_kitchen256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-kitchen-256x256-030000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-kitchen',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_apartment': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_apartment256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-apartment-256x256-060000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-bedroom-livingroom-diningroom-kitchen',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_churchoutdoor': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_churchoutdoor256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-churchoutdoor-256x256-030000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-churchoutdoor',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_tower': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_tower256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-tower-256x256-030000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-tower',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_bridge': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_bridge256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-bridge-256x256-025000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-bridge',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_restaurant': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_restaurant256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-restaurant-256x256-050000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-restaurant',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_classroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_classroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-classroom-256x256-050000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-classroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_conferenceroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_conferenceroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-conferenceroom-256x256-050000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-conferenceroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },

    # StyleGAN2 Official.
    'stylegan2_ffhq': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_ffhq1024_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-ffhq-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'ffhq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 1024,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },

    'stylegan2_ffhq_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_ffhq1024_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-ffhq-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2_disc',
        'dataset_name': 'ffhq',
        'resolution': 1024,
        'd_architecture_type': 'resnet',
        'fused_modulate': True,
    },

    'stylegan2_church': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_church256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-church-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'lsun-church',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },
    'stylegan2_church_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_church256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-church-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2_disc',
        'dataset_name': 'church',
        'resolution': 256,
        'd_architecture_type': 'resnet',
        'fused_modulate': True,
    },
    'stylegan2_cat': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_cat256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-cat-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'lsun-cat',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },
    'stylegan2_horse': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_horse256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-horse-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'lsun-horse',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },
    'stylegan2_horse_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_horse256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-horse-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2_disc',
        'dataset_name': 'horse',
        'resolution': 256,
        'd_architecture_type': 'resnet',
        'fused_modulate': True,
    },
    'stylegan2_car': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_car512_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-car-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'lsun-car',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 512,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },
    'biggandeep128_imagenet':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-128-pytorch_model.bin'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'biggan-deep-128-tensorflow_model.bin'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-128-pytorch_config.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'biggan_tf_official'),
        'gan_type': 'biggandeep',
        'dataset_name': 'imagenet',
        'z_space_dim': 128,
        'resolution': 128,
    },
    'biggandeep256_imagenet':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-256-pytorch_model.bin'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'biggan-deep-256-tensorflow_model.bin'),
        'config_path':  os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-256-pytorch_config.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'biggan_tf_official'),
        'gan_type': 'biggandeep',
        'dataset_name': 'imagenet',
        'z_space_dim': 128,
        'resolution': 256,
    },
    'biggandeep512_imagenet':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-512-pytorch_model.bin'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'biggan-deep-512-tensorflow_model.bin'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-512-pytorch_config.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'biggan_tf_official'),
        'gan_type': 'biggandeep',
        'dataset_name': 'imagenet',
        'z_space_dim': 128,
        'resolution': 512,
    },
    'bigganshallow_imagenet':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'BigGAN', 'G_ema.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'generators', 'BigGAN', 'G_ema.pth'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'BigGAN', 'generator_config.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'biggan_tf_official'),
        'gan_type': 'bigganshallow',
        'dataset_name': 'imagenet',
        'z_space_dim': 120,
        'resolution': 128,
    },
    'sngan_anime':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'SN_Anime', 'generator.pt'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'generators', 'SN_Anime', 'generator.pt'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'SN_Anime', 'args.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'sngan_tf_official'),
        'gan_type': 'sngan',
        'dataset_name': 'anime',
        'z_space_dim': 128,
        'resolution': 64,
    },
    'sngan_mnist':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'SN_MNIST', 'generator.pt'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'generators', 'SN_MNIST', 'generator.pt'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'SN_MNIST', 'args.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'sngan_tf_official'),
        'gan_type': 'sngan',
        'dataset_name': 'mnist',
        'z_space_dim': 128,
        'resolution': 32,
    }
}
# pylint: enable=line-too-long

# Settings for StyleGAN.
STYLEGAN_TRUNCATION_PSI = 0.7  # 1.0 means no truncation
STYLEGAN_TRUNCATION_LAYERS = 8  # 0 means no truncation
STYLEGAN_RANDOMIZE_NOISE = False

# Settings for StyleGAN2.
STYLEGAN2_TRUNCATION_PSI = 0.7  # 1.0 means no truncation
STYLEGAN2_TRUNCATION_LAYERS = 8  # 0 means no truncation
STYLEGAN2_RANDOMIZE_NOISE = False

# Settings for model running.
USE_CUDA = True

MAX_IMAGES_ON_DEVICE = 4

MAX_IMAGES_ON_RAM = 800



# python 3.7
"""Contains basic configurations for models used in this project.

Please download the public released models from the following repositories
OR train your own models, and then put them into the folder
`pretrain/tensorflow`.

PGGAN: https://github.com/tkarras/progressive_growing_of_gans
StyleGAN: https://github.com/NVlabs/stylegan
StyleGAN2: https://github.com/NVlabs/stylegan2

NOTE: Any new model should be registered in `MODEL_POOL` before used.
"""

import os
import copy
import json

BASE_DIR = os.path.dirname(os.path.relpath(__file__))

MODEL_DIR = os.path.join(BASE_DIR, 'pretrain')
PTH_MODEL_DIR = 'pytorch'
TF_MODEL_DIR = 'tensorflow'



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"

if not os.path.exists(os.path.join(MODEL_DIR, PTH_MODEL_DIR)):
  os.makedirs(os.path.join(MODEL_DIR, PTH_MODEL_DIR))


# pylint: disable=line-too-long
MODEL_POOL = {
    # PGGAN Official.
    'pggan_celebahq': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_celebahq1024_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-celebahq-1024x1024.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'celebahq',
        'z_space_dim': 512,
        'resolution': 1024,
        'fused_scale': False,
    },
    'pggan_celebahq_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_celebahq1024_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-celebahq-1024x1024.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan_disc',
        'dataset_name': 'celebahq',
        'resolution': 1024,
        'fused_scale': False,
    },
    'pggan_bedroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bedroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bedroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bedroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bedroom_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bedroom256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bedroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan_disc',
        'dataset_name': 'lsun-bedroom',
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_livingroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_livingroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-livingroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-livingroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_diningroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_diningroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-dining_room-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-diningroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_kitchen': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_kitchen256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-kitchen-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-kitchen',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_churchoutdoor': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_churchoutdoor256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-churchoutdoor-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-churchoutdoor',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_churchoutdoor_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_churchoutdoor256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-churchoutdoor-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan_disc',
        'dataset_name': 'lsun-churchoutdoor',
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_tower': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_tower256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-tower-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-tower',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bridge': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bridge256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bridge-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bridge',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_restaurant': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_restaurant256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-restaurant-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-restaurant',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_classroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_classroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-classroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-classroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_conferenceroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_conferenceroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-conferenceroom-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-conferenceroom',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_person': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_person256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-person-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-person',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_cat': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_cat256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-cat-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-cat',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_dog': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_dog256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-dog-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-dog',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bird': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bird256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bird-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bird',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_horse': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_horse256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-horse-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-horse',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_horse_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_horse256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-horse-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan_disc',
        'dataset_name': 'lsun-horse',
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_sheep': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_sheep256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-sheep-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-sheep',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_cow': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_cow256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-cow-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-cow',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_car': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_car256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-car-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-car',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bicycle': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bicycle256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bicycle-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bicycle',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_motorbike': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_motorbike256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-motorbike-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-motorbike',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bus': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bus256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bus-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bus',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_train': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_train256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-train-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-train',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_boat': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_boat256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-boat-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-boat',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_airplane': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_airplane256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-airplane-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-airplane',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_bottle': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_bottle256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-bottle-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-bottle',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_chair': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_chair256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-chair-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-chair',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_pottedplant': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_pottedplant256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-pottedplant-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-pottedplant',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_tvmonitor': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_tvmonitor256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-tvmonitor-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-tvmonitor',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_diningtable': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_diningtable256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-diningtable-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-diningtable',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },
    'pggan_sofa': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'pggan_sofa256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2018iclr-lsun-sofa-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'pggan_tf_official'),
        'gan_type': 'pggan',
        'dataset_name': 'lsun-sofa',
        'z_space_dim': 512,
        'resolution': 256,
        'fused_scale': False,
    },

    # StyleGAN Official.
    'stylegan_ffhq': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_ffhq1024_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-ffhq-1024x1024.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'ffhq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 1024,
        'fused_scale': 'auto',
    },
    'stylegan_celebahq': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_celebahq1024_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-celebahq-1024x1024.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'celebahq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 1024,
        'fused_scale': 'auto',
    },
    'stylegan_bedroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_bedroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-bedrooms-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-bedroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_cat': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_cat256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-cats-256x256.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-cat',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_car': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_car512_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'karras2019stylegan-cars-512x384.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-car',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 512,
        'fused_scale': 'auto',
    },

    # StyleGAN Self-Training.
    'stylegan_ffhq256': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_ffhq256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-ffhq-256x256-025000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'ffhq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_ffhq512': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_ffhq512_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-ffhq-512x512-025000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'ffhq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 512,
        'fused_scale': 'auto',
    },
    'stylegan_livingroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_livingroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-livingroom-256x256-030000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-livingroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_diningroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_diningroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-diningroom-256x256-025000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-diningroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_kitchen': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_kitchen256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-kitchen-256x256-030000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-kitchen',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_apartment': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_apartment256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-apartment-256x256-060000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-bedroom-livingroom-diningroom-kitchen',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_churchoutdoor': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_churchoutdoor256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-churchoutdoor-256x256-030000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-churchoutdoor',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_tower': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_tower256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-tower-256x256-030000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-tower',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_bridge': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_bridge256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-bridge-256x256-025000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-bridge',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_restaurant': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_restaurant256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-restaurant-256x256-050000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-restaurant',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_classroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_classroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-classroom-256x256-050000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-classroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },
    'stylegan_conferenceroom': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan_conferenceroom256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan-conferenceroom-256x256-050000.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan_tf_official'),
        'gan_type': 'stylegan',
        'dataset_name': 'lsun-conferenceroom',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'fused_scale': 'auto',
    },

    # StyleGAN2 Official.
    'stylegan2_ffhq': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_ffhq1024_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-ffhq-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'ffhq',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 1024,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },

    'stylegan2_ffhq_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_ffhq1024_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-ffhq-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2_disc',
        'dataset_name': 'ffhq',
        'resolution': 1024,
        'd_architecture_type': 'resnet',
        'fused_modulate': True,
    },

    'stylegan2_church': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_church256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-church-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'lsun-church',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },
    'stylegan2_church_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_church256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-church-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2_disc',
        'dataset_name': 'church',
        'resolution': 256,
        'd_architecture_type': 'resnet',
        'fused_modulate': True,
    },
    'stylegan2_cat': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_cat256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-cat-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'lsun-cat',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },
    'stylegan2_horse': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_horse256_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-horse-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'lsun-horse',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 256,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },
    'stylegan2_horse_disc': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_horse256_discriminator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-horse-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2_disc',
        'dataset_name': 'horse',
        'resolution': 256,
        'd_architecture_type': 'resnet',
        'fused_modulate': True,
    },
    'stylegan2_car': {
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'stylegan2_car512_generator.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'stylegan2-car-config-f.pkl'),
        'tf_code_path': os.path.join(BASE_DIR, 'stylegan2_tf_official'),
        'gan_type': 'stylegan2',
        'dataset_name': 'lsun-car',
        'z_space_dim': 512,
        'w_space_dim': 512,
        'resolution': 512,
        'g_architecture_type': 'skip',
        'fused_modulate': True,
    },
    'biggandeep128_imagenet':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-128-pytorch_model.bin'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'biggan-deep-128-tensorflow_model.bin'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-128-pytorch_config.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'biggan_tf_official'),
        'gan_type': 'biggandeep',
        'dataset_name': 'imagenet',
        'z_space_dim': 128,
        'resolution': 128,
    },
    'biggandeep256_imagenet':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-256-pytorch_model.bin'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'biggan-deep-256-tensorflow_model.bin'),
        'config_path':  os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-256-pytorch_config.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'biggan_tf_official'),
        'gan_type': 'biggandeep',
        'dataset_name': 'imagenet',
        'z_space_dim': 128,
        'resolution': 256,
    },
    'biggandeep512_imagenet':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-512-pytorch_model.bin'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'biggan-deep-512-tensorflow_model.bin'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'biggan-deep-512-pytorch_config.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'biggan_tf_official'),
        'gan_type': 'biggandeep',
        'dataset_name': 'imagenet',
        'z_space_dim': 128,
        'resolution': 512,
    },
    'bigganshallow_imagenet':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'BigGAN', 'G_ema.pth'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'generators', 'BigGAN', 'G_ema.pth'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'BigGAN', 'generator_config.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'biggan_tf_official'),
        'gan_type': 'bigganshallow',
        'dataset_name': 'imagenet',
        'z_space_dim': 120,
        'resolution': 128,
    },
    'sngan_anime':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'SN_Anime', 'generator.pt'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'generators', 'SN_Anime', 'generator.pt'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'SN_Anime', 'args.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'sngan_tf_official'),
        'gan_type': 'sngan',
        'dataset_name': 'anime',
        'z_space_dim': 128,
        'resolution': 64,
    },
    'sngan_mnist':{
        'weight_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'SN_MNIST', 'generator.pt'),
        'tf_weight_path': os.path.join(MODEL_DIR, TF_MODEL_DIR, 'generators', 'SN_MNIST', 'generator.pt'),
        'config_path': os.path.join(MODEL_DIR, PTH_MODEL_DIR, 'generators', 'SN_MNIST', 'args.json'),
        'tf_code_path': os.path.join(BASE_DIR, 'sngan_tf_official'),
        'gan_type': 'sngan',
        'dataset_name': 'mnist',
        'z_space_dim': 128,
        'resolution': 32,
    }
}
# pylint: enable=line-too-long

# Settings for StyleGAN.
STYLEGAN_TRUNCATION_PSI = 0.7  # 1.0 means no truncation
STYLEGAN_TRUNCATION_LAYERS = 8  # 0 means no truncation
STYLEGAN_RANDOMIZE_NOISE = False

# Settings for StyleGAN2.
STYLEGAN2_TRUNCATION_PSI = 0.7  # 1.0 means no truncation
STYLEGAN2_TRUNCATION_LAYERS = 8  # 0 means no truncation
STYLEGAN2_RANDOMIZE_NOISE = False

# Settings for model running.
USE_CUDA = True

MAX_IMAGES_ON_DEVICE = 4

MAX_IMAGES_ON_RAM = 800



# python 3.7 Contains basic configurations for models used in this project.  Please download the public released models from the following repositories OR train your own models, and then put them into the folder `pretrain/tensorflow`.  PGGAN: https://github.com/tkarras/progressive_growing_of_gans StyleGAN: https://github.com/NVlabs/stylegan StyleGAN2: https://github.com/NVlabs/stylegan2  NOTE: Any new model should be registered in `MODEL_POOL` before used. 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) Constructs BigGANConfig. Constructs a `BigGANConfig` from a Python dictionary of parameters. Constructs a `BigGANConfig` from a json file of parameters. Serializes this instance to a Python dictionary. Serializes this instance to a JSON string. # pylint: disable=line-too-long # PGGAN Official. # StyleGAN Official. # StyleGAN Self-Training. # StyleGAN2 Official. # pylint: enable=line-too-long # Settings for StyleGAN. # 1.0 means no truncation # 0 means no truncation # Settings for StyleGAN2. # 1.0 means no truncation # 0 means no truncation # Settings for model running.