From d1f2e3e47bb878bf623a42f9ac80325f9396fcb2 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Wed, 29 Jun 2022 11:43:30 +0000 Subject: [PATCH 1/2] up --- src/diffusers/models/resnet.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/diffusers/models/resnet.py b/src/diffusers/models/resnet.py index 462f87d8d850..0b5262907a53 100644 --- a/src/diffusers/models/resnet.py +++ b/src/diffusers/models/resnet.py @@ -167,8 +167,8 @@ def forward(self, x): # class GlideUpsample(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. #""" +# 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): # super().__init__() @@ -193,8 +193,8 @@ def forward(self, x): # class LDMUpsample(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. #""" +# 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__() From f15ab901a04a5cff71950ff2c837005f050edeae Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Wed, 29 Jun 2022 11:46:23 +0000 Subject: [PATCH 2/2] fix comments --- src/diffusers/models/unet_glide.py | 107 -------------- src/diffusers/models/unet_grad_tts.py | 20 --- src/diffusers/models/unet_ldm.py | 105 -------------- src/diffusers/models/unet_rl.py | 31 ----- .../models/unet_sde_score_estimation.py | 131 ------------------ 5 files changed, 394 deletions(-) diff --git a/src/diffusers/models/unet_glide.py b/src/diffusers/models/unet_glide.py index c9f06a2d80b4..477c1768ae4e 100644 --- a/src/diffusers/models/unet_glide.py +++ b/src/diffusers/models/unet_glide.py @@ -1,5 +1,3 @@ -from abc import abstractmethod - import torch import torch.nn as nn import torch.nn.functional as F @@ -96,16 +94,6 @@ def zero_module(module): return module -# 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 TimestepEmbedSequential(nn.Sequential, TimestepBlock): """ A sequential module that passes timestep embeddings to the children that support it as an extra input. @@ -122,101 +110,6 @@ def forward(self, x, emb, encoder_out=None): return 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, swish=1.0), -# nn.Identity(), -# conv_nd(dims, channels, self.out_channels, 3, padding=1), -# ) -# -# self.updown = up or down -# -# if up: -# self.h_upd = Upsample(channels, use_conv=False, dims=dims) -# self.x_upd = Upsample(channels, use_conv=False, dims=dims) -# elif down: -# self.h_upd = Downsample(channels, use_conv=False, dims=dims, padding=1, name="op") -# self.x_upd = Downsample(channels, use_conv=False, dims=dims, padding=1, name="op") -# 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, swish=0.0 if use_scale_shift_norm else 1.0), -# nn.SiLU() if use_scale_shift_norm else nn.Identity(), -# 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. #""" -# 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 = torch.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 - - class GlideUNetModel(ModelMixin, ConfigMixin): """ The full UNet model with attention and timestep embedding. diff --git a/src/diffusers/models/unet_grad_tts.py b/src/diffusers/models/unet_grad_tts.py index cb6122c897cd..e6dabe6dd41c 100644 --- a/src/diffusers/models/unet_grad_tts.py +++ b/src/diffusers/models/unet_grad_tts.py @@ -36,26 +36,6 @@ def forward(self, x, mask): return output * mask -# class ResnetBlock(torch.nn.Module): -# def __init__(self, dim, dim_out, time_emb_dim, groups=8): -# super(ResnetBlock, self).__init__() -# self.mlp = torch.nn.Sequential(Mish(), torch.nn.Linear(time_emb_dim, dim_out)) -# -# self.block1 = Block(dim, dim_out, groups=groups) -# self.block2 = Block(dim_out, dim_out, groups=groups) -# if dim != dim_out: -# self.res_conv = torch.nn.Conv2d(dim, dim_out, 1) -# else: -# self.res_conv = torch.nn.Identity() -# -# def forward(self, x, mask, time_emb): -# h = self.block1(x, mask) -# h += self.mlp(time_emb).unsqueeze(-1).unsqueeze(-1) -# h = self.block2(h, mask) -# output = h + self.res_conv(x * mask) -# return output - - class Residual(torch.nn.Module): def __init__(self, fn): super(Residual, self).__init__() diff --git a/src/diffusers/models/unet_ldm.py b/src/diffusers/models/unet_ldm.py index 650b262c9626..0571013d9d87 100644 --- a/src/diffusers/models/unet_ldm.py +++ b/src/diffusers/models/unet_ldm.py @@ -1,5 +1,4 @@ import math -from abc import abstractmethod from inspect import isfunction import numpy as np @@ -328,7 +327,6 @@ def normalization(channels, swish=0.0): return GroupNorm32(num_channels=channels, num_groups=32, swish=swish) -## go class AttentionPool2d(nn.Module): """ Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py @@ -359,16 +357,6 @@ def forward(self, x): return x[:, :, 0] -# 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 TimestepEmbedSequential(nn.Sequential, TimestepBlock): """ A sequential module that passes timestep embeddings to the children that support it as an extra input. @@ -385,99 +373,6 @@ def forward(self, x, emb, context=None): return x -# class A_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, use_conv=False, dims=dims) -# self.x_upd = Upsample(channels, use_conv=False, dims=dims) -# elif down: -# self.h_upd = Downsample(channels, use_conv=False, dims=dims, padding=1, name="op") -# self.x_upd = Downsample(channels, use_conv=False, dims=dims, padding=1, name="op") -# 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): -# 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 = torch.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 -# - - class QKVAttention(nn.Module): """ A module which performs QKV attention and splits in a different order. diff --git a/src/diffusers/models/unet_rl.py b/src/diffusers/models/unet_rl.py index c4e59b1001df..2dea6f21069b 100644 --- a/src/diffusers/models/unet_rl.py +++ b/src/diffusers/models/unet_rl.py @@ -73,37 +73,6 @@ def forward(self, x): return self.block(x) -# class ResidualTemporalBlock(nn.Module): -# def __init__(self, inp_channels, out_channels, embed_dim, horizon, kernel_size=5): -# super().__init__() -# -# self.blocks = nn.ModuleList( -# [ -# Conv1dBlock(inp_channels, out_channels, kernel_size), -# Conv1dBlock(out_channels, out_channels, kernel_size), -# ] -# ) -# -# self.time_mlp = nn.Sequential( -# nn.Mish(), -# nn.Linear(embed_dim, out_channels), -# RearrangeDim(), -# Rearrange("batch t -> batch t 1"), -# ) -# -# self.residual_conv = ( -# nn.Conv1d(inp_channels, out_channels, 1) if inp_channels != out_channels else nn.Identity() -# ) -# -# def forward(self, x, t): -# """ -# x : [ batch_size x inp_channels x horizon ] t : [ batch_size x embed_dim ] returns: out : [ batch_size x # -out_channels x horizon ] #""" -# out = self.blocks[0](x) + self.time_mlp(t) -# out = self.blocks[1](out) -# return out + self.residual_conv(x) - - class TemporalUNet(ModelMixin, ConfigMixin): # (nn.Module): def __init__( self, diff --git a/src/diffusers/models/unet_sde_score_estimation.py b/src/diffusers/models/unet_sde_score_estimation.py index c03b513ca6ff..8acf337268e6 100644 --- a/src/diffusers/models/unet_sde_score_estimation.py +++ b/src/diffusers/models/unet_sde_score_estimation.py @@ -491,137 +491,6 @@ def forward(self, x): return x -# class ResnetBlockDDPMpp(nn.Module): -# """ResBlock adapted from DDPM.""" -# -# def __init__( -# self, -# act, -# in_ch, -# out_ch=None, -# temb_dim=None, -# conv_shortcut=False, -# dropout=0.1, -# skip_rescale=False, -# init_scale=0.0, -# ): -# super().__init__() -# out_ch = out_ch if out_ch else in_ch -# self.GroupNorm_0 = nn.GroupNorm(num_groups=min(in_ch // 4, 32), num_channels=in_ch, eps=1e-6) -# self.Conv_0 = conv3x3(in_ch, out_ch) -# if temb_dim is not None: -# self.Dense_0 = nn.Linear(temb_dim, out_ch) -# self.Dense_0.weight.data = default_init()(self.Dense_0.weight.data.shape) -# nn.init.zeros_(self.Dense_0.bias) -# self.GroupNorm_1 = nn.GroupNorm(num_groups=min(out_ch // 4, 32), num_channels=out_ch, eps=1e-6) -# self.Dropout_0 = nn.Dropout(dropout) -# self.Conv_1 = conv3x3(out_ch, out_ch, init_scale=init_scale) -# if in_ch != out_ch: -# if conv_shortcut: -# self.Conv_2 = conv3x3(in_ch, out_ch) -# else: -# self.NIN_0 = NIN(in_ch, out_ch) -# -# self.skip_rescale = skip_rescale -# self.act = act -# self.out_ch = out_ch -# self.conv_shortcut = conv_shortcut -# -# def forward(self, x, temb=None): -# h = self.act(self.GroupNorm_0(x)) -# h = self.Conv_0(h) -# if temb is not None: -# h += self.Dense_0(self.act(temb))[:, :, None, None] -# h = self.act(self.GroupNorm_1(h)) -# h = self.Dropout_0(h) -# h = self.Conv_1(h) -# if x.shape[1] != self.out_ch: -# if self.conv_shortcut: -# x = self.Conv_2(x) -# else: -# x = self.NIN_0(x) -# if not self.skip_rescale: -# return x + h -# else: -# return (x + h) / np.sqrt(2.0) - - -# class ResnetBlockBigGANpp(nn.Module): -# def __init__( -# self, -# act, -# in_ch, -# out_ch=None, -# temb_dim=None, -# up=False, -# down=False, -# dropout=0.1, -# fir=False, -# fir_kernel=(1, 3, 3, 1), -# skip_rescale=True, -# init_scale=0.0, -# ): -# super().__init__() -# -# out_ch = out_ch if out_ch else in_ch -# self.GroupNorm_0 = nn.GroupNorm(num_groups=min(in_ch // 4, 32), num_channels=in_ch, eps=1e-6) -# self.up = up -# self.down = down -# self.fir = fir -# self.fir_kernel = fir_kernel -# -# self.Conv_0 = conv3x3(in_ch, out_ch) -# if temb_dim is not None: -# self.Dense_0 = nn.Linear(temb_dim, out_ch) -# self.Dense_0.weight.data = default_init()(self.Dense_0.weight.shape) -# nn.init.zeros_(self.Dense_0.bias) -# -# self.GroupNorm_1 = nn.GroupNorm(num_groups=min(out_ch // 4, 32), num_channels=out_ch, eps=1e-6) -# self.Dropout_0 = nn.Dropout(dropout) -# self.Conv_1 = conv3x3(out_ch, out_ch, init_scale=init_scale) -# if in_ch != out_ch or up or down: -# self.Conv_2 = conv1x1(in_ch, out_ch) -# -# self.skip_rescale = skip_rescale -# self.act = act -# self.in_ch = in_ch -# self.out_ch = out_ch -# -# def forward(self, x, temb=None): -# h = self.act(self.GroupNorm_0(x)) -# -# if self.up: -# if self.fir: -# h = upsample_2d(h, self.fir_kernel, factor=2) -# x = upsample_2d(x, self.fir_kernel, factor=2) -# else: -# h = naive_upsample_2d(h, factor=2) -# x = naive_upsample_2d(x, factor=2) -# elif self.down: -# if self.fir: -# h = downsample_2d(h, self.fir_kernel, factor=2) -# x = downsample_2d(x, self.fir_kernel, factor=2) -# else: -# h = naive_downsample_2d(h, factor=2) -# x = naive_downsample_2d(x, factor=2) -# -# h = self.Conv_0(h) -# Add bias to each feature map conditioned on the time embedding -# if temb is not None: -# h += self.Dense_0(self.act(temb))[:, :, None, None] -# h = self.act(self.GroupNorm_1(h)) -# h = self.Dropout_0(h) -# h = self.Conv_1(h) -# -# if self.in_ch != self.out_ch or self.up or self.down: -# x = self.Conv_2(x) -# -# if not self.skip_rescale: -# return x + h -# else: -# return (x + h) / np.sqrt(2.0) - - class NCSNpp(ModelMixin, ConfigMixin): """NCSN++ model"""