lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-02-21 15:23:58 +00:00
Code Issues 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

77
backend/nn/cnets/cldm.py
View File

@@ -17,7 +17,6 @@ class ControlNet(nn.Module):
dims=2,
num_classes=None,
use_checkpoint=False,
dtype=torch.float32,
num_heads=-1,
num_head_channels=-1,
num_heads_upsample=-1,
@@ -35,29 +34,27 @@ class ControlNet(nn.Module):
adm_in_channels=None,
transformer_depth_middle=None,
transformer_depth_output=None,
device=None,
dtype=None,
**kwargs,
):
super().__init__()
assert use_spatial_transformer == True, "use_spatial_transformer has to be true"
assert 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:
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)
assert use_spatial_transformer
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'
assert 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.in_channels = in_channels
self.model_channels = model_channels
@@ -66,12 +63,10 @@ class ControlNet(nn.Module):
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")
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)
@@ -84,7 +79,6 @@ class ControlNet(nn.Module):
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
@@ -92,24 +86,23 @@ class ControlNet(nn.Module):
time_embed_dim = model_channels * 4
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.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 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(
nn.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device),
nn.Linear(adm_in_channels, time_embed_dim),
nn.SiLU(),
nn.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
nn.Linear(time_embed_dim, time_embed_dim),
)
)
else:
@@ -118,28 +111,28 @@ class ControlNet(nn.Module):
self.input_blocks = nn.ModuleList(
[
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(
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(),
conv_nd(dims, 16, 16, 3, padding=1, dtype=self.dtype, device=device),
conv_nd(dims, 16, 16, 3, padding=1),
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(),
conv_nd(dims, 32, 32, 3, padding=1, dtype=self.dtype, device=device),
conv_nd(dims, 32, 32, 3, padding=1),
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(),
conv_nd(dims, 96, 96, 3, padding=1, dtype=self.dtype, device=device),
conv_nd(dims, 96, 96, 3, padding=1),
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(),
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
@@ -157,8 +150,6 @@ class ControlNet(nn.Module):
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
dtype=self.dtype,
device=device,
)
]
ch = mult * model_channels
@@ -180,11 +171,11 @@ class ControlNet(nn.Module):
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, dtype=self.dtype, device=device
use_checkpoint=use_checkpoint
)
)
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
input_block_chans.append(ch)
if level != len(channel_mult) - 1:
@@ -200,18 +191,16 @@ class ControlNet(nn.Module):
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
down=True,
dtype=self.dtype,
device=device,
)
if resblock_updown
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
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
self._feature_size += ch
@@ -229,15 +218,13 @@ class ControlNet(nn.Module):
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
dtype=self.dtype,
device=device,
)]
if transformer_depth_middle >= 0:
mid_block += [
SpatialTransformer( # always uses a self-attn
SpatialTransformer(
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, dtype=self.dtype, device=device
use_checkpoint=use_checkpoint
),
ResBlock(
ch,
@@ -246,15 +233,13 @@ class ControlNet(nn.Module):
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
dtype=self.dtype,
device=device,
)]
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
def make_zero_conv(self, channels, dtype=None, device=None):
return TimestepEmbedSequential(conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device))
def make_zero_conv(self, channels):
return TimestepEmbedSequential(conv_nd(self.dims, channels, channels, 1, padding=0))
def forward(self, x, hint, timesteps, context, y=None, **kwargs):
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)

4
modules_forge/supported_controlnet.py
View File

@@ -116,8 +116,8 @@ class ControlNetPatcher(ControlModelPatcher):
controlnet_config.pop("out_channels")
controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
with using_forge_operations():
control_model = cldm.ControlNet(**controlnet_config)
with using_forge_operations(dtype=unet_dtype):
control_model = cldm.ControlNet(**controlnet_config).to(dtype=unet_dtype)
if pth:
if 'difference' in controlnet_data: