coreml-community/ControlNet-v1-1-Annotators-cpu

1# Copyright (c) Facebook, Inc. and its affiliates.2 3from copy import deepcopy4import fvcore.nn.weight_init as weight_init5import torch6from torch import nn7from torch.nn import functional as F8 9from .batch_norm import get_norm10from .blocks import DepthwiseSeparableConv2d11from .wrappers import Conv2d12 13 14class ASPP(nn.Module):15    """16    Atrous Spatial Pyramid Pooling (ASPP).17    """18 19    def __init__(20        self,21        in_channels,22        out_channels,23        dilations,24        *,25        norm,26        activation,27        pool_kernel_size=None,28        dropout: float = 0.0,29        use_depthwise_separable_conv=False,30    ):31        """32        Args:33            in_channels (int): number of input channels for ASPP.34            out_channels (int): number of output channels.35            dilations (list): a list of 3 dilations in ASPP.36            norm (str or callable): normalization for all conv layers.37                See :func:`layers.get_norm` for supported format. norm is38                applied to all conv layers except the conv following39                global average pooling.40            activation (callable): activation function.41            pool_kernel_size (tuple, list): the average pooling size (kh, kw)42                for image pooling layer in ASPP. If set to None, it always43                performs global average pooling. If not None, it must be44                divisible by the shape of inputs in forward(). It is recommended45                to use a fixed input feature size in training, and set this46                option to match this size, so that it performs global average47                pooling in training, and the size of the pooling window stays48                consistent in inference.49            dropout (float): apply dropout on the output of ASPP. It is used in50                the official DeepLab implementation with a rate of 0.1:51                https://github.com/tensorflow/models/blob/21b73d22f3ed05b650e85ac50849408dd36de32e/research/deeplab/model.py#L532  # noqa52            use_depthwise_separable_conv (bool): use DepthwiseSeparableConv2d53                for 3x3 convs in ASPP, proposed in :paper:`DeepLabV3+`.54        """55        super(ASPP, self).__init__()56        assert len(dilations) == 3, "ASPP expects 3 dilations, got {}".format(len(dilations))57        self.pool_kernel_size = pool_kernel_size58        self.dropout = dropout59        use_bias = norm == ""60        self.convs = nn.ModuleList()61        # conv 1x162        self.convs.append(63            Conv2d(64                in_channels,65                out_channels,66                kernel_size=1,67                bias=use_bias,68                norm=get_norm(norm, out_channels),69                activation=deepcopy(activation),70            )71        )72        weight_init.c2_xavier_fill(self.convs[-1])73        # atrous convs74        for dilation in dilations:75            if use_depthwise_separable_conv:76                self.convs.append(77                    DepthwiseSeparableConv2d(78                        in_channels,79                        out_channels,80                        kernel_size=3,81                        padding=dilation,82                        dilation=dilation,83                        norm1=norm,84                        activation1=deepcopy(activation),85                        norm2=norm,86                        activation2=deepcopy(activation),87                    )88                )89            else:90                self.convs.append(91                    Conv2d(92                        in_channels,93                        out_channels,94                        kernel_size=3,95                        padding=dilation,96                        dilation=dilation,97                        bias=use_bias,98                        norm=get_norm(norm, out_channels),99                        activation=deepcopy(activation),100                    )101                )102                weight_init.c2_xavier_fill(self.convs[-1])103        # image pooling104        # We do not add BatchNorm because the spatial resolution is 1x1,105        # the original TF implementation has BatchNorm.106        if pool_kernel_size is None:107            image_pooling = nn.Sequential(108                nn.AdaptiveAvgPool2d(1),109                Conv2d(in_channels, out_channels, 1, bias=True, activation=deepcopy(activation)),110            )111        else:112            image_pooling = nn.Sequential(113                nn.AvgPool2d(kernel_size=pool_kernel_size, stride=1),114                Conv2d(in_channels, out_channels, 1, bias=True, activation=deepcopy(activation)),115            )116        weight_init.c2_xavier_fill(image_pooling[1])117        self.convs.append(image_pooling)118 119        self.project = Conv2d(120            5 * out_channels,121            out_channels,122            kernel_size=1,123            bias=use_bias,124            norm=get_norm(norm, out_channels),125            activation=deepcopy(activation),126        )127        weight_init.c2_xavier_fill(self.project)128 129    def forward(self, x):130        size = x.shape[-2:]131        if self.pool_kernel_size is not None:132            if size[0] % self.pool_kernel_size[0] or size[1] % self.pool_kernel_size[1]:133                raise ValueError(134                    "`pool_kernel_size` must be divisible by the shape of inputs. "135                    "Input size: {} `pool_kernel_size`: {}".format(size, self.pool_kernel_size)136                )137        res = []138        for conv in self.convs:139            res.append(conv(x))140        res[-1] = F.interpolate(res[-1], size=size, mode="bilinear", align_corners=False)141        res = torch.cat(res, dim=1)142        res = self.project(res)143        res = F.dropout(res, self.dropout, training=self.training) if self.dropout > 0 else res144        return res145