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

1# -*- coding: utf-8 -*-2# Copyright (c) Facebook, Inc. and its affiliates.3 4import fvcore.nn.weight_init as weight_init5from torch import nn6 7from .batch_norm import FrozenBatchNorm2d, get_norm8from .wrappers import Conv2d9 10 11"""12CNN building blocks.13"""14 15 16class CNNBlockBase(nn.Module):17    """18    A CNN block is assumed to have input channels, output channels and a stride.19    The input and output of `forward()` method must be NCHW tensors.20    The method can perform arbitrary computation but must match the given21    channels and stride specification.22 23    Attribute:24        in_channels (int):25        out_channels (int):26        stride (int):27    """28 29    def __init__(self, in_channels, out_channels, stride):30        """31        The `__init__` method of any subclass should also contain these arguments.32 33        Args:34            in_channels (int):35            out_channels (int):36            stride (int):37        """38        super().__init__()39        self.in_channels = in_channels40        self.out_channels = out_channels41        self.stride = stride42 43    def freeze(self):44        """45        Make this block not trainable.46        This method sets all parameters to `requires_grad=False`,47        and convert all BatchNorm layers to FrozenBatchNorm48 49        Returns:50            the block itself51        """52        for p in self.parameters():53            p.requires_grad = False54        FrozenBatchNorm2d.convert_frozen_batchnorm(self)55        return self56 57 58class DepthwiseSeparableConv2d(nn.Module):59    """60    A kxk depthwise convolution + a 1x1 convolution.61 62    In :paper:`xception`, norm & activation are applied on the second conv.63    :paper:`mobilenet` uses norm & activation on both convs.64    """65 66    def __init__(67        self,68        in_channels,69        out_channels,70        kernel_size=3,71        padding=1,72        dilation=1,73        *,74        norm1=None,75        activation1=None,76        norm2=None,77        activation2=None,78    ):79        """80        Args:81            norm1, norm2 (str or callable): normalization for the two conv layers.82            activation1, activation2 (callable(Tensor) -> Tensor): activation83                function for the two conv layers.84        """85        super().__init__()86        self.depthwise = Conv2d(87            in_channels,88            in_channels,89            kernel_size=kernel_size,90            padding=padding,91            dilation=dilation,92            groups=in_channels,93            bias=not norm1,94            norm=get_norm(norm1, in_channels),95            activation=activation1,96        )97        self.pointwise = Conv2d(98            in_channels,99            out_channels,100            kernel_size=1,101            bias=not norm2,102            norm=get_norm(norm2, out_channels),103            activation=activation2,104        )105 106        # default initialization107        weight_init.c2_msra_fill(self.depthwise)108        weight_init.c2_msra_fill(self.pointwise)109 110    def forward(self, x):111        return self.pointwise(self.depthwise(x))112