Spaces:
Sleeping
Sleeping
| #taken from: https://github.com/lllyasviel/ControlNet | |
| #and modified | |
| import torch | |
| import torch as th | |
| import torch.nn as nn | |
| from ..ldm.modules.diffusionmodules.util import ( | |
| zero_module, | |
| timestep_embedding, | |
| ) | |
| from ..ldm.modules.attention import SpatialTransformer | |
| from ..ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample | |
| from ..ldm.util import exists | |
| import comfy.ops | |
| class ControlledUnetModel(UNetModel): | |
| #implemented in the ldm unet | |
| pass | |
| class ControlNet(nn.Module): | |
| def __init__( | |
| self, | |
| image_size, | |
| in_channels, | |
| model_channels, | |
| hint_channels, | |
| num_res_blocks, | |
| dropout=0, | |
| channel_mult=(1, 2, 4, 8), | |
| conv_resample=True, | |
| dims=2, | |
| num_classes=None, | |
| use_checkpoint=False, | |
| dtype=torch.float32, | |
| num_heads=-1, | |
| num_head_channels=-1, | |
| num_heads_upsample=-1, | |
| use_scale_shift_norm=False, | |
| resblock_updown=False, | |
| use_new_attention_order=False, | |
| use_spatial_transformer=False, # custom transformer support | |
| transformer_depth=1, # custom transformer support | |
| context_dim=None, # custom transformer support | |
| n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model | |
| legacy=True, | |
| disable_self_attentions=None, | |
| num_attention_blocks=None, | |
| disable_middle_self_attn=False, | |
| use_linear_in_transformer=False, | |
| adm_in_channels=None, | |
| transformer_depth_middle=None, | |
| transformer_depth_output=None, | |
| attn_precision=None, | |
| device=None, | |
| operations=comfy.ops.disable_weight_init, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| assert use_spatial_transformer == True, "use_spatial_transformer has to be true" | |
| if use_spatial_transformer: | |
| assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...' | |
| if context_dim is not None: | |
| assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...' | |
| # from omegaconf.listconfig import ListConfig | |
| # if type(context_dim) == ListConfig: | |
| # context_dim = list(context_dim) | |
| if num_heads_upsample == -1: | |
| num_heads_upsample = num_heads | |
| if num_heads == -1: | |
| assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' | |
| if num_head_channels == -1: | |
| assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' | |
| self.dims = dims | |
| self.image_size = image_size | |
| self.in_channels = in_channels | |
| self.model_channels = model_channels | |
| if isinstance(num_res_blocks, int): | |
| self.num_res_blocks = len(channel_mult) * [num_res_blocks] | |
| else: | |
| if len(num_res_blocks) != len(channel_mult): | |
| raise ValueError("provide num_res_blocks either as an int (globally constant) or " | |
| "as a list/tuple (per-level) with the same length as channel_mult") | |
| self.num_res_blocks = num_res_blocks | |
| if disable_self_attentions is not None: | |
| # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not | |
| assert len(disable_self_attentions) == len(channel_mult) | |
| if num_attention_blocks is not None: | |
| assert len(num_attention_blocks) == len(self.num_res_blocks) | |
| assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks)))) | |
| transformer_depth = transformer_depth[:] | |
| self.dropout = dropout | |
| self.channel_mult = channel_mult | |
| self.conv_resample = conv_resample | |
| self.num_classes = num_classes | |
| self.use_checkpoint = use_checkpoint | |
| self.dtype = dtype | |
| self.num_heads = num_heads | |
| self.num_head_channels = num_head_channels | |
| self.num_heads_upsample = num_heads_upsample | |
| self.predict_codebook_ids = n_embed is not None | |
| time_embed_dim = model_channels * 4 | |
| self.time_embed = nn.Sequential( | |
| operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device), | |
| ) | |
| if self.num_classes is not None: | |
| if isinstance(self.num_classes, int): | |
| self.label_emb = nn.Embedding(num_classes, time_embed_dim) | |
| elif self.num_classes == "continuous": | |
| print("setting up linear c_adm embedding layer") | |
| self.label_emb = nn.Linear(1, time_embed_dim) | |
| elif self.num_classes == "sequential": | |
| assert adm_in_channels is not None | |
| self.label_emb = nn.Sequential( | |
| nn.Sequential( | |
| operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device), | |
| ) | |
| ) | |
| else: | |
| raise ValueError() | |
| self.input_blocks = nn.ModuleList( | |
| [ | |
| TimestepEmbedSequential( | |
| operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device) | |
| ) | |
| ] | |
| ) | |
| self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels, operations=operations, dtype=self.dtype, device=device)]) | |
| self.input_hint_block = TimestepEmbedSequential( | |
| operations.conv_nd(dims, hint_channels, 16, 3, padding=1, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 16, 16, 3, padding=1, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 16, 32, 3, padding=1, stride=2, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 32, 32, 3, padding=1, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 32, 96, 3, padding=1, stride=2, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 96, 96, 3, padding=1, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 96, 256, 3, padding=1, stride=2, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 256, model_channels, 3, padding=1, dtype=self.dtype, device=device) | |
| ) | |
| self._feature_size = model_channels | |
| input_block_chans = [model_channels] | |
| ch = model_channels | |
| ds = 1 | |
| for level, mult in enumerate(channel_mult): | |
| for nr in range(self.num_res_blocks[level]): | |
| layers = [ | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| out_channels=mult * model_channels, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| dtype=self.dtype, | |
| device=device, | |
| operations=operations, | |
| ) | |
| ] | |
| ch = mult * model_channels | |
| num_transformers = transformer_depth.pop(0) | |
| if num_transformers > 0: | |
| if num_head_channels == -1: | |
| dim_head = ch // num_heads | |
| else: | |
| num_heads = ch // num_head_channels | |
| dim_head = num_head_channels | |
| if legacy: | |
| #num_heads = 1 | |
| dim_head = ch // num_heads if use_spatial_transformer else num_head_channels | |
| if exists(disable_self_attentions): | |
| disabled_sa = disable_self_attentions[level] | |
| else: | |
| disabled_sa = False | |
| if not exists(num_attention_blocks) or nr < num_attention_blocks[level]: | |
| layers.append( | |
| SpatialTransformer( | |
| ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim, | |
| disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer, | |
| use_checkpoint=use_checkpoint, attn_precision=attn_precision, dtype=self.dtype, device=device, operations=operations | |
| ) | |
| ) | |
| self.input_blocks.append(TimestepEmbedSequential(*layers)) | |
| self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)) | |
| self._feature_size += ch | |
| input_block_chans.append(ch) | |
| if level != len(channel_mult) - 1: | |
| out_ch = ch | |
| self.input_blocks.append( | |
| TimestepEmbedSequential( | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| out_channels=out_ch, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| down=True, | |
| dtype=self.dtype, | |
| device=device, | |
| operations=operations | |
| ) | |
| if resblock_updown | |
| else Downsample( | |
| ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations | |
| ) | |
| ) | |
| ) | |
| ch = out_ch | |
| input_block_chans.append(ch) | |
| self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)) | |
| ds *= 2 | |
| self._feature_size += ch | |
| if num_head_channels == -1: | |
| dim_head = ch // num_heads | |
| else: | |
| num_heads = ch // num_head_channels | |
| dim_head = num_head_channels | |
| if legacy: | |
| #num_heads = 1 | |
| dim_head = ch // num_heads if use_spatial_transformer else num_head_channels | |
| mid_block = [ | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| dtype=self.dtype, | |
| device=device, | |
| operations=operations | |
| )] | |
| if transformer_depth_middle >= 0: | |
| mid_block += [SpatialTransformer( # always uses a self-attn | |
| ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim, | |
| disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer, | |
| use_checkpoint=use_checkpoint, attn_precision=attn_precision, dtype=self.dtype, device=device, operations=operations | |
| ), | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| dtype=self.dtype, | |
| device=device, | |
| operations=operations | |
| )] | |
| self.middle_block = TimestepEmbedSequential(*mid_block) | |
| self.middle_block_out = self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device) | |
| self._feature_size += ch | |
| def make_zero_conv(self, channels, operations=None, dtype=None, device=None): | |
| return TimestepEmbedSequential(operations.conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device)) | |
| def forward(self, x, hint, timesteps, context, y=None, **kwargs): | |
| t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype) | |
| emb = self.time_embed(t_emb) | |
| guided_hint = self.input_hint_block(hint, emb, context) | |
| out_output = [] | |
| out_middle = [] | |
| hs = [] | |
| if self.num_classes is not None: | |
| assert y.shape[0] == x.shape[0] | |
| emb = emb + self.label_emb(y) | |
| h = x | |
| for module, zero_conv in zip(self.input_blocks, self.zero_convs): | |
| if guided_hint is not None: | |
| h = module(h, emb, context) | |
| h += guided_hint | |
| guided_hint = None | |
| else: | |
| h = module(h, emb, context) | |
| out_output.append(zero_conv(h, emb, context)) | |
| h = self.middle_block(h, emb, context) | |
| out_middle.append(self.middle_block_out(h, emb, context)) | |
| return {"middle": out_middle, "output": out_output} | |