Official implementation of Aurora
# python3.7
"""Collects all models."""

from .clip_model import CLIPModel
from .text_generator import Text2ImageGenerator

__all__ = ['build_model']

_MODELS = {
    'CLIPModel': CLIPModel,
    'Text2ImageGenerator': Text2ImageGenerator
}


def build_model(model_type, **kwargs):
    """Builds a model based on its class type.

    Args:
        model_type: Class type to which the model belongs, which is case
            sensitive.
        **kwargs: Additional arguments to build the model.

    Raises:
        ValueError: If the `model_type` is not supported.
    """
    if model_type not in _MODELS:
        raise ValueError(f'Invalid model type: `{model_type}`!\n'
                         f'Types allowed: {list(_MODELS)}.')
    return _MODELS[model_type](**kwargs)

models/__init__.py
zhujiapeng/Aurora
# python3.7
"""Contains the implementation of text embedding code used in our text2image
generation. We used the pre-trained CLIP model, and the paper can be found
here: https://arxiv.org/abs/2103.00020. Also, part of codes are borrow from
open_clip (https://github.com/mlfoundations/open_clip).
"""
from itertools import repeat
import collections.abc
from dataclasses import dataclass
import os
from collections import OrderedDict
from typing import Tuple, Union, Optional, Callable, Sequence, List
import json
import re
from pathlib import Path
from copy import deepcopy
import numpy as np
from pkg_resources import packaging
from PIL import Image

import torch.nn as nn
import torch
import torch.nn.functional as F
import torch.distributed as dist
from torch.utils.checkpoint import checkpoint
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
from huggingface_hub import hf_hub_download

from utils.misc import download_url
from utils.misc import get_cache_dir
from third_party.clip_official.simple_tokenizer import SimpleTokenizer as _Tokenizer

try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

__all__ = ['CLIPModel']

# pylint: disable=arguments-renamed
# pylint: disable=line-too-long
# pylint: disable=global-statement
# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring


class GatherLayer(torch.autograd.Function):
    """
    Gather tensors from all workers with support for backward propagation:
    This implementation does not cut the gradients as torch.distributed.all_gather does.
    """

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        output = [torch.zeros_like(input) for _ in range(dist.get_world_size())]
        dist.all_gather(output, input)
        return tuple(output)

    @staticmethod
    def backward(ctx, *grads):
        input, = ctx.saved_tensors
        grad_out = torch.zeros_like(input)
        grad_out[:] = grads[dist.get_rank()]
        return grad_out

@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())]
    dist.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output


def _pcfg(url='', hf_hub='', mean=None, std=None):
    return dict(
        url=url,
        hf_hub=hf_hub,
        mean=mean,
        std=std,
    )

_VITB32 = dict(
    openai=_pcfg(
        'https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt'),
    laion400m_e31=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt'),
    laion400m_e32=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt'),
    laion2b_e16=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-laion2b_e16-af8dbd0c.pth'),
    laion2b_s34b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-laion2B-s34B-b79K/'),
    # DataComp-M models
    datacomp_m_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-DataComp.M-s128M-b4K/'),
    commonpool_m_clip_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.clip-s128M-b4K/'),
    commonpool_m_laion_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.laion-s128M-b4K/'),
    commonpool_m_image_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.image-s128M-b4K/'),
    commonpool_m_text_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.text-s128M-b4K/'),
    commonpool_m_basic_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.basic-s128M-b4K/'),
    commonpool_m_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M-s128M-b4K/'),
    # DataComp-S models
    datacomp_s_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-DataComp.S-s13M-b4K/'),
    commonpool_s_clip_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.clip-s13M-b4K/'),
    commonpool_s_laion_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.laion-s13M-b4K/'),
    commonpool_s_image_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.image-s13M-b4K/'),
    commonpool_s_text_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.text-s13M-b4K/'),
    commonpool_s_basic_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.basic-s13M-b4K/'),
    commonpool_s_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S-s13M-b4K/'),
)

_VITB32_quickgelu = dict(
    openai=_pcfg(
        'https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt'),
    laion400m_e31=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt'),
    laion400m_e32=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt'),
)

_VITB16 = dict(
    openai=_pcfg(
        'https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt'),
    laion400m_e31=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e31-00efa78f.pt'),
    laion400m_e32=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e32-55e67d44.pt'),
    laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-laion2B-s34B-b88K/'),
    # DataComp-L models
    datacomp_l_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-DataComp.L-s1B-b8K/'),
    commonpool_l_clip_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.clip-s1B-b8K/'),
    commonpool_l_laion_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.laion-s1B-b8K/'),
    commonpool_l_image_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.image-s1B-b8K/'),
    commonpool_l_text_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.text-s1B-b8K/'),
    commonpool_l_basic_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.basic-s1B-b8K/'),
    commonpool_l_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L-s1B-b8K/'),
)

_VITB16_PLUS_240 = dict(
    laion400m_e31=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e31-8fb26589.pt'),
    laion400m_e32=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e32-699c4b84.pt'),
)

_VITL14 = dict(
    openai=_pcfg(
        'https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt'),
    laion400m_e31=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e31-69988bb6.pt'),
    laion400m_e32=_pcfg(
        'https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e32-3d133497.pt'),
    laion2b_s32b_b82k=_pcfg(
        hf_hub='laion/CLIP-ViT-L-14-laion2B-s32B-b82K/',
        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    # DataComp-XL models
    datacomp_xl_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90K/'),
    commonpool_xl_clip_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-L-14-CommonPool.XL.clip-s13B-b90K/'),
    commonpool_xl_laion_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-L-14-CommonPool.XL.laion-s13B-b90K/'),
    commonpool_xl_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-L-14-CommonPool.XL-s13B-b90K/'),
)

_VITL14_336 = dict(
    openai=_pcfg(
        'https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt'),
)

_VITH14 = dict(
    laion2b_s32b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-H-14-laion2B-s32B-b79K/'),
)

_VITg14 = dict(
    laion2b_s12b_b42k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s12B-b42K/'),
    laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s34B-b88K/'),
)

_VITbigG14 = dict(
    laion2b_s39b_b160k=_pcfg(hf_hub='laion/CLIP-ViT-bigG-14-laion2B-39B-b160k/'),
)

_coca_VITB32 = dict(
    laion2b_s13b_b90k=_pcfg(hf_hub='laion/CoCa-ViT-B-32-laion2B-s13B-b90k/'),
    mscoco_finetuned_laion2b_s13b_b90k=_pcfg(hf_hub='laion/mscoco_finetuned_CoCa-ViT-B-32-laion2B-s13B-b90k/')
)

_coca_VITL14 = dict(
    laion2b_s13b_b90k=_pcfg(hf_hub='laion/CoCa-ViT-L-14-laion2B-s13B-b90k/'),
    mscoco_finetuned_laion2b_s13b_b90k=_pcfg(hf_hub='laion/mscoco_finetuned_CoCa-ViT-L-14-laion2B-s13B-b90k/')
)


_PRETRAINED = {
    'ViT-B-32': _VITB32,
    'ViT-B-32-quickgelu': _VITB32_quickgelu,
    'ViT-B-16': _VITB16,
    'ViT-B-16-plus-240': _VITB16_PLUS_240,
    'ViT-L-14': _VITL14,
    'ViT-L-14-336': _VITL14_336,
    'ViT-H-14': _VITH14,
    'ViT-g-14': _VITg14,
    'ViT-bigG-14': _VITbigG14,
    'coca_ViT-B-32': _coca_VITB32,
    'coca_ViT-L-14': _coca_VITL14,
}

HF_HUB_PREFIX = 'hf-hub:'
_MODEL_CONFIG_PATHS = [Path(__file__).parent / 'clip_configs/']
_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)

_tokenizer = _Tokenizer()

def _natural_key(string_):
    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]


def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable):
            return x
        return tuple(repeat(x, n))
    return parse


to_1tuple = _ntuple(1)
to_2tuple = _ntuple(2)
to_3tuple = _ntuple(3)
to_4tuple = _ntuple(4)
to_ntuple = lambda n, x: _ntuple(n)(x)


def _rescan_model_configs():
    global _MODEL_CONFIGS

    config_ext = ('.json',)
    config_files = []
    for config_path in _MODEL_CONFIG_PATHS:
        if config_path.is_file() and config_path.suffix in config_ext:
            config_files.append(config_path)
        elif config_path.is_dir():
            for ext in config_ext:
                config_files.extend(config_path.glob(f'*{ext}'))

    for cf in config_files:
        with open(cf, 'r') as f:
            model_cfg = json.load(f)
            if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
                _MODEL_CONFIGS[cf.stem] = model_cfg

    _MODEL_CONFIGS = {k: v for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))}


_rescan_model_configs()  # initial populate of model config registry


def list_pretrained_models_by_tag(tag: str):
    """ return all models having the specified pretrain tag """
    models = []
    tag = _clean_tag(tag)
    for k in _PRETRAINED.keys():
        if tag in _PRETRAINED[k]:
            models.append(k)
    return models


def list_pretrained_tags_by_model(model: str):
    """ return all pretrain tags for the specified model architecture """
    tags = []
    if model in _PRETRAINED:
        tags.extend(_PRETRAINED[model].keys())
    return tags


def list_models():
    """ enumerate available model architectures based on config files """
    return list(_MODEL_CONFIGS.keys())


def add_model_config(path):
    """ add model config path or file and update registry """
    if not isinstance(path, Path):
        path = Path(path)
    _MODEL_CONFIG_PATHS.append(path)
    _rescan_model_configs()


def get_model_config(model_name):
    if model_name in _MODEL_CONFIGS:
        return deepcopy(_MODEL_CONFIGS[model_name])
    else:
        return None


def download_pretrained(cfg,
                        path=None,
                        force_hf_hub=False,
                        revision=None):
    target = ''
    if not cfg:
        return target

    url = cfg.get('url', '')
    hf_hub = cfg.get('hf_hub', '')
    if hf_hub and force_hf_hub:
        # use HF hub even if url exists
        url = ''

    if url:
        expected_sha256 = None
        if 'mlfoundations' in url:
            expected_sha256 = os.path.splitext(url)[0].split('-')[-1]
        target, hash_check = download_url(url, path=path, sha256=expected_sha256)
        if hash_check is False:
            print(f'Hash check failed! The remote file from URL '
                  f'`{url}` may be changed, or the downloading is '
                  f'interrupted. The loaded perceptual model may '
                  f'have unexpected behavior.')
    elif hf_hub:
        # we assume the hf_hub entries in pretrained config combine model_id + filename in
        # 'org/model_name/filename.pt' form. To specify just the model id w/o filename and
        # use 'open_clip_pytorch_model.bin' default, there must be a trailing slash 'org/model_name/'.
        model_id, filename = os.path.split(hf_hub)
        if filename:
            target = download_hf_model(model_id,
                                       path=path,
                                       filename=filename,
                                       revision=revision)
        else:
            target = download_hf_model(model_id,
                                       path=path,
                                       revision=revision)

    return target


def _clean_tag(tag: str):
    # normalize pretrained tags
    return tag.lower().replace('-', '_')


def get_pretrained_cfg(model, tag):
    if model not in _PRETRAINED:
        return {}
    model_pretrained = _PRETRAINED[model]
    return model_pretrained.get(_clean_tag(tag), {})


def get_pretrained_url(model, tag):
    cfg = get_pretrained_cfg(model, _clean_tag(tag))
    return cfg.get('url', '')


def list_openai_models():
    """Returns the names of available CLIP models"""
    return list_pretrained_models_by_tag('openai')


def _convert_image_to_rgb(image):
    return image.convert('RGB')


def _transfor_ori(img_size):
    return Compose([
        Resize(img_size, interpolation=BICUBIC),
        CenterCrop(img_size),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073),
                  (0.26862954, 0.26130258, 0.27577711)),
    ])


def _transfor_mod(img_size):
    return Compose([
        Resize(img_size, interpolation=BICUBIC),
        Normalize((0.48145466, 0.4578275, 0.40821073),
                  (0.26862954, 0.26130258, 0.27577711)),
    ])



def download_hf_model(model_id,
                      path=None,
                      filename='open_clip_pytorch_model.bin',
                      library_name='open_clip',
                      version='2.20.0',
                      revision='85b38ccd4a98128c6e4f7428c45ab8d71660ab9c'):
    """Download file from Hugging Face.

    Args:
        model_id: A user or an organization name and a repo name separated
            by a `/`.
        path: Path (directory) to save the downloaded file. If set as `None`,
            the cache directory will be used. Please see `get_cache_dir()` for
            more details. (default: None)
        filename: The name of the file in the repo.
        library_name: The name of the library to which the object corresponds.
        version: The version of the library.
        revision: An optional Git revision id which can be a branch name, a tag,
            or a commit hash.

    Returns:
        Local path (string) of file or if networking is off, last version of
        file cached on disk.
    """
    # Handle file path.
    if path is None:
        path = get_cache_dir()
    save_path = hf_hub_download(model_id,
                                filename,
                                revision=revision,
                                library_name=library_name,
                                library_version=version,
                                cache_dir=path)
    return save_path

@dataclass
class CLIPVisionCfg:
    layers: Union[Tuple[int, int, int, int], int] = 12
    width: int = 768
    head_width: int = 64
    mlp_ratio: float = 4.0
    patch_size: int = 16
    image_size: Union[Tuple[int, int], int] = 224

    ls_init_value: Optional[float] = None  # layer scale initial value
    patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
    input_patchnorm: bool = False  # whether to use dual patchnorm - would only apply the input layernorm on each patch, as post-layernorm already exist in original clip vit design
    global_average_pool: bool = False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
    attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer
    n_queries: int = 256  # n_queries for attentional pooler
    attn_pooler_heads: int = 8  # n heads for attentional_pooling
    output_tokens: bool = False

    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
    timm_proj_bias: bool = False  # enable bias final projection
    timm_drop: float = 0.  # head dropout
    timm_drop_path: Optional[float] = None  # backbone stochastic depth


@dataclass
class CLIPTextCfg:
    context_length: int = 77
    vocab_size: int = 49408
    width: int = 512
    heads: int = 8
    layers: int = 12
    ls_init_value: Optional[float] = None  # layer scale initial value
    hf_model_name: str = None
    hf_tokenizer_name: str = None
    hf_model_pretrained: bool = True
    proj: str = 'mlp'
    pooler_type: str = 'mean_pooler'
    embed_cls: bool = False
    pad_id: int = 0
    output_tokens: bool = False


class LayerNormFp32(nn.LayerNorm):
    """Subclass torch's LayerNorm to handle fp16
      (by casting to float32 and back)."""

    def forward(self, x: torch.Tensor):
        orig_type = x.dtype
        x = F.layer_norm(x.to(torch.float32),
                         self.normalized_shape,
                         self.weight,
                         self.bias,
                         self.eps)
        return x.to(orig_type)


class LayerNorm(nn.LayerNorm):
    """Subclass torch's LayerNorm (with cast back to input dtype)."""

    def forward(self, x: torch.Tensor):
        orig_type = x.dtype
        x = F.layer_norm(x,
                         self.normalized_shape,
                         self.weight,
                         self.bias,
                         self.eps)
        return x.to(orig_type)


class QuickGELU(nn.Module):
    # NOTE This is slower than nn.GELU or nn.SiLU and uses more GPU memory
    def forward(self, x: torch.Tensor):
        return x * torch.sigmoid(1.702 * x)


class LayerScale(nn.Module):
    def __init__(self, dim, init_values=1e-5, inplace=False):
        super().__init__()
        self.inplace = inplace
        self.gamma = nn.Parameter(init_values * torch.ones(dim))

    def forward(self, x):
        return x.mul_(self.gamma) if self.inplace else x * self.gamma


class PatchDropout(nn.Module):
    """
    https://arxiv.org/abs/2212.00794
    """

    def __init__(self, prob, exclude_first_token=True):
        super().__init__()
        assert 0 <= prob < 1.
        self.prob = prob
        self.exclude_first_token = exclude_first_token  # exclude CLS token

    def forward(self, x):
        if not self.training or self.prob == 0.:
            return x

        if self.exclude_first_token:
            cls_tokens, x = x[:, :1], x[:, 1:]
        else:
            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])

        batch = x.size()[0]
        num_tokens = x.size()[1]

        batch_indices = torch.arange(batch)
        batch_indices = batch_indices[..., None]

        keep_prob = 1 - self.prob
        num_patches_keep = max(1, int(num_tokens * keep_prob))

        rand = torch.randn(batch, num_tokens)
        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices

        x = x[batch_indices, patch_indices_keep]

        if self.exclude_first_token:
            x = torch.cat((cls_tokens, x), dim=1)

        return x


class Attention(nn.Module):
    def __init__(
            self,
            dim,
            num_heads=8,
            qkv_bias=True,
            scaled_cosine=False,
            scale_heads=False,
            logit_scale_max=np.log(1. / 0.01),
            attn_drop=0.,
            proj_drop=0.
    ):
        super().__init__()
        self.scaled_cosine = scaled_cosine
        self.scale_heads = scale_heads
        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.logit_scale_max = logit_scale_max

        # keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
        self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
        if qkv_bias:
            self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
        else:
            self.in_proj_bias = None

        if self.scaled_cosine:
            self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
        else:
            self.logit_scale = None
        self.attn_drop = nn.Dropout(attn_drop)
        if self.scale_heads:
            self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
        else:
            self.head_scale = None
        self.out_proj = nn.Linear(dim, dim)
        self.out_drop = nn.Dropout(proj_drop)

    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
        L, N, C = x.shape
        q, k, v = F.linear(x, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
        q = q.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
        k = k.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
        v = v.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)

        if self.logit_scale is not None:
            attn = torch.bmm(F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2))
            logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
            attn = attn.view(N, self.num_heads, L, L) * logit_scale
            attn = attn.view(-1, L, L)
        else:
            q = q * self.scale
            attn = torch.bmm(q, k.transpose(-1, -2))

        if attn_mask is not None:
            if attn_mask.dtype == torch.bool:
                new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
                new_attn_mask.masked_fill_(attn_mask, float('-inf'))
                attn_mask = new_attn_mask
            attn += attn_mask

        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = torch.bmm(attn, v)
        if self.head_scale is not None:
            x = x.view(N, self.num_heads, L, C) * self.head_scale
            x = x.view(-1, L, C)
        x = x.transpose(0, 1).reshape(L, N, C)
        x = self.out_proj(x)
        x = self.out_drop(x)
        return x


class AttentionalPooler(nn.Module):
    def __init__(
            self,
            d_model: int,
            context_dim: int,
            n_head: int = 8,
            n_queries: int = 256,
            norm_layer: Callable = LayerNorm
    ):
        super().__init__()
        self.query = nn.Parameter(torch.randn(n_queries, d_model))
        self.attn = nn.MultiheadAttention(d_model, n_head, kdim=context_dim, vdim=context_dim)
        self.ln_q = norm_layer(d_model)
        self.ln_k = norm_layer(context_dim)

    def forward(self, x: torch.Tensor):
        x = self.ln_k(x).permute(1, 0, 2)  # NLD -> LND
        N = x.shape[1]
        q = self.ln_q(self.query)
        out = self.attn(self._repeat(q, N), x, x, need_weights=False)[0]
        return out.permute(1, 0, 2)  # LND -> NLD

    def _repeat(self, query, N: int):
        return query.unsqueeze(1).repeat(1, N, 1)


class ResidualAttentionBlock(nn.Module):
    def __init__(
            self,
            d_model: int,
            n_head: int,
            mlp_ratio: float = 4.0,
            ls_init_value: float = None,
            act_layer: Callable = nn.GELU,
            norm_layer: Callable = LayerNorm,
            is_cross_attention: bool = False,
    ):
        super().__init__()

        self.ln_1 = norm_layer(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_head)
        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
        if is_cross_attention:
            self.ln_1_kv = norm_layer(d_model)

        self.ln_2 = norm_layer(d_model)
        mlp_width = int(d_model * mlp_ratio)
        self.mlp = nn.Sequential(OrderedDict([
            ('c_fc', nn.Linear(d_model, mlp_width)),
            ('gelu', act_layer()),
            ('c_proj', nn.Linear(mlp_width, d_model))
        ]))
        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()

    def attention(
            self,
            q_x: torch.Tensor,
            k_x: Optional[torch.Tensor] = None,
            v_x: Optional[torch.Tensor] = None,
            attn_mask: Optional[torch.Tensor] = None,
    ):
        k_x = k_x if k_x is not None else q_x
        v_x = v_x if v_x is not None else q_x

        attn_mask = attn_mask.to(q_x.dtype) if attn_mask is not None else None
        return self.attn(
            q_x, k_x, v_x, need_weights=False, attn_mask=attn_mask
        )[0]

    def forward(
            self,
            q_x: torch.Tensor,
            k_x: Optional[torch.Tensor] = None,
            v_x: Optional[torch.Tensor] = None,
            attn_mask: Optional[torch.Tensor] = None,
    ):
        k_x = self.ln_1_kv(k_x) if hasattr(self, 'ln_1_kv') and k_x is not None else None
        v_x = self.ln_1_kv(v_x) if hasattr(self, 'ln_1_kv') and v_x is not None else None

        x = q_x + self.ls_1(self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, attn_mask=attn_mask))
        x = x + self.ls_2(self.mlp(self.ln_2(x)))
        return x


class Transformer(nn.Module):
    def __init__(
            self,
            width: int,
            layers: int,
            heads: int,
            mlp_ratio: float = 4.0,
            ls_init_value: float = None,
            act_layer: Callable = nn.GELU,
            norm_layer: Callable = LayerNorm,
    ):
        super().__init__()
        self.width = width
        self.layers = layers
        self.grad_checkpointing = False

        self.resblocks = nn.ModuleList([
            ResidualAttentionBlock(
                width, heads, mlp_ratio, ls_init_value=ls_init_value, act_layer=act_layer, norm_layer=norm_layer)
            for _ in range(layers)
        ])

    def get_cast_dtype(self) -> torch.dtype:
        if hasattr(self.resblocks[0].mlp.c_fc, 'int8_original_dtype'):
            return self.resblocks[0].mlp.c_fc.int8_original_dtype
        return self.resblocks[0].mlp.c_fc.weight.dtype

    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
        for r in self.resblocks:
            if self.grad_checkpointing and not torch.jit.is_scripting():
                # TODO: handle kwargs https://github.com/pytorch/pytorch/issues/79887#issuecomment-1161758372
                x = checkpoint(r, x, None, None, attn_mask)
            else:
                x = r(x, attn_mask=attn_mask)
        return x


class VisionTransformer(nn.Module):
    output_tokens: torch.jit.Final[bool]

    def __init__(
            self,
            image_size: int,
            patch_size: int,
            width: int,
            layers: int,
            heads: int,
            mlp_ratio: float,
            ls_init_value: float = None,
            global_average_pool: bool = False,
            attentional_pool: bool = False,
            n_queries: int = 256,
            attn_pooler_heads: int = 8,
            output_dim: int = 512,
            patch_dropout: float = 0.,
            input_patchnorm: bool = False,
            act_layer: Callable = nn.GELU,
            norm_layer: Callable = LayerNorm,
            output_tokens: bool = False
    ):
        super().__init__()
        self.output_tokens = output_tokens
        # self.image_size = image_size
        # self.patch_size = patch_size
        image_height, image_width = self.image_size = to_2tuple(image_size)
        patch_height, patch_width = self.patch_size = to_2tuple(patch_size)
        self.grid_size = (image_height // patch_height, image_width // patch_width)
        self.output_dim = output_dim

        # whether to layernorm each patch, as done in dual patchnorm paper - https://arxiv.org/abs/2302.01327v1
        self.input_patchnorm = input_patchnorm

        if input_patchnorm:
            patch_input_dim = patch_height * patch_width * 3
            self.patchnorm_pre_ln = LayerNorm(patch_input_dim)
            self.conv1 = nn.Linear(patch_input_dim, width)
        else:
            self.patchnorm_pre_ln = nn.Identity()
            self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)

        # class embeddings and positional embeddings
        scale = width ** -0.5
        self.class_embedding = nn.Parameter(scale * torch.randn(width))
        self.positional_embedding = nn.Parameter(scale * torch.randn(self.grid_size[0] * self.grid_size[1] + 1, width))

        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()

        self.ln_pre = norm_layer(width)
        self.transformer = Transformer(
            width,
            layers,
            heads,
            mlp_ratio,
            ls_init_value=ls_init_value,
            act_layer=act_layer,
            norm_layer=norm_layer,
        )

        self.global_average_pool = global_average_pool
        if attentional_pool:
            self.attn_pool = AttentionalPooler(output_dim, width, n_head=attn_pooler_heads, n_queries=n_queries)
            self.ln_post = norm_layer(output_dim)
            self.proj = nn.Parameter(scale * torch.randn(output_dim, output_dim))
        else:
            self.attn_pool = None
            self.ln_post = norm_layer(width)
            self.proj = nn.Parameter(scale * torch.randn(width, output_dim))

        self.init_parameters()

    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
        for param in self.parameters():
            param.requires_grad = False

        if unlocked_groups != 0:
            groups = [
                [
                    self.conv1,
                    self.class_embedding,
                    self.positional_embedding,
                    self.ln_pre,
                ],
                *self.transformer.resblocks[:-1],
                [
                    self.transformer.resblocks[-1],
                    self.ln_post,
                ],
                self.proj,
            ]

            def _unlock(x):
                if isinstance(x, Sequence):
                    for g in x:
                        _unlock(g)
                else:
                    if isinstance(x, torch.nn.Parameter):
                        x.requires_grad = True
                    else:
                        for p in x.parameters():
                            p.requires_grad = True

            _unlock(groups[-unlocked_groups:])

    def init_parameters(self):
        # FIXME OpenAI CLIP did not define an init for the VisualTransformer
        # TODO experiment if default PyTorch init, below, or alternate init is best.

        # nn.init.normal_(self.class_embedding, std=self.scale)
        # nn.init.normal_(self.positional_embedding, std=self.scale)
        #
        # proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
        # attn_std = self.transformer.width ** -0.5
        # fc_std = (2 * self.transformer.width) ** -0.5
        # for block in self.transformer.resblocks:
        #     nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
        #     nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
        #     nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
        #     nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
        #
        # if self.text_projection is not None:
        #     nn.init.normal_(self.text_projection, std=self.scale)
        pass

    @torch.jit.ignore
    def set_grad_checkpointing(self, enable=True):
        self.transformer.grad_checkpointing = enable

    def _global_pool(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        if self.global_average_pool:
            return x.mean(dim=1), x
        else:
            return x[:, 0], x[:, 1:]

    def forward(self, x: torch.Tensor):

        # to patches - whether to use dual patchnorm - https://arxiv.org/abs/2302.01327v1
        if self.input_patchnorm:
            # einops - rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)')
            x = x.reshape(x.shape[0], x.shape[1], self.grid_size[0], self.patch_size[0], self.grid_size[1], self.patch_size[1])
            x = x.permute(0, 2, 4, 1, 3, 5)
            x = x.reshape(x.shape[0], self.grid_size[0] * self.grid_size[1], -1)
            x = self.patchnorm_pre_ln(x)
            x = self.conv1(x)
        else:
            x = self.conv1(x)  # shape = [*, width, grid, grid]
            x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
            x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]

        # class embeddings and positional embeddings
        x = torch.cat(
            [self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
             x], dim=1)  # shape = [*, grid ** 2 + 1, width]
        x = x + self.positional_embedding.to(x.dtype)

        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
        x = self.patch_dropout(x)
        x = self.ln_pre(x)

        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # LND -> NLD

        if self.attn_pool is not None:
            x = self.attn_pool(x)
            x = self.ln_post(x)
            pooled, tokens = self._global_pool(x)
        else:
            pooled, tokens = self._global_pool(x)
            pooled = self.ln_post(pooled)

        if self.proj is not None:
            pooled = pooled @ self.proj

        if self.output_tokens:
            return pooled, tokens

        return pooled


class TextTransformer(nn.Module):
    output_tokens: torch.jit.Final[bool]

    def __init__(
            self,
            context_length: int = 77,
            vocab_size: int = 49408,
            width: int = 512,
            heads: int = 8,
            layers: int = 12,
            ls_init_value: float = None,
            output_dim: int = 512,
            act_layer: Callable = nn.GELU,
            norm_layer: Callable = LayerNorm,
            embed_cls: bool = False,
            pad_id: int = 0,
            output_tokens: bool = False,
    ):
        super().__init__()
        self.output_tokens = output_tokens
        self.num_pos = self.context_length = context_length
        self.vocab_size = vocab_size
        self.width = width
        self.output_dim = output_dim
        self.heads = heads
        self.pad_id = pad_id

        self.text_projection = nn.Parameter(torch.empty(width, output_dim))

        if embed_cls:
            self.cls_emb = nn.Parameter(torch.empty(width))
            self.num_pos += 1
        else:
            self.cls_emb = None

        self.token_embedding = nn.Embedding(vocab_size, width)
        self.positional_embedding = nn.Parameter(torch.empty(self.num_pos, width))
        self.transformer = Transformer(
            width=width,
            layers=layers,
            heads=heads,
            ls_init_value=ls_init_value,
            act_layer=act_layer,
            norm_layer=norm_layer,
        )
        self.ln_final = norm_layer(width)

        self.register_buffer('attn_mask', self.build_attention_mask(), persistent=False)

        self.init_parameters()

    def init_parameters(self):
        nn.init.normal_(self.token_embedding.weight, std=0.02)
        nn.init.normal_(self.positional_embedding, std=0.01)
        if self.cls_emb is not None:
            nn.init.normal_(self.cls_emb, std=0.01)

        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
        attn_std = self.transformer.width ** -0.5
        fc_std = (2 * self.transformer.width) ** -0.5
        for block in self.transformer.resblocks:
            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)

        if self.text_projection is not None:
            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)

    @torch.jit.ignore
    def set_grad_checkpointing(self, enable=True):
        self.transformer.grad_checkpointing = enable

    def build_attention_mask(self):
        # lazily create causal attention mask, with full attention between the tokens
        # pytorch uses additive attention mask; fill with -inf
        mask = torch.empty(self.num_pos, self.num_pos)
        mask.fill_(float('-inf'))
        mask.triu_(1)  # zero out the lower diagonal
        return mask

    def build_cls_mask(self, text, cast_dtype: torch.dtype):
        cls_mask = (text != self.pad_id).unsqueeze(1)
        cls_mask = F.pad(cls_mask, (1, 0, cls_mask.shape[2], 0), value=1.0)
        additive_mask = torch.empty(cls_mask.shape, dtype=cast_dtype, device=cls_mask.device)
        additive_mask.fill_(0)
        additive_mask.masked_fill_(~cls_mask, float('-inf'))
        additive_mask = torch.repeat_interleave(additive_mask, self.heads, 0)
        return additive_mask

    def _repeat(self, t, N: int):
        return t.reshape(1, 1, -1).repeat(N, 1, 1)

    def forward(self, text):
        cast_dtype = self.transformer.get_cast_dtype()
        seq_len = text.shape[1]

        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
        attn_mask = self.attn_mask
        if self.cls_emb is not None:
            seq_len += 1
            x = torch.cat([x, self._repeat(self.cls_emb, x.shape[0])], dim=1)
            cls_mask = self.build_cls_mask(text, cast_dtype)
            attn_mask = attn_mask[None, :seq_len, :seq_len] + cls_mask[:, :seq_len, :seq_len]

        x = x + self.positional_embedding[:seq_len].to(cast_dtype)
        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x, attn_mask=attn_mask)
        x = x.permute(1, 0, 2)  # LND -> NLD

        # x.shape = [batch_size, n_ctx, transformer.width]
        # take features from the eot embedding (eot_token is the highest number in each sequence)
        if self.cls_emb is not None:
            pooled, tokens = x[:, -1], x[:, :-1]
            pooled = self.ln_final(pooled)
        else:
            x = self.ln_final(x)
            pooled, tokens = x[torch.arange(x.shape[0]), text.argmax(dim=-1)], x

        if self.text_projection is not None:
            pooled = pooled @ self.text_projection

        if self.output_tokens:
            return pooled, tokens

        return pooled


def get_cast_dtype(precision: str):
    cast_dtype = None
    if precision == 'bf16':
        cast_dtype = torch.bfloat16
    elif precision == 'fp16':
        cast_dtype = torch.float16
    return cast_dtype


def get_input_dtype(precision: str):
    input_dtype = None
    if precision in ('bf16', 'pure_bf16'):
        input_dtype = torch.bfloat16
    elif precision in ('fp16', 'pure_fp16'):
        input_dtype = torch.float16
    return input_dtype


def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
    """Convert applicable model parameters to low-precision (bf16 or fp16)"""

    def _convert_weights(l):
        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
            l.weight.data = l.weight.data.to(dtype)
            if l.bias is not None:
                l.bias.data = l.bias.data.to(dtype)

        if isinstance(l, (nn.MultiheadAttention, Attention)):
            for attr in [*[f'{s}_proj_weight' for s in ['in', 'q', 'k', 'v']], 'in_proj_bias', 'bias_k', 'bias_v']:
                tensor = getattr(l, attr)
                if tensor is not None:
                    tensor.data = tensor.data.to(dtype)

        if isinstance(l, (CLIP, TextTransformer)):
            # convert text nn.Parameter projections
            attr = getattr(l, 'text_projection', None)
            if attr is not None:
                attr.data = attr.data.to(dtype)

        if isinstance(l, VisionTransformer):
            # convert vision nn.Parameter projections
            attr = getattr(l, 'proj', None)
            if attr is not None:
                attr.data = attr.data.to(dtype)

    model.apply(_convert_weights)


convert_weights_to_fp16 = convert_weights_to_lp


def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
    # Rescale the grid of position embeddings when loading from state_dict
    old_pos_embed = state_dict.get('visual.positional_embedding', None)
    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
        return
    grid_size = to_2tuple(model.visual.grid_size)
    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
    if new_seq_len == old_pos_embed.shape[0]:
        return

    if extra_tokens:
        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
    else:
        pos_emb_tok, pos_emb_img = None, old_pos_embed
    old_grid_size = to_2tuple(int(np.sqrt(len(pos_emb_img))))

    print('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
    pos_emb_img = F.interpolate(
        pos_emb_img,
        size=grid_size,
        mode=interpolation,
        antialias=antialias,
        align_corners=False,
    )
    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
    if pos_emb_tok is not None:
        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
    else:
        new_pos_embed = pos_emb_img
    state_dict['visual.positional_embedding'] = new_pos_embed


def load_state_dict(checkpoint_path: str, map_location='cpu'):
    checkpoint = torch.load(checkpoint_path, map_location=map_location)
    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
        state_dict = checkpoint['state_dict']
    else:
        state_dict = checkpoint
    if next(iter(state_dict.items()))[0].startswith('module'):
        state_dict = {k[7:]: v for k, v in state_dict.items()}
    return state_dict


def load_checkpoint(model, checkpoint_path, strict=True):
    state_dict = load_state_dict(checkpoint_path)
    # detect old format and make compatible with new format
    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
        state_dict = convert_to_custom_text_state_dict(state_dict)
    resize_pos_embed(state_dict, model)
    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
    return incompatible_keys


def _build_vision_tower(
        embed_dim: int,
        vision_cfg: CLIPVisionCfg,
        quick_gelu: bool = False,
        cast_dtype: Optional[torch.dtype] = None
):
    if isinstance(vision_cfg, dict):
        vision_cfg = CLIPVisionCfg(**vision_cfg)

    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
    # memory efficient in recent PyTorch releases (>= 1.10).
    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
    act_layer = QuickGELU if quick_gelu else nn.GELU
    vision_heads = vision_cfg.width // vision_cfg.head_width
    norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
    visual = VisionTransformer(
        image_size=vision_cfg.image_size,
        patch_size=vision_cfg.patch_size,
        width=vision_cfg.width,
        layers=vision_cfg.layers,
        heads=vision_heads,
        mlp_ratio=vision_cfg.mlp_ratio,
        ls_init_value=vision_cfg.ls_init_value,
        patch_dropout=vision_cfg.patch_dropout,
        input_patchnorm=vision_cfg.input_patchnorm,
        global_average_pool=vision_cfg.global_average_pool,
        attentional_pool=vision_cfg.attentional_pool,
        n_queries=vision_cfg.n_queries,
        attn_pooler_heads=vision_cfg.attn_pooler_heads,
        output_tokens=vision_cfg.output_tokens,
        output_dim=embed_dim,
        act_layer=act_layer,
        norm_layer=norm_layer,
    )

    return visual


def _build_text_tower(
        embed_dim: int,
        text_cfg: CLIPTextCfg,
        quick_gelu: bool = False,
        cast_dtype: Optional[torch.dtype] = None,
):
    if isinstance(text_cfg, dict):
        text_cfg = CLIPTextCfg(**text_cfg)

    act_layer = QuickGELU if quick_gelu else nn.GELU
    norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm

    text = TextTransformer(
        context_length=text_cfg.context_length,
        vocab_size=text_cfg.vocab_size,
        width=text_cfg.width,
        heads=text_cfg.heads,
        layers=text_cfg.layers,
        ls_init_value=text_cfg.ls_init_value,
        output_dim=embed_dim,
        embed_cls=text_cfg.embed_cls,
        output_tokens=text_cfg.output_tokens,
        pad_id=text_cfg.pad_id,
        act_layer=act_layer,
        norm_layer=norm_layer,
    )

    return text


def convert_to_custom_text_state_dict(state_dict: dict):
    if 'text_projection' in state_dict:
        # old format state_dict, move text tower -> .text
        new_state_dict = {}
        for k, v in state_dict.items():
            if any(k.startswith(p) for p in (
                'text_projection',
                'positional_embedding',
                'token_embedding',
                'transformer',
                'ln_final',
            )):
                k = 'text.' + k
            new_state_dict[k] = v
        return new_state_dict
    return state_dict


def build_model_from_openai_state_dict(state_dict,
                                       quick_gelu=True,
                                       cast_dtype=torch.float16,):
    vit = 'visual.proj' in state_dict

    if vit:
        vision_width = state_dict['visual.conv1.weight'].shape[0]
        vision_layers = len(
            [k for k in state_dict.keys() if k.startswith('visual.') and k.endswith('.attn.in_proj_weight')])
        vision_patch_size = state_dict['visual.conv1.weight'].shape[-1]
        grid_size = round((state_dict['visual.positional_embedding'].shape[0] - 1) ** 0.5)
        image_size = vision_patch_size * grid_size
    else:
        counts: list = [
            len(set(k.split('.')[2] for k in state_dict if k.startswith(f'visual.layer{b}'))) for b in [1, 2, 3, 4]]
        vision_layers = tuple(counts)
        vision_width = state_dict['visual.layer1.0.conv1.weight'].shape[0]
        output_width = round((state_dict['visual.attnpool.positional_embedding'].shape[0] - 1) ** 0.5)
        vision_patch_size = None
        assert output_width ** 2 + 1 == state_dict['visual.attnpool.positional_embedding'].shape[0]
        image_size = output_width * 32

    embed_dim = state_dict['text_projection'].shape[1]
    context_length = state_dict['positional_embedding'].shape[0]
    vocab_size = state_dict['token_embedding.weight'].shape[0]
    transformer_width = state_dict['ln_final.weight'].shape[0]
    transformer_heads = transformer_width // 64
    transformer_layers = len(set(k.split('.')[2] for k in state_dict if k.startswith('transformer.resblocks')))

    vision_cfg = CLIPVisionCfg(
        layers=vision_layers,
        width=vision_width,
        patch_size=vision_patch_size,
        image_size=image_size,
    )
    text_cfg = CLIPTextCfg(
        context_length=context_length,
        vocab_size=vocab_size,
        width=transformer_width,
        heads=transformer_heads,
        layers=transformer_layers,
    )
    model = CLIP(
        embed_dim,
        vision_cfg=vision_cfg,
        text_cfg=text_cfg,
        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
        cast_dtype=cast_dtype,
    )

    for key in ['input_resolution', 'context_length', 'vocab_size']:
        state_dict.pop(key, None)

    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
    model.load_state_dict(state_dict)
    return model.eval()


def load_openai_model(name,
                      precision=None,
                      device=None,
                      cache_dir=None):
    """Load CLIP model from OpenAI.

    Args:
        name: A model name listed by `clip.available_models()`, or the path to
              a model checkpoint containing the state_dict.
        precision: Model precision, if None defaults to 'fp32'
            if device == 'cpu' else 'fp16'.
        device: The device to put the loaded model.
        cache_dir : The directory to cache the downloaded model weights.

    Returns:
        model: The CLIP model
        preprocess:  A torchvision transform that converts a PIL image into a
            tensor that the returned model can take as its input.
    """
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if precision is None:
        precision = 'fp32' if device == 'cpu' else 'fp16'
    model_url = get_pretrained_url(name, 'openai')
    if model_url:
        expected_sha256 = None
        if 'openaipublic' in model_url:
            expected_sha256 = model_url.split('/')[-2]
        model_path, hash_check = download_url(model_url,
                                              path=cache_dir,
                                              sha256=expected_sha256)
        if hash_check is False:
            print(f'Hash check failed! The remote file from URL '
                  f'`{model_url}` may be changed, or the downloading is '
                  f'interrupted. The loaded perceptual model may '
                  f'have unexpected behavior.')
    elif os.path.isfile(name):
        model_path = name
    else:
        raise RuntimeError(f'Model {name} not found; available models '
                           f'= {list_openai_models()}')

    try:
        # loading JIT archive
        model = torch.jit.load(model_path, map_location='cpu').eval()
        state_dict = None
    except RuntimeError:
        # loading saved state dict
        state_dict = torch.load(model_path, map_location='cpu')

    # Build a non-jit model from the OpenAI jitted model state dict
    cast_dtype = get_cast_dtype(precision)
    try:
        model = build_model_from_openai_state_dict(state_dict or model.state_dict(), cast_dtype=cast_dtype)
    except KeyError:
        sd = {k[7:]: v for k, v in state_dict['state_dict'].items()}
        model = build_model_from_openai_state_dict(sd, cast_dtype=cast_dtype)

    # model from OpenAI state dict is in manually cast fp16 mode, must be converted for AMP/fp32/bf16 use
    model = model.to(device)
    # FIXME support pure fp16/bf16 precision modes
    if precision != 'fp16':
        model.float()
        if precision == 'bf16':
            # for bf16, convert back to low-precision
            convert_weights_to_lp(model, dtype=torch.bfloat16)

    # add mean / std attributes for consistency with OpenCLIP models
    model.visual.image_mean = OPENAI_DATASET_MEAN
    model.visual.image_std = OPENAI_DATASET_STD
    return model


class CLIP(nn.Module):
    output_dict: torch.jit.Final[bool]

    def __init__(
            self,
            embed_dim: int,
            vision_cfg: CLIPVisionCfg,
            text_cfg: CLIPTextCfg,
            quick_gelu: bool = False,
            cast_dtype: Optional[torch.dtype] = None,
            output_dict: bool = False,
    ):
        super().__init__()
        self.output_dict = output_dict
        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)

        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
        self.transformer = text.transformer
        self.context_length = text.context_length
        self.vocab_size = text.vocab_size
        self.token_embedding = text.token_embedding
        self.positional_embedding = text.positional_embedding
        self.ln_final = text.ln_final
        self.text_projection = text.text_projection
        self.register_buffer('attn_mask', text.attn_mask, persistent=False)

        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))

    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)

    @torch.jit.ignore
    def set_grad_checkpointing(self, enable=True):
        self.visual.set_grad_checkpointing(enable)
        self.transformer.grad_checkpointing = enable

    def encode_image(self, image, normalize=True):
        features = self.visual(image)
        if normalize:
            features = F.normalize(features, dim=-1)
        return  features

    def encode_text(self, text, normalize=True):
        cast_dtype = self.transformer.get_cast_dtype()
        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
        x = x + self.positional_embedding.to(cast_dtype)
        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x, attn_mask=self.attn_mask)
        x = x.permute(1, 0, 2)  # LND -> NLD
        x_ = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
        # take features from the eot embedding (eot_token is the highest number in each sequence)
        x = x_[torch.arange(x_.shape[0]), text.argmax(dim=-1)] @ self.text_projection
        if normalize:
            x = F.normalize(x, dim=-1)
        return x, x_

    def forward(self, image=None, text=None):
        image_features = self.encode_image(image, normalize=True) if image is not None else None
        pooled_out, out = self.encode_text(text, normalize=True) if text is not None else None
        return {
                'image_features': image_features,
                'text_features_pooled': pooled_out,
                'text_features': out,
                'logit_scale': self.logit_scale.exp()
            }


class CLIPModel(nn.Module):
    """CLIP model used for extracting the image and text features."""
    def __init__(self,
                 model_name='ViT-L/14',
                 pretrained='openai',
                 max_length=77,
                 precision='fp32',
                 cache_dir=None,
                 device='cuda',
                 jit=False,
                 force_quick_gelu=False,
                 force_patch_dropout=None,
                 force_image_size=None,
                 output_dict=None,
                 require_pretrained=False,
                 freeze_clip=True,
                 revision='85b38ccd4a98128c6e4f7428c45ab8d71660ab9c'):
        super().__init__()

        self.device = device
        self.freeze_clip = freeze_clip
        self.max_length = max_length

        has_hf_hub_prefix = model_name.startswith(HF_HUB_PREFIX)
        if has_hf_hub_prefix:
            model_id = model_name[len(HF_HUB_PREFIX):]
            self.checkpoint_path = download_hf_model(model_id,
                                                     path=cache_dir,
                                                     revision=revision)
            config_path = download_hf_model(model_id,
                                            path=cache_dir,
                                            filename='open_clip_config.json',
                                            revision=revision)
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            self.pretrained_cfg = config['preprocess_cfg']
            self.model_cfg = config['model_cfg']
        else:
            # for callers using old naming with / in ViT names
            model_name = model_name.replace('/', '-')
            self.checkpoint_path = None
            self.pretrained_cfg = {}
            self.model_cfg = None

        if pretrained and pretrained.lower() == 'openai':
            print(f'Loading pretrained {model_name} from OpenAI.')
            self.model = load_openai_model(model_name,
                                           precision=precision,
                                           device=device,
                                           cache_dir=cache_dir)
        else:
            self.model_cfg = self.model_cfg or get_model_config(model_name)
            if self.model_cfg is not None:
                print(f'Loaded {model_name} model config.')
            else:
                print(f'Model config for {model_name} not found; '
                      f'available models {list_models()}.')
                raise RuntimeError(f'Model config for {model_name} not found.')

            if force_quick_gelu:
                # override for use of QuickGELU on non-OpenAI transformer models
                self.model_cfg['quick_gelu'] = True

            if force_patch_dropout is not None:
                # override the default patch dropout value
                self.model_cfg['vision_cfg']['patch_dropout'] = force_patch_dropout

            if force_image_size is not None:
                # override model config's image size
                self.model_cfg['vision_cfg']['image_size'] = force_image_size

            # cast_dtype set for fp16 and bf16 (manual mixed-precision),
            # not set for 'amp' or 'pure' modes
            cast_dtype = get_cast_dtype(precision)
            self.model = CLIP(**self.model_cfg, cast_dtype=cast_dtype)
            if precision in ('fp16', 'bf16'):
                dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
                # manual mixed precision that matches original OpenAI behaviour
                self.model.to(device=device)
                convert_weights_to_lp(self.model, dtype=dtype)
            elif precision in ('pure_fp16', 'pure_bf16'):
                dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
                self.model.to(device=device, dtype=dtype)
            else:
                self.model.to(device=device)

            pretrained_loaded = False
            if pretrained:
                checkpoint_path = ''
                pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
                if pretrained_cfg:
                    checkpoint_path = download_pretrained(pretrained_cfg,
                                                          path=cache_dir,
                                                          revision=revision)
                elif os.path.exists(pretrained):
                    checkpoint_path = pretrained

                if checkpoint_path:
                    print(f'Loading pretrained {model_name} weights ({pretrained}).')
                    load_checkpoint(self.model, checkpoint_path)
                else:
                    error_str = (
                        f'Pretrained weights ({pretrained}) not found for model {model_name}.'
                        f'Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
                    print(error_str)
                    raise RuntimeError(error_str)
                pretrained_loaded = True
            elif has_hf_hub_prefix:
                print(f'Loading pretrained {model_name} weights ({pretrained}).')
                load_checkpoint(self.model, checkpoint_path)
                pretrained_loaded = True

            if require_pretrained and not pretrained_loaded:
                # callers of create_model_from_pretrained always expect pretrained weights
                raise RuntimeError(
                    f'Pretrained weights were required for (model: '
                    f'{model_name}, pretrained: {pretrained}) but not loaded.')

            # set image / mean metadata from pretrained_cfg if available, or use default
            self.model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
            self.model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD

        self.preprocess = _transfor_ori(self.model.visual.image_size[0])
        self.preprocess_mod = _transfor_mod(self.model.visual.image_size[0])

        if output_dict and hasattr(self.model, 'output_dict'):
            self.model.output_dict = True

        if jit:
            self.model = torch.jit.script(self.model)

        if freeze_clip:
            self.freeze()

        self.model.to(self.device)

    def freeze(self):
        """Freeze the parameters of the clip."""
        self.model = self.model.eval()
        for param in self.parameters():
            param.requires_grad = False

    @staticmethod
    def tokenize(texts, context_length=77, truncate=True):
        """
        Returns the tokenized representation of given input string(s)

        Parameters
        ----------
        texts : Union[str, List[str]]
            An input string or a list of input strings to tokenize

        context_length : int
            The context length to use; all CLIP models use 77 as the context length

        truncate: bool
            Whether to truncate the text in case its encoding is longer than the context length

        Returns
        -------
        A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
        We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
        """
        if isinstance(texts, str):
            texts = [texts]

        sot_token = _tokenizer.encoder['<|startoftext|>']
        eot_token = _tokenizer.encoder['<|endoftext|>']
        all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
        if packaging.version.parse(torch.__version__) < packaging.version.parse('1.8.0'):
            result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
        else:
            result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)

        for i, tokens in enumerate(all_tokens):
            if len(tokens) > context_length:
                if truncate:
                    tokens = tokens[:context_length]
                    tokens[-1] = eot_token
                else:
                    raise RuntimeError(f'Input {texts[i]} is too long for context length {context_length}')
            result[i, :len(tokens)] = torch.tensor(tokens)

        return result


    def encode_image(self, images, normalize=True, is_preprocess=False):
        if is_preprocess:
            images = self.preprocess(images)
        if images.ndim == 3:
            images = images.unsqueeze(0)
        images = images.to(self.device)
        image_features = self.model.encode_image(images, normalize=normalize)
        return image_features


    def encode_text(self, text, normalize=True, is_tokenize=False):
        if is_tokenize:
            if not isinstance(text, list):
                text = [text]
            tokens = self.tokenize(text, context_length=self.max_length)
        else:
            tokens = text
        eot_ind = tokens.argmax(dim=-1)
        tokens = tokens.to(self.device)
        pooled_out, out = self.model.encode_text(tokens, normalize=normalize)
        return pooled_out, out, eot_ind


    def forward(self, text, image):
        results = self.model.forward(image=image, text=text)
        return results

    def clip_loss(self, text, image, real_image=None, img_clip_loss_weight=0.0):
        image = self.preprocess_mod(image)
        img_feats = self.encode_image(image, is_preprocess=False)
        txt_feats, _, _ = self.encode_text(text, is_tokenize=True)
        img_feats = torch.cat(GatherLayer.apply(img_feats), dim=0)
        txt_feats = torch.cat(GatherLayer.apply(txt_feats), dim=0)
        if real_image is not None:
            with torch.no_grad():
                real_image = self.preprocess_mod(real_image)
                real_img_feats = self.encode_image(real_image, is_preprocess=False)
                real_img_feats = concat_all_gather(real_img_feats)

        # cosine similarity as logits
        logit_scale = self.model.logit_scale
        logit_scale = logit_scale.exp()
        logits_per_image = logit_scale * img_feats @ txt_feats.t()
        logits_per_text = logits_per_image.t()
        # logits_per_image = torch.diagonal(logits_per_image, 0)
        # similarity = 1 - logits_per_image / 100

        batch_size = img_feats.shape[0]
        labels = torch.arange(batch_size, device=img_feats.device).long()
        total_loss = (F.cross_entropy(logits_per_image, labels) +
                      F.cross_entropy(logits_per_text, labels)) / 2

        if real_image is not None:
            logits = logit_scale * img_feats @ real_img_feats.t()
            img_clip_loss = F.cross_entropy(logits, labels)
            total_loss = total_loss + img_clip_loss_weight * img_clip_loss
        return total_loss


    def clip_score(self, text, image):
        image = self.preprocess_mod(image)
        tokens = self.tokenize(text).to(self.device)
        results = self(text=tokens, image=image)
        img_features = results['image_features']
        txt_features = results['text_features_pooled']
        score = F.relu((img_features * txt_features).sum(axis=-1))
        return score

models/clip_model.py
zhujiapeng/Aurora
# python3.7
"""Contains the implementation of the text2image generator.
"""
from collections import OrderedDict
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from einops import rearrange, repeat, reduce

from third_party.stylegan3_official_ops import bias_act
from third_party.stylegan2_official_ops import upfirdn2d
from third_party.stylegan3_official_ops import conv2d_gradfix
from .utils.ops import all_gather
# import torch.utils.checkpoint as checkpoint

__all__ = ['Text2ImageGenerator']

# Resolutions allowed.
_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256]

# pylint: disable=missing-function-docstring
# pylint: disable=arguments-renamed

def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


def checkpoint(func, inputs, params, flag):
    """
    Evaluate a function without caching intermediate activations, allowing for
    reduced memory at the expense of extra compute in the backward pass.

    :param func: the function to evaluate.
    :param inputs: the argument sequence to pass to `func`.
    :param params: a sequence of parameters `func` depends on but does not
                   explicitly take as arguments.
    :param flag: if False, disable gradient checkpointing.
    """
    if flag:
        args = tuple(inputs) + tuple(params)
        return CheckpointFunction.apply(func, len(inputs), *args)
    return func(*inputs)


class Text2ImageGenerator(nn.Module):
    """Defines the generator network of a text to images.

    NOTE: The synthesized images are with `RGB` channel order and pixel range
    [-1, 1].
    Settings for the mapping network:

    (1) z_dim: Dimension of the input latent space, Z. (default: 128)
    (2) w_dim: Dimension of the output latent space, W. (default: 1024)
    (3) repeat_w: Repeat w-code for different layers. (default: True)
    (4) normalize_z: Whether to normalize the z-code. (default: True)
    (5) mapping_layers: Number of layers of the mapping network. (default: 4)
    (6) mapping_fmaps: Number of hidden channels of the mapping network.
        (default: 1024)
    (7) mapping_use_wscale: Whether to use weight scaling for the mapping
        network. (default: True)
    (8) mapping_wscale_gain: The factor to control weight scaling for the
        mapping network. (default: 1.0)
    (9) mapping_lr_mul: Learning rate multiplier for the mapping network.
        (default: 0.01)

    Settings for conditional generation (class label):

    (1) use_class_label: Whether to use class label condition. (default False)
    (2) label_dim:  Dimension of the additional label for conditional
        generation. In one-hot conditioning case, it is equal to the number of
        classes. If set to 0, conditioning training will be disabled.
        (default: 0)
    (3) embedding_dim: Dimension of the embedding space, if needed.
        (default: 512)
    (4) embedding_bias: Whether to add bias to embedding learning.
        (default: True)
    (5) embedding_use_wscale: Whether to use weight scaling for embedding
        learning. (default: True)
    (6) embedding_wscale_gian: The factor to control weight scaling for
        embedding. (default: 1.0)
    (7) embedding_lr_mul: Learning rate multiplier for the embedding
        learning. (default: 1.0)
    (8) normalize_class_embed: Whether to normalize the class embedding.
        (default: True)

    Settings for conditional generation (text):

    (1) use_text_cond: Whether to use text condition. (default True)
    (2) num_layers_text_enc: Number of transformer blocks will be used in the
        trainable text encoder (default: 4).
    (3) weight_name: Weight name for the pre-trained clip model.
    (4) dropout: Dropout ratio for in the project out layer in transformer.
        (default: 0.)
    (5) text_use_wscale: Whether to use weight scaling for the trainable text
        network. (default: True)
    (6) text_wscale_gain: The factor to control weight scaling for the
        trainable text network. (default: 1.0)
    (7) text_lr_mul: Learning rate multiplier for the trainable text network.
        (default: 0.01)
    (8) use_transformer_clip: Whether to use the clip model from `transformers`
        library, otherwise, using clip from  `open_clip`. (default: True)
    (9) trainable_head: Whether to use trainable head in text network.
        (default: True)
    (10) context_dim: Dimension of the text embedding. (default: 768)
    (11) normalize_text_cond: Whether to normalize the text embedding.
        (default: True)

    Settings for the synthesis network:

    (1) resolution: The resolution of the output image. (default: -1)
    (2) init_res: The initial resolution to start with convolution. (default: 4)
    (3) image_channels: Number of channels of the output image. (default: 3)
    (4) num_adaptive_kernels: Number of adaptive conv kernels for each
        resolution blocks. (default: {4:1, 8:1, 16:2, 32:4, 64:8})
    (5) num_block_per_res: Number of synthesis blocks per resolution.
        (default: {4:3, 8:3, 16: 3, 32: 2, 64: 2})
    (6) attn_resolutions: Attention add on which resolution.
        (default: (8, 16, 32))
    (7) attn_depth: Attention depth for each resolution. (default: {8:2,
        16:2, 32:1})
    (8) attn_ch_factor: Attention dimension multiplier. (default 1.0)
    (9) attn_gain: Residual gain on attention. (default: 0.3)
    (10) zero_out: Whether to zero the weights for the attention project out
        layer. (default: True)
    (11) final_tanh: Whether to use `tanh` to control the final pixel range.
        (default: False)
    (12) fourier_feat: Whether to use a fourier feature in the first
        convolutional layer. (default: True)
    (13) demodulate: Whether to perform style demodulation. (default: True)
    (14) use_wscale: Whether to use weight scaling. (default: True)
    (15) wscale_gain: The factor to control weight scaling. (default: 1.0)
    (16) lr_mul: Learning rate multiplier for the synthesis network.
         (default: 1.0)
    (17) fmaps_base: Factor to control number of feature maps for each layer.
         (default: 32 << 9)
    (18) fmaps_max: Maximum number of feature maps in each layer.
          (default: 1600)
    (19) filter_kernel: Kernel used for filtering (e.g., downsampling).
         (default: (1, 3, 3, 1))
    (20) conv_clamp: A threshold to clamp the output of convolution layers to
         avoid overflow under FP16 training. (default: None)
    (21) num_heads: Number of head in multi-head attention. (default 8)
    (22) head_dim: Dimension of each head. (default 64)
    (23) l2_attention: Whether or not to use l2 distance to compute the
         similarity. (default True)
    (24) use_checkpoint: if True, use gradient checkpointing on this module.
        (default False) NOTE: In current GAN training, we can not use this
        function since the second derivative need the intermediate results.
    (25) eps: A small value to avoid divide overflow. (default: 1e-8)
    (26) mtm: Whether to apply modulated transformation module. (default: False)
    (27) num_experts: Number of experts for MoE. (default: {4:-1,8:-1,16:-1,32:-1,64:-1})

    Runtime settings:

    (1) w_moving_decay: Decay factor for updating `w_avg`, which is used for
        training only. Set `None` to disable. (default: None)
    (2) sync_w_avg: Synchronizing the stats of `w_avg` across replicas. If set
        as `True`, the stats will be more accurate, yet the speed maybe a little
        bit slower. (default: False)
    (3) trunc_psi: Truncation psi, set `None` to disable. (default: None)
    (4) trunc_layers: Number of layers to perform truncation. (default: None)
    (5) impl: Implementation mode of some particular ops, e.g., `filtering`,
        `bias_act`, etc. `cuda` means using the official CUDA implementation
        from StyleGAN2, while `ref` means using the native PyTorch ops.
        (default: `cuda`)
    """

    def __init__(self,
                 # Settings for mapping network.
                 z_dim=128,
                 w_dim=1024,
                 repeat_w=True,
                 normalize_z=True,
                 mapping_layers=4,
                 mapping_fmaps=1024,
                 mapping_use_wscale=True,
                 mapping_wscale_gain=1.0,
                 mapping_lr_mul=0.01,
                 # Settings for conditional generation (class label).
                 use_class_label=False,
                 label_dim=0,
                 embedding_dim=512,
                 embedding_bias=True,
                 embedding_use_wscale=True,
                 embedding_wscale_gain=1.0,
                 embedding_lr_mul=1.0,
                 normalize_class_embed=True,
                 # Settings for conditional generation (text).
                 use_text_cond=True,
                 use_w_cond=False,
                 num_layers_text_enc=4,
                 dropout=0.0,
                 text_use_wscale=True,
                 text_wscale_gain=1.0,
                 text_lr_mul=0.01,
                 clip_out_dim=768,
                 context_dim=768,
                 normalize_text_cond=True,
                 # Settings for synthesis network.
                 resolution=-1,
                 init_res=4,
                 image_channels=3,
                 num_adaptive_kernels={},
                 num_block_per_res={},
                 attn_resolutions=(8, 16, 32),
                 attn_depth={},
                 attn_ch_factor=1.0,
                 mask_self=False,
                 residual_gain=1.0,
                 attn_gain=1.0,
                 text_head_gain=1.0,
                 zero_out=True,
                 final_tanh=False,
                 fourier_feat=True,
                 demodulate=True,
                 use_wscale=True,
                 wscale_gain=1.0,
                 lr_mul=1.0,
                 fmaps_base=32 << 9,
                 fmaps_max=1600,
                 filter_kernel=(1, 3, 3, 1),
                 conv_clamp=None,
                 num_heads=8,
                 head_dim=None,
                 l2_attention=True,
                 tie=True,
                 scale_in=True,
                 include_ff=True,
                 glu=False,
                 use_checkpoint=False,
                 checkpoint_res=None,
                 eps=1e-8,
                 mtm=False,
                 num_experts={},
                 ms_training_res=['4', '8', '16', '32', '64'],
                 skip_connection=False):
        """Initializes with basic settings.

        Raises:
            ValueError: If the `resolution` is not supported, or `architecture`
                is not supported.
        """
        super().__init__()
        if resolution not in _RESOLUTIONS_ALLOWED:
            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')

        assert ((use_text_cond and not use_class_label) or
                (not use_text_cond and use_class_label)), ('Condition '
                'information can not from both the text and class label.')

        # assert ((label_dim > 0 and not context_dim) or
        #         (label_dim <= 0 and context_dim)), ('Class label '
        #         'and context dimension should be just give one.')

        self.z_dim = z_dim
        self.w_dim = w_dim
        self.repeat_w = repeat_w
        self.normalize_z = normalize_z
        self.mapping_layers = mapping_layers
        self.mapping_fmaps = mapping_fmaps
        self.mapping_use_wscale = mapping_use_wscale
        self.mapping_wscale_gain = mapping_wscale_gain
        self.mapping_lr_mul = mapping_lr_mul

        self.use_class_label=use_class_label
        self.label_dim = label_dim
        self.embedding_dim = embedding_dim
        self.embedding_bias = embedding_bias
        self.embedding_use_wscale = embedding_use_wscale
        self.embedding_wscale_gain = embedding_wscale_gain
        self.embedding_lr_mul = embedding_lr_mul
        self.normalize_class_embed = normalize_class_embed

        self.use_text_cond = use_text_cond
        self.use_w_cond = use_w_cond
        self.num_layers_text_enc = num_layers_text_enc
        self.dropout = dropout
        self.text_use_wscale = text_use_wscale
        self.text_wscale_gain = text_wscale_gain
        self.text_lr_mul = text_lr_mul
        self.context_dim = context_dim
        self.normalize_text_cond = normalize_text_cond

        self.resolution = resolution
        self.init_res = init_res
        self.image_channels = image_channels
        self.num_adaptive_kernel = num_adaptive_kernels
        self.num_block_per_res = num_block_per_res
        self.attn_depth = attn_depth
        if attn_resolutions is None:
            self.attn_resolutions = list(attn_depth.keys())
        else:
            self.attn_resolutions = attn_resolutions

        self.mask_self = mask_self
        self.residual_gain = residual_gain
        self.attn_ch_factor = attn_ch_factor
        self.attn_gain = attn_gain
        self.zero_out = zero_out
        self.final_tanh = final_tanh
        self.fourier_feat = fourier_feat
        self.demodulate = demodulate
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.fmaps_base = fmaps_base
        self.fmaps_max = fmaps_max
        self.filter_kernel = filter_kernel
        self.conv_clamp = conv_clamp
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.use_checkpoint = use_checkpoint
        if checkpoint_res is None:
            self.checkpoint_res = self.attn_resolutions
        else:
            self.checkpoint_res = checkpoint_res
        self.l2_attention = l2_attention
        self.tie = tie
        self.scale_in = scale_in
        self.include_ff = include_ff
        self.glu = glu
        self.eps = eps
        self.mtm = mtm
        self.num_experts = num_experts
        self.ms_training_res = ms_training_res
        self.skip_connection = skip_connection

        # Dimension of latent space, which is convenient for sampling.
        self.latent_dim = (z_dim,)

        # Number of synthesis (convolutional) layers.
        num_block_per_res_keys = num_block_per_res.keys()
        self.num_layers = 0
        self.layer_id_to_res = {}
        for res_ in num_block_per_res_keys:
            if res_ in self.attn_resolutions:
                self.num_layers += num_block_per_res[res_] * 2
            else:
                self.num_layers += num_block_per_res[res_]
            self.layer_id_to_res[self.num_layers - 1] = res_
        print(f'Number of layers: {self.num_layers}')
        print(f'layer_id_to_res: {self.layer_id_to_res}')

        if use_text_cond:
            self.text_head = TextHead(in_dim=clip_out_dim,
                                      context_dim=context_dim,
                                      num_heads=num_heads,
                                      num_layers=num_layers_text_enc,
                                      dropout=dropout,
                                      use_wscale=text_use_wscale,
                                      wscale_gain=text_wscale_gain,
                                      lr_mul=text_lr_mul,
                                      use_checkpoint=use_checkpoint,
                                      mlp_ratio=4,
                                      l2_attention=l2_attention,
                                      scale_in=scale_in,
                                      mask_self=mask_self,
                                      attn_gain=text_head_gain)

        self.mapping = MappingNetwork(
                           input_dim=z_dim,
                           output_dim=w_dim,
                           num_outputs=self.num_layers,
                           repeat_output=repeat_w,
                           normalize_z=normalize_z,
                           num_layers=mapping_layers,
                           hidden_dim=mapping_fmaps,
                           use_wscale=mapping_use_wscale,
                           wscale_gain=mapping_wscale_gain,
                           lr_mul=mapping_lr_mul,
                           use_class_label=use_class_label,
                           label_dim=label_dim,
                           embedding_dim=embedding_dim,
                           embedding_bias=embedding_bias,
                           embedding_use_wscale=embedding_use_wscale,
                           embedding_wscale_gain=embedding_wscale_gain,
                           embedding_lr_mul=embedding_lr_mul,
                           normalize_class_embed=normalize_class_embed,
                           use_text_cond=use_text_cond,
                           context_dim=context_dim,
                           normalize_text_cond=normalize_text_cond,
                           eps=eps)

        # This is used for truncation trick.
        if self.repeat_w:
            self.register_buffer('w_avg', torch.zeros(w_dim))
        else:
            self.register_buffer('w_avg', torch.zeros(self.num_layers * w_dim))

        self.synthesis = SynthesisNetwork(
                            resolution=resolution,
                            init_res=init_res,
                            w_dim=w_dim,
                            num_adaptive_kernels=num_adaptive_kernels,
                            num_block_per_res=num_block_per_res,
                            image_channels=image_channels,
                            attn_resolutions=self.attn_resolutions,
                            attn_depth=attn_depth,
                            attn_ch_factor=attn_ch_factor,
                            layer_id_to_res=self.layer_id_to_res,
                            fourier_feat=fourier_feat,
                            demodulate=demodulate,
                            use_wscale=use_wscale,
                            wscale_gain=wscale_gain,
                            lr_mul=lr_mul,
                            fmaps_base=fmaps_base,
                            fmaps_max=fmaps_max,
                            filter_kernel=filter_kernel,
                            conv_clamp=conv_clamp,
                            use_text_cond=use_text_cond,
                            use_w_cond=use_w_cond,
                            context_dim=context_dim,
                            num_heads=num_heads,
                            head_dim=head_dim,
                            use_checkpoint=use_checkpoint,
                            checkpoint_res=self.checkpoint_res,
                            l2_attention=l2_attention,
                            tie=tie,
                            scale_in=scale_in,
                            include_ff=include_ff,
                            glu=glu,
                            mask_self=mask_self,
                            residual_gain=residual_gain,
                            attn_gain=attn_gain,
                            final_tanh=final_tanh,
                            zero_out=zero_out,
                            dropout=dropout,
                            eps=eps,
                            mtm=mtm,
                            num_experts=num_experts,
                            ms_training_res=ms_training_res,
                            skip_connection=skip_connection)

    def forward(self,
                z,
                label=None,
                context=None,
                eot_ind=None,
                w_moving_decay=None,
                sync_w_avg=False,
                trunc_psi=None,
                trunc_layers=None,
                impl='cuda'):

        if self.use_text_cond:
            global_text, local_text = self.text_head(context, eot_ind=eot_ind)
            mapping_results = self.mapping(z,
                                           label=None,
                                           context=global_text)
        else:
            local_text = None
            mapping_results = self.mapping(z,
                                           label=label,
                                           context=None)

        w = mapping_results['w']
        if self.training and w_moving_decay is not None:
            if sync_w_avg:
                batch_w_avg = all_gather(w.detach()).mean(dim=0)
            else:
                batch_w_avg = w.detach().mean(dim=0)
            self.w_avg.copy_(batch_w_avg.lerp(self.w_avg, w_moving_decay))

        wp = mapping_results.pop('wp')
        if not self.training:
            trunc_psi = 1.0 if trunc_psi is None else trunc_psi
            trunc_layers = 0 if trunc_layers is None else trunc_layers
            if trunc_psi < 1.0 and trunc_layers > 0:
                w_avg = self.w_avg.reshape(1, -1, self.w_dim)[:, :trunc_layers]
                wp[:, :trunc_layers] = w_avg.lerp(
                    wp[:, :trunc_layers], trunc_psi)

        synthesis_results = self.synthesis(wp, context=local_text)

        return {**mapping_results, **synthesis_results}


class TextHead(nn.Module):
    """Define the trainable text head.

    Trainable head is build at the top of a text encoder (i.g., CLIP text
    encoder). Namely, the text description is firstly send to a pre-trained text
    encoder to extract the features. Then this trainable head is used as an
    adaptor between the pre-trained text encoder and learned generator. Besides,
    the features of the trainable head is categorized into two groups. One is
    global features and the other is the local features, The last token (`end of
    text` or `EOT`) of embedding serves as the global feature, while the
    remaining tokens are considered as the local features. Specifically, the
    output of the trainable head has the shape of [B, C, dim], where B is the
    batch size, C is the max token length (usually C equals to 77), and dim is
    the embedding dimension. If an input description has token length `len` and
    features `feats`. Then, the global feature is feats[:, len, :] and the
    local feature is torch.cat([feats[1, :len-1, :], feats[1, len:, :]], 1).
    """

    def __init__(self,
                 in_dim,
                 context_dim,
                 num_heads,
                 num_layers,
                 dropout,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 use_checkpoint,
                 mlp_ratio,
                 l2_attention,
                 scale_in,
                 mask_self,
                 attn_gain):
        super().__init__()
        self.context_dim = context_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.dropout = dropout
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.use_checkpoint = use_checkpoint
        self.mlp_ratio = mlp_ratio
        self.l2_attention = l2_attention
        self.scale_in = scale_in
        self.mask_self = mask_self
        self.align_ch = in_dim != context_dim
        self.attn_gain = attn_gain

        if self.align_ch:
            self.align_dim = DenseLayer(in_channels=in_dim,
                                        out_channels=context_dim,
                                        add_bias=True,
                                        init_bias=1.0,
                                        use_wscale=use_wscale,
                                        wscale_gain=wscale_gain,
                                        lr_mul=lr_mul,
                                        activation_type='linear')

        self.transformer = TextTransformer(width=context_dim,
                                           num_heads=num_heads,
                                           num_layers=num_layers,
                                           dropout=dropout,
                                           add_bias=True,
                                           init_bias=0.0,
                                           use_wscale=use_wscale,
                                           wscale_gain=wscale_gain,
                                           lr_mul=lr_mul,
                                           use_checkpoint=use_checkpoint,
                                           mlp_ratio=4,
                                           l2_attention=l2_attention,
                                           scale_in=scale_in,
                                           mask_self=mask_self,
                                           attn_gain=attn_gain)

    def forward(self, text_feats, eot_ind):
        if self.align_ch:
            text_feats = self.align_dim(text_feats)
        bs, emb_length, dim = text_feats.shape
        device = text_feats.device

        text_feats = text_feats.permute(1, 0, 2)
        text_feats = self.transformer(text_feats)
        text_feats = text_feats.permute(1, 0, 2)
        length_ind = eot_ind[:, None, None].repeat(1, 1, dim).to(device)
        mask = torch.ones_like(text_feats, device=device)
        zero_fill = torch.zeros_like(text_feats, device=device)
        mask = torch.scatter(mask, dim=1, index=length_ind, src=zero_fill)
        local_feats = text_feats[mask.bool()].view(bs, emb_length-1, dim)
        global_feats = torch.gather(text_feats, dim=1, index=length_ind)
        # global_feats = text_feats[torch.arange(text_feats.shape[0]), eot_ind]
        # global_feats = global_feats.unsqueeze(1)

        return global_feats[:, 0], local_feats


class MappingNetwork(nn.Module):
    """Implements the latent space mapping network.

    Basically, this network executes several dense layers in sequence, and also
    support add text embedding if needed.
    """

    def __init__(self,
                 input_dim,
                 output_dim,
                 num_outputs,
                 repeat_output,
                 normalize_z,
                 num_layers,
                 hidden_dim,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 use_class_label,
                 label_dim,
                 embedding_dim,
                 embedding_bias,
                 embedding_use_wscale,
                 embedding_wscale_gain,
                 embedding_lr_mul,
                 normalize_class_embed,
                 use_text_cond,
                 context_dim,
                 normalize_text_cond,
                 eps):
        super().__init__()

        self.input_dim = input_dim
        self.output_dim = output_dim
        self.num_outputs = num_outputs
        self.repeat_output = repeat_output
        self.normalize_z = normalize_z
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul

        self.use_class_label = use_class_label
        self.label_dim = label_dim
        self.embedding_dim = embedding_dim
        self.embedding_bias = embedding_bias
        self.embedding_use_wscale = embedding_use_wscale
        self.embedding_wscale_gain = embedding_wscale_gain
        self.embedding_lr_mul = embedding_lr_mul
        self.normalize_class_embed = normalize_class_embed

        self.use_text_cond = use_text_cond
        self.context_dim = context_dim
        self.normalize_text_cond = normalize_text_cond

        self.eps = eps

        self.norm = PixelNormLayer(dim=1, eps=eps)

        if self.use_text_cond:
            input_dim = input_dim + context_dim

        if self.use_class_label:
            input_dim = input_dim + embedding_dim
            self.embedding = DenseLayer(in_channels=label_dim,
                                        out_channels=embedding_dim,
                                        add_bias=embedding_bias,
                                        init_bias=0.0,
                                        use_wscale=embedding_use_wscale,
                                        wscale_gain=embedding_wscale_gain,
                                        lr_mul=embedding_lr_mul,
                                        activation_type='linear')

        if num_outputs is not None and not repeat_output:
            output_dim = output_dim * num_outputs

        for i in range(num_layers):
            in_channels = (input_dim if i == 0 else hidden_dim)
            out_channels = (output_dim if i == (num_layers - 1) else hidden_dim)
            self.add_module(f'dense{i}',
                            DenseLayer(in_channels=in_channels,
                                       out_channels=out_channels,
                                       add_bias=True,
                                       init_bias=0.0,
                                       use_wscale=use_wscale,
                                       wscale_gain=wscale_gain,
                                       lr_mul=lr_mul,
                                       activation_type='lrelu'))

    def forward(self, z, label=None, context=None):
        if z.ndim != 2 or z.shape[1] != self.input_dim:
            raise ValueError(f'Input latent code should be with shape '
                             f'[batch_size, input_dim], where '
                             f'`input_dim` equals to {self.input_dim}!\n'
                             f'But `{z.shape}` is received!')

        if self.use_text_cond and context is None:
            raise ValueError(f'Text condition is needed! but got {context}')
        if self.use_class_label and label is None:
            raise ValueError(f'Class condition is needed! but got {label}')

        if self.normalize_z:
            z = self.norm(z)

        if self.use_text_cond:
            if (context.ndim != 2 or
                context.shape != (z.shape[0], self.context_dim)):
                raise ValueError(f'Input text condition should be with shape '
                                 f'[batch_size, context_dim], where '
                                 f'`batch_size` equals to that of '
                                 f'latent codes ({z.shape[0]}) and '
                                 f'`context_dim` equals to {self.context_dim}! '
                                 f'But `{context.shape}` is received!')
            if self.normalize_text_cond:
                context = self.norm(context)
            w = torch.cat((z, context), dim=1)
        elif self.use_class_label:
            if label.ndim != 2 or label.shape != (z.shape[0], self.label_dim):
                raise ValueError(f'Input label should be with shape '
                                 f'[batch_size, label_dim], where '
                                 f'`batch_size` equals to that of '
                                 f'latent codes ({z.shape[0]}) and '
                                 f'`label_dim` equals to {self.label_dim}!'
                                 f'But `{label.shape}` is received!')
            label = label.to(dtype=torch.float32)
            embedding = self.embedding(label)
            if self.normalize_class_embed:
                embedding = self.norm(embedding)
            w = torch.cat((z, embedding), dim=1)
        else:
            w = z

        for i in range(self.num_layers):
            w = getattr(self, f'dense{i}')(w)

        wp = None
        if self.num_outputs is not None:
            if self.repeat_output:
                wp = w.unsqueeze(1).repeat((1, self.num_outputs, 1))
            else:
                wp = w.reshape(-1, self.num_outputs, self.output_dim)

        results = {
            'z': z,
            'global_text': context,
            'label': label,
            'w': w,
            'wp': wp,
        }
        return results


