libcom 0.1.8__py3-none-any.whl

libcom/image_harmonization/source/src/lbm/trainer/utils.py

@@ -0,0 +1,193 @@
+import logging
+import os
+import re
+import time
+from typing import Dict, List, Literal, Optional, Tuple
+
+import torch
+
+
+class StateDictAdapter:
+    """
+    StateDictAdapter for adapting the state dict of a model to a checkpoint state dict.
+
+    This class will iterate over all keys in the checkpoint state dict and filter them by a list of regex keys.
+    For each matching key, the class will adapt the checkpoint state dict to the model state dict.
+    Depending on the target size, the class will add missing blocks or cut the block.
+    When adding missing blocks, the class will use a strategy to fill the missing blocks: either adding zeros or normal random values.
+
+    Example:
+
+    ```
+    adapter = StateDictAdapter()
+    new_state_dict = adapter(
+        model_state_dict=model.state_dict(),
+        checkpoint_state_dict=state_dict,
+        regex_keys=[
+            r"class_embedding.linear_1.weight",
+            r"conv_in.weight",
+            r"(down_blocks|up_blocks)\.\d+\.attentions\.\d+\.transformer_blocks\.\d+\.attn\d+\.(to_k|to_v)\.weight",
+            r"mid_block\.attentions\.\d+\.transformer_blocks\.\d+\.attn\d+\.(to_k|to_v)\.weight"
+        ]
+    )
+    ```
+
+    Args:
+        model_state_dict (Dict[str, torch.Tensor]): The model state dict.
+        checkpoint_state_dict (Dict[str, torch.Tensor]): The checkpoint state dict.
+        regex_keys (Optional[List[str]]): A list of regex keys to adapt the checkpoint state dict. Defaults to None.
+            Passing a list of regex will drastically reduce the latency.
+            If None, all keys in the checkpoint state dict will be adapted.
+        strategy (Literal["zeros", "normal"], optional): The strategy to fill the missing blocks. Defaults to "normal".
+
+    """
+
+    def _create_block(
+        self,
+        shape: List[int],
+        strategy: Literal["zeros", "normal"],
+        input: torch.Tensor = None,
+    ):
+        if strategy == "zeros":
+            return torch.zeros(shape)
+        elif strategy == "normal":
+            if input is not None:
+                mean = input.mean().item()
+                std = input.std().item()
+                return torch.randn(shape) * std + mean
+            else:
+                return torch.randn(shape)
+        else:
+            raise ValueError(f"Unknown strategy {strategy}")
+
+    def __call__(
+        self,
+        model_state_dict: Dict[str, torch.Tensor],
+        checkpoint_state_dict: Dict[str, torch.Tensor],
+        regex_keys: Optional[List[str]] = None,
+        strategy: Literal["zeros", "normal"] = "normal",
+    ):
+        start = time.perf_counter()
+        # if no regex keys are provided, we use all keys in the model state dict
+        if regex_keys is None:
+            regex_keys = list(model_state_dict.keys())
+
+        # iterate over all keys in the checkpoint state dict
+        for checkpoint_key in list(checkpoint_state_dict.keys()):
+            # iterate over all regex keys
+            for regex_key in regex_keys:
+                if re.match(regex_key, checkpoint_key):
+                    dst_shape = model_state_dict[checkpoint_key].shape
+                    src_shape = checkpoint_state_dict[checkpoint_key].shape
+
+                    ## Sizes adapter
+                    # if length of shapes are different, we need to unsqueeze or squeeze the tensor
+                    if len(dst_shape) != len(src_shape):
+                        # in the case [a] vs [a, b] -> unsqueeze [a, 1]
+                        if len(src_shape) == 1:
+                            checkpoint_state_dict[checkpoint_key] = (
+                                checkpoint_state_dict[checkpoint_key].unsqueeze(1)
+                            )
+                            logging.info(
+                                f"Unsqueeze {checkpoint_key}: {src_shape} -> {checkpoint_state_dict[checkpoint_key].shape}"
+                            )
+                        # in the case [a, b] vs [a] -> squeeze [a]
+                        elif len(dst_shape) == 1:
+                            checkpoint_state_dict[checkpoint_key] = (
+                                checkpoint_state_dict[checkpoint_key][:, 0]
+                            )
+                            logging.info(
+                                f"Squeeze {checkpoint_key}: {src_shape} -> {checkpoint_state_dict[checkpoint_key].shape}"
+                            )
+                        # in the other cases, raise an error
+                        else:
+                            raise ValueError(
+                                f"Shapes of {checkpoint_key} are different: {dst_shape} != {src_shape}"
+                            )
+
+                        # update the shapes
+                        dst_shape = model_state_dict[checkpoint_key].shape
+                        src_shape = checkpoint_state_dict[checkpoint_key].shape
+                        assert len(dst_shape) == len(
+                            src_shape
+                        ), f"Shapes of {checkpoint_key} are different: {dst_shape} != {src_shape}"
+
+                    ## Shapes adapter
+                    # modify the checkpoint state dict only if the shapes are different
+                    if dst_shape != src_shape:
+                        # create a copy of the tensor
+                        tmp = torch.clone(checkpoint_state_dict[checkpoint_key])
+
+                        # iterate over all dimensions
+                        for i in range(len(dst_shape)):
+                            if dst_shape[i] != src_shape[i]:
+                                diff = dst_shape[i] - src_shape[i]
+
+                                # if the difference is greater than 0, we need to add missing blocks
+                                if diff > 0:
+                                    missing_shape = list(tmp.shape)
+                                    missing_shape[i] = diff
+                                    missing = self._create_block(
+                                        shape=missing_shape,
+                                        strategy=strategy,
+                                        input=tmp,
+                                    )
+                                    tmp = torch.cat((tmp, missing), dim=i)
+                                    logging.info(
+                                        f"Adapting {checkpoint_key} with strategy:{strategy} from shape {src_shape} to {dst_shape}"
+                                    )
+                                # if the difference is less than 0, we need to cut the block
+                                else:
+                                    tmp = tmp.narrow(i, 0, dst_shape[i])
+                                    logging.info(
+                                        f"Adapting {checkpoint_key} by narrowing from shape {src_shape} to {dst_shape}"
+                                    )
+
+                        checkpoint_state_dict[checkpoint_key] = tmp
+        end = time.perf_counter()
+        logging.info(f"StateDictAdapter took {end-start:.2f} seconds")
+        return checkpoint_state_dict
+
+
+class StateDictRenamer:
+    """
+    StateDictRenamer for renaming keys in a checkpoint state dict.
+    This class will iterate over all keys in the checkpoint state dict and rename them according to a rename dict.
+
+    Example:
+
+        ```
+        renamer = StateDictRenamer()
+        new_state_dict = renamer(
+            checkpoint_state_dict=state_dict,
+            rename_dict={
+                "add_embedding.linear_1.weight": "class_embedding.linear_1.weight",
+                "add_embedding.linear_1.bias": "class_embedding.linear_1.bias",
+                "add_embedding.linear_2.weight": "class_embedding.linear_2.weight",
+                "add_embedding.linear_2.bias": "class_embedding.linear_2.bias",
+            }
+        )
+        ```
+
+    Args:
+
+        checkpoint_state_dict (Dict[str, torch.Tensor]): The checkpoint state dict.
+        rename_dict (Dict[str, str]): The dictionary mapping the old keys to new keys
+    """
+
+    def __call__(
+        self,
+        checkpoint_state_dict: Dict[str, torch.Tensor],
+        rename_dict: Dict[str, str],
+    ) -> Dict[str, torch.Tensor]:
+        for old_key, new_key in rename_dict.items():
+            if old_key not in checkpoint_state_dict:
+                logging.warning(f"Key {old_key} not found in checkpoint state dict")
+                continue
+            else:
+                assert (
+                    new_key not in checkpoint_state_dict
+                ), f"Key {new_key} already exists in checkpoint state dict"
+                checkpoint_state_dict[new_key] = checkpoint_state_dict.pop(old_key)
+                logging.info(f"Renaming {old_key} to {new_key}")
+        return checkpoint_state_dict

libcom/inharmonious_region_localization/__init__.py

@@ -0,0 +1,4 @@
+# change to your lib name
+from .inharmonious_region_localization import InharmoniousLocalizationModel
+
+__all__ = ['InharmoniousLocalizationModel']

libcom/inharmonious_region_localization/inharmonious_region_localization.py

@@ -0,0 +1,97 @@
+import torch
+import torchvision
+from libcom.utils.model_download import download_pretrained_model
+from libcom.utils.process_image import *
+from libcom.utils.environment import *
+import os
+import torchvision.transforms as transforms
+from .source.madis_net import *
+
+cur_dir   = os.path.dirname(os.path.abspath(__file__))
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+model_set = ['IHDRNet'] 
+
+class InharmoniousLocalizationModel:
+    """
+    Inharmonious region localization model.
+
+    Args:
+        device (str | torch.device): gpu id
+        model_type (str): predefined model type
+        kwargs (dict): other parameters for building model
+
+    Examples:
+        >>> from libcom import InharmoniousLocalizationModel
+        >>> import cv2
+        >>> net = InharmoniousLocalizationModel(device=0)
+        >>> comp_img1  = '../tests/source/composite/comp1_MadisNet.png'
+        >>> inharmonious_localization1 = net(comp_img1)
+        >>> comp_img2  = '../tests/source/composite/comp2_MadisNet.png'
+        >>> inharmonious_localization2 = net(comp_img2)
+        >>> cv2.imwrite('../docs/_static/image/inharmonious_localization_result1.jpg', np.concatenate([cv2.resize(cv2.imread(comp_img1),(256,256)), inharmonious_localization1],axis=1))
+        >>> cv2.imwrite('../docs/_static/image/inharmonious_localization_result2.jpg', np.concatenate([cv2.resize(cv2.imread(comp_img2),(256,256)), inharmonious_localization2],axis=1))
+
+
+    Expected result:
+
+    .. image:: _static/image/inharmonious_localization_result3_4.jpg
+
+    """
+    def __init__(self, device=0, model_type='IHDRNet', **kwargs):
+        assert model_type in model_set, f'Not implementation for {model_type}'
+        self.model_type = model_type
+        self.option = kwargs
+        weight_path_g = os.path.join(model_dir, 'pretrained_models', 'Inharmonious_G.pth')
+        download_pretrained_model(weight_path_g)
+        weight_path_ihdrnet = os.path.join(model_dir, 'pretrained_models', 'IHDRNet.pth')
+        download_pretrained_model(weight_path_ihdrnet)
+        self.device = check_gpu_device(device)
+        self.build_pretrained_model(weight_path_g, weight_path_ihdrnet)
+        self.build_data_transformer()
+
+    def build_pretrained_model(self, weight_path_g, weight_path_ihdrnet):
+        model = MadisNet()
+        model.g.load_state_dict(torch.load(weight_path_g, weights_only=True)['state_dict'])
+        model.ihdrnet.load_state_dict(torch.load(weight_path_ihdrnet, weights_only=True)['state_dict'])
+        self.MadisNet_model = model.to(self.device).eval()
+        
+    def build_data_transformer(self):
+        self.transformer = transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    (0.485, 0.456, 0.406),
+                    (0.229, 0.224, 0.225)
+                )
+            ])
+    
+    def inputs_preprocess(self, composite_image):
+        img = read_image_opencv(composite_image)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = cv2.resize(img, (256,256))
+        img = self.transformer(img).float().to(self.device).unsqueeze(0)
+        return img
+    
+    def outputs_postprocess(self, outputs):
+        if len(outputs.shape) == 4:
+            outputs = outputs.squeeze(0)
+        outputs = (torch.clamp(255.0 * outputs.permute(1, 2, 0), 0, 255)).cpu().numpy()
+        outputs = cv2.cvtColor(outputs, cv2.COLOR_RGB2BGR)
+        return outputs
+    
+
+    @torch.no_grad()
+    def __call__(self, composite_image):
+        """
+        Given a composite image, predict the mask of the inharmonious region.
+
+        Args:
+            composite_image (str | numpy.ndarray): The path to composite image or the compposite image in ndarray form.
+
+        Returns:
+            inharmonious_mask (np.array): The inharmonious mask.
+        
+        """
+        img = self.inputs_preprocess(composite_image)
+        outputs = self.MadisNet_model(img)[0]
+        preds = self.outputs_postprocess(outputs)
+        return preds

libcom/inharmonious_region_localization/source/blocks.py

@@ -0,0 +1,368 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+import math
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=nn.BatchNorm2d):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        if downsample is None and inplanes != planes:
+            self.downsample = nn.Conv2d(inplanes, planes, 1,1,0)
+        else:
+            self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+# Used for spatial attention
+class BasicConv(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
+        super(BasicConv, self).__init__()
+        self.out_channels = out_planes
+        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+        self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
+        self.relu = nn.ReLU() if relu else None
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn is not None:
+            x = self.bn(x)
+        if self.relu is not None:
+            x = self.relu(x)
+        return x
+
+
+#  Central Difference Convolutional Network
+class Conv2d_cd(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
+                 padding=1, dilation=1, groups=1, bias=False, theta=0.7, use_bn=True, actv='relu'):
+
+        super(Conv2d_cd, self).__init__() 
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+        self.theta = theta
+        self.bn = nn.BatchNorm2d(out_channels)
+        if actv == 'relu':
+            self.actv = nn.ReLU(True)
+        elif actv == 'lrelu':
+            self.actv = nn.LeakyReLU(0.2)
+        elif actv == 'tanh':
+            self.actv = nn.Tanh()
+        elif actv == 'sigmoid':
+            self.actv = nn.Sigmoid()
+        elif actv == 'elu':
+            self.actv = nn.ELU()
+        else:
+            self.actv = nn.Identity()
+
+    def forward(self, x):
+        out_normal = self.conv(x)
+
+        if math.fabs(self.theta - 0.0) < 1e-6:
+            diff =  out_normal 
+        else:
+            #pdb.set_trace()
+            [C_out,C_in, kernel_size,kernel_size] = self.conv.weight.shape
+            kernel_diff = self.conv.weight.sum(2).sum(2)
+            kernel_diff = kernel_diff[:, :, None, None]
+            out_diff = F.conv2d(input=x, weight=kernel_diff, bias=self.conv.bias, stride=self.conv.stride, padding=0, groups=self.conv.groups)
+            
+            diff =  out_normal - self.theta * out_diff
+        out = self.bn(diff)
+        out = self.actv(out)
+        return out
+
+class PartialConv2d(nn.Module):
+    def __init__(self, conv_module):
+        super(PartialConv2d, self).__init__()
+        # whether the mask is multi-channel or not
+        self.multi_channel = False
+        self.return_mask = True
+        self.kernel_size = conv_module.kernel_size
+        self.stride = conv_module.stride
+        self.padding = conv_module.padding
+        self.dilation = conv_module.dilation
+        self.bias = conv_module.bias
+        self.in_channels = conv_module.in_channels
+        self.out_channels = conv_module.out_channels
+
+        self.conv = conv_module
+
+        self.weight_maskUpdater = torch.ones(1, 1, self.kernel_size[0], self.kernel_size[1])
+
+        self.slide_winsize = self.weight_maskUpdater.shape[1] * self.weight_maskUpdater.shape[2] * \
+                             self.weight_maskUpdater.shape[3]
+
+        self.last_size = (None, None, None, None)
+        self.update_mask = None
+        self.mask_ratio = None
+
+    def forward(self, input, mask_in=None):
+        assert len(input.shape) == 4
+        if mask_in is not None or self.last_size != tuple(input.shape):
+            self.last_size = tuple(input.shape)
+
+            with torch.no_grad():
+                if self.weight_maskUpdater.type() != input.type():
+                    self.weight_maskUpdater = self.weight_maskUpdater.to(input)
+
+                if mask_in is None:
+                    # if mask is not provided, create a mask
+                    mask = torch.ones(1, 1, input.data.shape[2], input.data.shape[3]).to(input)
+                else:
+                    mask = mask_in
+
+                self.update_mask = F.conv2d(mask, self.weight_maskUpdater, bias=None, stride=self.stride,
+                                            padding=self.padding, dilation=self.dilation, groups=1)
+
+                self.mask_ratio = self.slide_winsize / (self.update_mask + 1e-8)
+                self.update_mask = torch.clamp(self.update_mask, 0, 1)
+                self.mask_ratio = torch.mul(self.mask_ratio, self.update_mask)
+
+
+        raw_out = self.conv.forward(torch.mul(input, mask) if mask_in is not None else input)
+
+        if self.bias is not None:
+            bias_view = self.bias.view(1, self.out_channels, 1, 1)
+            output = torch.mul(raw_out - bias_view, self.mask_ratio) + bias_view
+            output = torch.mul(output, self.update_mask)
+        else:
+            output = torch.mul(raw_out, self.mask_ratio)
+
+        if self.return_mask:
+            return output, self.update_mask
+        else:
+            return output
+
+class ResNetBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride = 1, bias=False, actv='relu'):
+        super(ResNetBlock, self).__init__()
+        self.bn1 = nn.BatchNorm2d(int(in_channels*(1.5)))
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.conv1 = nn.Conv2d(in_channels, int(in_channels*(1.5)), 3,1,1, bias=False)
+        self.conv2 = nn.Conv2d(int(in_channels*(1.5)), out_channels, 3,1,1, bias=False)
+        self.adapt = nn.Conv2d(in_channels, out_channels, 1,1,0, bias=False) if in_channels != out_channels else nn.Identity()
+        if actv == 'relu':
+            self.actv = nn.ReLU(True)
+        elif actv == 'lrelu':
+            self.actv = nn.LeakyReLU(0.2)
+        elif actv == 'tanh':
+            self.actv = nn.Tanh()
+        elif actv == 'sigmoid':
+            self.actv = nn.Sigmoid()
+        elif actv == 'elu':
+            self.actv = nn.ELU() 
+
+    def forward(self, x):
+        res = x
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+
+        res = self.adapt(x)
+        out = res + x
+        return self.actv(out)
+
+
+
+
+
+class Conv2dBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, ks, st, padding=0, dilation=1,
+                 norm='none', activation='relu', pad_type='zero', 
+                 use_bias=True, activation_first=False):
+        super(Conv2dBlock, self).__init__()
+        self.use_bias = use_bias
+        self.activation_first = activation_first
+        # initialize padding
+        if pad_type == 'reflect':
+            self.pad = nn.ReflectionPad2d(padding)
+        elif pad_type == 'replicate':
+            self.pad = nn.ReplicationPad2d(padding)
+        elif pad_type == 'zero':
+            self.pad = nn.ZeroPad2d(padding)
+        elif pad_type == 'none':
+            self.pad = None
+        else:
+            assert 0, "Unsupported padding type: {}".format(pad_type)
+
+        # initialize normalization
+        norm_dim = out_dim
+        if norm == 'bn':
+            self.norm = nn.BatchNorm2d(norm_dim)
+        elif norm == 'in':
+            self.norm = nn.InstanceNorm2d(norm_dim)
+        elif norm == 'adain':
+            self.norm = AdaptiveInstanceNorm2d(norm_dim)
+        elif norm == 'none':
+            self.norm = None
+        else:
+            assert 0, "Unsupported normalization: {}".format(norm)
+
+        # initialize activation
+        if activation == 'relu':
+            if activation_first == True:
+                self.activation = nn.ReLU(inplace=False)
+            else:
+                self.activation = nn.ReLU(inplace=True)
+        elif activation == 'lrelu':
+            self.activation = nn.LeakyReLU(0.2, inplace=True)
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'none':
+            self.activation = None
+        else:
+            assert 0, "Unsupported activation: {}".format(activation)
+
+        self.conv = nn.Conv2d(in_dim, out_dim, ks, st, dilation=dilation, bias=self.use_bias)
+
+    def forward(self, x):
+        if self.activation_first:
+            if self.activation:
+                x = self.activation(x)
+            if self.pad is not None:
+                x = self.conv(self.pad(x))
+            else:
+                x = self.conv(x)
+            
+            if self.norm:
+                x = self.norm(x)
+        else:
+            if self.pad is not None:
+                x = self.conv(self.pad(x))
+            else:
+                x = self.conv(x)
+            
+            if self.norm:
+                x = self.norm(x)
+            if self.activation:
+                x = self.activation(x)
+        return x
+
+class Swish(nn.Module):
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+class SwishImplementation(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, i):
+        result = i * torch.sigmoid(i)
+        ctx.save_for_backward(i)
+        return result
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        i = ctx.saved_variables[0]
+        sigmoid_i = torch.sigmoid(i)
+        return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i)))
+
+
+class MemoryEfficientSwish(nn.Module):
+    def forward(self, x):
+        return SwishImplementation.apply(x)
+
+
+
+class AdaptiveInstanceNorm2d(nn.Module):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1):
+        super(AdaptiveInstanceNorm2d, self).__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.weight = None
+        self.bias = None
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+
+    def forward(self, x):
+        assert self.weight is not None and \
+               self.bias is not None, "Please assign AdaIN weight first"
+        b, c = x.size(0), x.size(1)
+        running_mean = self.running_mean.repeat(b)
+        running_var = self.running_var.repeat(b)
+        x_reshaped = x.contiguous().view(1, b * c, *x.size()[2:])
+        out = F.batch_norm(
+            x_reshaped, running_mean, running_var, self.weight, self.bias,
+            True, self.momentum, self.eps)
+        return out.view(b, c, *x.size()[2:])
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(' + str(self.num_features) + ')'