Check num_control_type against ControlNetUnionInput, ControlNetUnionInputProMax

Author: hlky

src/diffusers/models/controlnet_union.py

@@ -725,6 +725,20 @@ def forward(
                 If `return_dict` is `True`, a [`~models.controlnet.ControlNetOutput`] is returned, otherwise a tuple is
                 returned where the first element is the sample tensor.
         """
+        if not isinstance(controlnet_cond, (ControlNetUnionInput, ControlNetUnionInputProMax)):
+            raise ValueError(
+                "Expected type of `controlnet_cond` to be one of `ControlNetUnionInput` or `ControlNetUnionInputProMax`"
+            )
+        if len(controlnet_cond) != self.config.num_control_type:
+            if isinstance(controlnet_cond, ControlNetUnionInput):
+                raise ValueError(
+                    f"Expected num_control_type {self.config.num_control_type}, got {len(controlnet_cond)}. Try `ControlNetUnionInputProMax`."
+                )
+            elif isinstance(controlnet_cond, ControlNetUnionInputProMax):
+                raise ValueError(
+                    f"Expected num_control_type {self.config.num_control_type}, got {len(controlnet_cond)}. Try `ControlNetUnionInput`."
+                )
+
         # check channel order
         channel_order = self.config.controlnet_conditioning_channel_order
 

src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py

@@ -1400,6 +1400,20 @@ def __call__(
 
         controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet
 
+        if not isinstance(control_image_list, (ControlNetUnionInput, ControlNetUnionInputProMax)):
+            raise ValueError(
+                "Expected type of `control_image_list` to be one of `ControlNetUnionInput` or `ControlNetUnionInputProMax`"
+            )
+        if len(control_image_list) != controlnet.config.num_control_type:
+            if isinstance(control_image_list, ControlNetUnionInput):
+                raise ValueError(
+                    f"Expected num_control_type {controlnet.config.num_control_type}, got {len(control_image_list)}. Try `ControlNetUnionInputProMax`."
+                )
+            elif isinstance(control_image_list, ControlNetUnionInputProMax):
+                raise ValueError(
+                    f"Expected num_control_type {controlnet.config.num_control_type}, got {len(control_image_list)}. Try `ControlNetUnionInput`."
+                )
+
         # align format for control guidance
         if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
             control_guidance_start = len(control_guidance_end) * [control_guidance_start]

src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py

@@ -1232,6 +1232,20 @@ def __call__(
 
         controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet
 
+        if not isinstance(image_list, (ControlNetUnionInput, ControlNetUnionInputProMax)):
+            raise ValueError(
+                "Expected type of `image_list` to be one of `ControlNetUnionInput` or `ControlNetUnionInputProMax`"
+            )
+        if len(image_list) != controlnet.config.num_control_type:
+            if isinstance(image_list, ControlNetUnionInput):
+                raise ValueError(
+                    f"Expected num_control_type {controlnet.config.num_control_type}, got {len(image_list)}. Try `ControlNetUnionInputProMax`."
+                )
+            elif isinstance(image_list, ControlNetUnionInputProMax):
+                raise ValueError(
+                    f"Expected num_control_type {controlnet.config.num_control_type}, got {len(image_list)}. Try `ControlNetUnionInput`."
+                )
+
         # align format for control guidance
         if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
             control_guidance_start = len(control_guidance_end) * [control_guidance_start]

src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py

@@ -1323,6 +1323,20 @@ def __call__(
 
         controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet
 
+        if not isinstance(control_image_list, (ControlNetUnionInput, ControlNetUnionInputProMax)):
+            raise ValueError(
+                "Expected type of `control_image_list` to be one of `ControlNetUnionInput` or `ControlNetUnionInputProMax`"
+            )
+        if len(control_image_list) != controlnet.config.num_control_type:
+            if isinstance(control_image_list, ControlNetUnionInput):
+                raise ValueError(
+                    f"Expected num_control_type {controlnet.config.num_control_type}, got {len(control_image_list)}. Try `ControlNetUnionInputProMax`."
+                )
+            elif isinstance(control_image_list, ControlNetUnionInputProMax):
+                raise ValueError(
+                    f"Expected num_control_type {controlnet.config.num_control_type}, got {len(control_image_list)}. Try `ControlNetUnionInput`."
+                )
+
         # align format for control guidance
         if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
             control_guidance_start = len(control_guidance_end) * [control_guidance_start]