class SynthesisNetwork(nn.Module):
    """Implements the image synthesis network.

    Basically, this network executes several convolutional layers in sequence.
    """
    def __init__(self,
                 resolution,
                 init_res,
                 w_dim,
                 num_adaptive_kernels,
                 num_block_per_res,
                 image_channels,
                 attn_resolutions,
                 attn_depth,
                 attn_ch_factor,
                 layer_id_to_res,
                 fourier_feat,
                 demodulate,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 fmaps_base,
                 fmaps_max,
                 filter_kernel,
                 conv_clamp,
                 use_text_cond,
                 use_w_cond,
                 context_dim,
                 num_heads,
                 head_dim,
                 use_checkpoint,
                 checkpoint_res,
                 l2_attention,
                 tie,
                 scale_in,
                 include_ff,
                 glu,
                 mask_self,
                 residual_gain,
                 attn_gain,
                 final_tanh,
                 zero_out,
                 dropout,
                 eps,
                 mtm,
                 num_experts,
                 ms_training_res,
                 skip_connection):
        super().__init__()

        self.init_res = init_res
        self.init_res_log2 = int(np.log2(init_res))
        self.resolution = resolution
        self.final_res_log2 = int(np.log2(resolution))
        self.w_dim = w_dim
        self.num_adaptive_kernels = num_adaptive_kernels
        self.num_block_per_res = num_block_per_res
        self.attn_ch_factor = attn_ch_factor
        self.image_channels = image_channels
        self.attn_resolutions = attn_resolutions
        self.attn_depth = attn_depth
        self.fourier_feat = fourier_feat
        self.demodulate = demodulate
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.fmaps_base = fmaps_base
        self.fmaps_max = fmaps_max
        self.filter_kernel = filter_kernel
        self.final_tanh = final_tanh
        self.conv_clamp = conv_clamp
        self.eps = eps
        self.mtm = mtm
        self.num_experts = num_experts
        self.ms_training_res = ms_training_res
        self.skip_connection = skip_connection

        # hyper-parameter on the attention.
        self.use_text_cond = use_text_cond
        self.use_w_cond = use_w_cond
        self.context_dim = context_dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.use_checkpoint = use_checkpoint
        self.checkpoint_res = checkpoint_res
        self.l2_attention = l2_attention
        self.tie = tie
        self.scale_in = scale_in
        self.include_ff = include_ff
        self.glu = glu
        self.attn_gain = attn_gain
        self.zero_out = zero_out
        self.dropout = dropout

        self.res_level = []
        for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1):
            self.res_level.append(2**res_log2)

        assert len(num_adaptive_kernels) == len(self.res_level), (f'For each '
                f'resolution level, must assign a corresponding number, but '
                f'got resolution: {self.res_level}, and number of kernel list: '
                f'{self.num_kernels}')
        assert len(self.res_level) >= len(attn_resolutions), (f'Attention '
                f'resolution list must less than or equal to the resolution '
                f'level, but get resolution: {self.res_level}, and attention '
                f'resolution: {self.attn_resolutions}')

        self.layer_id_to_res = layer_id_to_res

        self.net = nn.ModuleList()
        for res in self.res_level:
            in_channels = self.get_nf(res // 2)
            out_channels = self.get_nf(res)
            num_blocks = num_block_per_res[str(res)]

            # Layer with initial resolution, i.e., 4
            if res == self.init_res:
                for i in range(num_blocks):
                    if i == 0:  # First block at resolution 4x4.
                        # First conv layer
                        if fourier_feat:
                            layer0 = FourierInputLayer(w_dim=w_dim,
                                                       size=res,
                                                       channels=in_channels,
                                                       sampling_rate=res,
                                                       bandwidth=1)
                        else:
                            layer0 = ConstInputLayer(init_res=res,
                                                     channels=in_channels)

                        # Second conv layer
                        layer1 = AdaptiveModulateConvLayer(
                                    in_channels=in_channels,
                                    out_channels=out_channels,
                                    modulate_dim=w_dim,
                                    kernel_size=3,
                                    num_kernels=num_adaptive_kernels[str(res)],
                                    add_bias=True,
                                    scale_factor=1,
                                    filter_kernel=None,
                                    demodulate=demodulate,
                                    use_wscale=use_wscale,
                                    wscale_gain=wscale_gain,
                                    lr_mul=lr_mul,
                                    activation_type='lrelu',
                                    conv_clamp=conv_clamp,
                                    eps=eps,
                                    mtm=mtm)
                        self.net.append(ResolutionSequential(*[layer0, layer1]))
                    else: # The rest blocks at resolution 4x4.
                        layer = ResidualConvBlock(
                                    in_channels=out_channels,
                                    out_channels=out_channels,
                                    modulate_dim=w_dim,
                                    num_kernels=num_adaptive_kernels[str(res)],
                                    upsample=False,
                                    residual_gain=residual_gain,
                                    filter_kernel=filter_kernel,
                                    demodulate=demodulate,
                                    use_wscale=use_wscale,
                                    wscale_gain=wscale_gain,
                                    lr_mul=lr_mul,
                                    conv_clamp=conv_clamp,
                                    eps=eps,
                                    mtm=mtm)
                        self.net.append(ResolutionSequential(layer))

                    if str(res) in attn_resolutions:
                        layer = ResidualAttnBlock(
                                            depth=attn_depth[str(res)],
                                            in_channels=out_channels,
                                            res=res,
                                            context_dim=context_dim,
                                            attn_ch_factor=attn_ch_factor,
                                            num_heads=num_heads,
                                            head_dim=head_dim,
                                            w_dim=w_dim,
                                            group_channels=32,
                                            demodulate=demodulate,
                                            use_text_cond=use_text_cond,
                                            use_w_cond=use_w_cond,
                                            use_wscale=use_wscale,
                                            wscale_gain=wscale_gain,
                                            lr_mul=lr_mul,
                                            conv_clamp=conv_clamp,
                                            use_checkpoint=use_checkpoint,
                                            checkpoint_res=checkpoint_res,
                                            l2_attention=l2_attention,
                                            tie=tie,
                                            scale_in=scale_in,
                                            include_ff=include_ff,
                                            mask_self=mask_self,
                                            attn_gain=attn_gain,
                                            dropout=dropout,
                                            mlp_ratio=4,
                                            add_bias=True,
                                            init_bias=0.0,
                                            glu=glu,
                                            zero_out=zero_out,
                                            eps=eps,
                                            num_experts=num_experts[str(res)])
                        self.net.append(ResolutionSequential(layer))

            else: # Layers with rest resolution i.e., res >= 8
                for i in range(num_blocks):
                    if i == 0:
                        upsample = True
                    else:
                        upsample = False

                    layer = ResidualConvBlock(
                                in_channels=in_channels,
                                out_channels=out_channels,
                                modulate_dim=w_dim,
                                num_kernels=num_adaptive_kernels[str(res)],
                                upsample=upsample,
                                residual_gain=residual_gain,
                                filter_kernel=filter_kernel,
                                demodulate=demodulate,
                                use_wscale=use_wscale,
                                wscale_gain=wscale_gain,
                                lr_mul=lr_mul,
                                conv_clamp=conv_clamp,
                                eps=eps,
                                mtm=mtm and res < 32)
                    self.net.append(ResolutionSequential(layer))

                    if str(res) in attn_resolutions:
                        layer = ResidualAttnBlock(
                                          depth=attn_depth[str(res)],
                                          in_channels=out_channels,
                                          res=res,
                                          context_dim=context_dim,
                                          attn_ch_factor=attn_ch_factor,
                                          num_heads=num_heads,
                                          head_dim=head_dim,
                                          w_dim=w_dim,
                                          group_channels=32,
                                          demodulate=demodulate,
                                          use_text_cond=use_text_cond,
                                          use_w_cond=use_w_cond,
                                          use_wscale=use_wscale,
                                          wscale_gain=wscale_gain,
                                          lr_mul=lr_mul,
                                          conv_clamp=conv_clamp,
                                          use_checkpoint=use_checkpoint,
                                          checkpoint_res=checkpoint_res,
                                          l2_attention=l2_attention,
                                          tie=tie,
                                          scale_in=scale_in,
                                          include_ff=include_ff,
                                          mask_self=mask_self,
                                          attn_gain=attn_gain,
                                          dropout=dropout,
                                          mlp_ratio=4,
                                          add_bias=True,
                                          init_bias=0.0,
                                          glu=glu,
                                          zero_out=zero_out,
                                          eps=eps,
                                          num_experts=num_experts[str(res)])
                        self.net.append(ResolutionSequential(layer))
                    in_channels = out_channels
            if use_text_cond: # Multi-scale output for text input.
                layer_name = f'output_res_{res}x{res}'
                self.add_module(
                    layer_name,
                    AdaptiveModulateConvLayer(
                                    in_channels=out_channels,
                                    out_channels=image_channels,
                                    modulate_dim=w_dim,
                                    kernel_size=1,
                                    num_kernels=num_adaptive_kernels[str(res)],
                                    add_bias=True,
                                    scale_factor=1,
                                    filter_kernel=None,
                                    demodulate=False,
                                    use_wscale=use_wscale,
                                    wscale_gain=wscale_gain,
                                    lr_mul=lr_mul,
                                    activation_type='linear',
                                    conv_clamp=conv_clamp,
                                    eps=eps))
            else: # Single output for the class label input.
                if res == self.resolution:
                    layer_name = f'output_res_{res}x{res}'
                    self.add_module(
                        layer_name,
                        AdaptiveModulateConvLayer(
                                    in_channels=out_channels,
                                    out_channels=image_channels,
                                    modulate_dim=w_dim,
                                    kernel_size=1,
                                    num_kernels=num_adaptive_kernels[str(res)],
                                    add_bias=True,
                                    scale_factor=1,
                                    filter_kernel=None,
                                    demodulate=False,
                                    use_wscale=use_wscale,
                                    wscale_gain=wscale_gain,
                                    lr_mul=lr_mul,
                                    activation_type='linear',
                                    conv_clamp=conv_clamp,
                                    eps=eps))

        # Used for upsampling output images for each resolution block for sum.
        if self.skip_connection:
            self.register_buffer(
                'filter', upfirdn2d.setup_filter(filter_kernel))


    def get_nf(self, res):
        """Gets number of feature maps according to the given resolution."""
        return min(self.fmaps_base // res, self.fmaps_max)


    def forward(self, wp, context):
        results = {'wp': wp,
                   'local_context': context}
        h = None
        img = None
        moe_losses = 0.0
        for idx, module in enumerate(self.net):
            h, moe_loss = module(h, wp[:, idx], context)
            moe_losses += moe_loss
            if idx in self.layer_id_to_res:
                res = self.layer_id_to_res[idx]
                if res in self.ms_training_res:
                    if self.use_text_cond:
                        out_layer = getattr(self, f'output_res_{res}x{res}')
                        y, _ = out_layer(h, wp[:, idx])
                        if self.final_tanh:
                            y = torch.tanh(y)
                        if self.skip_connection:
                            y = y.to(torch.float32)
                            if res == str(self.init_res):
                                img = y
                            else:
                                img = y + upfirdn2d.upsample2d(img, self.filter)
                        else:
                            img = y
                        results[f'img_res_{res}x{res}'] = img
                    else:
                        if int(res) == self.resolution:
                            out_layer = getattr(self, f'output_res_{res}x{res}')
                            img, _ = out_layer(h, wp[:, idx])
                            if self.final_tanh:
                                img = torch.tanh(img)
                            results[f'img_res_{res}x{res}'] = img
                    if int(res) == max([int(each) for each in self.ms_training_res]):
                        break
        results['image'] = img
        results['moe_loss'] = moe_losses
        return results


class ResolutionSequential(nn.Sequential):
    """A sequential module that passes modulate code `w` and `text` as
       the condition.
    """

    def forward(self, x, w, context=None):
        moe_loss = 0.0
        for layer in self:
            if isinstance(layer, (FourierInputLayer, ConstInputLayer)):
                x = layer(w)
            elif isinstance(layer, (ResidualConvBlock,
                                    AdaptiveModulateConvLayer)):
                x, _ = layer(x, w)
            elif isinstance(layer, ResidualAttnBlock):
                x, _, moe_loss = layer(x, w, context)
            else:
                x = layer(x)
        return x, moe_loss


class ResidualConvBlock(nn.Module):
    """Residual block that that consists of two convolutional layers. If the

    dimension of the input and output features changed, an extra conv layer will
    be added.
    """
    def __init__(self,
                 in_channels,
                 out_channels,
                 modulate_dim,
                 num_kernels,
                 upsample,
                 residual_gain,
                 filter_kernel,
                 demodulate,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 conv_clamp,
                 eps,
                 mtm):
        """Initializes with layer settings.

        Args:
            in_channels: Number of channels of projection in.
            out_channels: Number of channels of projection out.
            w_dim: Dimension of W space for style modulation.
            num_kernels: Number of adaptive kernels.
            upsample: Whether or not to do the upsampling.
            residual_gain: residual gain for the residual connection.
            filter_kernel: Kernel used for filtering.
            demodulate: Whether to perform style demodulation.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            conv_clamp: A threshold to clamp the output of convolution layers to
                avoid overflow under FP16 training.
            eps: A small value to avoid divide overflow.
            mtm: Whether to apply modulated tranformation module
        """
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modulate_dim = modulate_dim
        self.num_kernels = num_kernels
        self.upsample = upsample
        self.residual_gain = residual_gain
        self.filter_kernel = filter_kernel
        self.demodulate = demodulate
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.conv_clamp = conv_clamp
        self.eps = eps
        self.mtm = mtm

        if upsample:
            scale_factor = 2
        else:
            scale_factor = 1
            filter_kernel = None

        self.conv1 = AdaptiveModulateConvLayer(
                                    in_channels=in_channels,
                                    out_channels=out_channels,
                                    modulate_dim=modulate_dim,
                                    kernel_size=3,
                                    num_kernels=num_kernels,
                                    add_bias=True,
                                    scale_factor=scale_factor,
                                    filter_kernel=filter_kernel,
                                    demodulate=demodulate,
                                    use_wscale=use_wscale,
                                    wscale_gain=wscale_gain,
                                    lr_mul=lr_mul,
                                    activation_type='lrelu',
                                    conv_clamp=conv_clamp,
                                    eps=eps)

        self.conv2 = AdaptiveModulateConvLayer(
                                    in_channels=out_channels,
                                    out_channels=out_channels,
                                    modulate_dim=modulate_dim,
                                    kernel_size=3,
                                    num_kernels=num_kernels,
                                    add_bias=True,
                                    scale_factor=1,
                                    filter_kernel=None,
                                    demodulate=demodulate,
                                    use_wscale=use_wscale,
                                    wscale_gain=wscale_gain,
                                    lr_mul=lr_mul,
                                    activation_type='lrelu',
                                    conv_clamp=conv_clamp,
                                    eps=eps,
                                    mtm=mtm) # using mtm

        if scale_factor == 1 and in_channels == out_channels:
            self.skip_connection = nn.Identity()
        else:
            self.skip_connection = ConvLayer(in_channels=in_channels,
                                             out_channels=out_channels,
                                             kernel_size=1,
                                             add_bias=False,
                                             scale_factor=scale_factor,
                                             filter_kernel=filter_kernel,
                                             use_wscale=use_wscale,
                                             wscale_gain=wscale_gain,
                                             lr_mul=lr_mul,
                                             activation_type='linear',
                                             conv_clamp=None)

    def forward(self, x, modulate_code):
        h, _ = self.conv1(x, modulate_code)
        h, style = self.conv2(h, modulate_code)
        residual = self.skip_connection(x)
        return (h + residual) * self.residual_gain, style


class ConstInputLayer(nn.Module):
    """Implements the input layer to start convolution with.

    Basically, this block starts from a const input, which is with shape
    `(channels, init_res, init_res)`.
    """

    def __init__(self, init_res, channels):
        super().__init__()
        self.const = nn.Parameter(torch.randn(1, channels, init_res, init_res))

    def forward(self, w):
        x = self.const.repeat(w.shape[0], 1, 1, 1)
        return x


class FourierInputLayer(nn.Module):
    """Implements the input layer with Fourier features."""
    def __init__(self, w_dim, size, channels, sampling_rate, bandwidth):
        super().__init__()

        self.w_dim = w_dim
        self.channels = channels
        self.size = np.broadcast_to(np.asarray(size), [2])
        self.sampling_rate = sampling_rate
        self.bandwidth = bandwidth

        # Draw random frequencies from uniform 2D disc.
        freqs = torch.randn([self.channels, 2])
        radii = freqs.square().sum(dim=1, keepdim=True).sqrt()
        freqs /= radii * radii.square().exp().pow(0.25)
        freqs *= bandwidth
        phases = torch.rand([self.channels]) - 0.5

        # Setup parameters and buffers.
        self.weight = nn.Parameter(torch.randn([self.channels, self.channels]))
        self.affine = DenseLayer(w_dim, 4, init_weight=0, init_bias=[1,0,0,0], add_bias=True, use_wscale=True, wscale_gain=1.0, lr_mul=1.0, activation_type='linear')

        self.register_buffer('transform', torch.eye(3, 3)) # User-specified inverse transform wrt. resulting image.
        self.register_buffer('freqs', freqs)
        self.register_buffer('phases', phases)


    def extra_repr(self):
        return (f'w_dim={self.w_dim}, '
                f'channels={self.channels}, '
                f'sampling_rate={self.sampling_rate}, '
                f'bandwidth={self.bandwidth}, '
                f'size={self.size}, ')

    def forward(self, w):

        # Introduce batch dimension.
        transforms = self.transform.unsqueeze(0) # [batch, row, col]
        freqs = self.freqs.unsqueeze(0) # [batch, channel, xy]
        phases = self.phases.unsqueeze(0) # [batch, channel]

        # Apply learned transformation.
        t = self.affine(w) # t = (r_c, r_s, t_x, t_y)
        t = t / t[:, :2].norm(dim=1, keepdim=True) # t' = (r'_c, r'_s, t'_x, t'_y)
        m_r = torch.eye(3, device=w.device).unsqueeze(0).repeat([w.shape[0], 1, 1]) # Inverse rotation wrt. resulting image.
        m_r[:, 0, 0] = t[:, 0]  # r'_c
        m_r[:, 0, 1] = -t[:, 1] # r'_s
        m_r[:, 1, 0] = t[:, 1]  # r'_s
        m_r[:, 1, 1] = t[:, 0]  # r'_c
        m_t = torch.eye(3, device=w.device).unsqueeze(0).repeat([w.shape[0], 1, 1]) # Inverse translation wrt. resulting image.
        m_t[:, 0, 2] = -t[:, 2] # t'_x
        m_t[:, 1, 2] = -t[:, 3] # t'_y
        transforms = m_r @ m_t @ transforms # First rotate resulting image, then translate, and finally apply user-specified transform.

        # Transform frequencies.
        phases = phases + (freqs @ transforms[:, :2, 2:]).squeeze(2)
        freqs = freqs @ transforms[:, :2, :2]

        # Dampen out-of-band frequencies that may occur due to the user-specified transform.
        #amplitudes = (1 - (freqs.norm(dim=2) - self.bandwidth) / (self.sampling_rate / 2 - self.bandwidth)).clamp(0, 1)

        # Construct sampling grid.
        theta = torch.eye(2, 3, device=w.device)
        theta[0, 0] = 0.5 * self.size[0] / self.sampling_rate
        theta[1, 1] = 0.5 * self.size[1] / self.sampling_rate
        grids = torch.nn.functional.affine_grid(theta.unsqueeze(0), [1, 1, self.size[1], self.size[0]], align_corners=False)

        # Compute Fourier features.
        x = (grids.unsqueeze(3) @ freqs.permute(0, 2, 1).unsqueeze(1).unsqueeze(2)).squeeze(3) # [batch, height, width, channel]
        x = x + phases.unsqueeze(1).unsqueeze(2)
        x = torch.sin(x * (np.pi * 2))
        #x = x * amplitudes.unsqueeze(1).unsqueeze(2)

        # Apply trainable mapping.
        weight = self.weight / np.sqrt(self.channels)
        x = x @ weight.t()

        # Ensure correct shape.
        x = x.permute(0, 3, 1, 2) # [batch, channel, height, width]
        return x.contiguous()


class ResidualAttnBlock(nn.Module):
    """Implements the attention used in the generator.

    Attention is composed by several `AttentionBlock`
    """
    def __init__(self,
                 depth,
                 in_channels,
                 res,
                 context_dim,
                 attn_ch_factor,
                 num_heads,
                 head_dim,
                 w_dim,
                 group_channels,
                 demodulate,
                 use_text_cond,
                 use_w_cond,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 conv_clamp,
                 use_checkpoint,
                 checkpoint_res,
                 l2_attention,
                 tie,
                 scale_in,
                 include_ff,
                 mask_self,
                 attn_gain,
                 dropout,
                 mlp_ratio,
                 add_bias,
                 init_bias,
                 glu,
                 zero_out,
                 eps,
                 num_experts):
        """Initializes with layer settings.

        Args:
            depth: Number of attention block used in this layer.
            in_channels: Number of channels of projection out.
            res:
            context_dim: Number of channels of the context tensor.
            attn_ch_factor: Attention dimension multiplier
            num_heads: Number of head of multi-head attention.
            head_dim: Dimension of each attention head.
            w_dim: Dimension of W space for style modulation.
            group_channel: Group number of group norm.
            demodulate: Whether to perform style demodulation.
            use_text_cond: Whether to involve the text condition.
            use_w_cond: Whether to involve the wp as the condition.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            conv_clamp: A threshold to clamp the output of convolution layers to
                avoid overflow under FP16 training.
            use_checkpoint: if True, use gradient checkpointing on this module.
            checkpoint_res:
            l2_attention: Whether or not to use l2 distance to compute the
                similarity.
            tie:
            scale_in:
            include_ff:
            mask_self:
            attn_gain: residual gain for the self attention.
            dropout: Dropout ratio.
            mlp_ratio:
            add_bias: Whether to add bias onto the fully-connected result.
            init_bias: The initial bias value before training.
            glu: use glu/geglu as activation, true: glu, false: geglu.
            zero_out: Whether or not to init zero weight at the out layer.
            eps: A small value to avoid divide overflow.
            num_experts: Number of experts for MoE
        """
        super().__init__()
        self.depth = depth
        self.in_channels = in_channels
        self.res = res
        self.context_dim = context_dim
        self.attn_ch_factor = attn_ch_factor
        self.num_heads = num_heads
        if head_dim is None:
            self.head_dim = in_channels // num_heads
            assert self.head_dim * num_heads == in_channels
        else:
            self.head_dim = head_dim
        self.w_dim = w_dim
        self.group_channels = group_channels
        self.demodulate = demodulate
        self.use_text_cond = use_text_cond
        self.use_w_cond = use_w_cond
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.conv_clamp = conv_clamp
        self.use_checkpoint = use_checkpoint
        self.checkpoint_res = checkpoint_res
        self.l2_attention = l2_attention
        self.tie = tie
        self.scale_in = scale_in
        self.include_ff = include_ff
        self.attn_gain = attn_gain
        self.dropout = dropout
        self.mlp_ratio = mlp_ratio
        self.add_bias = add_bias
        self.init_bias = init_bias
        self.glu = glu
        self.zero_out = zero_out
        self.eps = eps
        self.num_experts = num_experts

        inner_dim = int(num_heads * self.head_dim)

        self.norm = GroupNorm(group_channels, in_channels)

        self.proj_in = ConvLayer(in_channels=in_channels,
                                 out_channels=inner_dim,
                                 kernel_size=1,
                                 add_bias=False,
                                 scale_factor=1,
                                 filter_kernel=None,
                                 use_wscale=use_wscale,
                                 wscale_gain=wscale_gain,
                                 lr_mul=lr_mul,
                                 activation_type='linear',
                                 conv_clamp=conv_clamp)

        self.attn_blocks = nn.ModuleList(
            [AttentionBlock(in_channels=inner_dim,
                            res=res,
                            context_dim=context_dim,
                            num_heads=num_heads,
                            head_dim=self.head_dim,
                            w_dim=w_dim,
                            demodulate=demodulate,
                            use_text_cond=use_text_cond,
                            use_w_cond=use_w_cond,
                            attn_gain=attn_gain,
                            use_wscale=use_wscale,
                            wscale_gain=wscale_gain,
                            lr_mul=lr_mul,
                            conv_clamp=conv_clamp,
                            use_checkpoint=use_checkpoint,
                            checkpoint_res=checkpoint_res,
                            l2_attention=l2_attention,
                            tie=tie,
                            scale_in=scale_in,
                            include_ff=include_ff,
                            mask_self=mask_self,
                            dropout=dropout,
                            mlp_ratio=mlp_ratio,
                            add_bias=add_bias,
                            init_bias=init_bias,
                            glu=glu,
                            eps=eps,
                            num_experts=num_experts)
                for _ in range(depth)])

        self.proj_out = ConvLayer(in_channels=inner_dim,
                                  out_channels=in_channels,
                                  kernel_size=1,
                                  add_bias=False,
                                  scale_factor=1,
                                  filter_kernel=None,
                                  use_wscale=use_wscale,
                                  wscale_gain=wscale_gain,
                                  lr_mul=lr_mul,
                                  activation_type='linear',
                                  conv_clamp=conv_clamp)
        if zero_out:
            self.proj_out = zero_module(self.proj_out)

    def extra_repr(self):
        return (f'depth={self.depth}, '
                f'in_channels={self.in_channels}, '
                f'num_heads={self.num_heads}, '
                f'context_dim={self.context_dim}, '
                f'head_dim={self.head_dim}, '
                f'group_channels={self.group_channels}, '
                f'w_dim={self.w_dim}')

    def forward(self, x, w, context=None):
        moe_losses = 0.0
        h = self.norm(x)
        h = self.proj_in(h)
        for block in self.attn_blocks:
            h, style, moe_loss = block(h, w, context=context)
            moe_losses += moe_loss
        h = self.proj_out(h)
        return (x + h) * self.attn_gain, style, moe_losses


class FeedForward(nn.Module):
    """Implements of the feed-forward layer in attention layer.

    """
    def __init__(self,
                 in_dim,
                 mlp_ratio,
                 dropout,
                 w_dim,
                 add_bias,
                 init_bias,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 glu,
                 num_experts=-1):
        """Initializes with layer settings.

        Args:
            in_dim: Number of channels of the input tensor.
            mlp_ratio: A factor used in increasing the dimension of MLP.
            dropout: Dropout ratio.
            w_dim: Dimension of W space for style modulation.
            add_bias: Whether to add bias onto the fully-connected result.
            init_bias: The initial bias value before training.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            glu: True: Use glu, otherwise, use geglu.
            num_experts: Number of involved experts. Disable in default.
        """
        super().__init__()
        self.in_dim = in_dim
        self.mlp_ratio = mlp_ratio
        self.dropout = dropout
        self.add_bias = add_bias
        self.init_bias = init_bias
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.glu = glu
        self.num_experts = num_experts

        if self.num_experts < 2: # one ffn layer
            if glu:
                self.project_in = nn.Sequential(
                        DenseLayer(in_channels=in_dim,
                                   out_channels=in_dim * mlp_ratio,
                                   add_bias=add_bias,
                                   init_bias=init_bias,
                                   use_wscale=use_wscale,
                                   wscale_gain=wscale_gain,
                                   lr_mul=lr_mul,
                                   activation_type='linear'),
                        nn.GELU())
            else:
                self.project_in = GEGLU(in_channels=in_dim,
                                        out_channels=in_dim * mlp_ratio,
                                        add_bias=add_bias,
                                        init_bias=init_bias,
                                        use_wscale=use_wscale,
                                        wscale_gain=wscale_gain,
                                        lr_mul=lr_mul)

            self.proj_out = DenseLayer(in_channels=in_dim * mlp_ratio,
                                       out_channels=in_dim,
                                       add_bias=add_bias,
                                       init_bias=init_bias,
                                       use_wscale=use_wscale,
                                       wscale_gain=wscale_gain,
                                       lr_mul=lr_mul,
                                       activation_type='linear')

            self.net = nn.Sequential(
                            self.project_in,
                            nn.Dropout(dropout),
                            self.proj_out)
        else:
            self.net = nn.ModuleList([])
            # build each expert
            for _ in range(self.num_experts):
                if glu:
                    project_in = nn.Sequential(
                            DenseLayer(in_channels=in_dim,
                                       out_channels=in_dim * mlp_ratio,
                                       add_bias=add_bias,
                                       init_bias=init_bias,
                                       use_wscale=use_wscale,
                                       wscale_gain=wscale_gain,
                                       lr_mul=lr_mul,
                                       activation_type='linear'),
                            nn.GELU())
                else:
                    project_in = GEGLU(in_channels=in_dim,
                                            out_channels=in_dim * mlp_ratio,
                                            add_bias=add_bias,
                                            init_bias=init_bias,
                                            use_wscale=use_wscale,
                                            wscale_gain=wscale_gain,
                                            lr_mul=lr_mul)

                proj_out = DenseLayer(in_channels=in_dim * mlp_ratio,
                                           out_channels=in_dim,
                                           add_bias=add_bias,
                                           init_bias=init_bias,
                                           use_wscale=use_wscale,
                                           wscale_gain=wscale_gain,
                                           lr_mul=lr_mul,
                                           activation_type='linear')

                net = nn.Sequential(
                                project_in,
                                nn.Dropout(dropout),
                                proj_out)
                self.net.append(net)
            # build routing layer
            self.routing = AdaptiveModulateConvLayer(
                                    in_channels=in_dim,
                                    out_channels=self.num_experts,
                                    modulate_dim=w_dim,
                                    kernel_size=1,
                                    num_kernels=1,
                                    add_bias=add_bias,
                                    scale_factor=1,
                                    filter_kernel=None,
                                    demodulate=True,
                                    use_wscale=use_wscale,
                                    wscale_gain=wscale_gain,
                                    lr_mul=lr_mul,
                                    activation_type='lrelu',
                                    conv_clamp=None,
                                    eps=1e-8,
                                    mtm=False)
            # hyper-parameters for MoE
            self.jitter_noise = 0.0 # currently we disable the representation jittering
            self.expert_capacity = 262144 # 512*512

    def router(self, hidden_states, w):
        # float32 is used to ensure stability. See the discussion of "selective precision" in
        # https://arxiv.org/abs/2101.03961.
        original_dtype = hidden_states.dtype
        assert(original_dtype==torch.float32)

        if self.jitter_noise > 0: # question: diable this jittering when inference??
            # Get the lower and upper bound of the uniform distribution
            # Adapted from: https://stackoverflow.com/questions/44328530/how-to-get-a-uniform-distribution-in-a-range-r1-r2-in-pytorch
            distrib_lower_bound = 1.0 - self.jitter_noise
            distrib_upper_bound = 1.0 + self.jitter_noise

            uniform_distrib = torch.rand(hidden_states.shape, device=hidden_states.device, dtype=original_dtype)
            uniform_distrib = uniform_distrib * (distrib_lower_bound - distrib_upper_bound)

            uniform_distrib = uniform_distrib + distrib_upper_bound
            # Multiply the token inputs by the uniform distribution - adding some noise
            hidden_states *= uniform_distrib

        # Shape: [num_groups, tokens_per_group, num_experts]
        height = int(np.sqrt(hidden_states.shape[1]))
        width = int(np.sqrt(hidden_states.shape[1]))
        assert(height*width==hidden_states.shape[1])
        hidden_states = rearrange(hidden_states, 'b (h w) c -> b c h w', h=height, w=width)
        router_logits = self.routing(hidden_states, w)[0]
        router_logits = rearrange(router_logits, 'b c h w -> b (h w) c')


        # Apply Softmax and cast back to the original `dtype`
        router_probs = nn.functional.softmax(router_logits, dim=-1)


        expert_index = torch.argmax(router_probs, dim=-1)
        expert_index = torch.nn.functional.one_hot(expert_index, num_classes=self.num_experts)

        # Mask tokens outside expert capacity. Sum over each sequence
        token_priority = torch.cumsum(expert_index, dim=-2)
        # mask if the token routed to to the expert will overflow
        expert_capacity_mask = token_priority <= self.expert_capacity
        expert_index = expert_index * expert_capacity_mask

        router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(-1)
        return expert_index, router_probs, router_logits

    def cal_moe_loss(self, router_logits, expert_index):
        z_loss = router_z_loss_func(router_logits)
        router_probs = nn.Softmax(dim=-1)(router_logits)
        aux_loss = load_balancing_loss_func(router_probs, expert_index)
        loss = 0.001 * z_loss + 0.001 * aux_loss
        return loss

    def forward(self, x, w=None):
        if self.num_experts < 2:
            return self.net(x), 0.0

        # Step 1: Get the router_mask from the router as wel as the probabilities
        router_mask, router_probs, router_logits = self.router(x, w)
        expert_index = torch.argmax(router_mask, dim=-1)

        # The routers introduced might not always map all the tokens, to a router, which means that some hidden states
        # can be unchanged from one layer to another. That is why the hidden states are cloned before updating only the seleced ones.
        next_states = x.clone()
        for idx, expert in enumerate(self.net):
            token_indices = router_mask[:, :, idx].bool()
            next_states[token_indices] = expert(x[token_indices])

        x = router_probs * next_states
        moe_loss = self.cal_moe_loss(router_logits, expert_index)
        return x, moe_loss


class AttentionBlock(nn.Module):
    """Implements the attention block used in the `Attention`.

    Attention block composed by self-attention and cross-attention.
    Self-attention allows spatial positions to attend to each other, while the
    cross-attention is used to involve the text condition information. Note that
    cross-attention only be involved when perform the text to image generation.
    """
    def __init__(self,
                 in_channels,
                 res,
                 context_dim,
                 num_heads,
                 head_dim,
                 w_dim,
                 demodulate,
                 use_text_cond,
                 use_w_cond,
                 attn_gain,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 conv_clamp,
                 use_checkpoint,
                 checkpoint_res,
                 l2_attention,
                 tie,
                 scale_in,
                 include_ff,
                 mask_self,
                 dropout,
                 mlp_ratio,
                 add_bias,
                 init_bias,
                 glu,
                 eps,
                 num_experts):
        """Initializes with layer settings.

        Args:
            in_channels: Number of channels of projection out.
            res:
            context_dim: Number of channels of the context tensor.
            num_heads: Number of head of multi-head attention.
            head_dim: Dimension of each head.
            w_dim: Dimension of W space for style modulation.
            demodulate: Whether to perform style demodulation.
            use_text_cond: Whether to involve the text condition.
            use_w_cond:
            attn_gain:
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            conv_clamp: A threshold to clamp the output of convolution layers to
                avoid overflow under FP16 training.
            use_checkpoint: if True, use gradient checkpointing on this module.
            checkpoint_res:
            l2_attention: Whether or not to use l2 distance to compute the
                similarity.
            dropout: Dropout ratio.
            mlp_ratio: Ratio for the attention multipler.
            add_bias: Whether to add bias onto the fully-connected result.
            init_bias: The initial bias value before training.
            glu: use glu/geglu as activation, true: glu, false: geglu.
            eps: A small value to avoid divide overflow.
            num_experts: Numer of experts for MoE
        """
        super().__init__()
        self.in_channels = in_channels
        self.res = res
        self.context_dim = context_dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.w_dim = w_dim
        self.demodulate = demodulate
        self.use_text_cond = use_text_cond
        self.use_w_cond = use_w_cond
        self.attn_gain = attn_gain
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.conv_clamp = conv_clamp
        self.use_checkpoint = use_checkpoint
        self.checkpoint_res = checkpoint_res
        self.l2_attention = l2_attention
        self.tie = tie
        self.scale_in = scale_in
        self.include_ff = include_ff
        self.dropout = dropout
        self.mlp_ratio = mlp_ratio
        self.add_bias = add_bias
        self.init_bias = init_bias
        self.glu = glu
        self.eps = eps
        self.num_experts = num_experts

        if self.use_text_cond or self.use_w_cond:
            assert context_dim, ('Must given the context dimension when '
                                 '`use_text_cond` or `use_w_cond` is true.')

        self.norm1 = LayerNorm(in_channels)
        self.self_attn = SelfAttentionLayer(in_channels=in_channels,
                                            num_heads=num_heads,
                                            head_dim=head_dim,
                                            w_dim=w_dim,
                                            demodulate=demodulate,
                                            use_wscale=use_wscale,
                                            wscale_gain=wscale_gain,
                                            lr_mul=lr_mul,
                                            use_checkpoint=use_checkpoint,
                                            checkpoint_res=checkpoint_res,
                                            mask_self=mask_self,
                                            l2_attention=l2_attention,
                                            tie=tie,
                                            scale_in=scale_in,
                                            eps=eps)
        if use_text_cond or use_w_cond:
            self.norm2 = LayerNorm(in_channels)
            if use_w_cond:
                self.align_dim = DenseLayer(in_channels=w_dim,
                                            out_channels=context_dim,
                                            add_bias=True,
                                            init_bias=1.0,
                                            use_wscale=use_wscale,
                                            wscale_gain=wscale_gain,
                                            lr_mul=lr_mul,
                                            activation_type='linear')

            self.cross_attn = CrossAttentionLayer(query_dim=in_channels,
                                                  context_dim=context_dim,
                                                  num_heads=num_heads,
                                                  head_dim=head_dim,
                                                  dropout=0.,
                                                  use_wscale=use_wscale,
                                                  wscale_gain=wscale_gain,
                                                  lr_mul=lr_mul,
                                                  use_checkpoint=use_checkpoint,
                                                  checkpoint_res=checkpoint_res,
                                                  l2_attention=l2_attention,
                                                  scale_in=scale_in,
                                                  mask_self=mask_self)
        if include_ff:
            self.norm3 = LayerNorm(in_channels)
            self.ff = FeedForward(in_dim=in_channels,
                                mlp_ratio=mlp_ratio,
                                dropout=0.,
                                w_dim=w_dim,
                                add_bias=add_bias,
                                init_bias=init_bias,
                                use_wscale=use_wscale,
                                wscale_gain=wscale_gain,
                                lr_mul=lr_mul,
                                glu=glu,
                                num_experts=num_experts)

    def extra_repr(self) -> str:
        return (f'in_channels={self.in_channels}, '
                f'num_heads={self.num_heads}, '
                f'context_dim={self.context_dim}, '
                f'head_dim={self.head_dim}, '
                f'w_dim={self.w_dim}')

    def forward(self, x, w, context=None):
        moe_loss = 0.0
        residual = x
        _, _, height, width = x.shape
        x = rearrange(x, 'b c h w -> b (h w) c')
        x = self.norm1(x)
        x = rearrange(x, 'b (h w) c -> b c h w', h=height, w=width)
        x, style = self.self_attn(x, w)
        #x = (x + residual) * self.attn_gain
        x = x + residual

        x = rearrange(x, 'b c h w -> b (h w) c')
        if self.use_text_cond or self.use_w_cond:
            contexts = []
            if self.use_text_cond:
                assert context is not None
                contexts.append(context)
            if self.use_w_cond:
                w_ = self.align_dim(w)
                contexts.append(w_[:, None])
            contexts = torch.cat(contexts, dim=1)
            #x = (self.cross_attn(self.norm2(x), contexts) + x) * self.attn_gain
            x = self.cross_attn(self.norm2(x), contexts) + x

        if self.include_ff:
            out, moe_loss = self.ff(self.norm3(x), w)
            #x = (out + x) * self.attn_gain
            x = out + x

        x = rearrange(x, 'b (h w) c -> b c h w', h=height, w=width)

        return x, style, moe_loss


class SelfAttentionLayer(nn.Module):
    """
    Implements the self-attention layer.

    Supporting the similarity computed using dot product or l2 norm.

    TODO: Scale down the l2 distance logit to match the unit normal distribution
    when init weight, reduce residual gain from attention.
    """

    def __init__(
        self,
        in_channels,
        num_heads,
        head_dim,
        w_dim,
        demodulate,
        use_wscale,
        wscale_gain,
        lr_mul,
        use_checkpoint,
        checkpoint_res,
        mask_self,
        l2_attention,
        tie,
        scale_in,
        eps):
        """Initializes with layer settings.

        Args:
            in_channels: Number of channels of projection out.
            num_heads: Number of head of multi-head attention.
            head_dim: Dimension of the each head.
            w_dim: Dimension of W space for style modulation.
            demodulate: Whether to perform style demodulation.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            use_checkpoint: if True, use gradient checkpointing on this module.
            checkpoint_res:
            mask_self: Whether to omit the attention to self.
            l2_attention: Whether or not to use l2 distance to compute the
                similarity.
            eps: A small value to avoid divide overflow.
        """
        super().__init__()
        self.in_channels = in_channels
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.w_dim = w_dim
        self.demodulate = demodulate
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul

        self.use_checkpoint = use_checkpoint
        self.checkpoint_res = checkpoint_res
        self.l2_attention = l2_attention
        self.tie = tie
        self.scale_in = scale_in
        self.mask_self = mask_self
        self.eps = eps

        self.scale = self.head_dim ** -0.5

        if self.l2_attention:
            if self.tie:
                self.to_qk = DenseLayer(in_channels=in_channels,
                                        out_channels=in_channels,
                                        add_bias=False,
                                        init_bias=0.,
                                        use_wscale=use_wscale,
                                        wscale_gain=wscale_gain,
                                        lr_mul=lr_mul,
                                        activation_type='linear')
            else:
                self.to_qk = DenseLayer(in_channels=in_channels,
                                        out_channels=in_channels * 2,
                                        add_bias=False,
                                        init_bias=0.,
                                        use_wscale=use_wscale,
                                        wscale_gain=wscale_gain,
                                        lr_mul=lr_mul,
                                        activation_type='linear')

            self.to_v = DenseLayer(in_channels=in_channels,
                                   out_channels=in_channels,
                                   add_bias=False,
                                   init_bias=0.,
                                   use_wscale=use_wscale,
                                   wscale_gain=wscale_gain,
                                   lr_mul=lr_mul,
                                   activation_type='linear')
        else:
            self.to_qkv = DenseLayer(in_channels=in_channels,
                                     out_channels=in_channels * 3,
                                     add_bias=False,
                                     init_bias=0.,
                                     use_wscale=use_wscale,
                                     wscale_gain=wscale_gain,
                                     lr_mul=lr_mul,
                                     activation_type='linear')

        self.attention = QKVAttention(n_heads=self.num_heads,
                                      scale=self.scale,
                                      l2_attention=l2_attention,
                                      scale_in=scale_in,
                                      mask_self=mask_self)

        self.proj_out = AdaptiveModulateConvLayer(in_channels=in_channels,
                                                  out_channels=in_channels,
                                                  modulate_dim=w_dim,
                                                  kernel_size=1,
                                                  num_kernels=1,
                                                  add_bias=True,
                                                  scale_factor=1,
                                                  filter_kernel=None,
                                                  demodulate=demodulate,
                                                  use_wscale=use_wscale,
                                                  wscale_gain=wscale_gain,
                                                  lr_mul=lr_mul,
                                                  activation_type='linear',
                                                  conv_clamp=None,
                                                  eps=eps)

    def extra_repr(self):
        return (f'in_channels={self.in_channels}, '
                f'num_heads={self.num_heads}, '
                f'head_dim={self.head_dim}, '
                f'w_dim={self.w_dim}, '
                f'wscale_gain={self.wscale_gain:.3f}, '
                f'lr_mul={self.lr_mul:.3f}')

    def forward(self, x, w):
        right_res = str(x.shape[-1]) in self.checkpoint_res
        do_checkpoint = right_res and self.use_checkpoint
        return checkpoint(self._forward,
                          (x, w),
                          self.parameters(),
                          do_checkpoint)

    def _forward(self, x, w):
        _, _, height, width = x.shape
        h = rearrange(x, 'b c h w -> b (h w) c')
        if self.l2_attention:
            qk = self.to_qk(h) # b, x*y, c (num_heads * head_dim)
            qk = rearrange(qk, 'b x (h d) -> (b h) x d', h=self.num_heads)
            v = self.to_v(h)
            v = rearrange(v, 'b x (h d) -> (b h) x d', h=self.num_heads)
            if self.tie:
                qkv = torch.cat([qk, qk, v], dim=2)
            else:
                qkv = torch.cat([qk, v], dim=2)
        else:
            qkv = self.to_qkv(h)
            qkv = rearrange(qkv, 'b x (h d) -> (b h) x d', h=self.num_heads)
        out = self.attention(qkv)

        out = rearrange(out, '(b h) (x y) d -> b (h d) x y',
                        x=height, y=width, h=self.num_heads)
        out, style = self.proj_out(out, w)

        return out, style


class QKVAttention(nn.Module):
    """
    QKV attention module that support similarity computed using dot product and

    L2 norm.
    """
    def __init__(self,
                 n_heads,
                 scale,
                 l2_attention=True,
                 scale_in=True,
                 mask_self=False,
                 mask_value='-inf'):
        """Initializes with layer settings.

        Args:
            n_head: Number of heads for the multi-head attention.
            scale: The scale for the attention.
            l2_attention: Whether or not use l2 distance to compute similarity.
            mask_self: Whether to omit the attention to self.
            mask_value: Default value to fill in the distance matrix.
        """
        super().__init__()
        self.n_heads = n_heads
        self.scale = scale
        self.l2_attention = l2_attention
        self.scale_in = scale_in
        self.mask_value = mask_value
        self.mask_self = mask_self

    def extra_repr(self) -> str:
        return (f'num_heads={self.n_heads}, '
                f'scale={self.scale}, '
                f'mask_val={self.mask_value}')

    def forward(self, qkv):
        """
        Apply QKV attention.

        :param qkv: an [N x (HW) x C*3] tensor of Qs, Ks, and Vs.
        :return: an [N x (HW) x C] tensor after attention.
        """
        q, k, v = qkv.chunk(3, dim=2)

        if self.scale_in:
            q = q * self.scale
            k = k * self.scale

        if self.l2_attention:
            attn_logits = - torch.cdist(q, k, p=2)
            if self.mask_self:
                mask = torch.eye(attn_logits.shape[1],
                                 device=attn_logits.device,
                                 dtype=torch.bool)
                attn_logits = attn_logits.masked_fill(mask,
                                                      float(self.mask_value))
        else:
            attn_logits = torch.einsum('b t c, b s c -> b t s', q, k)

        if not self.scale_in:
            attn_logits = attn_logits * self.scale

        attn = torch.softmax(attn_logits, dim=-1)
        a = torch.einsum('b t s, b s c -> b t c', attn, v)
        return a


class CrossAttentionLayer(nn.Module):
    """Implements the cross-attention layer.

    NOTE: This layer support compute similarity using l2 distance. However, the
    query and key can not share the same weight as ViTGAN dose since query comes
    from image feature while key comes from text condition.

    TODO: Need to make sure whether this layer has a residual connection.
    """
    def __init__(self,
                 query_dim,
                 context_dim,
                 num_heads,
                 head_dim,
                 dropout,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 use_checkpoint,
                 checkpoint_res,
                 l2_attention,
                 scale_in,
                 mask_self):
        """Initializes with layer settings.

        Args:
            query_dim: Number of channels of the query tensor.
            context_dim: Number of channels of the context tensor.
            num_heads: Number of attention heads.
            head_dim: Dimension of each head.
            dropout: Dropout ratio.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            use_checkpoint: if True, use gradient checkpointing on this module.
            checkpoint_res:
            l2_attention: Whether or not to use l2 distance to compute the
                similarity.
            scale_in:
            mask_self:
        """
        super().__init__()
        self.query_dim = query_dim
        self.context_dim = context_dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.dropout = dropout
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.use_checkpoint = use_checkpoint
        self.checkpoint_res = checkpoint_res
        self.l2_attention = l2_attention
        self.scale_in = scale_in
        self.mask_self = mask_self

        inner_dim = head_dim * num_heads
        self.scale = head_dim ** -0.5

        self.to_q = DenseLayer(in_channels=query_dim,
                               out_channels=inner_dim,
                               add_bias=False,
                               init_bias=0.,
                               use_wscale=use_wscale,
                               wscale_gain=wscale_gain,
                               lr_mul=lr_mul,
                               activation_type='linear')

        self.to_kv = DenseLayer(in_channels=context_dim,
                                out_channels=inner_dim * 2,
                                add_bias=False,
                                init_bias=0.,
                                use_wscale=use_wscale,
                                wscale_gain=wscale_gain,
                                lr_mul=lr_mul,
                                activation_type='linear')

        self.proj_out = nn.Sequential(
            DenseLayer(in_channels=inner_dim,
                       out_channels=query_dim,
                       add_bias=True,
                       init_bias=0.,
                       use_wscale=use_wscale,
                       wscale_gain=wscale_gain,
                       lr_mul=lr_mul,
                       activation_type='linear'),
            nn.Dropout(dropout)
        )

    def extra_repr(self):
        return (f'query_dim={self.query_dim}, '
                f'context_dim={self.context_dim}, '
                f'num_heads={self.num_heads}, '
                f'head_dim={self.head_dim}, '
                f'dropout={self.dropout}, '
                f'wscale_gain={self.wscale_gain:.3f}, '
                f'lr_mul={self.lr_mul:.3f}')

    def forward(self, x, context):
        return checkpoint(self._forward,
                          (x, context),
                          self.parameters(),
                          self.use_checkpoint)

    def _forward(self, x, context):
        q = self.to_q(x) # b, x*y, c (num_heads * head_dim)
        kv = self.to_kv(context)
        k, v = kv.chunk(2, dim=2)
        q, k, v = map(lambda t: rearrange(t, 'b x (h d) -> (b h) x d',
                                          h=self.num_heads), (q, k, v))
        if self.scale_in:
            q = q * self.scale
            k = k * self.scale

        if self.l2_attention:
            attn_logits =  - torch.cdist(q, k, p=2)
            if self.mask_self:
                mask = torch.eye(attn_logits.shape[1],
                                 device=attn_logits.device,
                                 dtype=torch.bool)
                attn_logits = attn_logits.masked_fill(mask, float('-inf'))
        else:
            attn_logits = torch.einsum('b i d, b j d -> b i j', q, k)

        if not self.scale_in:
            attn_logits = attn_logits * self.scale

        attn = torch.softmax(attn_logits, dim=-1)
        out = torch.einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h=self.num_heads)
        out = self.proj_out(out)

        return out


class ResidualSparseAttnBlock(nn.Module):
    """Implements the sparse attention used in the generator.

    Attention is composed by several `SparseAttentionBlock`
    """
    def __init__(self) -> None:
        super().__init__()

    def forward(self):
        pass


class ResidualDownsampleAttnBlock(nn.Module):
    """Implements the downsampling attention used in the generator.

    Attention is composed by several `DownsampleAttentionBlock`
    """
    def __init__(self) -> None:
        super().__init__()

    def forward(self):
        pass


class ResidualLongNetAttnBlock(nn.Module):
    """Implements the longnet attention used in the generator.

    Attention is composed by several `LongNetAttentionBlock`
    """
    def __init__(self) -> None:
        super().__init__()

    def forward(self):
        pass



class MultiHeadAttention(nn.Module):
    """Implements the multi-head attention layer with linear layer."""
    def __init__(self,
                 in_dim,
                 num_heads,
                 dropout,
                 add_bias,
                 init_bias,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 use_checkpoint,
                 l2_attention,
                 scale_in,
                 mask_self):
        """Initializes with layer settings.

        Args:
            in_dim: Number of channels of the input tensor.
            num_heads: Number of attention heads.
            dropout: Dropout ratio.
            add_bias: Whether to add bias onto the fully-connected result.
            init_bias: The initial bias value before training.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            use_checkpoint: if True, use gradient checkpointing on this module.
            l2_attention: Whether or not to use l2 distance to compute the
                similarity.
            scale_in:
            mask_self:
        """
        super().__init__()
        self.in_dim = in_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.add_bias = add_bias
        self.init_bias = init_bias
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.use_checkpoint = use_checkpoint
        self.l2_attention = l2_attention
        self.scale_in = scale_in
        self.mask_self = mask_self

        self.head_dim = in_dim // num_heads

        self.scale = self.head_dim ** -0.5

        self.to_qkv = DenseLayer(in_channels=in_dim,
                                 out_channels=in_dim * 3,
                                 add_bias=add_bias,
                                 init_bias=init_bias,
                                 use_wscale=use_wscale,
                                 wscale_gain=wscale_gain,
                                 lr_mul=lr_mul,
                                 activation_type='linear')

        self.to_out = nn.Sequential(
            DenseLayer(in_channels=in_dim,
                       out_channels=in_dim,
                       add_bias=add_bias,
                       init_bias=init_bias,
                       use_wscale=use_wscale,
                       wscale_gain=wscale_gain,
                       lr_mul=lr_mul,
                       activation_type='linear'),
            nn.Dropout(dropout)
        )

    def extra_repr(self):
        return (f'in_dim={self.in_dim}, '
                f'num_heads={self.num_heads}, '
                f'head_dim={self.head_dim}, '
                f'init_bias={self.init_bias}, '
                f'dropout={self.dropout}, '
                f'wscale_gain={self.wscale_gain:.3f}, '
                f'lr_mul={self.lr_mul:.3f}, '
                f'l2_attantion={self.l2_attention}, '
                f'scale_in={self.scale_in}, '
                f'mask_self={self.mask_self}')

    def forward(self, x):
        return checkpoint(self._forward,
                          (x,),
                          self.parameters(),
                          self.use_checkpoint)

    def _forward(self, x):
        qkv = self.to_qkv(x)
        q, k, v = qkv.chunk(3, dim=2)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d',
                                          h=self.num_heads), (q, k, v))
        if self.scale_in:
            q = q * self.scale
            k = k * self.scale

        if self.l2_attention:
            attn_logits = - torch.cdist(q, k, p=2)
            if self.mask_self:
                mask = torch.eye(attn_logits.shape[1],
                                 device=attn_logits.device,
                                 dtype=torch.bool)
                attn_logits = attn_logits.masked_fill(mask, float('-inf'))
        else:
            attn_logits = torch.einsum('b i d, b j d -> b i j', q, k)

        if not self.scale_in:
            attn_logits = attn_logits * self.scale_in

        # attention, what we cannot get enough of
        attn = torch.softmax(attn_logits, dim=-1)

        out = torch.einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h=self.num_heads)
        out = self.to_out(out)

        return out


class ResidualAttentionBlock(nn.Module):
    """Residual attention block in transformer."""
    def __init__(self,
                 in_dim,
                 num_heads,
                 dropout,
                 add_bias,
                 init_bias,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 use_checkpoint,
                 mlp_ratio,
                 l2_attention,
                 scale_in,
                 mask_self,
                 attn_gain):
        """Initializes with layer settings.

        Args:
            in_dim: Number of channels of the input tensor.
            num_heads: Number of attention heads.
            dropout: Dropout ratio.
            add_bias: Whether to add bias onto the fully-connected result.
            init_bias: The initial bias value before training.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            use_checkpoint: if True, use gradient checkpointing on this module.
            mlp_ratio: A factor used in increasing the dimension of MLP.
            l2_attention:
            scale_in:
            mask_self:
            attn_gain
        """
        super().__init__()
        self.attn_gain = attn_gain

        self.attn = MultiHeadAttention(in_dim=in_dim,
                                       num_heads=num_heads,
                                       dropout=dropout,
                                       add_bias=add_bias,
                                       init_bias=init_bias,
                                       use_wscale=use_wscale,
                                       wscale_gain=wscale_gain,
                                       lr_mul=lr_mul,
                                       use_checkpoint=use_checkpoint,
                                       l2_attention=l2_attention,
                                       scale_in=scale_in,
                                       mask_self=mask_self)

        self.c_fc = DenseLayer(in_channels=in_dim,
                               out_channels=in_dim * mlp_ratio,
                               add_bias=add_bias,
                               init_bias=init_bias,
                               use_wscale=use_wscale,
                               wscale_gain=wscale_gain,
                               lr_mul=lr_mul,
                               activation_type='linear')

        self.c_proj = DenseLayer(in_channels=in_dim * mlp_ratio,
                                 out_channels=in_dim,
                                 add_bias=add_bias,
                                 init_bias=init_bias,
                                 use_wscale=use_wscale,
                                 wscale_gain=wscale_gain,
                                 lr_mul=lr_mul,
                                 activation_type='linear')

        self.ln_1 = LayerNorm(in_dim)
        self.ln_2 = LayerNorm(in_dim)

        self.mlp = nn.Sequential(OrderedDict([('c_fc', self.c_fc),
                                              ('gelu', QuickGELU()),
                                              ('c_proj', self.c_proj)]))

    def forward(self, x):
        x = (x + self.attn(self.ln_1(x))) * self.attn_gain
        x = (x + self.mlp(self.ln_2(x))) * self.attn_gain
        return x


class TextTransformer(nn.Module):
    """TextTransformer implementation."""
    def __init__(self,
                 width,
                 num_heads,
                 num_layers,
                 dropout,
                 add_bias,
                 init_bias,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 use_checkpoint,
                 mlp_ratio,
                 l2_attention,
                 scale_in,
                 mask_self,
                 attn_gain):
        """Initializes with layer settings.

        Args:
            width: Total dimension of the model.
            num_heads: Number of parallel attention heads.
            num_layers: Number of transformer blocks.
            dropout: Dropout ratio.
            add_bias: Whether to add bias onto the fully-connected result.
            init_bias: The initial bias value before training.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            use_checkpoint: if True, use gradient checkpoint
            activation_type: Type of activation.
            mlp_ratio: Ratio used in increases the dimension of MLP.
            l2_attention:
            scale_in:
            mask_self:
        """
        super().__init__()
        self.width = width
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.dropout = dropout
        self.add_bias = add_bias
        self.init_bias = init_bias
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.use_checkpoint = use_checkpoint
        self.mlp_ratio = mlp_ratio
        self.l2_attention = l2_attention
        self.scale_in = scale_in
        self.mask_self = mask_self
        self.attn_gain = attn_gain

        net_blocks = []
        for _ in range(self.num_layers):
            net_blocks.append(
                ResidualAttentionBlock(in_dim=width,
                                       num_heads=num_heads,
                                       dropout=dropout,
                                       add_bias=add_bias,
                                       init_bias=init_bias,
                                       use_wscale=use_wscale,
                                       wscale_gain=wscale_gain,
                                       lr_mul=lr_mul,
                                       use_checkpoint=use_checkpoint,
                                       mlp_ratio=mlp_ratio,
                                       l2_attention=l2_attention,
                                       scale_in=scale_in,
                                       mask_self=mask_self,
                                       attn_gain=attn_gain))

        self.res_blocks = nn.Sequential(*net_blocks)

    def extra_repr(self) -> str:
        return (f'width={self.width}, '
                f'num_heads={self.num_heads}, '
                f'num_layers={self.num_layers}, '
                f'dropout={self.dropout}, '
                f'bias={self.add_bias}, '
                f'init_bias={self.init_bias}, '
                f'use_wscale={self.use_wscale}, '
                f'wscale_gain={self.wscale_gain:.3f}, '
                f'lr_mul={self.lr_mul:.3f}, '
                f'mlp_ration={self.mlp_ratio}')

    def forward(self, x):
        if self.num_layers == 0:
            return x
        return self.res_blocks(x)


class AdaptiveModulateConvLayer(nn.Module):
    """Implements the adaptive convolutional layer with style modulation.

    When num_kernels equals to 1, it is the same as the conventional
    ModulateConvLayer.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 modulate_dim,
                 kernel_size,
                 num_kernels,
                 add_bias,
                 scale_factor,
                 filter_kernel,
                 demodulate,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 activation_type,
                 conv_clamp,
                 eps,
                 mtm=False):
        """Initializes with layer settings.

        Args:
            in_channels: Number of channels of the input tensor.
            out_channels: Number of channels of the output tensor.
            modulate_dim: Dimension of the input modulate code.
            kernel_size: Size of the convolutional kernels.
            num_kernels: Number of convolution filters used.
            add_bias: Whether to add bias onto the convolutional result.
            scale_factor: Scale factor for upsampling.
            filter_kernel: Kernel used for filtering.
            demodulate: Whether to perform style demodulation.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            activation_type: Type of activation.
            conv_clamp: A threshold to clamp the output of convolution layers to
                avoid overflow under FP16 training.
            eps: A small value to avoid divide overflow.
            mtm: Whether to apply modulated tranformation module
        """
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modulate_dim = modulate_dim
        self.kernel_size = kernel_size
        self.num_kernels = num_kernels
        self.add_bias = add_bias
        self.scale_factor = scale_factor
        self.filter_kernel = filter_kernel
        self.demodulate = demodulate
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.activation_type = activation_type
        self.conv_clamp = conv_clamp
        self.eps = eps
        self.mtm = mtm
        self.adaptive = self.num_kernels > 1

        # Set up weight.
        weight_shape = (num_kernels, out_channels, in_channels,
                        kernel_size, kernel_size)
        fan_in = kernel_size * kernel_size * in_channels
        wscale = wscale_gain / np.sqrt(fan_in)
        if use_wscale:
            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
            self.wscale = wscale * lr_mul
        else:
            self.weight = nn.Parameter(
                    torch.randn(*weight_shape) * wscale / lr_mul)
            self.wscale = lr_mul

        # Set up bias.
        if add_bias:
            self.bias = nn.Parameter(torch.zeros(out_channels))
            self.bscale = lr_mul
        else:
            self.bias = None
        self.act_gain = bias_act.activation_funcs[activation_type].def_gain

        # Set up style.
        self.style = DenseLayer(in_channels=modulate_dim,
                                out_channels=in_channels,
                                add_bias=True,
                                init_bias=1.0,
                                use_wscale=use_wscale,
                                wscale_gain=wscale_gain,
                                lr_mul=lr_mul,
                                activation_type='linear')
        if self.adaptive:
            # Set up weight selection
            self.filter_weight = DenseLayer(in_channels=modulate_dim,
                                            out_channels=num_kernels,
                                            add_bias=True,
                                            init_bias=1.0,
                                            use_wscale=use_wscale,
                                            wscale_gain=wscale_gain,
                                            lr_mul=lr_mul,
                                            activation_type='linear')
        if scale_factor > 1:
            assert filter_kernel is not None
            self.register_buffer(
                'filter', upfirdn2d.setup_filter(filter_kernel))
            fh, fw = self.filter.shape
            self.filter_padding = (
                kernel_size // 2 + (fw + scale_factor - 1) // 2,
                kernel_size // 2 + (fw - scale_factor) // 2,
                kernel_size // 2 + (fh + scale_factor - 1) // 2,
                kernel_size // 2 + (fh - scale_factor) // 2)

        if self.mtm:
            self.mtm_style = DenseLayer(in_channels=modulate_dim,
                                        out_channels=in_channels,
                                        add_bias=True,
                                        init_bias=1.0,
                                        use_wscale=use_wscale,
                                        wscale_gain=wscale_gain,
                                        lr_mul=lr_mul,
                                        activation_type='linear')
            mtm_out_channels = 3 * kernel_size * kernel_size
            self.mtm_weight = nn.Parameter(torch.randn([mtm_out_channels, in_channels, kernel_size, kernel_size]))
            self.mtm_bias = nn.Parameter(torch.zeros([1, mtm_out_channels, 1, 1]))
            self.mtm_scales = nn.Parameter(torch.zeros([1, mtm_out_channels, 1, 1]))

    def extra_repr(self):
        return (f'in_ch={self.in_channels}, '
                f'out_ch={self.out_channels}, '
                f'modulate_dim={self.modulate_dim}, '
                f'ksize={self.kernel_size}, '
                f'num_kernel={self.num_kernels}, '
                f'wscale_gain={self.wscale_gain:.3f}, '
                f'bias={self.add_bias}, '
                f'lr_mul={self.lr_mul:.3f}, '
                f'upsample={self.scale_factor}, '
                f'upsample_filter={self.filter_kernel}, '
                f'demodulate={self.demodulate}, '
                f'act={self.activation_type}, '
                f'mtm={self.mtm}, ')

    def get_offsets(self, x, modulate_code):
        dtype = x.dtype
        N, _, H, W = x.shape
        mtm_weight = self.mtm_weight

        out_ch, in_ch, kh, kw = mtm_weight.shape
        mtm_style = self.mtm_style(modulate_code)
        wscale = 1.0 / np.sqrt(in_ch * kh * kw)
        if not self.demodulate:
            _style = mtm_style * wscale  # Equivalent to scaling weight.
        else:
            _style = mtm_style

        # Pre-normalize inputs.
        if self.demodulate:
            mtm_weight = mtm_weight.unsqueeze(0)
            mtm_weight = (mtm_weight *
                      mtm_weight.square().mean(dim=(2, 3, 4), keepdim=True).rsqrt())
            _style = _style * _style.square().mean().rsqrt()

        mtm_weight = mtm_weight * _style.reshape(N, 1, in_ch, 1, 1)  # modulation

        if self.demodulate:
            decoef = (mtm_weight.square().sum(dim=(2, 3, 4)) + self.eps).rsqrt()
            mtm_weight = mtm_weight * decoef.reshape(N, out_ch, 1, 1, 1)  # demodulation

        # Fuse `conv` and `style modulation` as one op, using group convolution.
        x = x.reshape(1, N * in_ch, H, W)
        w = mtm_weight.reshape(N * out_ch, in_ch, kh, kw).to(dtype)
        offset = conv2d_gradfix.conv2d(
                x, w, stride=1, padding=kh//2, groups=N)
        offset = offset.reshape(N, -1, *offset.shape[2:])
        mtm_bias = self.mtm_bias.to(dtype)
        mtm_scales = self.mtm_scales.to(dtype)
        offset = offset * mtm_scales + mtm_bias
        offset_x, offset_y, offsets_mask = torch.chunk(offset, 3, dim=1)
        offsets = torch.cat([offset_x, offset_y], 1)
        offsets_mask = offsets_mask.sigmoid() * 2 # range in [0, 2]
        return offsets, offsets_mask


    def forward(self, x, modulate_code):
        dtype = x.dtype
        N, C, H, W = x.shape

        # calculate offsets
        if self.mtm:
            offsets, offsets_mask = self.get_offsets(x, modulate_code)

        if self.adaptive:
            # Affine on `modulate_code` to get the weight for each filter.
            filter_weights  = self.filter_weight(modulate_code)
            filter_weights = torch.softmax(filter_weights, dim=-1)
            filter_weights = filter_weights[:, :, None, None, None, None]
            weight = repeat(self.weight, '... -> b ...', b=N)
            weight = reduce(weight * filter_weights, 'b n ... -> b ...', 'sum')
        else:
            weight = self.weight

        _, out_ch, in_ch, kh, kw = weight.shape
        assert in_ch == C

        # Affine on `modulate_code`.
        style = self.style(modulate_code)
        if not self.demodulate:
            _style = style * self.wscale  # Equivalent to scaling weight.
        else:
            _style = style

        # Pre-normalize inputs.
        if self.demodulate:
            weight = (weight *
                      weight.square().mean(dim=(2, 3, 4), keepdim=True).rsqrt())
            _style = _style * _style.square().mean().rsqrt()

        weight = weight * _style.reshape(N, 1, in_ch, 1, 1)  # modulation

        if self.demodulate:
            decoef = (weight.square().sum(dim=(2, 3, 4)) + self.eps).rsqrt()
            weight = weight * decoef.reshape(N, out_ch, 1, 1, 1)  # demodulation

        # Fuse `conv` and `style modulation` as one op, using group convolution.
        x = x.reshape(1, N * in_ch, H, W)
        w = weight.reshape(N * out_ch, in_ch, kh, kw).to(dtype)

        if self.scale_factor == 1:  # Native convolution without upsampling.
            padding = self.kernel_size // 2
            if self.mtm: # Perform modulated transformation module
                offsets = offsets.reshape(1, -1, *offsets.shape[2:])
                offsets_mask = offsets_mask.reshape(1, -1, *offsets_mask.shape[2:])
                x = torchvision.ops.deform_conv2d(x, offsets, w,
                    mask=offsets_mask, padding=padding)
            else:
                x = conv2d_gradfix.conv2d(
                    x, w, stride=1, padding=padding, groups=N)
        else:  # Convolution with upsampling.
            up = self.scale_factor
            f = self.filter
            # When kernel size = 1, use filtering function for upsampling.
            if self.kernel_size == 1:
                padding = self.filter_padding
                x = conv2d_gradfix.conv2d(
                    x, w, stride=1, padding=0, groups=N)
                x = upfirdn2d.upfirdn2d(
                    x, f, up=up, padding=padding, gain=up ** 2)
            # When kernel size != 1, use stride convolution for upsampling.
            else:
                # Following codes are borrowed from
                # https://github.com/NVlabs/stylegan2-ada-pytorch
                px0, px1, py0, py1 = self.filter_padding
                px0 = px0 - (kw - 1)
                px1 = px1 - (kw - up)
                py0 = py0 - (kh - 1)
                py1 = py1 - (kh - up)
                pxt = max(min(-px0, -px1), 0)
                pyt = max(min(-py0, -py1), 0)
                if N == 1:
                    w = w.transpose(0, 1)
                else:
                    w = w.reshape(N, out_ch, in_ch, kh, kw)
                    w = w.transpose(1, 2)
                    w = w.reshape(N * in_ch, out_ch, kh, kw)
                padding = (pyt, pxt)
                x = conv2d_gradfix.conv_transpose2d(
                    x, w, stride=up, padding=padding, groups=N)
                padding = (px0 + pxt, px1 + pxt, py0 + pyt, py1 + pyt)
                x = upfirdn2d.upfirdn2d(
                    x, f, up=1, padding=padding, gain=up ** 2)

        x = x.reshape(N, out_ch, x.shape[2], x.shape[3])

        bias = None
        if self.bias is not None:
            bias = self.bias.to(dtype)
            if self.bscale != 1.0:
                bias = bias * self.bscale

        if self.activation_type == 'linear':  # Shortcut for output layer.
            x = bias_act.bias_act(
                x, bias, act='linear', clamp=self.conv_clamp)
        else:
            act_gain = self.act_gain
            act_clamp = None
            if self.conv_clamp is not None:
                act_clamp = self.conv_clamp
            x = bias_act.bias_act(x, bias,
                                  act=self.activation_type,
                                  gain=act_gain,
                                  clamp=act_clamp)

        assert x.dtype == dtype
        assert style.dtype == torch.float32
        return x, style


class ConvLayer(nn.Module):
    """Implements the convolutional layer.

    If upsampling is needed (i.e., `scale_factor = 2`), the feature map will
    be filtered with `filter_kernel` after convolution. This layer will only be
    used for skip connection in `resnet` architecture.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 add_bias,
                 scale_factor,
                 filter_kernel,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 activation_type,
                 conv_clamp):
        """Initializes with layer settings.

        Args:
            in_channels: Number of channels of the input tensor.
            out_channels: Number of channels of the output tensor.
            kernel_size: Size of the convolutional kernels.
            add_bias: Whether to add bias onto the convolutional result.
            scale_factor: Scale factor for upsampling.
            filter_kernel: Kernel used for filtering.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            activation_type: Type of activation.
            conv_clamp: A threshold to clamp the output of convolution layers to
                avoid overflow under FP16 training.
        """
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.add_bias = add_bias
        self.scale_factor = scale_factor
        self.filter_kernel = filter_kernel
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.activation_type = activation_type
        self.conv_clamp = conv_clamp

        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
        fan_in = kernel_size * kernel_size * in_channels
        wscale = wscale_gain / np.sqrt(fan_in)
        if use_wscale:
            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
            self.wscale = wscale * lr_mul
        else:
            self.weight = nn.Parameter(
                torch.randn(*weight_shape) * wscale / lr_mul)
            self.wscale = lr_mul

        if add_bias:
            self.bias = nn.Parameter(torch.zeros(out_channels))
            self.bscale = lr_mul
        else:
            self.bias = None
        self.act_gain = bias_act.activation_funcs[activation_type].def_gain

        if scale_factor > 1:
            assert filter_kernel is not None
            self.register_buffer(
                'filter', upfirdn2d.setup_filter(filter_kernel))
            fh, fw = self.filter.shape
            self.filter_padding = (
                kernel_size // 2 + (fw + scale_factor - 1) // 2,
                kernel_size // 2 + (fw - scale_factor) // 2,
                kernel_size // 2 + (fh + scale_factor - 1) // 2,
                kernel_size // 2 + (fh - scale_factor) // 2)

    def extra_repr(self):
        return (f'in_ch={self.in_channels}, '
                f'out_ch={self.out_channels}, '
                f'ksize={self.kernel_size}, '
                f'wscale_gain={self.wscale_gain:.3f}, '
                f'bias={self.add_bias}, '
                f'lr_mul={self.lr_mul:.3f}, '
                f'upsample={self.scale_factor}, '
                f'upsample_filter={self.filter_kernel}, '
                f'act={self.activation_type}, '
                f'clamp={self.conv_clamp}')

    def forward(self, x):
        dtype = x.dtype

        weight = self.weight
        if self.wscale != 1.0:
            weight = weight * self.wscale
        bias = None
        if self.bias is not None:
            bias = self.bias.to(dtype)
            if self.bscale != 1.0:
                bias = bias * self.bscale

        if self.scale_factor == 1:  # Native convolution without upsampling.
            padding = self.kernel_size // 2
            x = conv2d_gradfix.conv2d(
                x, weight.to(dtype), stride=1, padding=padding)
        else:  # Convolution with upsampling.
            up = self.scale_factor
            f = self.filter
            # When kernel size = 1, use filtering function for upsampling.
            if self.kernel_size == 1:
                padding = self.filter_padding
                x = conv2d_gradfix.conv2d(
                    x, weight.to(dtype), stride=1, padding=0)
                x = upfirdn2d.upfirdn2d(
                    x, f, up=up, padding=padding, gain=up ** 2)
            # When kernel size != 1, use transpose convolution for upsampling.
            else:
                # Following codes are borrowed from
                # https://github.com/NVlabs/stylegan2-ada-pytorch
                px0, px1, py0, py1 = self.filter_padding
                kh, kw = weight.shape[2:]
                px0 = px0 - (kw - 1)
                px1 = px1 - (kw - up)
                py0 = py0 - (kh - 1)
                py1 = py1 - (kh - up)
                pxt = max(min(-px0, -px1), 0)
                pyt = max(min(-py0, -py1), 0)
                weight = weight.transpose(0, 1)
                padding = (pyt, pxt)
                x = conv2d_gradfix.conv_transpose2d(
                    x, weight.to(dtype), stride=up, padding=padding)
                padding = (px0 + pxt, px1 + pxt, py0 + pyt, py1 + pyt)
                x = upfirdn2d.upfirdn2d(
                    x, f, up=1, padding=padding, gain=up ** 2)

        act_gain = self.act_gain
        act_clamp = None
        if self.conv_clamp is not None:
            act_clamp = self.conv_clamp
        x = bias_act.bias_act(x, bias,
                              act=self.activation_type,
                              gain=act_gain,
                              clamp=act_clamp)

        assert x.dtype == dtype
        return x


class QuickGELU(nn.Module):
    """
    Applies GELU approximation that is fast but somewhat inaccurate.

    See: https://github.com/hendrycks/GELUs.
    """
    def forward(self, x):
        return x * torch.sigmoid(1.702 * x)


class GEGLU(nn.Module):
    """ A variant of the gated linear unit activation function
        from https://arxiv.org/abs/2002.05202.
    """
    def __init__(self,
                 in_channels,
                 out_channels,
                 add_bias,
                 init_bias,
                 use_wscale,
                 wscale_gain,
                 lr_mul):
        """Initializes with layer settings.

        Args:
            in_channels: Number of channels of the input tensor.
            out_channels: Number of channels of the output tensor.
            add_bias: Whether to add bias onto the fully-connected result.
            init_bias: The initial bias value before training.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
        """
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.add_bias = add_bias
        self.init_bias = init_bias
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul

        self.proj = DenseLayer(in_channels=in_channels,
                               out_channels=out_channels * 2,
                               add_bias=add_bias,
                               init_bias=init_bias,
                               use_wscale=use_wscale,
                               wscale_gain=wscale_gain,
                               lr_mul=lr_mul,
                               activation_type='linear')

    def forward(self, x):
        x, gate = self.proj(x).chunk(2, dim=-1)
        return x * F.gelu(gate)


class PixelNormLayer(nn.Module):
    """Implements pixel-wise feature vector normalization layer."""

    def __init__(self, dim, eps):
        super().__init__()
        self.dim = dim
        self.eps = eps

    def extra_repr(self):
        return f'dim={self.dim}, epsilon={self.eps}'

    def forward(self, x):
        scale = (x.square().mean(dim=self.dim, keepdim=True) + self.eps).rsqrt()
        return x * scale


class LayerNorm(nn.LayerNorm):
    """Subclass torch's LayerNorm to handle fp16."""
    def __init__(self, embedding_dim):
        super().__init__(embedding_dim)
        self.embedding_dim = embedding_dim

    def extra_repr(self):
        return f'embedding_dim={self.embedding_dim}'

    def forward(self, x):
        orig_type = x.dtype
        ret = super().forward(x.type(torch.float32))
        return ret.type(orig_type)


class GroupNorm(nn.GroupNorm):
    """Implement of the group norm 32"""
    def __init__(self, num_groups, num_channels):
        super().__init__(num_groups, num_channels)
        self.num_groups = num_groups
        self.num_channels = num_channels

    def extra_repr(self):
        return (f'num_group={self.num_groups}, '
                f'num_channels={self.num_channels}')

    def forward(self, x):
        return super().forward(x.float()).type(x.dtype)


class DenseLayer(nn.Module):
    """Implements the dense layer."""

    def __init__(self,
                 in_channels,
                 out_channels,
                 add_bias,
                 init_bias,
                 use_wscale,
                 wscale_gain,
                 lr_mul,
                 activation_type,
                 init_weight=1.0):
        """Initializes with layer settings.

        Args:
            in_channels: Number of channels of the input tensor.
            out_channels: Number of channels of the output tensor.
            add_bias: Whether to add bias onto the fully-connected result.
            init_bias: The initial bias value before training.
            use_wscale: Whether to use weight scaling.
            wscale_gain: Gain factor for weight scaling.
            lr_mul: Learning multiplier for both weight and bias.
            activation_type: Type of activation.
        """
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.add_bias = add_bias
        self.init_bias = init_bias
        self.use_wscale = use_wscale
        self.wscale_gain = wscale_gain
        self.lr_mul = lr_mul
        self.activation_type = activation_type

        weight_shape = (out_channels, in_channels)
        wscale = wscale_gain / np.sqrt(in_channels)
        if use_wscale:
            self.weight = nn.Parameter(torch.randn(*weight_shape) * init_weight / lr_mul)
            self.wscale = wscale * lr_mul
        else:
            self.weight = nn.Parameter(
                torch.randn(*weight_shape) * wscale * init_weight / lr_mul)
            self.wscale = lr_mul

        if add_bias:
            init_bias = np.broadcast_to(np.asarray(init_bias, dtype=np.float32), [out_channels])
            self.bias = nn.Parameter(torch.from_numpy(init_bias / lr_mul))
            self.bscale = lr_mul
        else:
            self.bias = None

    def extra_repr(self):
        return (f'in_ch={self.in_channels}, '
                f'out_ch={self.out_channels}, '
                f'wscale_gain={self.wscale_gain:.3f}, '
                f'bias={self.add_bias}, '
                f'init_bias={self.init_bias}, '
                f'lr_mul={self.lr_mul:.3f}, '
                f'act={self.activation_type}')

    def forward(self, x):
        dtype = x.dtype

        weight = self.weight.to(dtype) * self.wscale
        bias = None
        if self.bias is not None:
            bias = self.bias.to(dtype)
            if self.bscale != 1.0:
                bias = bias * self.bscale

        # Fast pass for two dimension input.
        if x.dim() == 2 and bias is not None:
            x = torch.addmm(bias.unsqueeze(0), x, weight.t())
            if self.activation_type != 'linear':
                x = bias_act.bias_act(x,
                                      b=None,
                                      dim=1,
                                      act=self.activation_type)
        else:
            x = x.matmul(weight.t())
            x = bias_act.bias_act(x, bias,
                                  dim=2,
                                  act=self.activation_type)

        assert x.dtype == dtype
        return x


class CheckpointFunction(torch.autograd.Function):
    """CheckpointFunction """
    @staticmethod
    def forward(ctx, run_function, length, *args):
        ctx.run_function = run_function
        ctx.input_tensors = list(args[:length])
        ctx.input_params = list(args[length:])
        with torch.no_grad():
            output_tensors = ctx.run_function(*ctx.input_tensors)
        return output_tensors

    @staticmethod
    def backward(ctx, *output_grads):
        ctx.input_tensors = [x.detach().requires_grad_(True)
                             for x in ctx.input_tensors]
        with torch.enable_grad():
            # Fixes a bug where the first op in run_function modifies the
            # Tensor storage in place, which is not allowed for detach()'d
            # Tensors.
            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
            output_tensors = ctx.run_function(*shallow_copies)
        input_grads = torch.autograd.grad(
            output_tensors,
            ctx.input_tensors + ctx.input_params,
            output_grads,
            allow_unused=True,
        )
        del ctx.input_tensors
        del ctx.input_params
        del output_tensors
        return (None, None) + input_grads


# list two loss for MoE training
def router_z_loss_func(router_logits):
    r"""
    Compute the router z-loss implemented in PyTorch.

    The router z-loss was introduced in [Designing Effective Sparse Expert Models](https://arxiv.org/abs/2202.08906).
    It encourages router logits to remain small in an effort to improve stability.

    Args:
        router_logits (`float`):
            Input logits of shape [batch_size, sequence_length, num_experts]

    Returns:
        Scalar router z-loss.
    """
    num_groups, tokens_per_group, _ = router_logits.shape
    log_z = torch.logsumexp(router_logits, dim=-1)
    z_loss = log_z**2
    return torch.sum(z_loss) / (num_groups * tokens_per_group)


def load_balancing_loss_func(router_probs, expert_indices):
    r"""
    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.

    See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
    experts is too unbalanced.

    Args:
        router_probs (`torch.Tensor`):
            Probability assigned to each expert per token. Shape: [batch_size, seqeunce_length, num_experts].
        expert_indices (`torch.Tensor`):
            Indices tensor of shape [batch_size, seqeunce_length] identifying the selected expert for a given token.

    Returns:
        The auxiliary loss.
    """
    num_experts = router_probs.shape[-1]

    # cast the expert indices to int64, otherwise one-hot encoding will fail
    if expert_indices.dtype != torch.int64:
        expert_indices = expert_indices.to(torch.int64)

    if len(expert_indices.shape) == 2:
        expert_indices = expert_indices.unsqueeze(2)

    expert_mask = torch.nn.functional.one_hot(expert_indices, num_experts)

    # For a given token, determine if it was routed to a given expert.
    expert_mask = torch.max(expert_mask, axis=-2).values

    # cast to float32 otherwise mean will fail
    expert_mask = expert_mask.to(torch.float32)
    tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)

    router_prob_per_group_and_expert = torch.mean(router_probs, axis=-2)
    return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert) * (num_experts**2)

models/text_generator.py
zhujiapeng/Aurora

# python3.7
"""Contains the code to interpolate text embedding using a pre-trained models.
"""
import os
import argparse
import numpy as np
from tqdm import tqdm
import torch

from models import build_model
from utils.image_utils import postprocess_image
from utils.visualizers import HtmlVisualizer


def run_text_head(G, context, eot_ind=None):
    """Run the text head"""
    with torch.no_grad():
        global_text, local_text = G.text_head(context, eot_ind=eot_ind)
    return global_text, local_text


def run_mapping(G, z, global_text):
    """Run mapping network of the generator."""
    with torch.no_grad():
        mapping_results = G.mapping(z,
                                    label=None,
                                    context=global_text)
    return mapping_results['wp']


def run_synthesize(G, wp, local_text):
    """Run synthesis network of the generator."""
    with torch.no_grad():
        res = G.synthesis(wp, context=local_text)
    return res


def read_text(text_path):
    """Prepare snapshot text that will be used for evaluation."""
    print(f'Loading text from {text_path}')
    with open(text_path) as f:
        text = [line.strip() for line in f.readlines()]
    return text


def parse_float(arg):
    """Parse float number in string."""
    if not arg:
        return None
    arg = arg.split(',')
    arg = [float(i) for i in arg]
    return arg


def to_numpy(data):
    """Converts the input data to `numpy.ndarray`."""
    if isinstance(data, (int, float)):
        return np.array(data)
    if isinstance(data, np.ndarray):
        return data
    if isinstance(data, torch.Tensor):
        return data.detach().cpu().numpy()
    raise TypeError(f'Not supported data type `{type(data)}` for '
                    f'converting to `numpy.ndarray`!')


def linear_interpolate(src_code,
                       dst_code,
                       layer_index=None,
                       steps=7):
    """Interpolate between the latent code and boundary."""
    assert (len(src_code.shape) == 3 and len(dst_code.shape) == 3 and
            src_code.shape[0] == 1 and dst_code.shape[0] == 1 and
            src_code.shape[1] == dst_code.shape[1])
    if not layer_index:
        layer_index = list(range(src_code.shape[1]))
    linspace = np.linspace(0.0, 1.0, steps)
    linspace = linspace.reshape([-1, 1, 1]).astype(np.float32)
    inter_code = src_code + linspace * (dst_code - src_code)
    is_inter = np.zeros(inter_code.shape, dtype=bool)
    is_inter[:, layer_index, :] = True
    inter_code = np.where(is_inter, inter_code, src_code)
    return inter_code


def parse_args():
    """Parses arguments."""

    parser = argparse.ArgumentParser()
    parser.add_argument('weight_path', type=str, default='',
                        help='Path to the pre-trained models.')
    parser.add_argument('--src_prompt', type=str, default='',
                        help='The src text prompt, support reading from a file '
                             'or just given a text prompt.')
    parser.add_argument('--dst_prompt', type=str, default='',
                        help='The dst text prompt, support reading from a file '
                             'or just given a text prompt.')
    parser.add_argument('--num_src', type=int, default=10,
                        help='Number of src images.')
    parser.add_argument('--num_dst', type=int, default=10,
                        help='Number of dst images.')
    parser.add_argument('--resolution', type=int, default=64,
                        help='Resolution of the model output.')
    parser.add_argument('--results_dir', type=str, default='work_dirs/inter_res',
                        help='Results directory.')
    parser.add_argument('--seed', type=int, default=4,
                        help='Random seed.')
    parser.add_argument('--trunc_layers', type=int, default=None,
                        help='Number of layers to perform truncation.')
    parser.add_argument('--trunc_val', type=float, default=0.2,
                        help='Default value for truncation.')
    parser.add_argument('--loop_mapping', type=int, default=0,
                        help='Loop number for getting average for wp.')
    parser.add_argument('--save_name', type=str, default='0',
                        help='Name to help save the file.')
    parser.add_argument('--num_z', type=int, default=3,
                        help='Number of z for each text prompt.')
    parser.add_argument('--inter_layers', type=str, default=None,
                        help='The layers will be interpolated')
    parser.add_argument('--inter_step', type=int, default=9,
                        help='Number of interpolation steps. (default: 9)')
    parser.add_argument('--inter_g', action='store_true',
                        help='Whether to interp global code.')
    return parser.parse_args()


