cxllm-vendors/stable-diffusion-webui-forge
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix controlnet

Browse Source 
#961
This commit is contained in:
layerdiffusion
2024-08-07 19:21:05 -07:00
parent 692a0b2422
commit 002341af5b
2 changed files with 33 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/nn/cnets/cldm.py
+31 -46
View File
@@ -17,7 +17,6 @@ class ControlNet(nn.Module):
dims=2, dims=2,
num_classes=None, num_classes=None,
use_checkpoint=False, use_checkpoint=False,
dtype=torch.float32,
num_heads=-1, num_heads=-1,
num_head_channels=-1, num_head_channels=-1,
num_heads_upsample=-1, num_heads_upsample=-1,
@@ -35,29 +34,27 @@ class ControlNet(nn.Module):
adm_in_channels=None, adm_in_channels=None,
transformer_depth_middle=None, transformer_depth_middle=None,
transformer_depth_output=None, transformer_depth_output=None,
device=None, dtype=None,
**kwargs, **kwargs,
): ):
super().__init__() super().__init__()
assert use_spatial_transformer == True, "use_spatial_transformer has to be true" assert use_spatial_transformer
if use_spatial_transformer: if use_spatial_transformer:
assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...' assert context_dim is not None
if context_dim is not None: if context_dim is not None:
assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...' assert use_spatial_transformer
# from omegaconf.listconfig import ListConfig
# if type(context_dim) == ListConfig:
# context_dim = list(context_dim)
if num_heads_upsample == -1: if num_heads_upsample == -1:
num_heads_upsample = num_heads num_heads_upsample = num_heads
if num_heads == -1: if num_heads == -1:
assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' assert num_head_channels != -1
if num_head_channels == -1: if num_head_channels == -1:
assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' assert num_heads != -1
self.dtype = dtype
self.dims = dims self.dims = dims
self.in_channels = in_channels self.in_channels = in_channels
self.model_channels = model_channels self.model_channels = model_channels
@@ -66,12 +63,10 @@ class ControlNet(nn.Module):
self.num_res_blocks = len(channel_mult) * [num_res_blocks] self.num_res_blocks = len(channel_mult) * [num_res_blocks]
else: else:
if len(num_res_blocks) != len(channel_mult): if len(num_res_blocks) != len(channel_mult):
raise ValueError("provide num_res_blocks either as an int (globally constant) or " 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")
"as a list/tuple (per-level) with the same length as channel_mult")
self.num_res_blocks = num_res_blocks self.num_res_blocks = num_res_blocks
if disable_self_attentions is not None: 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) assert len(disable_self_attentions) == len(channel_mult)
if num_attention_blocks is not None: if num_attention_blocks is not None:
assert len(num_attention_blocks) == len(self.num_res_blocks) assert len(num_attention_blocks) == len(self.num_res_blocks)
@@ -84,7 +79,6 @@ class ControlNet(nn.Module):
self.conv_resample = conv_resample self.conv_resample = conv_resample
self.num_classes = num_classes self.num_classes = num_classes
self.use_checkpoint = use_checkpoint self.use_checkpoint = use_checkpoint
self.dtype = dtype
self.num_heads = num_heads self.num_heads = num_heads
self.num_head_channels = num_head_channels self.num_head_channels = num_head_channels
self.num_heads_upsample = num_heads_upsample self.num_heads_upsample = num_heads_upsample
@@ -92,24 +86,23 @@ class ControlNet(nn.Module):
time_embed_dim = model_channels * 4 time_embed_dim = model_channels * 4
self.time_embed = nn.Sequential( self.time_embed = nn.Sequential(
nn.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device), nn.Linear(model_channels, time_embed_dim),
nn.SiLU(), nn.SiLU(),
nn.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device), nn.Linear(time_embed_dim, time_embed_dim),
) )
if self.num_classes is not None: if self.num_classes is not None:
if isinstance(self.num_classes, int): if isinstance(self.num_classes, int):
self.label_emb = nn.Embedding(num_classes, time_embed_dim) self.label_emb = nn.Embedding(num_classes, time_embed_dim)
elif self.num_classes == "continuous": elif self.num_classes == "continuous":
print("setting up linear c_adm embedding layer")
self.label_emb = nn.Linear(1, time_embed_dim) self.label_emb = nn.Linear(1, time_embed_dim)
elif self.num_classes == "sequential": elif self.num_classes == "sequential":
assert adm_in_channels is not None assert adm_in_channels is not None
self.label_emb = nn.Sequential( self.label_emb = nn.Sequential(
nn.Sequential( nn.Sequential(
nn.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device), nn.Linear(adm_in_channels, time_embed_dim),
nn.SiLU(), nn.SiLU(),
nn.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device), nn.Linear(time_embed_dim, time_embed_dim),
) )
) )
else: else:
@@ -118,28 +111,28 @@ class ControlNet(nn.Module):
self.input_blocks = nn.ModuleList( self.input_blocks = nn.ModuleList(
[ [
TimestepEmbedSequential( TimestepEmbedSequential(
nn.Conv2d(in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device) nn.Conv2d(in_channels, model_channels, 3, padding=1)
) )
] ]
) )
self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels, dtype=self.dtype, device=device)]) self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels)])
self.input_hint_block = TimestepEmbedSequential( self.input_hint_block = TimestepEmbedSequential(
conv_nd(dims, hint_channels, 16, 3, padding=1, dtype=self.dtype, device=device), conv_nd(dims, hint_channels, 16, 3, padding=1),
nn.SiLU(), nn.SiLU(),
conv_nd(dims, 16, 16, 3, padding=1, dtype=self.dtype, device=device), conv_nd(dims, 16, 16, 3, padding=1),
nn.SiLU(), nn.SiLU(),
conv_nd(dims, 16, 32, 3, padding=1, stride=2, dtype=self.dtype, device=device), conv_nd(dims, 16, 32, 3, padding=1, stride=2),
nn.SiLU(), nn.SiLU(),
conv_nd(dims, 32, 32, 3, padding=1, dtype=self.dtype, device=device), conv_nd(dims, 32, 32, 3, padding=1),
nn.SiLU(), nn.SiLU(),
conv_nd(dims, 32, 96, 3, padding=1, stride=2, dtype=self.dtype, device=device), conv_nd(dims, 32, 96, 3, padding=1, stride=2),
nn.SiLU(), nn.SiLU(),
conv_nd(dims, 96, 96, 3, padding=1, dtype=self.dtype, device=device), conv_nd(dims, 96, 96, 3, padding=1),
nn.SiLU(), nn.SiLU(),
conv_nd(dims, 96, 256, 3, padding=1, stride=2, dtype=self.dtype, device=device), conv_nd(dims, 96, 256, 3, padding=1, stride=2),
nn.SiLU(), nn.SiLU(),
conv_nd(dims, 256, model_channels, 3, padding=1, dtype=self.dtype, device=device) conv_nd(dims, 256, model_channels, 3, padding=1)
) )
self._feature_size = model_channels self._feature_size = model_channels
@@ -157,8 +150,6 @@ class ControlNet(nn.Module):
dims=dims, dims=dims,
use_checkpoint=use_checkpoint, use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm, use_scale_shift_norm=use_scale_shift_norm,
dtype=self.dtype,
device=device,
) )
] ]
ch = mult * model_channels ch = mult * model_channels
@@ -180,11 +171,11 @@ class ControlNet(nn.Module):
SpatialTransformer( SpatialTransformer(
ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim, ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer, disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
use_checkpoint=use_checkpoint, dtype=self.dtype, device=device use_checkpoint=use_checkpoint
) )
) )
self.input_blocks.append(TimestepEmbedSequential(*layers)) self.input_blocks.append(TimestepEmbedSequential(*layers))
self.zero_convs.append(self.make_zero_conv(ch, dtype=self.dtype, device=device)) self.zero_convs.append(self.make_zero_conv(ch))
self._feature_size += ch self._feature_size += ch
input_block_chans.append(ch) input_block_chans.append(ch)
if level != len(channel_mult) - 1: if level != len(channel_mult) - 1:
@@ -200,18 +191,16 @@ class ControlNet(nn.Module):
use_checkpoint=use_checkpoint, use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm, use_scale_shift_norm=use_scale_shift_norm,
down=True, down=True,
dtype=self.dtype,
device=device,
) )
if resblock_updown if resblock_updown
else Downsample( else Downsample(
ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device ch, conv_resample, dims=dims, out_channels=out_ch
) )
) )
) )
ch = out_ch ch = out_ch
input_block_chans.append(ch) input_block_chans.append(ch)
self.zero_convs.append(self.make_zero_conv(ch, dtype=self.dtype, device=device)) self.zero_convs.append(self.make_zero_conv(ch))
ds *= 2 ds *= 2
self._feature_size += ch self._feature_size += ch
@@ -229,15 +218,13 @@ class ControlNet(nn.Module):
dims=dims, dims=dims,
use_checkpoint=use_checkpoint, use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm, use_scale_shift_norm=use_scale_shift_norm,
dtype=self.dtype,
device=device,
)] )]
if transformer_depth_middle >= 0: if transformer_depth_middle >= 0:
mid_block += [ mid_block += [
SpatialTransformer( # always uses a self-attn SpatialTransformer(
ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim, 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, disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
use_checkpoint=use_checkpoint, dtype=self.dtype, device=device use_checkpoint=use_checkpoint
), ),
ResBlock( ResBlock(
ch, ch,
@@ -246,15 +233,13 @@ class ControlNet(nn.Module):
dims=dims, dims=dims,
use_checkpoint=use_checkpoint, use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm, use_scale_shift_norm=use_scale_shift_norm,
dtype=self.dtype,
device=device,
)] )]
self.middle_block = TimestepEmbedSequential(*mid_block) self.middle_block = TimestepEmbedSequential(*mid_block)
self.middle_block_out = self.make_zero_conv(ch, dtype=self.dtype, device=device) self.middle_block_out = self.make_zero_conv(ch)
self._feature_size += ch self._feature_size += ch
def make_zero_conv(self, channels, dtype=None, device=None): def make_zero_conv(self, channels):
return TimestepEmbedSequential(conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device)) return TimestepEmbedSequential(conv_nd(self.dims, channels, channels, 1, padding=0))
def forward(self, x, hint, timesteps, context, y=None, **kwargs): def forward(self, x, hint, timesteps, context, y=None, **kwargs):
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype) t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
modules_forge/supported_controlnet.py
+2 -2
View File
@@ -116,8 +116,8 @@ class ControlNetPatcher(ControlModelPatcher):
controlnet_config.pop("out_channels") controlnet_config.pop("out_channels")
controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1] controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
with using_forge_operations(): with using_forge_operations(dtype=unet_dtype):
control_model = cldm.ControlNet(**controlnet_config) control_model = cldm.ControlNet(**controlnet_config).to(dtype=unet_dtype)
if pth: if pth:
if 'difference' in controlnet_data: if 'difference' in controlnet_data: