ICML2022/resefa

1# python3.72"""Collects all models."""3 4from .pggan_generator import PGGANGenerator5from .pggan_discriminator import PGGANDiscriminator6from .stylegan_generator import StyleGANGenerator7from .stylegan_discriminator import StyleGANDiscriminator8from .stylegan2_generator import StyleGAN2Generator9from .stylegan2_discriminator import StyleGAN2Discriminator10from .stylegan3_generator import StyleGAN3Generator11from .ghfeat_encoder import GHFeatEncoder12from .perceptual_model import PerceptualModel13from .inception_model import InceptionModel14 15__all__ = ['build_model']16 17_MODELS = {18    'PGGANGenerator': PGGANGenerator,19    'PGGANDiscriminator': PGGANDiscriminator,20    'StyleGANGenerator': StyleGANGenerator,21    'StyleGANDiscriminator': StyleGANDiscriminator,22    'StyleGAN2Generator': StyleGAN2Generator,23    'StyleGAN2Discriminator': StyleGAN2Discriminator,24    'StyleGAN3Generator': StyleGAN3Generator,25    'GHFeatEncoder': GHFeatEncoder,26    'PerceptualModel': PerceptualModel.build_model,27    'InceptionModel': InceptionModel.build_model28}29 30 31def build_model(model_type, **kwargs):32    """Builds a model based on its class type.33 34    Args:35        model_type: Class type to which the model belongs, which is case36            sensitive.37        **kwargs: Additional arguments to build the model.38 39    Raises:40        ValueError: If the `model_type` is not supported.41    """42    if model_type not in _MODELS:43        raise ValueError(f'Invalid model type: `{model_type}`!\n'44                         f'Types allowed: {list(_MODELS)}.')45    return _MODELS[model_type](**kwargs)46