Unverified Commit ef3844d3 authored by Hans Brouwer's avatar Hans Brouwer Committed by GitHub
Browse files

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

support ControlNet models with a different hint_channels value (e.g. TemporalNet2)
parent 3ebbaf7c
...@@ -93,6 +93,7 @@ class ControlNetModel(ModelMixin, ConfigMixin): ...@@ -93,6 +93,7 @@ class ControlNetModel(ModelMixin, ConfigMixin):
def __init__( def __init__(
self, self,
in_channels: int = 4, in_channels: int = 4,
conditioning_channels: int = 3,
flip_sin_to_cos: bool = True, flip_sin_to_cos: bool = True,
freq_shift: int = 0, freq_shift: int = 0,
down_block_types: Tuple[str] = ( down_block_types: Tuple[str] = (
...@@ -185,6 +186,7 @@ class ControlNetModel(ModelMixin, ConfigMixin): ...@@ -185,6 +186,7 @@ class ControlNetModel(ModelMixin, ConfigMixin):
self.controlnet_cond_embedding = ControlNetConditioningEmbedding( self.controlnet_cond_embedding = ControlNetConditioningEmbedding(
conditioning_embedding_channels=block_out_channels[0], conditioning_embedding_channels=block_out_channels[0],
block_out_channels=conditioning_embedding_out_channels, block_out_channels=conditioning_embedding_out_channels,
conditioning_channels=conditioning_channels,
) )
self.down_blocks = nn.ModuleList([]) self.down_blocks = nn.ModuleList([])
......
...@@ -286,6 +286,7 @@ def create_unet_diffusers_config(original_config, image_size: int, controlnet=Fa ...@@ -286,6 +286,7 @@ def create_unet_diffusers_config(original_config, image_size: int, controlnet=Fa
"use_linear_projection": use_linear_projection, "use_linear_projection": use_linear_projection,
"class_embed_type": class_embed_type, "class_embed_type": class_embed_type,
"projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim,
"conditioning_channels": unet_params.hint_channels,
} }
if not controlnet: if not controlnet:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment