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):
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 @@ class ControlNetModel(ModelMixin, ConfigMixin):
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([])
......
......@@ -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:
......
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