<reponame>open-mmlab/MMGEN-FaceStylor
import random

import mmcv
import torch
from mmgen.models.architectures.common import get_module_device
from mmgen.models.architectures.stylegan.generator_discriminator_v2 import \
    StyleGANv2Generator
from mmgen.models.builder import MODULES


@MODULES.register_module()
class SwapStyleGANv2Generator(StyleGANv2Generator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def swap_forward(self,
                     styles,
                     num_batches=-1,
                     return_noise=False,
                     return_latents=False,
                     inject_index=None,
                     truncation=1,
                     truncation_latent=None,
                     input_is_latent=False,
                     injected_noise=None,
                     randomize_noise=True,
                     swap=False,
                     swap_layer_tensor=None,
                     swap_layer_num=1):
        """Forward function incoporated with layer-swapping.

        This function has been integrated with the truncation trick. Please
        refer to the usage of `truncation` and `truncation_latent`.

        Args:
            styles (torch.Tensor | list[torch.Tensor] | callable | None): In
                StyleGAN2, you can provide noise tensor or latent tensor. Given
                a list containing more than one noise or latent tensors, style
                mixing trick will be used in training. Of course, You can
                directly give a batch of noise through a ``torch.Tensor`` or
                offer a callable function to sample a batch of noise data.
                Otherwise, the ``None`` indicates to use the default noise
                sampler.
            num_batches (int, optional): The number of batch size.
                Defaults to 0.
            return_noise (bool, optional): If True, ``noise_batch`` will be
                returned in a dict with ``fake_img``. Defaults to False.
            return_latents (bool, optional): If True, ``latent`` will be
                returned in a dict with ``fake_img``. Defaults to False.
            inject_index (int | None, optional): The index number for mixing
                style codes. Defaults to None.
            truncation (float, optional): Truncation factor. Give value less
                than 1., the truncation trick will be adopted. Defaults to 1.
            truncation_latent (torch.Tensor, optional): Mean truncation latent.
                Defaults to None.
            input_is_latent (bool, optional): If `True`, the input tensor is
                the latent tensor. Defaults to False.
            injected_noise (torch.Tensor | None, optional): Given a tensor, the
                random noise will be fixed as this input injected noise.
                Defaults to None.
            randomize_noise (bool, optional): If `False`, images are sampled
                with the buffered noise tensor injected to the style conv
                block. Defaults to True.
            swap (bool, optional): Whether use layer Swapping.
                Defaults to False.
            swap_layer_tensor (torch.Tensor, optional): Tensor used to replace
                median feature in swap forward. Defaults to None.
            swap_layer_num (int, optional): The index of swap layer.
                Defaults to 1.

        Returns:
            torch.Tensor | dict: Generated image tensor or dictionary \
                containing more data.
        """
        # receive noise and conduct sanity check.
        if isinstance(styles, torch.Tensor):
            assert styles.shape[1] == self.style_channels
            styles = [styles]
        elif mmcv.is_seq_of(styles, torch.Tensor):
            for t in styles:
                assert t.shape[-1] == self.style_channels
        # receive a noise generator and sample noise.
        elif callable(styles):
            device = get_module_device(self)
            noise_generator = styles
            assert num_batches > 0
            if self.default_style_mode == 'mix' and random.random(
            ) < self.mix_prob:
                styles = [
                    noise_generator((num_batches, self.style_channels))
                    for _ in range(2)
                ]
            else:
                styles = [noise_generator((num_batches, self.style_channels))]
            styles = [s.to(device) for s in styles]
        # otherwise, we will adopt default noise sampler.
        else:
            device = get_module_device(self)
            assert num_batches > 0 and not input_is_latent
            if self.default_style_mode == 'mix' and random.random(
            ) < self.mix_prob:
                styles = [
                    torch.randn((num_batches, self.style_channels))
                    for _ in range(2)
                ]
            else:
                styles = [torch.randn((num_batches, self.style_channels))]
            styles = [s.to(device) for s in styles]

        if not input_is_latent:
            noise_batch = styles
            styles = [self.style_mapping(s) for s in styles]
        else:
            noise_batch = None

        if injected_noise is None:
            if randomize_noise:
                injected_noise = [None] * self.num_injected_noises
            else:
                injected_noise = [
                    getattr(self, f'injected_noise_{i}')
                    for i in range(self.num_injected_noises)
                ]
        # use truncation trick
        if truncation < 1:
            style_t = []
            # calculate truncation latent on the fly
            if truncation_latent is None and not hasattr(
                    self, 'truncation_latent'):
                self.truncation_latent = self.get_mean_latent()
                truncation_latent = self.truncation_latent
            elif truncation_latent is None and hasattr(self,
                                                       'truncation_latent'):
                truncation_latent = self.truncation_latent

            for style in styles:
                style_t.append(truncation_latent + truncation *
                               (style - truncation_latent))

            styles = style_t
        # no style mixing
        if len(styles) < 2:
            inject_index = self.num_latents

            if styles[0].ndim < 3:
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)

            else:
                latent = styles[0]
        # style mixing
        else:
            if inject_index is None:
                inject_index = random.randint(1, self.num_latents - 1)

            latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = styles[1].unsqueeze(1).repeat(
                1, self.num_latents - inject_index, 1)

            latent = torch.cat([latent, latent2], 1)

        # 4x4 stage
        out = self.constant_input(latent)
        out = self.conv1(out, latent[:, 0], noise=injected_noise[0])
        skip = self.to_rgb1(out, latent[:, 1])

        _index = 1

        # 8x8 ---> higher resolutions
        for up_conv, conv, noise1, noise2, to_rgb in zip(
                self.convs[::2], self.convs[1::2], injected_noise[1::2],
                injected_noise[2::2], self.to_rgbs):
            out = up_conv(out, latent[:, _index], noise=noise1)
            out = conv(out, latent[:, _index + 1], noise=noise2)
            # swap layer
            if swap and _index == (2 * swap_layer_num - 1):
                save_swap_layer = out
                if swap_layer_tensor is not None:
                    out = swap_layer_tensor

            skip = to_rgb(out, latent[:, _index + 2], skip)
            _index += 2

        # make sure the output image is torch.float32 to avoid RunTime Error
        # in other modules
        img = skip.to(torch.float32)

        if return_latents or return_noise:
            output_dict = dict(fake_img=img,
                               latent=latent,
                               inject_index=inject_index,
                               noise_batch=noise_batch,
                               save_swap_layer=save_swap_layer)
            return output_dict

        return img, save_swap_layer

import random

import mmcv
import torch
from mmgen.models.architectures.common import get_module_device
from mmgen.models.architectures.stylegan.generator_discriminator_v2 import \
    StyleGANv2Generator
from mmgen.models.builder import MODULES


@MODULES.register_module()
class SwapStyleGANv2Generator(StyleGANv2Generator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def swap_forward(self,
                     styles,
                     num_batches=-1,
                     return_noise=False,
                     return_latents=False,
                     inject_index=None,
                     truncation=1,
                     truncation_latent=None,
                     input_is_latent=False,
                     injected_noise=None,
                     randomize_noise=True,
                     swap=False,
                     swap_layer_tensor=None,
                     swap_layer_num=1):
        """Forward function incoporated with layer-swapping.

        This function has been integrated with the truncation trick. Please
        refer to the usage of `truncation` and `truncation_latent`.

        Args:
            styles (torch.Tensor | list[torch.Tensor] | callable | None): In
                StyleGAN2, you can provide noise tensor or latent tensor. Given
                a list containing more than one noise or latent tensors, style
                mixing trick will be used in training. Of course, You can
                directly give a batch of noise through a ``torch.Tensor`` or
                offer a callable function to sample a batch of noise data.
                Otherwise, the ``None`` indicates to use the default noise
                sampler.
            num_batches (int, optional): The number of batch size.
                Defaults to 0.
            return_noise (bool, optional): If True, ``noise_batch`` will be
                returned in a dict with ``fake_img``. Defaults to False.
            return_latents (bool, optional): If True, ``latent`` will be
                returned in a dict with ``fake_img``. Defaults to False.
            inject_index (int | None, optional): The index number for mixing
                style codes. Defaults to None.
            truncation (float, optional): Truncation factor. Give value less
                than 1., the truncation trick will be adopted. Defaults to 1.
            truncation_latent (torch.Tensor, optional): Mean truncation latent.
                Defaults to None.
            input_is_latent (bool, optional): If `True`, the input tensor is
                the latent tensor. Defaults to False.
            injected_noise (torch.Tensor | None, optional): Given a tensor, the
                random noise will be fixed as this input injected noise.
                Defaults to None.
            randomize_noise (bool, optional): If `False`, images are sampled
                with the buffered noise tensor injected to the style conv
                block. Defaults to True.
            swap (bool, optional): Whether use layer Swapping.
                Defaults to False.
            swap_layer_tensor (torch.Tensor, optional): Tensor used to replace
                median feature in swap forward. Defaults to None.
            swap_layer_num (int, optional): The index of swap layer.
                Defaults to 1.

        Returns:
            torch.Tensor | dict: Generated image tensor or dictionary \
                containing more data.
        """
        # receive noise and conduct sanity check.
        if isinstance(styles, torch.Tensor):
            assert styles.shape[1] == self.style_channels
            styles = [styles]
        elif mmcv.is_seq_of(styles, torch.Tensor):
            for t in styles:
                assert t.shape[-1] == self.style_channels
        # receive a noise generator and sample noise.
        elif callable(styles):
            device = get_module_device(self)
            noise_generator = styles
            assert num_batches > 0
            if self.default_style_mode == 'mix' and random.random(
            ) < self.mix_prob:
                styles = [
                    noise_generator((num_batches, self.style_channels))
                    for _ in range(2)
                ]
            else:
                styles = [noise_generator((num_batches, self.style_channels))]
            styles = [s.to(device) for s in styles]
        # otherwise, we will adopt default noise sampler.
        else:
            device = get_module_device(self)
            assert num_batches > 0 and not input_is_latent
            if self.default_style_mode == 'mix' and random.random(
            ) < self.mix_prob:
                styles = [
                    torch.randn((num_batches, self.style_channels))
                    for _ in range(2)
                ]
            else:
                styles = [torch.randn((num_batches, self.style_channels))]
            styles = [s.to(device) for s in styles]

        if not input_is_latent:
            noise_batch = styles
            styles = [self.style_mapping(s) for s in styles]
        else:
            noise_batch = None

        if injected_noise is None:
            if randomize_noise:
                injected_noise = [None] * self.num_injected_noises
            else:
                injected_noise = [
                    getattr(self, f'injected_noise_{i}')
                    for i in range(self.num_injected_noises)
                ]
        # use truncation trick
        if truncation < 1:
            style_t = []
            # calculate truncation latent on the fly
            if truncation_latent is None and not hasattr(
                    self, 'truncation_latent'):
                self.truncation_latent = self.get_mean_latent()
                truncation_latent = self.truncation_latent
            elif truncation_latent is None and hasattr(self,
                                                       'truncation_latent'):
                truncation_latent = self.truncation_latent

            for style in styles:
                style_t.append(truncation_latent + truncation *
                               (style - truncation_latent))

            styles = style_t
        # no style mixing
        if len(styles) < 2:
            inject_index = self.num_latents

            if styles[0].ndim < 3:
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)

            else:
                latent = styles[0]
        # style mixing
        else:
            if inject_index is None:
                inject_index = random.randint(1, self.num_latents - 1)

            latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = styles[1].unsqueeze(1).repeat(
                1, self.num_latents - inject_index, 1)

            latent = torch.cat([latent, latent2], 1)

        # 4x4 stage
        out = self.constant_input(latent)
        out = self.conv1(out, latent[:, 0], noise=injected_noise[0])
        skip = self.to_rgb1(out, latent[:, 1])

        _index = 1

        # 8x8 ---> higher resolutions
        for up_conv, conv, noise1, noise2, to_rgb in zip(
                self.convs[::2], self.convs[1::2], injected_noise[1::2],
                injected_noise[2::2], self.to_rgbs):
            out = up_conv(out, latent[:, _index], noise=noise1)
            out = conv(out, latent[:, _index + 1], noise=noise2)
            # swap layer
            if swap and _index == (2 * swap_layer_num - 1):
                save_swap_layer = out
                if swap_layer_tensor is not None:
                    out = swap_layer_tensor

            skip = to_rgb(out, latent[:, _index + 2], skip)
            _index += 2

        # make sure the output image is torch.float32 to avoid RunTime Error
        # in other modules
        img = skip.to(torch.float32)

        if return_latents or return_noise:
            output_dict = dict(fake_img=img,
                               latent=latent,
                               inject_index=inject_index,
                               noise_batch=noise_batch,
                               save_swap_layer=save_swap_layer)
            return output_dict

        return img, save_swap_layer
Forward function incoporated with layer-swapping.          This function has been integrated with the truncation trick. Please         refer to the usage of `truncation` and `truncation_latent`.          Args:             styles (torch.Tensor | list[torch.Tensor] | callable | None): In                 StyleGAN2, you can provide noise tensor or latent tensor. Given                 a list containing more than one noise or latent tensors, style                 mixing trick will be used in training. Of course, You can                 directly give a batch of noise through a ``torch.Tensor`` or                 offer a callable function to sample a batch of noise data.                 Otherwise, the ``None`` indicates to use the default noise                 sampler.             num_batches (int, optional): The number of batch size.                 Defaults to 0.             return_noise (bool, optional): If True, ``noise_batch`` will be                 returned in a dict with ``fake_img``. Defaults to False.             return_latents (bool, optional): If True, ``latent`` will be                 returned in a dict with ``fake_img``. Defaults to False.             inject_index (int | None, optional): The index number for mixing                 style codes. Defaults to None.             truncation (float, optional): Truncation factor. Give value less                 than 1., the truncation trick will be adopted. Defaults to 1.             truncation_latent (torch.Tensor, optional): Mean truncation latent.                 Defaults to None.             input_is_latent (bool, optional): If `True`, the input tensor is                 the latent tensor. Defaults to False.             injected_noise (torch.Tensor | None, optional): Given a tensor, the                 random noise will be fixed as this input injected noise.                 Defaults to None.             randomize_noise (bool, optional): If `False`, images are sampled                 with the buffered noise tensor injected to the style conv                 block. Defaults to True.             swap (bool, optional): Whether use layer Swapping.                 Defaults to False.             swap_layer_tensor (torch.Tensor, optional): Tensor used to replace                 median feature in swap forward. Defaults to None.             swap_layer_num (int, optional): The index of swap layer.                 Defaults to 1.          Returns:             torch.Tensor | dict: Generated image tensor or dictionary \                 containing more data. # receive noise and conduct sanity check. # receive a noise generator and sample noise. # otherwise, we will adopt default noise sampler. # use truncation trick # calculate truncation latent on the fly # no style mixing # style mixing # 4x4 stage # 8x8 ---> higher resolutions # swap layer # make sure the output image is torch.float32 to avoid RunTime Error # in other modules