Source code for mmagic.models.editors.global_local.gl_dilation

# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
from mmcv.cnn import ConvModule
from mmengine.model import BaseModule

from mmagic.models.archs import SimpleGatedConvModule
from mmagic.registry import MODELS


@MODELS.register_module()
class GLDilationNeck(BaseModule):
    """Dilation Backbone used in Global&Local model.

    This implementation follows:
    Globally and locally Consistent Image Completion

    Args:
        in_channels (int): Channel number of input feature.
        conv_type (str): The type of conv module. In DeepFillv1 model, the
            `conv_type` should be 'conv'. In DeepFillv2 model, the `conv_type`
            should be 'gated_conv'.
        norm_cfg (dict): Config dict to build norm layer.
        act_cfg (dict): Config dict for activation layer, "relu" by default.
        kwargs (keyword arguments).
    """
    _conv_type = dict(conv=ConvModule, gated_conv=SimpleGatedConvModule)

    def __init__(self,
                 in_channels=256,
                 conv_type='conv',
                 norm_cfg=None,
                 act_cfg=dict(type='ReLU'),
                 **kwargs):
        super().__init__()
        conv_module = self._conv_type[conv_type]
        dilation_convs_ = []
        for i in range(4):
            dilation_ = int(2**(i + 1))
            dilation_convs_.append(
                conv_module(
                    in_channels,
                    in_channels,
                    kernel_size=3,
                    padding=dilation_,
                    dilation=dilation_,
                    stride=1,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg,
                    **kwargs))
        self.dilation_convs = nn.Sequential(*dilation_convs_)

    def forward(self, x):
        """Forward Function.

        Args:
            x (torch.Tensor): Input tensor with shape of (n, c, h, w).

        Returns:
            torch.Tensor: Output tensor with shape of (n, c, h', w').
        """
        x = self.dilation_convs(x)
        return x