def main():
    """Main function."""
    args = parse_args()
    if torch.cuda.is_available():
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')
    assert args.src_prompt
    assert args.dst_prompt
    if os.path.exists(args.src_prompt):
        src_prompt = read_text(args.src_prompt)
    else:
        src_prompt = [args.src_prompt]
    if os.path.exists(args.dst_prompt):
        dst_prompt = read_text(args.dst_prompt)
    else:
        dst_prompt = [args.dst_prompt]
    num_src = min(args.num_src, len(src_prompt))
    num_dst = min(args.num_dst, len(dst_prompt))

    clip_config = {'model_name':'ViT-L-14',
                   'pretrained':'openai',
                   'freeze_clip': True}
    clip = build_model('CLIPModel', **clip_config)

    g_config = {'resolution': args.resolution,
                'image_channels': 3,
                'init_res': 4,
                'z_dim': 128,
                'w_dim': 1024,
                'mapping_fmaps': 1024,
                'label_dim': 0,
                'context_dim': 1024,
                'clip_out_dim': 768,
                'head_dim': 64,
                'embedding_dim': 1024,
                'use_text_cond': True,
                'num_layers_text_enc': 4,
                'use_w_cond': False,
                'use_class_label': False,
                'mapping_layers': 4,
                'fmaps_base': 16384,
                'fmaps_max': 1600,
                'num_adaptive_kernels': {"4":1,"8":1,"16":2,"32":4,"64":8},
                'num_block_per_res': {"4":3,"8":3,"16":3,"32":2,"64":2},
                'attn_resolutions': ['8', '16', '32', '64'],
                'attn_depth': {"8":2,"16":2,"32":2,"64":1},
                'attn_ch_factor': 1,
                'attn_gain': 0.3,
                'residual_gain': 0.4,
                'text_head_gain': 1.0,
                'zero_out': True,
                'fourier_feat': True,
                'l2_attention': True,
                'tie': False,
                'scale_in': False,
                'include_ff': True,
                'use_checkpoint': False,
                'checkpoint_res': ['8', '16', '32'],
                'mask_self': False,
                'conv_clamp': None,
                'mtm': True,
                'num_experts': {"8":4,"16":8,"32":16,"64":16},
                'ms_training_res': ['4','8','16','32','64'],
                'skip_connection': True}

    G = build_model('Text2ImageGenerator', **g_config)
    checkpoint = torch.load(args.weight_path, map_location='cpu')
    if 'generator_smooth' in checkpoint:
        print('Loading checkpoint from generator smooth!')
        G.load_state_dict(checkpoint['generator_smooth'])
    else:
        print('Loading checkpoint from generator!')
        G.load_state_dict(checkpoint['generator'])
    G = G.eval().to(device)

    trunc_layers = args.trunc_layers
    if not trunc_layers:
        trunc_layers = G.num_layers

    w_avg = G.w_avg.reshape(1, -1, G.w_dim)[:, :trunc_layers]
    num_rows = args.num_z * num_src * num_dst
    visualizer_syn = HtmlVisualizer(image_size=args.resolution)
    visualizer_syn.reset(num_rows=num_rows, num_cols=args.inter_step)
    head = [f'Step: {i:02d}' for i in range(1, args.inter_step + 1)]
    visualizer_syn.set_headers(head)
    torch.manual_seed(args.seed)
    os.makedirs(args.results_dir, exist_ok=True)

    for z_idx in tqdm(range(args.num_z)):
        latent_z = torch.randn((1, *G.latent_dim), device=device)
        zs_avg = torch.randn((args.loop_mapping, *G.latent_dim), device=device)
        for src_idx in range(num_src):
            src_text = src_prompt[src_idx]
            _, src_enc_text, src_eot_ind = clip.encode_text(text=src_text,
                                                            is_tokenize=True)
            src_g_text, src_l_text = run_text_head(G,
                                                   src_enc_text,
                                                   eot_ind=src_eot_ind)
            for dst_idx in range(num_dst):
                dst_text = dst_prompt[dst_idx]
                _, dst_enc_text, dst_eot_ind = clip.encode_text(text=dst_text,
                                                                is_tokenize=True)
                dst_g_text, dst_l_text = run_text_head(G,
                                                       dst_enc_text,
                                                       eot_ind=dst_eot_ind)
                inter_g_codes = linear_interpolate(
                                   src_code=to_numpy(src_g_text)[:, np.newaxis],
                                   dst_code=to_numpy(dst_g_text)[:, np.newaxis],
                                   steps=args.inter_step)
                inter_l_codes = linear_interpolate(
                                    src_code=to_numpy(src_l_text),
                                    dst_code=to_numpy(dst_l_text),
                                    steps=args.inter_step)
                inter_g_codes = torch.from_numpy(inter_g_codes).to(device)
                inter_l_codes = torch.from_numpy(inter_l_codes).to(device)
                row_ind = (z_idx * num_src + src_idx) * num_dst + dst_idx
                for idx in range(inter_l_codes.shape[0]):
                    if args.loop_mapping > 0:
                        sum_wp = 0
                        for i in range(args.loop_mapping):
                            if args.inter_g:
                                tmp_res = run_mapping(G,
                                                      zs_avg[i:i+1],
                                                      inter_g_codes[idx:idx+1, 0])
                            else:
                                tmp_res = run_mapping(G,
                                                      zs_avg[i:i+1],
                                                      src_g_text)
                            sum_wp += tmp_res
                        avg_wp = sum_wp / args.loop_mapping
                        avg_wp = avg_wp[:, :trunc_layers]
                    if args.inter_g:
                        wp_i = run_mapping(G,
                                           latent_z,
                                           inter_g_codes[idx:idx+1, 0])
                    else:
                        wp_i = run_mapping(G, latent_z, src_g_text)
                    wp_i[:, :trunc_layers] = w_avg.lerp(wp_i[:, :trunc_layers], args.trunc_val)
                    if args.loop_mapping > 0:
                        wp_i[:, :trunc_layers] = avg_wp.lerp(wp_i[:, :trunc_layers], args.trunc_val)
                    img_inter = run_synthesize(G,
                                               wp_i,
                                               inter_l_codes[idx:idx+1])
                    img_inter = postprocess_image(to_numpy(img_inter['image']))
                    if idx == 0:
                        text = src_text
                    elif idx == args.inter_step - 1:
                        text = dst_text
                    else:
                        text = 'inter_res'
                    visualizer_syn.set_cell(row_ind,
                                            idx,
                                            text=text,
                                            image=img_inter[0])
    # Save result.
    save_name_syn = f'inter_seed_{args.seed}_{args.save_name}.html'
    save_path_syn = os.path.join(args.results_dir, save_name_syn)
    visualizer_syn.save(save_path_syn)


if __name__ == '__main__':
    main()

run_interpolate.py
zhujiapeng/Aurora

# python3.7
"""Contains the code to synthesize images from a pre-trained models.
"""
import os
import argparse
from tqdm import tqdm

import torch
from models import build_model
from utils.image_utils import postprocess_image, save_image
from utils.visualizers import HtmlVisualizer


def run_mapping(G, z, context, eot_ind=None):
    """Run mapping network of the generator."""
    with torch.no_grad():
        global_text, local_text = G.text_head(context, eot_ind=eot_ind)
        mapping_results = G.mapping(z,
                                    label=None,
                                    context=global_text)
    return mapping_results['wp'], local_text


def run_synthesize(G, wp, local_text):
    """Run synthesis network of the generator."""
    with torch.no_grad():
        res = G.synthesis(wp, context=local_text)
    return res


def read_text(text_path):
    """Prepare snapshot text that will be used for evaluation."""
    print(f'Loading text from {text_path}')
    with open(text_path) as f:
        text = [line.strip() for line in f.readlines()]
    return text


def parse_float(arg):
    """Parse float number in string."""
    if not arg:
        return None
    arg = arg.split(',')
    arg = [float(i) for i in arg]
    return arg


def parse_args():
    """Parses arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument('weight_path', type=str, default='',
                        help='Path to the pre-trained models.')
    parser.add_argument('text_prompt', type=str, default='',
                        help='The text prompt, support reading from a file '
                             'or just given a text prompt.')
    parser.add_argument('--batch_size', type=int, default=1,
                        help='Batch size.')
    parser.add_argument('--syn_num', type=int, default=100,
                        help='Number of synthesized images.')
    parser.add_argument('--resolution', type=int, default=64,
                        help='Resolution of the model output.')
    parser.add_argument('--results_dir', type=str, default='work_dirs/syn_res',
                        help='Results directory.')
    parser.add_argument('--seed', type=int, default=4,
                        help='Random seed.')
    parser.add_argument('--trunc_layers', type=int, default=None,
                        help='Number of layers to perform truncation.')
    parser.add_argument('--loop_mapping', type=int, default=16,
                        help='Loop number for getting average for wp.')
    parser.add_argument('--save_name', type=str, default='0',
                        help='Name to help save the file.')
    parser.add_argument('--num_z', type=int, default=3,
                        help='Number of z for each text prompt.')
    parser.add_argument('--save_png', action='store_true',
                        help='Whether or not to save the synthesized images.')
    parser.add_argument('--trunc_vals', type=str, default='0,0.05,0.1,0.15,0.2',
                        help='Default values for truncation.')
    return parser.parse_args()


def main():
    """Main function."""
    args = parse_args()
    assert args.batch_size == 1, 'Current script only support bs equals to 1.'
    if os.path.exists(args.text_prompt):
        text_prompt = read_text(args.text_prompt)
    else:
        text_prompt = [args.text_prompt]
    syn_num = min(args.syn_num, len(text_prompt))
    if torch.cuda.is_available():
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')

    clip_config = {'model_name':'ViT-L-14',
                   'pretrained':'openai',
                   'freeze_clip': True}
    clip = build_model('CLIPModel', **clip_config)

    g_config = {'resolution': 64,
                'image_channels':3,
                'init_res': 4,
                'z_dim': 128,
                'w_dim': 1024,
                'mapping_fmaps': 1024,
                'label_dim': 0,
                'context_dim': 1024,
                'clip_out_dim': 768,
                'head_dim': 64,
                'embedding_dim': 1024,
                'use_text_cond': True,
                'num_layers_text_enc': 4,
                'use_w_cond': False,
                'use_class_label': False,
                'mapping_layers': 4,
                'fmaps_base': 16384,
                'fmaps_max': 1600,
                'num_adaptive_kernels': {"4":1,"8":1,"16":2,"32":4,"64":8},
                'num_block_per_res': {"4":3,"8":3,"16":3,"32":2,"64":2},
                'attn_resolutions': ['8', '16', '32', '64'],
                'attn_depth': {"8":2,"16":2,"32":2,"64":1},
                'attn_ch_factor': 1,
                'attn_gain': 0.3,
                'residual_gain': 0.4,
                'text_head_gain': 1.0,
                'zero_out': True,
                'fourier_feat': True,
                'l2_attention': True,
                'tie': False,
                'scale_in': False,
                'include_ff': True,
                'use_checkpoint': False,
                'checkpoint_res': ['8', '16', '32'],
                'mask_self': False,
                'conv_clamp': None,
                'mtm': True,
                'num_experts': {"8":4,"16":8,"32":16,"64":16},
                'ms_training_res': ['4','8','16','32','64'],
                'skip_connection': True}

    G = build_model('Text2ImageGenerator', **g_config)
    checkpoint = torch.load(args.weight_path, map_location='cpu')
    if 'generator_smooth' in checkpoint:
        print('Loading checkpoint from generator smooth!')
        G.load_state_dict(checkpoint['generator_smooth'])
    else:
        print('Loading checkpoint from generator!')
        G.load_state_dict(checkpoint['generator'])
    G = G.eval().to(device)

    trunc_vals = parse_float(args.trunc_vals)
    if not trunc_vals:
        trunc_vals = [0, 0.05, 0.1, 0.15, 0.2]
    trunc_layers = args.trunc_layers
    if not trunc_layers:
        trunc_layers = G.num_layers

    visualizer_syn = HtmlVisualizer(image_size=args.resolution)
    visualizer_syn.reset(num_rows=syn_num * args.num_z,
                         num_cols=len(trunc_vals) + 1)
    head = ['Number Z']
    head += [f'trunc_val_{val_}' for val_ in trunc_vals]
    visualizer_syn.set_headers(head)
    torch.manual_seed(args.seed)
    os.makedirs(args.results_dir, exist_ok=True)
    if args.save_png:
        os.makedirs(f'{args.results_dir}/images', exist_ok=True)
    w_avg = G.w_avg.reshape(1, -1, G.w_dim)[:, :args.trunc_layers]
    for idx in tqdm(range(syn_num)):
        text = text_prompt[idx]
        _, enc_text, eot_ind = clip.encode_text(text=text, is_tokenize=True)
        if args.loop_mapping > 0:
            sum_wp = 0
            for _ in range(args.loop_mapping):
                z = torch.randn((args.batch_size, *G.latent_dim), device=device)
                tmp_res, _ = run_mapping(G, z, enc_text, eot_ind=eot_ind)
                sum_wp += tmp_res
            avg_wp = sum_wp / args.loop_mapping
            avg_wp = avg_wp[:, :args.trunc_layers]
        for z_i in range(args.num_z):
            z = torch.randn((args.batch_size, *G.latent_dim), device=device)
            row_ind = idx * args.num_z + z_i
            visualizer_syn.set_cell(row_ind, 0, text=f'z_{z_i}')
            for col_idx, trunc_psi in enumerate(trunc_vals):
                wp, local_text = run_mapping(G, z, enc_text, eot_ind=eot_ind)
                wp[:, :args.trunc_layers] = w_avg.lerp(wp[:, :args.trunc_layers], trunc_psi)
                if args.loop_mapping > 0:
                    wp[:, :args.trunc_layers] = avg_wp.lerp(wp[:, :args.trunc_layers], trunc_psi)
                fake_results = run_synthesize(G, wp, local_text)
                syn_imgs = fake_results['image'].detach().cpu().numpy()
                syn_imgs = postprocess_image(syn_imgs)
                visualizer_syn.set_cell(row_ind,
                                        col_idx + 1,
                                        text=text,
                                        image=syn_imgs[0])
                if args.save_png:
                    prefix = f'{args.results_dir}/images/text_{idx:04d}_z_'
                    save_path0 = f'{prefix}{z_i:02d}_psi_{trunc_psi:-2.1f}.png'
                    save_image(save_path0, syn_imgs[0])

    # Save result.
    save_name_syn = f'syn_{syn_num:04d}_seed_{args.seed}_{args.save_name}.html'
    save_path_syn = os.path.join(args.results_dir, save_name_syn)
    visualizer_syn.save(save_path_syn)


if __name__ == '__main__':
    main()

run_synthesize.py
zhujiapeng/Aurora
# python3.7

# Copyright (c) 2021 OpenAI.  All rights reserved.

"""Text tokenizer.

Please refer to https://github.com/openai/CLIP
"""

# pylint: disable=line-too-long
# pylint: disable=inconsistent-quotes
# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring
# pylint: disable=bare-except
# pylint: disable=no-else-break

import gzip
import html
import os
from functools import lru_cache

import ftfy
import regex as re


@lru_cache()
def default_bpe():
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")


@lru_cache()
def bytes_to_unicode():
    """
    Returns list of utf-8 byte and a corresponding list of unicode strings.
    The reversible bpe codes work on unicode strings.
    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
    This is a signficant percentage of your normal, say, 32K bpe vocab.
    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
    And avoids mapping to whitespace/control characters the bpe code barfs on.
    """
    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8+n)
            n += 1
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))


def get_pairs(word):
    """Return set of symbol pairs in a word.
    Word is represented as tuple of symbols (symbols being variable-length strings).
    """
    pairs = set()
    prev_char = word[0]
    for char in word[1:]:
        pairs.add((prev_char, char))
        prev_char = char
    return pairs


def basic_clean(text):
    text = ftfy.fix_text(text)
    text = html.unescape(html.unescape(text))
    return text.strip()


def whitespace_clean(text):
    text = re.sub(r'\s+', ' ', text)
    text = text.strip()
    return text


class SimpleTokenizer(object):
    def __init__(self, bpe_path: str = default_bpe()):
        self.byte_encoder = bytes_to_unicode()
        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
        merges = merges[1:49152-256-2+1]
        merges = [tuple(merge.split()) for merge in merges]
        vocab = list(bytes_to_unicode().values())
        vocab = vocab + [v+'</w>' for v in vocab]
        for merge in merges:
            vocab.append(''.join(merge))
        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
        self.encoder = dict(zip(vocab, range(len(vocab))))
        self.decoder = {v: k for k, v in self.encoder.items()}
        self.bpe_ranks = dict(zip(merges, range(len(merges))))
        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)

    def bpe(self, token):
        if token in self.cache:
            return self.cache[token]
        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
        pairs = get_pairs(word)

        if not pairs:
            return token+'</w>'

        while True:
            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
            if bigram not in self.bpe_ranks:
                break
            first, second = bigram
            new_word = []
            i = 0
            while i < len(word):
                try:
                    j = word.index(first, i)
                    new_word.extend(word[i:j])
                    i = j
                except:
                    new_word.extend(word[i:])
                    break

                if word[i] == first and i < len(word)-1 and word[i+1] == second:
                    new_word.append(first+second)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_word = tuple(new_word)
            word = new_word
            if len(word) == 1:
                break
            else:
                pairs = get_pairs(word)
        word = ' '.join(word)
        self.cache[token] = word
        return word

    def encode(self, text):
        bpe_tokens = []
        text = whitespace_clean(basic_clean(text)).lower()
        for token in re.findall(self.pat, text):
            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
        return bpe_tokens

    def decode(self, tokens):
        text = ''.join([self.decoder[token] for token in tokens])
        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
        return text

# pylint: enable=line-too-long
# pylint: enable=inconsistent-quotes
# pylint: enable=missing-class-docstring
# pylint: enable=missing-function-docstring
# pylint: enable=bare-except
# pylint: enable=no-else-break

third_party/clip_official/simple_tokenizer.py
zhujiapeng/Aurora
# python3.7

# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

"""Custom replacement for convolution operators.

Operators in this file support arbitrarily high order gradients with zero
performance penalty. Please set `impl` as `cuda` to use faster customized
operators, OR as `ref` to use native `torch.nn.functional.conv2d` and
`torch.nn.functional.conv_transpose2d`.

Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
"""

# pylint: disable=redefined-builtin
# pylint: disable=arguments-differ
# pylint: disable=protected-access
# pylint: disable=line-too-long
# pylint: disable=global-statement
# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring
# pylint: disable=deprecated-module
# pylint: disable=wrong-import-order

import warnings
import contextlib
import torch

from distutils.version import LooseVersion

enabled = True                     # Enable the custom op by setting this to true.
weight_gradients_disabled = False  # Forcefully disable computation of gradients with respect to the weights.

@contextlib.contextmanager
def no_weight_gradients():
    global weight_gradients_disabled
    old = weight_gradients_disabled
    weight_gradients_disabled = True
    yield
    weight_gradients_disabled = old

#----------------------------------------------------------------------------

def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, impl='cuda'):
    if impl == 'cuda' and _should_use_custom_op(input):
        return _conv2d_gradfix(transpose=False, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=0, dilation=dilation, groups=groups).apply(input, weight, bias)
    return torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)

def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1, impl='cuda'):
    if impl == 'cuda' and _should_use_custom_op(input):
        return _conv2d_gradfix(transpose=True, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation).apply(input, weight, bias)
    return torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation)

#----------------------------------------------------------------------------

def _should_use_custom_op(input):
    assert isinstance(input, torch.Tensor)
    if (not enabled) or (not torch.backends.cudnn.enabled):
        return False
    if input.device.type != 'cuda':
        return False
    if LooseVersion(torch.__version__) >= LooseVersion('1.7.0'):
        return True
    warnings.warn(f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().')
    return False

def _tuple_of_ints(xs, ndim):
    xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
    assert len(xs) == ndim
    assert all(isinstance(x, int) for x in xs)
    return xs

#----------------------------------------------------------------------------

_conv2d_gradfix_cache = dict()

def _conv2d_gradfix(transpose, weight_shape, stride, padding, output_padding, dilation, groups):
    # Parse arguments.
    ndim = 2
    weight_shape = tuple(weight_shape)
    stride = _tuple_of_ints(stride, ndim)
    padding = _tuple_of_ints(padding, ndim)
    output_padding = _tuple_of_ints(output_padding, ndim)
    dilation = _tuple_of_ints(dilation, ndim)

    # Lookup from cache.
    key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
    if key in _conv2d_gradfix_cache:
        return _conv2d_gradfix_cache[key]

    # Validate arguments.
    assert groups >= 1
    assert len(weight_shape) == ndim + 2
    assert all(stride[i] >= 1 for i in range(ndim))
    assert all(padding[i] >= 0 for i in range(ndim))
    assert all(dilation[i] >= 0 for i in range(ndim))
    if not transpose:
        assert all(output_padding[i] == 0 for i in range(ndim))
    else: # transpose
        assert all(0 <= output_padding[i] < max(stride[i], dilation[i]) for i in range(ndim))

    # Helpers.
    common_kwargs = dict(stride=stride, padding=padding, dilation=dilation, groups=groups)
    def calc_output_padding(input_shape, output_shape):
        if transpose:
            return [0, 0]
        return [
            input_shape[i + 2]
            - (output_shape[i + 2] - 1) * stride[i]
            - (1 - 2 * padding[i])
            - dilation[i] * (weight_shape[i + 2] - 1)
            for i in range(ndim)
        ]

    # Forward & backward.
    class Conv2d(torch.autograd.Function):
        @staticmethod
        def forward(ctx, input, weight, bias):
            assert weight.shape == weight_shape
            if not transpose:
                output = torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)
            else: # transpose
                output = torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, output_padding=output_padding, **common_kwargs)
            ctx.save_for_backward(input, weight)
            return output

        @staticmethod
        def backward(ctx, grad_output):
            input, weight = ctx.saved_tensors
            grad_input = None
            grad_weight = None
            grad_bias = None

            if ctx.needs_input_grad[0]:
                p = calc_output_padding(input_shape=input.shape, output_shape=grad_output.shape)
                grad_input = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs).apply(grad_output, weight, None)
                assert grad_input.shape == input.shape

            if ctx.needs_input_grad[1] and not weight_gradients_disabled:
                grad_weight = Conv2dGradWeight.apply(grad_output, input)
                assert grad_weight.shape == weight_shape

            if ctx.needs_input_grad[2]:
                grad_bias = grad_output.sum([0, 2, 3])

            return grad_input, grad_weight, grad_bias

    # Gradient with respect to the weights.
    class Conv2dGradWeight(torch.autograd.Function):
        @staticmethod
        def forward(ctx, grad_output, input):
            op = torch._C._jit_get_operation('aten::cudnn_convolution_backward_weight' if not transpose else 'aten::cudnn_convolution_transpose_backward_weight')
            flags = [torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32]
            grad_weight = op(weight_shape, grad_output, input, padding, stride, dilation, groups, *flags)
            assert grad_weight.shape == weight_shape
            ctx.save_for_backward(grad_output, input)
            return grad_weight

        @staticmethod
        def backward(ctx, grad2_grad_weight):
            grad_output, input = ctx.saved_tensors
            grad2_grad_output = None
            grad2_input = None

            if ctx.needs_input_grad[0]:
                grad2_grad_output = Conv2d.apply(input, grad2_grad_weight, None)
                assert grad2_grad_output.shape == grad_output.shape

            if ctx.needs_input_grad[1]:
                p = calc_output_padding(input_shape=input.shape, output_shape=grad_output.shape)
                grad2_input = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs).apply(grad_output, grad2_grad_weight, None)
                assert grad2_input.shape == input.shape

            return grad2_grad_output, grad2_input

    _conv2d_gradfix_cache[key] = Conv2d
    return Conv2d

#----------------------------------------------------------------------------

# pylint: enable=redefined-builtin
# pylint: enable=arguments-differ
# pylint: enable=protected-access
# pylint: enable=line-too-long
# pylint: enable=global-statement
# pylint: enable=missing-class-docstring
# pylint: enable=missing-function-docstring
# pylint: enable=deprecated-module
# pylint: enable=wrong-import-order

third_party/stylegan2_official_ops/conv2d_gradfix.py
zhujiapeng/Aurora
# python3.7

# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

"""Custom replacement for `torch.nn.functional.grid_sample`.

This is useful for differentiable augmentation. This customized operator
supports arbitrarily high order gradients between the input and output. Only
works on 2D images and assumes `mode=bilinear`, `padding_mode=zeros`, and
`align_corners=False`.

Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
"""

# pylint: disable=redefined-builtin
# pylint: disable=arguments-differ
# pylint: disable=protected-access
# pylint: disable=line-too-long
# pylint: disable=missing-function-docstring
# pylint: disable=deprecated-module
# pylint: disable=wrong-import-order

import warnings
import torch
from distutils.version import LooseVersion

#----------------------------------------------------------------------------

enabled = True  # Enable the custom op by setting this to true.

#----------------------------------------------------------------------------

def grid_sample(input, grid, impl='cuda'):
    if impl == 'cuda' and _should_use_custom_op():
        return _GridSample2dForward.apply(input, grid)
    return torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False)

#----------------------------------------------------------------------------

def _should_use_custom_op():
    if not enabled:
        return False
    if LooseVersion(torch.__version__) >= LooseVersion('1.7.0'):
        return True
    warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().')
    return False

#----------------------------------------------------------------------------

class _GridSample2dForward(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, grid):
        assert input.ndim == 4
        assert grid.ndim == 4
        output = torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False)
        ctx.save_for_backward(input, grid)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, grid = ctx.saved_tensors
        grad_input, grad_grid = _GridSample2dBackward.apply(grad_output, input, grid)
        return grad_input, grad_grid

#----------------------------------------------------------------------------

class _GridSample2dBackward(torch.autograd.Function):
    @staticmethod
    def forward(ctx, grad_output, input, grid):
        op = torch._C._jit_get_operation('aten::grid_sampler_2d_backward')
        grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False)
        ctx.save_for_backward(grid)
        return grad_input, grad_grid

    @staticmethod
    def backward(ctx, grad2_grad_input, grad2_grad_grid):
        _ = grad2_grad_grid # unused
        grid, = ctx.saved_tensors
        grad2_grad_output = None
        grad2_input = None
        grad2_grid = None

        if ctx.needs_input_grad[0]:
            grad2_grad_output = _GridSample2dForward.apply(grad2_grad_input, grid)

        assert not ctx.needs_input_grad[2]
        return grad2_grad_output, grad2_input, grad2_grid

#----------------------------------------------------------------------------

# pylint: enable=redefined-builtin
# pylint: enable=arguments-differ
# pylint: enable=protected-access
# pylint: enable=line-too-long
# pylint: enable=missing-function-docstring
# pylint: enable=deprecated-module
# pylint: enable=wrong-import-order

third_party/stylegan2_official_ops/grid_sample_gradfix.py
zhujiapeng/Aurora
# python3.7

# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

"""Custom operators for efficient resampling of 2D images.

`upfirdn` means executing upsampling, FIR filtering, downsampling in sequence.

Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
"""

# pylint: disable=line-too-long
# pylint: disable=missing-class-docstring
# pylint: disable=bare-except
# pylint: disable=global-variable-not-assigned

import os
import warnings
import traceback
import numpy as np
import torch

from . import custom_ops
from . import misc
from . import conv2d_gradfix

#----------------------------------------------------------------------------

_inited = False
_plugin = None

def _init():
    global _inited, _plugin
    if not _inited:
        sources = ['upfirdn2d.cpp', 'upfirdn2d.cu']
        sources = [os.path.join(os.path.dirname(__file__), s) for s in sources]
        try:
            _plugin = custom_ops.get_plugin('upfirdn2d_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
        except:
            warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
    return _plugin is not None

def _parse_scaling(scaling):
    if isinstance(scaling, int):
        scaling = [scaling, scaling]
    assert isinstance(scaling, (list, tuple))
    assert all(isinstance(x, int) for x in scaling)
    sx, sy = scaling
    assert sx >= 1 and sy >= 1
    return sx, sy

def _parse_padding(padding):
    if isinstance(padding, int):
        padding = [padding, padding]
    assert isinstance(padding, (list, tuple))
    assert all(isinstance(x, int) for x in padding)
    if len(padding) == 2:
        padx, pady = padding
        padding = [padx, padx, pady, pady]
    padx0, padx1, pady0, pady1 = padding
    return padx0, padx1, pady0, pady1

def _get_filter_size(f):
    if f is None:
        return 1, 1
    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
    fw = f.shape[-1]
    fh = f.shape[0]
    with misc.suppress_tracer_warnings():
        fw = int(fw)
        fh = int(fh)
    misc.assert_shape(f, [fh, fw][:f.ndim])
    assert fw >= 1 and fh >= 1
    return fw, fh

#----------------------------------------------------------------------------

def setup_filter(f, device=torch.device('cpu'), normalize=True, flip_filter=False, gain=1, separable=None):
    r"""Convenience function to setup 2D FIR filter for `upfirdn2d()`.

    Args:
        f:           Torch tensor, numpy array, or python list of the shape
                     `[filter_height, filter_width]` (non-separable),
                     `[filter_taps]` (separable),
                     `[]` (impulse), or
                     `None` (identity).
        device:      Result device (default: cpu).
        normalize:   Normalize the filter so that it retains the magnitude
                     for constant input signal (DC)? (default: True).
        flip_filter: Flip the filter? (default: False).
        gain:        Overall scaling factor for signal magnitude (default: 1).
        separable:   Return a separable filter? (default: select automatically).

    Returns:
        Float32 tensor of the shape
        `[filter_height, filter_width]` (non-separable) or
        `[filter_taps]` (separable).
    """
    # Validate.
    if f is None:
        f = 1
    f = torch.as_tensor(f, dtype=torch.float32)
    assert f.ndim in [0, 1, 2]
    assert f.numel() > 0
    if f.ndim == 0:
        f = f[np.newaxis]

    # Separable?
    if separable is None:
        separable = (f.ndim == 1 and f.numel() >= 8)
    if f.ndim == 1 and not separable:
        f = f.ger(f)
    assert f.ndim == (1 if separable else 2)

    # Apply normalize, flip, gain, and device.
    if normalize:
        f /= f.sum()
    if flip_filter:
        f = f.flip(list(range(f.ndim)))
    f = f * (gain ** (f.ndim / 2))
    f = f.to(device=device)
    return f

#----------------------------------------------------------------------------

def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, impl='cuda'):
    r"""Pad, upsample, filter, and downsample a batch of 2D images.

    Performs the following sequence of operations for each channel:

    1. Upsample the image by inserting N-1 zeros after each pixel (`up`).

    2. Pad the image with the specified number of zeros on each side (`padding`).
       Negative padding corresponds to cropping the image.

    3. Convolve the image with the specified 2D FIR filter (`f`), shrinking it
       so that the footprint of all output pixels lies within the input image.

    4. Downsample the image by keeping every Nth pixel (`down`).

    This sequence of operations bears close resemblance to scipy.signal.upfirdn().
    The fused op is considerably more efficient than performing the same calculation
    using standard PyTorch ops. It supports gradients of arbitrary order.

    Args:
        x:           Float32/float64/float16 input tensor of the shape
                     `[batch_size, num_channels, in_height, in_width]`.
        f:           Float32 FIR filter of the shape
                     `[filter_height, filter_width]` (non-separable),
                     `[filter_taps]` (separable), or
                     `None` (identity).
        up:          Integer upsampling factor. Can be a single int or a list/tuple
                     `[x, y]` (default: 1).
        down:        Integer downsampling factor. Can be a single int or a list/tuple
                     `[x, y]` (default: 1).
        padding:     Padding with respect to the upsampled image. Can be a single number
                     or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
                     (default: 0).
        flip_filter: False = convolution, True = correlation (default: False).
        gain:        Overall scaling factor for signal magnitude (default: 1).
        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).

    Returns:
        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
    """
    assert isinstance(x, torch.Tensor)
    assert impl in ['ref', 'cuda']
    if impl == 'cuda' and x.device.type == 'cuda' and _init():
        return _upfirdn2d_cuda(up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain).apply(x, f)
    return _upfirdn2d_ref(x, f, up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain)

#----------------------------------------------------------------------------

@misc.profiled_function
def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
    """Slow reference implementation of `upfirdn2d()` using standard PyTorch ops.
    """
    # Validate arguments.
    assert isinstance(x, torch.Tensor) and x.ndim == 4
    if f is None:
        f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
    assert f.dtype == torch.float32 and not f.requires_grad
    batch_size, num_channels, in_height, in_width = x.shape
    upx, upy = _parse_scaling(up)
    downx, downy = _parse_scaling(down)
    padx0, padx1, pady0, pady1 = _parse_padding(padding)

    # Upsample by inserting zeros.
    x = x.reshape([batch_size, num_channels, in_height, 1, in_width, 1])
    x = torch.nn.functional.pad(x, [0, upx - 1, 0, 0, 0, upy - 1])
    x = x.reshape([batch_size, num_channels, in_height * upy, in_width * upx])

    # Pad or crop.
    x = torch.nn.functional.pad(x, [max(padx0, 0), max(padx1, 0), max(pady0, 0), max(pady1, 0)])
    x = x[:, :, max(-pady0, 0) : x.shape[2] - max(-pady1, 0), max(-padx0, 0) : x.shape[3] - max(-padx1, 0)]

    # Setup filter.
    f = f * (gain ** (f.ndim / 2))
    f = f.to(x.dtype)
    if not flip_filter:
        f = f.flip(list(range(f.ndim)))

    # Convolve with the filter.
    f = f[np.newaxis, np.newaxis].repeat([num_channels, 1] + [1] * f.ndim)
    if f.ndim == 4:
        x = conv2d_gradfix.conv2d(input=x, weight=f, groups=num_channels)
    else:
        x = conv2d_gradfix.conv2d(input=x, weight=f.unsqueeze(2), groups=num_channels)
        x = conv2d_gradfix.conv2d(input=x, weight=f.unsqueeze(3), groups=num_channels)

    # Downsample by throwing away pixels.
    x = x[:, :, ::downy, ::downx]
    return x

#----------------------------------------------------------------------------

_upfirdn2d_cuda_cache = dict()

def _upfirdn2d_cuda(up=1, down=1, padding=0, flip_filter=False, gain=1):
    """Fast CUDA implementation of `upfirdn2d()` using custom ops.
    """
    # Parse arguments.
    upx, upy = _parse_scaling(up)
    downx, downy = _parse_scaling(down)
    padx0, padx1, pady0, pady1 = _parse_padding(padding)

    # Lookup from cache.
    key = (upx, upy, downx, downy, padx0, padx1, pady0, pady1, flip_filter, gain)
    if key in _upfirdn2d_cuda_cache:
        return _upfirdn2d_cuda_cache[key]

    # Forward op.
    class Upfirdn2dCuda(torch.autograd.Function):
        @staticmethod
        def forward(ctx, x, f): # pylint: disable=arguments-differ
            assert isinstance(x, torch.Tensor) and x.ndim == 4
            if f is None:
                f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
            assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
            y = x
            if f.ndim == 2:
                y = _plugin.upfirdn2d(y, f, upx, upy, downx, downy, padx0, padx1, pady0, pady1, flip_filter, gain)
            else:
                y = _plugin.upfirdn2d(y, f.unsqueeze(0), upx, 1, downx, 1, padx0, padx1, 0, 0, flip_filter, np.sqrt(gain))
                y = _plugin.upfirdn2d(y, f.unsqueeze(1), 1, upy, 1, downy, 0, 0, pady0, pady1, flip_filter, np.sqrt(gain))
            ctx.save_for_backward(f)
            ctx.x_shape = x.shape
            return y

        @staticmethod
        def backward(ctx, dy): # pylint: disable=arguments-differ
            f, = ctx.saved_tensors
            _, _, ih, iw = ctx.x_shape
            _, _, oh, ow = dy.shape
            fw, fh = _get_filter_size(f)
            p = [
                fw - padx0 - 1,
                iw * upx - ow * downx + padx0 - upx + 1,
                fh - pady0 - 1,
                ih * upy - oh * downy + pady0 - upy + 1,
            ]
            dx = None
            df = None

            if ctx.needs_input_grad[0]:
                dx = _upfirdn2d_cuda(up=down, down=up, padding=p, flip_filter=(not flip_filter), gain=gain).apply(dy, f)

            assert not ctx.needs_input_grad[1]
            return dx, df

    # Add to cache.
    _upfirdn2d_cuda_cache[key] = Upfirdn2dCuda
    return Upfirdn2dCuda

#----------------------------------------------------------------------------

def filter2d(x, f, padding=0, flip_filter=False, gain=1, impl='cuda'):
    r"""Filter a batch of 2D images using the given 2D FIR filter.

    By default, the result is padded so that its shape matches the input.
    User-specified padding is applied on top of that, with negative values
    indicating cropping. Pixels outside the image are assumed to be zero.

    Args:
        x:           Float32/float64/float16 input tensor of the shape
                     `[batch_size, num_channels, in_height, in_width]`.
        f:           Float32 FIR filter of the shape
                     `[filter_height, filter_width]` (non-separable),
                     `[filter_taps]` (separable), or
                     `None` (identity).
        padding:     Padding with respect to the output. Can be a single number or a
                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
                     (default: 0).
        flip_filter: False = convolution, True = correlation (default: False).
        gain:        Overall scaling factor for signal magnitude (default: 1).
        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).

    Returns:
        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
    """
    padx0, padx1, pady0, pady1 = _parse_padding(padding)
    fw, fh = _get_filter_size(f)
    p = [
        padx0 + fw // 2,
        padx1 + (fw - 1) // 2,
        pady0 + fh // 2,
        pady1 + (fh - 1) // 2,
    ]
    return upfirdn2d(x, f, padding=p, flip_filter=flip_filter, gain=gain, impl=impl)

#----------------------------------------------------------------------------

def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl='cuda'):
    r"""Upsample a batch of 2D images using the given 2D FIR filter.

    By default, the result is padded so that its shape is a multiple of the input.
    User-specified padding is applied on top of that, with negative values
    indicating cropping. Pixels outside the image are assumed to be zero.

    Args:
        x:           Float32/float64/float16 input tensor of the shape
                     `[batch_size, num_channels, in_height, in_width]`.
        f:           Float32 FIR filter of the shape
                     `[filter_height, filter_width]` (non-separable),
                     `[filter_taps]` (separable), or
                     `None` (identity).
        up:          Integer upsampling factor. Can be a single int or a list/tuple
                     `[x, y]` (default: 1).
        padding:     Padding with respect to the output. Can be a single number or a
                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
                     (default: 0).
        flip_filter: False = convolution, True = correlation (default: False).
        gain:        Overall scaling factor for signal magnitude (default: 1).
        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).

    Returns:
        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
    """
    upx, upy = _parse_scaling(up)
    padx0, padx1, pady0, pady1 = _parse_padding(padding)
    fw, fh = _get_filter_size(f)
    p = [
        padx0 + (fw + upx - 1) // 2,
        padx1 + (fw - upx) // 2,
        pady0 + (fh + upy - 1) // 2,
        pady1 + (fh - upy) // 2,
    ]
    return upfirdn2d(x, f, up=up, padding=p, flip_filter=flip_filter, gain=gain*upx*upy, impl=impl)

#----------------------------------------------------------------------------

def downsample2d(x, f, down=2, padding=0, flip_filter=False, gain=1, impl='cuda'):
    r"""Downsample a batch of 2D images using the given 2D FIR filter.

    By default, the result is padded so that its shape is a fraction of the input.
    User-specified padding is applied on top of that, with negative values
    indicating cropping. Pixels outside the image are assumed to be zero.

    Args:
        x:           Float32/float64/float16 input tensor of the shape
                     `[batch_size, num_channels, in_height, in_width]`.
        f:           Float32 FIR filter of the shape
                     `[filter_height, filter_width]` (non-separable),
                     `[filter_taps]` (separable), or
                     `None` (identity).
        down:        Integer downsampling factor. Can be a single int or a list/tuple
                     `[x, y]` (default: 1).
        padding:     Padding with respect to the input. Can be a single number or a
                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
                     (default: 0).
        flip_filter: False = convolution, True = correlation (default: False).
        gain:        Overall scaling factor for signal magnitude (default: 1).
        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).

    Returns:
        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
    """
    downx, downy = _parse_scaling(down)
    padx0, padx1, pady0, pady1 = _parse_padding(padding)
    fw, fh = _get_filter_size(f)
    p = [
        padx0 + (fw - downx + 1) // 2,
        padx1 + (fw - downx) // 2,
        pady0 + (fh - downy + 1) // 2,
        pady1 + (fh - downy) // 2,
    ]
    return upfirdn2d(x, f, down=down, padding=p, flip_filter=flip_filter, gain=gain, impl=impl)

#----------------------------------------------------------------------------

# pylint: enable=line-too-long
# pylint: enable=missing-class-docstring
# pylint: enable=bare-except
# pylint: enable=global-variable-not-assigned

third_party/stylegan2_official_ops/upfirdn2d.py
zhujiapeng/Aurora
# python3.7

# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

"""Custom replacement for convolution operators.

Operators in this file support arbitrarily high order gradients with zero
performance penalty. Please set `impl` as `cuda` to use faster customized
operators, OR as `ref` to use native `torch.nn.functional.conv2d` and
`torch.nn.functional.conv_transpose2d`.

Please refer to https://github.com/NVlabs/stylegan3
"""

# pylint: disable=redefined-builtin
# pylint: disable=arguments-differ
# pylint: disable=protected-access
# pylint: disable=line-too-long
# pylint: disable=global-statement
# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring

import contextlib
import torch
from pkg_resources import parse_version

#----------------------------------------------------------------------------

enabled = True                      # Enable the custom op by setting this to true.
weight_gradients_disabled = False   # Forcefully disable computation of gradients with respect to the weights.
_use_pytorch_1_11_api = parse_version(torch.__version__) >= parse_version('1.11.0')

@contextlib.contextmanager
def no_weight_gradients(disable=True):
    global weight_gradients_disabled
    old = weight_gradients_disabled
    if disable:
        weight_gradients_disabled = True
    yield
    weight_gradients_disabled = old

#----------------------------------------------------------------------------

def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, impl='cuda'):
    if impl == 'cuda' and _should_use_custom_op(input):
        return _conv2d_gradfix(transpose=False, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=0, dilation=dilation, groups=groups).apply(input, weight, bias)
    return torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)

def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1, impl='cuda'):
    if impl == 'cuda' and _should_use_custom_op(input):
        return _conv2d_gradfix(transpose=True, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation).apply(input, weight, bias)
    return torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation)

#----------------------------------------------------------------------------

def _should_use_custom_op(input):
    assert isinstance(input, torch.Tensor)
    if (not enabled) or (not torch.backends.cudnn.enabled):
        return False
    if input.device.type != 'cuda':
        return False
    return True

def _tuple_of_ints(xs, ndim):
    xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
    assert len(xs) == ndim
    assert all(isinstance(x, int) for x in xs)
    return xs

#----------------------------------------------------------------------------

_conv2d_gradfix_cache = dict()
_null_tensor = torch.empty([0])

def _conv2d_gradfix(transpose, weight_shape, stride, padding, output_padding, dilation, groups):
    # Parse arguments.
    ndim = 2
    weight_shape = tuple(weight_shape)
    stride = _tuple_of_ints(stride, ndim)
    padding = _tuple_of_ints(padding, ndim)
    output_padding = _tuple_of_ints(output_padding, ndim)
    dilation = _tuple_of_ints(dilation, ndim)

    # Lookup from cache.
    key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
    if key in _conv2d_gradfix_cache:
        return _conv2d_gradfix_cache[key]

    # Validate arguments.
    assert groups >= 1
    assert len(weight_shape) == ndim + 2
    assert all(stride[i] >= 1 for i in range(ndim))
    assert all(padding[i] >= 0 for i in range(ndim))
    assert all(dilation[i] >= 0 for i in range(ndim))
    if not transpose:
        assert all(output_padding[i] == 0 for i in range(ndim))
    else: # transpose
        assert all(0 <= output_padding[i] < max(stride[i], dilation[i]) for i in range(ndim))

    # Helpers.
    common_kwargs = dict(stride=stride, padding=padding, dilation=dilation, groups=groups)
    def calc_output_padding(input_shape, output_shape):
        if transpose:
            return [0, 0]
        return [
            input_shape[i + 2]
            - (output_shape[i + 2] - 1) * stride[i]
            - (1 - 2 * padding[i])
            - dilation[i] * (weight_shape[i + 2] - 1)
            for i in range(ndim)
        ]

    # Forward & backward.
    class Conv2d(torch.autograd.Function):
        @staticmethod
        def forward(ctx, input, weight, bias):
            assert weight.shape == weight_shape
            ctx.save_for_backward(
                input if weight.requires_grad else _null_tensor,
                weight if input.requires_grad else _null_tensor,
            )
            ctx.input_shape = input.shape

            # Simple 1x1 convolution => cuBLAS (only on Volta, not on Ampere).
            if weight_shape[2:] == stride == dilation == (1, 1) and padding == (0, 0) and torch.cuda.get_device_capability(input.device) < (8, 0):
                a = weight.reshape(groups, weight_shape[0] // groups, weight_shape[1])
                b = input.reshape(input.shape[0], groups, input.shape[1] // groups, -1)
                c = (a.transpose(1, 2) if transpose else a) @ b.permute(1, 2, 0, 3).flatten(2)
                c = c.reshape(-1, input.shape[0], *input.shape[2:]).transpose(0, 1)
                c = c if bias is None else c + bias.unsqueeze(0).unsqueeze(2).unsqueeze(3)
                return c.contiguous(memory_format=(torch.channels_last if input.stride(1) == 1 else torch.contiguous_format))

            # General case => cuDNN.
            if transpose:
                return torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, output_padding=output_padding, **common_kwargs)
            return torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)

        @staticmethod
        def backward(ctx, grad_output):
            input, weight = ctx.saved_tensors
            input_shape = ctx.input_shape
            grad_input = None
            grad_weight = None
            grad_bias = None

            if ctx.needs_input_grad[0]:
                p = calc_output_padding(input_shape=input_shape, output_shape=grad_output.shape)
                op = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs)
                grad_input = op.apply(grad_output, weight, None)
                assert grad_input.shape == input_shape

            if ctx.needs_input_grad[1] and not weight_gradients_disabled:
                if _use_pytorch_1_11_api:
                    grad_weight = Conv2dGradWeight.apply(grad_output, input, weight)
                else:
                    grad_weight = Conv2dGradWeight.apply(grad_output, input)
                assert grad_weight.shape == weight_shape

            if ctx.needs_input_grad[2]:
                grad_bias = grad_output.sum([0, 2, 3])

            return grad_input, grad_weight, grad_bias

    # Gradient with respect to the weights.
    class Conv2dGradWeight(torch.autograd.Function):
        if _use_pytorch_1_11_api:
            @staticmethod
            def forward(ctx, grad_output, input, weight):
                ctx.save_for_backward(
                    grad_output if input.requires_grad else _null_tensor,
                    input if grad_output.requires_grad else _null_tensor,
                )
                ctx.grad_output_shape = grad_output.shape
                ctx.input_shape = input.shape

                # Simple 1x1 convolution => cuBLAS (on both Volta and Ampere).
                if weight_shape[2:] == stride == dilation == (1, 1) and padding == (0, 0):
                    a = grad_output.reshape(grad_output.shape[0], groups, grad_output.shape[1] // groups, -1).permute(1, 2, 0, 3).flatten(2)
                    b = input.reshape(input.shape[0], groups, input.shape[1] // groups, -1).permute(1, 2, 0, 3).flatten(2)
                    c = (b @ a.transpose(1, 2) if transpose else a @ b.transpose(1, 2)).reshape(weight_shape)
                    return c.contiguous(memory_format=(torch.channels_last if input.stride(1) == 1 else torch.contiguous_format))

                # General case => cuDNN.
                return torch.ops.aten.convolution_backward(grad_output=grad_output, input=input, weight=weight, bias_sizes=None, stride=stride, padding=padding, dilation=dilation, transposed=transpose, output_padding=output_padding, groups=groups, output_mask=[False, True, False])[1]
        else:
            @staticmethod
            def forward(ctx, grad_output, input):
                ctx.save_for_backward(
                    grad_output if input.requires_grad else _null_tensor,
                    input if grad_output.requires_grad else _null_tensor,
                )
                ctx.grad_output_shape = grad_output.shape
                ctx.input_shape = input.shape

                # Simple 1x1 convolution => cuBLAS (on both Volta and Ampere).
                if weight_shape[2:] == stride == dilation == (1, 1) and padding == (0, 0):
                    a = grad_output.reshape(grad_output.shape[0], groups, grad_output.shape[1] // groups, -1).permute(1, 2, 0, 3).flatten(2)
                    b = input.reshape(input.shape[0], groups, input.shape[1] // groups, -1).permute(1, 2, 0, 3).flatten(2)
                    c = (b @ a.transpose(1, 2) if transpose else a @ b.transpose(1, 2)).reshape(weight_shape)
                    return c.contiguous(memory_format=(torch.channels_last if input.stride(1) == 1 else torch.contiguous_format))

                # General case => cuDNN.
                name = 'aten::cudnn_convolution_transpose_backward_weight' if transpose else 'aten::cudnn_convolution_backward_weight'
                flags = [torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32]
                return torch._C._jit_get_operation(name)(weight_shape, grad_output, input, padding, stride, dilation, groups, *flags)

        @staticmethod
        def backward(ctx, grad2_grad_weight):
            grad_output, input = ctx.saved_tensors
            grad_output_shape = ctx.grad_output_shape
            input_shape = ctx.input_shape
            grad2_grad_output = None
            grad2_input = None

            if ctx.needs_input_grad[0]:
                grad2_grad_output = Conv2d.apply(input, grad2_grad_weight, None)
                assert grad2_grad_output.shape == grad_output_shape

            if ctx.needs_input_grad[1]:
                p = calc_output_padding(input_shape=input_shape, output_shape=grad_output_shape)
                op = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs)
                grad2_input = op.apply(grad_output, grad2_grad_weight, None)
                assert grad2_input.shape == input_shape

            return grad2_grad_output, grad2_input

    _conv2d_gradfix_cache[key] = Conv2d
    return Conv2d

#----------------------------------------------------------------------------

# pylint: enable=redefined-builtin
# pylint: enable=arguments-differ
# pylint: enable=protected-access
# pylint: enable=line-too-long
# pylint: enable=global-statement
# pylint: enable=missing-class-docstring
# pylint: enable=missing-function-docstring

third_party/stylegan3_official_ops/conv2d_gradfix.py
zhujiapeng/Aurora
# python3.7
"""Misc utility functions."""

import os
import hashlib

from pathlib import Path
from torch.hub import download_url_to_file

__all__ = [
    'REPO_NAME', 'Infix', 'print_and_execute', 'check_file_ext',
    'IMAGE_EXTENSIONS', 'VIDEO_EXTENSIONS', 'MEDIA_EXTENSIONS',
    'parse_file_format', 'set_cache_dir', 'get_cache_dir',
    'md5_update_from_file', 'md5_update_from_dir', 'md5', 'download_url'
]

REPO_NAME = 'Hammer'  # Name of the repository (project).


class Infix(object):
    """Helper class to create custom infix operators.

    When using it, make sure to put the operator between `<<` and `>>`.
    `<< INFIX_OP_NAME >>` should be considered as a whole operator.

    Examples:

    # Use `Infix` to create infix operators directly.
    add = Infix(lambda a, b: a + b)
    1 << add >> 2  # gives 3
    1 << add >> 2 << add >> 3  # gives 6

    # Use `Infix` as a decorator.
    @Infix
    def mul(a, b):
        return a * b
    2 << mul >> 4  # gives 8
    2 << mul >> 3 << mul >> 7  # gives 42
    """

    def __init__(self, function):
        self.function = function
        self.left_value = None

    def __rlshift__(self, left_value):  # override `<<` before `Infix` instance
        assert self.left_value is None  # make sure left is only called once
        self.left_value = left_value
        return self

    def __rshift__(self, right_value):  # override `>>` after `Infix` instance
        result = self.function(self.left_value, right_value)
        self.left_value = None  # reset to None
        return result


def print_and_execute(cmd):
    """Prints and executes a system command.

    Args:
        cmd: Command to be executed.
    """
    print(cmd)
    os.system(cmd)


def check_file_ext(filename, *ext_list):
    """Checks whether the given filename is with target extension(s).

    NOTE: If `ext_list` is empty, this function will always return `False`.

    Args:
        filename: Filename to check.
        *ext_list: A list of extensions.

    Returns:
        `True` if the filename is with one of extensions in `ext_list`,
        otherwise `False`.
    """
    if len(ext_list) == 0:
        return False
    ext_list = [ext if ext.startswith('.') else '.' + ext for ext in ext_list]
    ext_list = [ext.lower() for ext in ext_list]
    basename = os.path.basename(filename)
    ext = os.path.splitext(basename)[1].lower()
    return ext in ext_list


# File extensions regarding images (not including GIFs).
IMAGE_EXTENSIONS = (
    '.bmp', '.ppm', '.pgm', '.jpeg', '.jpg', '.jpe', '.jp2', '.png', '.webp',
    '.tiff', '.tif'
)
# File extensions regarding videos.
VIDEO_EXTENSIONS = (
    '.avi', '.mkv', '.mp4', '.m4v', '.mov', '.webm', '.flv', '.rmvb', '.rm',
    '.3gp'
)
# File extensions regarding media, i.e., images, videos, GIFs.
MEDIA_EXTENSIONS = ('.gif', *IMAGE_EXTENSIONS, *VIDEO_EXTENSIONS)


def parse_file_format(path):
    """Parses the file format of a given path.

    This function basically parses the file format according to its extension.
    It will also return `dir` is the given path is a directory.

    Parable file formats:

    - zip: with `.zip` extension.
    - tar: with `.tar` / `.tgz` / `.tar.gz` extension.
    - lmdb: a folder ending with `lmdb`.
    - txt: with `.txt` / `.text` extension, OR without extension (e.g. LICENSE).
    - json: with `.json` extension.
    - jpg: with `.jpeg` / `jpg` / `jpe` extension.
    - png: with `.png` extension.

    Args:
        path: The path to the file to parse format from.

    Returns:
        A lower-case string, indicating the file format, or `None` if the format
            cannot be successfully parsed.
    """
    # Handle directory.
    if os.path.isdir(path) or path.endswith('/'):
        if path.rstrip('/').lower().endswith('lmdb'):
            return 'lmdb'
        return 'dir'
    # Handle file.
    if os.path.isfile(path) and os.path.splitext(path)[1] == '':
        return 'txt'
    path = path.lower()
    if path.endswith('.tar.gz'):  # Cannot parse accurate extension.
        return 'tar'
    ext = os.path.splitext(path)[1]
    if ext == '.zip':
        return 'zip'
    if ext in ['.tar', '.tgz']:
        return 'tar'
    if ext in ['.txt', '.text']:
        return 'txt'
    if ext == '.json':
        return 'json'
    if ext in ['.jpeg', '.jpg', '.jpe']:
        return 'jpg'
    if ext == '.png':
        return 'png'
    # Unparsable.
    return None


_cache_dir = None


def set_cache_dir(directory=None):
    """Sets the global cache directory.

    The cache directory can be used to save some files that will be shared
    across jobs. The default cache directory is set as `~/.cache/`. This
    function can be used to redirect the cache directory. Or, users can use
    `None` to reset the cache directory back to default.

    Args:
        directory: The target directory used to cache files. If set as `None`,
            the cache directory will be reset back to default. (default: None)
    """
    assert directory is None or isinstance(directory, str), 'Invalid directory!'
    global _cache_dir  # pylint: disable=global-statement
    _cache_dir = directory


def get_cache_dir(use_repo_name=True):
    """Gets the global cache directory.

    The global cache directory is primarily set as `~/.cache/` by default, and
    can be redirected with `set_cache_dir()`.

    Args:
        use_repo_name: Whether to create a folder, named `REPO_NAME`, under
            `_cache_dir` as the actual cache directory. (default: True)

    Returns:
        A string, representing the global cache directory.
    """
    if _cache_dir is None:
        cache_dir = os.path.join(os.path.expanduser('~'), '.cache')
    else:
        cache_dir = _cache_dir
    if use_repo_name:
        return os.path.join(cache_dir, REPO_NAME)
    return cache_dir


def md5_update_from_file(filename, hash_obj):
    """Updates the `hash_obj` with a given file.

    Args:
        filename: Path to the file.
        hash_obj: The `hashlib.md5` object to be updated.

    Returns:
        An updated `hashlib.md5` object.
    """
    assert Path(filename).is_file()
    with open(str(filename), 'rb') as f:
        for chunk in iter(lambda: f.read(4096), b''):
            hash_obj.update(chunk)
    return hash_obj


def md5_update_from_dir(directory, hash_obj):
    """Updates the `hash_obj` with a given directory.

    Args:
        directory: Path to the directory.
        hash_obj: The `hashlib.md5` object to be updated.

    Returns:
        An updated `hashlib.md5` object.
    """
    assert Path(directory).is_dir()
    for path in sorted(Path(directory).iterdir()):
        hash_obj.update(path.name.encode())
        if path.is_file():
            hash_obj = md5_update_from_file(path, hash_obj)
        elif path.is_dir():
            hash_obj = md5_update_from_dir(path, hash_obj)
    return hash_obj


def md5(path):
    """Returns the MD5 of the given path in hex digest."""
    if Path(path).is_file():
        md5_hash = md5_update_from_file(path, hashlib.md5())
        return md5_hash.hexdigest()
    if Path(path).is_dir():
        md5_hash = md5_update_from_dir(path, hashlib.md5())
        return md5_hash.hexdigest()
    raise ValueError(f'Currently calculation of MD5 does not supported for '
                     f'`{path}`')


def download_url(url, path=None, filename=None, sha256=None):
    """Downloads file from URL.

    This function downloads a file from given URL, and executes Hash check if
    needed.

    Args:
        url: The URL to download file from.
        path: Path (directory) to save the downloaded file. If set as `None`,
            the cache directory will be used. Please see `get_cache_dir()` for
            more details. (default: None)
        filename: The name to save the file. If set as `None`, this name will be
            automatically parsed from the given URL. (default: None)
        sha256: The expected sha256 of the downloaded file. If set as `None`,
            the hash check will be skipped. Otherwise, this function will check
            whether the sha256 of the downloaded file matches this field.

    Returns:
        A two-element tuple, where the first term is the full path of the
            downloaded file, and the second term indicate the hash check result.
            `True` means hash check passes, `False` means hash check fails,
            while `None` means no hash check is executed.
    """
    # Handle file path.
    if path is None:
        path = get_cache_dir()
    if filename is None:
        filename = os.path.basename(url)
    save_path = os.path.join(path, filename)
    # Download file if needed.
    if not os.path.exists(save_path):
        print(f'Downloading URL `{url}` to path `{save_path}` ...')
        os.makedirs(path, exist_ok=True)
        download_url_to_file(url, save_path, hash_prefix=None, progress=True)
    # Check hash if needed.
    check_result = None
    if sha256 is not None:
        with open(save_path, 'rb') as f:
            file_hash = hashlib.sha256(f.read())
            check_result = (file_hash.hexdigest() == sha256)

    return save_path, check_result

utils/misc.py
zhujiapeng/Aurora