Support ControlNet models with different number of channels in control images (huggingface#3815)

JCBrouwer · web-flow · commit ef3844d3a835 · 2023-06-21T13:11:45.000+02:00
support ControlNet models with a different hint_channels value (e.g. TemporalNet2)
diff --git a/src/diffusers/models/controlnet.py b/src/diffusers/models/controlnet.py
@@ -93,6 +93,7 @@ class ControlNetModel(ModelMixin, ConfigMixin):
     def __init__(
         self,
         in_channels: int = 4,
+        conditioning_channels: int = 3,
         flip_sin_to_cos: bool = True,
         freq_shift: int = 0,
         down_block_types: Tuple[str] = (
@@ -185,6 +186,7 @@ def __init__(
         self.controlnet_cond_embedding = ControlNetConditioningEmbedding(
             conditioning_embedding_channels=block_out_channels[0],
             block_out_channels=conditioning_embedding_out_channels,
+            conditioning_channels=conditioning_channels,
         )
 
         self.down_blocks = nn.ModuleList([])
diff --git a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
@@ -286,6 +286,7 @@ def create_unet_diffusers_config(original_config, image_size: int, controlnet=Fa
         "use_linear_projection": use_linear_projection,
         "class_embed_type": class_embed_type,
         "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim,
+        "conditioning_channels": unet_params.hint_channels,
     }
 
     if not controlnet:

Original file line number	Diff line number	Diff line change
`@@ -286,6 +286,7 @@ def create_unet_diffusers_config(original_config, image_size: int, controlnet=Fa`
`286`	`286`	`"use_linear_projection": use_linear_projection,`
`287`	`287`	`"class_embed_type": class_embed_type,`
`288`	`288`	`"projection_class_embeddings_input_dim": projection_class_embeddings_input_dim,`
	`289`	`+ "conditioning_channels": unet_params.hint_channels,`
`289`	`290`	`}`
`290`	`291`
`291`	`292`	`if not controlnet:`