| """Code is adapted from https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/ldm/modules/diffusionmodules/openaimodel.py""" |
| from abc import abstractmethod |
| from functools import partial |
| import math |
| from typing import Iterable |
|
|
| import numpy as np |
| import torch as th |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .utils import ( |
| checkpoint, |
| conv_nd, |
| linear, |
| avg_pool_nd, |
| zero_module, |
| normalization, |
| ) |
|
|
|
|
| class TimestepBlock(nn.Module): |
| """ |
| Any module where forward() takes timestep embeddings as a second argument. |
| """ |
|
|
| @abstractmethod |
| def forward(self, x, emb): |
| """ |
| Apply the module to `x` given `emb` timestep embeddings. |
| """ |
|
|
|
|
| class Upsample(nn.Module): |
| """ |
| An upsampling layer with an optional convolution. |
| :param channels: channels in the inputs and outputs. |
| :param use_conv: a bool determining if a convolution is applied. |
| :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then |
| upsampling occurs in the inner-two dimensions. |
| """ |
|
|
| def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1): |
| super().__init__() |
| self.channels = channels |
| self.out_channels = out_channels or channels |
| self.use_conv = use_conv |
| self.dims = dims |
| if use_conv: |
| self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding) |
|
|
| def forward(self, x): |
| assert x.shape[1] == self.channels |
| if self.dims == 3: |
| x = F.interpolate( |
| x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest" |
| ) |
| else: |
| x = F.interpolate(x, scale_factor=2, mode="nearest") |
| if self.use_conv: |
| x = self.conv(x) |
| return x |
|
|
|
|
| class TransposedUpsample(nn.Module): |
| 'Learned 2x upsampling without padding' |
| def __init__(self, channels, out_channels=None, ks=5): |
| super().__init__() |
| self.channels = channels |
| self.out_channels = out_channels or channels |
|
|
| self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2) |
|
|
| def forward(self,x): |
| return self.up(x) |
|
|
|
|
| class Downsample(nn.Module): |
| """ |
| A downsampling layer with an optional convolution. |
| :param channels: channels in the inputs and outputs. |
| :param use_conv: a bool determining if a convolution is applied. |
| :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then |
| downsampling occurs in the inner-two dimensions. |
| """ |
|
|
| def __init__(self, channels, use_conv, dims=2, out_channels=None,padding=1): |
| super().__init__() |
| self.channels = channels |
| self.out_channels = out_channels or channels |
| self.use_conv = use_conv |
| self.dims = dims |
| stride = 2 if dims != 3 else (1, 2, 2) |
| if use_conv: |
| self.op = conv_nd( |
| dims, self.channels, self.out_channels, 3, stride=stride, padding=padding |
| ) |
| else: |
| assert self.channels == self.out_channels |
| self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) |
|
|
| def forward(self, x): |
| assert x.shape[1] == self.channels |
| return self.op(x) |
|
|
|
|
| class ResBlock(TimestepBlock): |
| """ |
| A residual block that can optionally change the number of channels. |
| :param channels: the number of input channels. |
| :param emb_channels: the number of timestep embedding channels. |
| :param dropout: the rate of dropout. |
| :param out_channels: if specified, the number of out channels. |
| :param use_conv: if True and out_channels is specified, use a spatial |
| convolution instead of a smaller 1x1 convolution to change the |
| channels in the skip connection. |
| :param dims: determines if the signal is 1D, 2D, or 3D. |
| :param use_checkpoint: if True, use gradient checkpointing on this module. |
| :param up: if True, use this block for upsampling. |
| :param down: if True, use this block for downsampling. |
| """ |
|
|
| def __init__( |
| self, |
| channels, |
| emb_channels, |
| dropout, |
| out_channels=None, |
| use_conv=False, |
| use_scale_shift_norm=False, |
| dims=2, |
| use_checkpoint=False, |
| up=False, |
| down=False, |
| ): |
| super().__init__() |
| self.channels = channels |
| self.emb_channels = emb_channels |
| self.dropout = dropout |
| self.out_channels = out_channels or channels |
| self.use_conv = use_conv |
| self.use_checkpoint = use_checkpoint |
| self.use_scale_shift_norm = use_scale_shift_norm |
|
|
| self.in_layers = nn.Sequential( |
| normalization(channels), |
| nn.SiLU(), |
| conv_nd(dims, channels, self.out_channels, 3, padding=1), |
| ) |
|
|
| self.updown = up or down |
|
|
| if up: |
| self.h_upd = Upsample(channels, False, dims) |
| self.x_upd = Upsample(channels, False, dims) |
| elif down: |
| self.h_upd = Downsample(channels, False, dims) |
| self.x_upd = Downsample(channels, False, dims) |
| else: |
| self.h_upd = self.x_upd = nn.Identity() |
|
|
| self.emb_layers = nn.Sequential( |
| nn.SiLU(), |
| linear( |
| emb_channels, |
| 2 * self.out_channels if use_scale_shift_norm else self.out_channels, |
| ), |
| ) |
| self.out_layers = nn.Sequential( |
| normalization(self.out_channels), |
| nn.SiLU(), |
| nn.Dropout(p=dropout), |
| zero_module( |
| conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) |
| ), |
| ) |
|
|
| if self.out_channels == channels: |
| self.skip_connection = nn.Identity() |
| elif use_conv: |
| self.skip_connection = conv_nd( |
| dims, channels, self.out_channels, 3, padding=1 |
| ) |
| else: |
| self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) |
|
|
| def forward(self, x, emb): |
| """ |
| Apply the block to a Tensor, conditioned on a timestep embedding. |
| :param x: an [N x C x ...] Tensor of features. |
| :param emb: an [N x emb_channels] Tensor of timestep embeddings. |
| :return: an [N x C x ...] Tensor of outputs. |
| """ |
| return checkpoint( |
| self._forward, (x, emb), self.parameters(), self.use_checkpoint |
| ) |
|
|
| def _forward(self, x, emb): |
| if self.updown: |
| in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] |
| h = in_rest(x) |
| h = self.h_upd(h) |
| x = self.x_upd(x) |
| h = in_conv(h) |
| else: |
| h = self.in_layers(x) |
| emb_out = self.emb_layers(emb).type(h.dtype) |
| while len(emb_out.shape) < len(h.shape): |
| emb_out = emb_out[..., None] |
| if self.use_scale_shift_norm: |
| out_norm, out_rest = self.out_layers[0], self.out_layers[1:] |
| scale, shift = th.chunk(emb_out, 2, dim=1) |
| h = out_norm(h) * (1 + scale) + shift |
| h = out_rest(h) |
| else: |
| h = h + emb_out |
| h = self.out_layers(h) |
| return self.skip_connection(x) + h |
|
|
