fix from_transformer() with extra conditioning channels (huggingface#…

…9364) * fix from_transformer() with extra conditioning channels * style fix --------- Co-authored-by: YiYi Xu <[email protected]> Co-authored-by: Álvaro Somoza <[email protected]>
tolgacangoz · Sep 9, 2024 · f28a8c2 · f28a8c2
1 parent 2c6a6c9
commit f28a8c2
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/diffusers/models/controlnet_sd3.py b/src/diffusers/models/controlnet_sd3.py
@@ -242,9 +242,12 @@ def _set_gradient_checkpointing(self, module, value=False):
             module.gradient_checkpointing = value
 
     @classmethod
-    def from_transformer(cls, transformer, num_layers=12, load_weights_from_transformer=True):
+    def from_transformer(
+        cls, transformer, num_layers=12, num_extra_conditioning_channels=1, load_weights_from_transformer=True
+    ):
         config = transformer.config
         config["num_layers"] = num_layers or config.num_layers
+        config["extra_conditioning_channels"] = num_extra_conditioning_channels
         controlnet = cls(**config)
 
         if load_weights_from_transformer: