zparadox/stable-video-diffusion

1import functools2 3import torch.nn as nn4 5from ..util import ActNorm6 7 8def weights_init(m):9    classname = m.__class__.__name__10    if classname.find("Conv") != -1:11        nn.init.normal_(m.weight.data, 0.0, 0.02)12    elif classname.find("BatchNorm") != -1:13        nn.init.normal_(m.weight.data, 1.0, 0.02)14        nn.init.constant_(m.bias.data, 0)15 16 17class NLayerDiscriminator(nn.Module):18    """Defines a PatchGAN discriminator as in Pix2Pix19    --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py20    """21 22    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):23        """Construct a PatchGAN discriminator24        Parameters:25            input_nc (int)  -- the number of channels in input images26            ndf (int)       -- the number of filters in the last conv layer27            n_layers (int)  -- the number of conv layers in the discriminator28            norm_layer      -- normalization layer29        """30        super(NLayerDiscriminator, self).__init__()31        if not use_actnorm:32            norm_layer = nn.BatchNorm2d33        else:34            norm_layer = ActNorm35        if (36            type(norm_layer) == functools.partial37        ):  # no need to use bias as BatchNorm2d has affine parameters38            use_bias = norm_layer.func != nn.BatchNorm2d39        else:40            use_bias = norm_layer != nn.BatchNorm2d41 42        kw = 443        padw = 144        sequence = [45            nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),46            nn.LeakyReLU(0.2, True),47        ]48        nf_mult = 149        nf_mult_prev = 150        for n in range(1, n_layers):  # gradually increase the number of filters51            nf_mult_prev = nf_mult52            nf_mult = min(2**n, 8)53            sequence += [54                nn.Conv2d(55                    ndf * nf_mult_prev,56                    ndf * nf_mult,57                    kernel_size=kw,58                    stride=2,59                    padding=padw,60                    bias=use_bias,61                ),62                norm_layer(ndf * nf_mult),63                nn.LeakyReLU(0.2, True),64            ]65 66        nf_mult_prev = nf_mult67        nf_mult = min(2**n_layers, 8)68        sequence += [69            nn.Conv2d(70                ndf * nf_mult_prev,71                ndf * nf_mult,72                kernel_size=kw,73                stride=1,74                padding=padw,75                bias=use_bias,76            ),77            norm_layer(ndf * nf_mult),78            nn.LeakyReLU(0.2, True),79        ]80 81        sequence += [82            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)83        ]  # output 1 channel prediction map84        self.main = nn.Sequential(*sequence)85 86    def forward(self, input):87        """Standard forward."""88        return self.main(input)89