# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # This work is made available under the Nvidia Source Code License-NC. # To view a copy of this license, check out LICENSE.md import warnings from types import SimpleNamespace import torch from torch import nn from torch.nn import Upsample as NearestUpsample from imaginaire.layers import Conv2dBlock, LinearBlock, Res2dBlock from imaginaire.generators.unit import ContentEncoder class Generator(nn.Module): r"""Improved MUNIT generator. Args: gen_cfg (obj): Generator definition part of the yaml config file. data_cfg (obj): Data definition part of the yaml config file. """ def __init__(self, gen_cfg, data_cfg): super().__init__() self.autoencoder_a = AutoEncoder(**vars(gen_cfg)) self.autoencoder_b = AutoEncoder(**vars(gen_cfg)) def forward(self, data, random_style=True, image_recon=True, latent_recon=True, cycle_recon=True, within_latent_recon=False): r"""In MUNIT's forward pass, it generates a content code and a style code from images in both domain. It then performs a within-domain reconstruction step and a cross-domain translation step. In within-domain reconstruction, it reconstructs an image using the content and style from the same image and optionally encodes the image back to the latent space. In cross-domain translation, it generates an translated image by mixing the content and style from images in different domains, and optionally encodes the image back to the latent space. Args: data (dict): Training data at the current iteration. - images_a (tensor): Images from domain A. - images_b (tensor): Images from domain B. random_style (bool): If ``True``, samples the style code from the prior distribution, otherwise uses the style code encoded from the input images in the other domain. image_recon (bool): If ``True``, also returns reconstructed images. latent_recon (bool): If ``True``, also returns reconstructed latent code during cross-domain translation. cycle_recon (bool): If ``True``, also returns cycle reconstructed images. within_latent_recon (bool): If ``True``, also returns reconstructed latent code during within-domain reconstruction. """ images_a = data['images_a'] images_b = data['images_b'] net_G_output = dict() # encode input images into content and style code content_a, style_a = self.autoencoder_a.encode(images_a) content_b, style_b = self.autoencoder_b.encode(images_b) # decode (within domain) if image_recon: images_aa = self.autoencoder_a.decode(content_a, style_a) images_bb = self.autoencoder_b.decode(content_b, style_b) net_G_output.update(dict(images_aa=images_aa, images_bb=images_bb)) # decode (cross domain) if random_style: # use randomly sampled style code style_a_rand = torch.randn_like(style_a) style_b_rand = torch.randn_like(style_b) else: # use style code encoded from the other domain style_a_rand = style_a style_b_rand = style_b images_ba = self.autoencoder_a.decode(content_b, style_a_rand) images_ab = self.autoencoder_b.decode(content_a, style_b_rand) # encode translated images into content and style code if latent_recon or cycle_recon: content_ba, style_ba = self.autoencoder_a.encode(images_ba) content_ab, style_ab = self.autoencoder_b.encode(images_ab) net_G_output.update(dict(content_ba=content_ba, style_ba=style_ba, content_ab=content_ab, style_ab=style_ab)) # encode reconstructed images into content and style code if image_recon and within_latent_recon: content_aa, style_aa = self.autoencoder_a.encode(images_aa) content_bb, style_bb = self.autoencoder_b.encode(images_bb) net_G_output.update(dict(content_aa=content_aa, style_aa=style_aa, content_bb=content_bb, style_bb=style_bb)) # cycle reconstruction if cycle_recon: images_aba = self.autoencoder_a.decode(content_ab, style_a) images_bab = self.autoencoder_b.decode(content_ba, style_b) net_G_output.update( dict(images_aba=images_aba, images_bab=images_bab)) # required outputs net_G_output.update(dict(content_a=content_a, content_b=content_b, style_a=style_a, style_b=style_b, style_a_rand=style_a_rand, style_b_rand=style_b_rand, images_ba=images_ba, images_ab=images_ab)) return net_G_output def inference(self, data, a2b=True, random_style=True): r"""MUNIT inference. Args: data (dict): Training data at the current iteration. - images_a (tensor): Images from domain A. - images_b (tensor): Images from domain B. a2b (bool): If ``True``, translates images from domain A to B, otherwise from B to A. random_style (bool): If ``True``, samples the style code from the prior distribution, otherwise uses the style code encoded from the input images in the other domain. """ if a2b: input_key = 'images_a' content_encode = self.autoencoder_a.content_encoder style_encode = self.autoencoder_b.style_encoder decode = self.autoencoder_b.decode else: input_key = 'images_b' content_encode = self.autoencoder_b.content_encoder style_encode = self.autoencoder_a.style_encoder decode = self.autoencoder_a.decode content_images = data[input_key] content = content_encode(content_images) if random_style: style_channels = self.autoencoder_a.style_channels style = torch.randn(content.size(0), style_channels, 1, 1, device=torch.device('cuda')) file_names = data['key'][input_key]['filename'] else: style_key = 'images_b' if a2b else 'images_a' assert style_key in data.keys(), \ "{} must be provided when 'random_style' " \ "is set to False".format(style_key) style_images = data[style_key] style = style_encode(style_images) file_names = \ [content_name + '_style_' + style_name for content_name, style_name in zip(data['key'][input_key]['filename'], data['key'][style_key]['filename'])] output_images = decode(content, style) return output_images, file_names class AutoEncoder(nn.Module): r"""Improved MUNIT autoencoder. Args: num_filters (int): Base filter numbers. max_num_filters (int): Maximum number of filters in the encoder. num_filters_mlp (int): Base filter number in the MLP module. latent_dim (int): Dimension of the style code. num_res_blocks (int): Number of residual blocks at the end of the content encoder. num_mlp_blocks (int): Number of layers in the MLP module. num_downsamples_style (int): Number of times we reduce resolution by 2x2 for the style image. num_downsamples_content (int): Number of times we reduce resolution by 2x2 for the content image. num_image_channels (int): Number of input image channels. content_norm_type (str): Type of activation normalization in the content encoder. style_norm_type (str): Type of activation normalization in the style encoder. decoder_norm_type (str): Type of activation normalization in the decoder. weight_norm_type (str): Type of weight normalization. decoder_norm_params (obj): Parameters of activation normalization in the decoder. If not ``None``, decoder_norm_params.__dict__ will be used as keyword arguments when initializing activation normalization. output_nonlinearity (str): Type of nonlinearity before final output, ``'tanh'`` or ``'none'``. pre_act (bool): If ``True``, uses pre-activation residual blocks. apply_noise (bool): If ``True``, injects Gaussian noise in the decoder. """ def __init__(self, num_filters=64, max_num_filters=256, num_filters_mlp=256, latent_dim=8, num_res_blocks=4, num_mlp_blocks=2, num_downsamples_style=4, num_downsamples_content=2, num_image_channels=3, content_norm_type='instance', style_norm_type='', decoder_norm_type='instance', weight_norm_type='', decoder_norm_params=SimpleNamespace(affine=False), output_nonlinearity='', pre_act=False, apply_noise=False, **kwargs): super().__init__() for key in kwargs: if key != 'type': warnings.warn( "Generator argument '{}' is not used.".format(key)) self.style_encoder = StyleEncoder(num_downsamples_style, num_image_channels, num_filters, latent_dim, 'reflect', style_norm_type, weight_norm_type, 'relu') self.content_encoder = ContentEncoder(num_downsamples_content, num_res_blocks, num_image_channels, num_filters, max_num_filters, 'reflect', content_norm_type, weight_norm_type, 'relu', pre_act) self.decoder = Decoder(num_downsamples_content, num_res_blocks, self.content_encoder.output_dim, num_image_channels, num_filters_mlp, 'reflect', decoder_norm_type, decoder_norm_params, weight_norm_type, 'relu', output_nonlinearity, pre_act, apply_noise) self.mlp = MLP(latent_dim, num_filters_mlp, num_filters_mlp, num_mlp_blocks, 'none', 'relu') self.style_channels = latent_dim def forward(self, images): r"""Reconstruct an image. Args: images (Tensor): Input images. Returns: images_recon (Tensor): Reconstructed images. """ content, style = self.encode(images) images_recon = self.decode(content, style) return images_recon def encode(self, images): r"""Encode an image to content and style code. Args: images (Tensor): Input images. Returns: (tuple): - content (Tensor): Content code. - style (Tensor): Style code. """ style = self.style_encoder(images) content = self.content_encoder(images) return content, style def decode(self, content, style): r"""Decode content and style code to an image. Args: content (Tensor): Content code. style (Tensor): Style code. Returns: images (Tensor): Output images. """ style = self.mlp(style) images = self.decoder(content, style) return images class StyleEncoder(nn.Module): r"""MUNIT style encoder. Args: num_downsamples (int): Number of times we reduce resolution by 2x2. num_image_channels (int): Number of input image channels. num_filters (int): Base filter numbers. style_channels (int): Dimension of the style code. padding_mode (string): Type of padding. activation_norm_type (str): Type of activation normalization. weight_norm_type (str): Type of weight normalization. nonlinearity (str): Type of nonlinear activation function. """ def __init__(self, num_downsamples, num_image_channels, num_filters, style_channels, padding_mode, activation_norm_type, weight_norm_type, nonlinearity): super().__init__() conv_params = dict(padding_mode=padding_mode, activation_norm_type=activation_norm_type, weight_norm_type=weight_norm_type, nonlinearity=nonlinearity, inplace_nonlinearity=True) model = [] model += [Conv2dBlock(num_image_channels, num_filters, 7, 1, 3, **conv_params)] for i in range(2): model += [Conv2dBlock(num_filters, 2 * num_filters, 4, 2, 1, **conv_params)] num_filters *= 2 for i in range(num_downsamples - 2): model += [Conv2dBlock(num_filters, num_filters, 4, 2, 1, **conv_params)] model += [nn.AdaptiveAvgPool2d(1)] model += [nn.Conv2d(num_filters, style_channels, 1, 1, 0)] self.model = nn.Sequential(*model) self.output_dim = num_filters def forward(self, x): r""" Args: x (tensor): Input image. """ return self.model(x) class Decoder(nn.Module): r"""Improved MUNIT decoder. The network consists of - $(num_res_blocks) residual blocks. - $(num_upsamples) residual blocks or convolutional blocks - output layer. Args: num_upsamples (int): Number of times we increase resolution by 2x2. num_res_blocks (int): Number of residual blocks. num_filters (int): Base filter numbers. num_image_channels (int): Number of input image channels. style_channels (int): Dimension of the style code. padding_mode (string): Type of padding. activation_norm_type (str): Type of activation normalization. activation_norm_params (obj): Parameters of activation normalization. If not ``None``, decoder_norm_params.__dict__ will be used as keyword arguments when initializing activation normalization. weight_norm_type (str): Type of weight normalization. nonlinearity (str): Type of nonlinear activation function. output_nonlinearity (str): Type of nonlinearity before final output, ``'tanh'`` or ``'none'``. pre_act (bool): If ``True``, uses pre-activation residual blocks. apply_noise (bool): If ``True``, injects Gaussian noise. """ def __init__(self, num_upsamples, num_res_blocks, num_filters, num_image_channels, style_channels, padding_mode, activation_norm_type, activation_norm_params, weight_norm_type, nonlinearity, output_nonlinearity, pre_act=False, apply_noise=False): super().__init__() adain_params = SimpleNamespace( activation_norm_type=activation_norm_type, activation_norm_params=activation_norm_params, cond_dims=style_channels) conv_params = dict(padding_mode=padding_mode, nonlinearity=nonlinearity, inplace_nonlinearity=True, apply_noise=apply_noise, weight_norm_type=weight_norm_type, activation_norm_type='adaptive', activation_norm_params=adain_params) # The order of operations in residual blocks. order = 'pre_act' if pre_act else 'CNACNA' # Residual blocks with AdaIN. self.decoder = nn.ModuleList() for _ in range(num_res_blocks): self.decoder += [Res2dBlock(num_filters, num_filters, **conv_params, order=order)] # Convolutional blocks with upsampling. for i in range(num_upsamples): self.decoder += [NearestUpsample(scale_factor=2)] self.decoder += [Conv2dBlock(num_filters, num_filters // 2, 5, 1, 2, **conv_params)] num_filters //= 2 self.decoder += [Conv2dBlock(num_filters, num_image_channels, 7, 1, 3, nonlinearity=output_nonlinearity, padding_mode=padding_mode)] def forward(self, x, style): r""" Args: x (tensor): Content embedding of the content image. style (tensor): Style embedding of the style image. """ for block in self.decoder: if getattr(block, 'conditional', False): x = block(x, style) else: x = block(x) return x class MLP(nn.Module): r"""The multi-layer perceptron (MLP) that maps Gaussian style code to a feature vector that is given as the conditional input to AdaIN. Args: input_dim (int): Number of channels in the input tensor. output_dim (int): Number of channels in the output tensor. latent_dim (int): Number of channels in the latent features. num_layers (int): Number of layers in the MLP. norm (str): Type of activation normalization. nonlinearity (str): Type of nonlinear activation function. """ def __init__(self, input_dim, output_dim, latent_dim, num_layers, norm, nonlinearity): super().__init__() model = [] model += [LinearBlock(input_dim, latent_dim, activation_norm_type=norm, nonlinearity=nonlinearity)] for i in range(num_layers - 2): model += [LinearBlock(latent_dim, latent_dim, activation_norm_type=norm, nonlinearity=nonlinearity)] model += [LinearBlock(latent_dim, output_dim, activation_norm_type=norm, nonlinearity=nonlinearity)] self.model = nn.Sequential(*model) def forward(self, x): r""" Args: x (tensor): Input image. """ return self.model(x.view(x.size(0), -1))