StableDiffusionControlNetInpaintPipeline - does not work with 4 channel Unets

[cryptexis]

Describe the bug

Hi,
I have tried to use an example from here https://huggingface.co/docs/diffusers/main/en/api/pipelines/controlnet (posted below). However, when a model is passed that has a Unet that expects 4 channels the pipeline breaks with the following error

RuntimeError: Given groups=1, weight of size [320, 9, 3, 3], expected input[2, 4, 64, 64] to have 9 channels, but got 4 channels instead

I raise this problem as it seems that in the source code it seems the handling of 4 vs 9 channel unet is handled
https://github.com/huggingface/diffusers/blob/v0.26.2-patch/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py#L1532

Reproduction

# !pip install transformers accelerate
import cv2
from PIL import Image
from diffusers.pipelines.stable_diffusion import safety_checker
from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDIMScheduler, DPMSolverMultistepScheduler
from diffusers.utils import load_image
import numpy as np
import torch


init_image = load_image(
    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png"
)
init_image = init_image.resize((512, 512))

generator = torch.Generator(device="cpu").manual_seed(1)

mask_image = load_image(
    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png"
)
mask_image = mask_image.resize((512, 512))


def make_canny_condition(image):
    image = np.array(image)
    image = cv2.Canny(image, 100, 200)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    image = Image.fromarray(image)
    return image


control_image = make_canny_condition(init_image)

controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/control_v11p_sd15_canny", torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetInpaintPipeline.from_single_file(
    "./models/majicmixRealistic_v7-inpainting.safetensors", controlnet=controlnet, torch_dtype=torch.float16,
    safety_checker = None,
    requires_safety_checker = False
)

pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, algorithm_type = 'sde-dpmsolver++', use_karras_sigmas = True)
pipe.enable_model_cpu_offload()


# generate image
image = pipe(
    "a handsome man with ray-ban sunglasses",
    num_inference_steps=20,
    generator=generator,
    image=init_image,
    mask_image=mask_image,
    control_image=control_image
    
).images[0]

Logs

No response

System Info

• diffusers version: 0.26.2
• Platform: Linux-5.15.0-1039-aws-x86_64-with-glibc2.31
• Python version: 3.11.7
• PyTorch version (GPU?): 2.1.2+cu121 (True)
• Huggingface_hub version: 0.20.3
• Transformers version: 4.37.2
• Accelerate version: 0.21.0
• xFormers version: 0.0.23.post1
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

@yiyixuxu @DN6 @sayakpaul @patrickvonplaten

[sayakpaul]

Where does "/models/majicmixRealistic_v7-inpainting.safetensors"come from? Additionally, could you check if theStableDiffusionInpaintPipeline` works as expected with this checkpoint?

[cryptexis]

@sayakpaul
apologies https://civitai.com/models/43331?modelVersionId=221343 this is where it comes from.

Sure will check and report

[cryptexis]

@sayakpaul

# !pip install transformers accelerate
from diffusers.pipelines.stable_diffusion import safety_checker
from diffusers import StableDiffusionInpaintPipeline, ControlNetModel, DDIMScheduler, DPMSolverMultistepScheduler
from diffusers.utils import load_image, make_image_grid
import numpy as np
import torch

init_image = load_image(
    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png"
)
init_image = init_image.resize((512, 512))

generator = torch.Generator(device="cpu").manual_seed(1)

mask_image = load_image(
    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png"
)
mask_image = mask_image.resize((512, 512))


pipe = StableDiffusionInpaintPipeline.from_single_file(
    "./models/majicmixRealistic_v7-inpainting.safetensors", torch_dtype=torch.float16,
    safety_checker = None,
    requires_safety_checker = False
)

pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, algorithm_type = 'sde-dpmsolver++', use_karras_sigmas = True)
pipe.enable_model_cpu_offload()


# generate image
image = pipe(
    "a handsome man with ray-ban sunglasses",
    num_inference_steps=20,
    generator=generator,
    image=init_image,
    mask_image=mask_image,    
).images[0]

make_image_grid([init_image, mask_image, image], 1, 3)

[sayakpaul]

Cc @yiyixuxu ^

[cryptexis]

@yiyixuxu

A bit more information for you.
Before I was getting the error when trying to run the pipeline. Then I have noticed that the exception was on the step where accelerate was calling old_forward method. I have upgraded accelerate to 0.27.

Now I am getting the same error however on a different step - when I am creating the pipeline.

Here is the error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[4], line 37
     32 control_image = make_canny_condition(init_image)
     34 controlnet = ControlNetModel.from_pretrained(
     35     "lllyasviel/control_v11p_sd15_canny", torch_dtype=torch.float16
     36 )
---> 37 pipe = StableDiffusionControlNetInpaintPipeline.from_single_file(
     38     "./models/majicmixRealistic_v7-inpainting.safetensors", 
     39     # controlnet=controlnet, 
     40     torch_dtype=torch.float16,
     41     safety_checker = None,
     42     requires_safety_checker = False
     43 )
     45 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, algorithm_type = 'sde-dpmsolver++', use_karras_sigmas = True)
     46 pipe.enable_model_cpu_offload()

File ~/.local/lib/python3.11/site-packages/huggingface_hub/utils/_validators.py:118, in validate_hf_hub_args.<locals>._inner_fn(*args, **kwargs)
    115 if check_use_auth_token:
    116     kwargs = smoothly_deprecate_use_auth_token(fn_name=fn.__name__, has_token=has_token, kwargs=kwargs)
--> 118 return fn(*args, **kwargs)

File ~/.conda/envs/cog/lib/python3.11/site-packages/diffusers/loaders/single_file.py:253, in FromSingleFileMixin.from_single_file(cls, pretrained_model_link_or_path, **kwargs)
    251     init_kwargs[name] = passed_class_obj[name]
    252 else:
--> 253     components = build_sub_model_components(
    254         init_kwargs,
    255         class_name,
    256         name,
    257         original_config,
    258         checkpoint,
    259         model_type=model_type,
    260         image_size=image_size,
    261         load_safety_checker=load_safety_checker,
    262         local_files_only=local_files_only,
    263         **kwargs,
    264     )
    265     if not components:
    266         continue

File ~/.conda/envs/cog/lib/python3.11/site-packages/diffusers/loaders/single_file.py:58, in build_sub_model_components(pipeline_components, pipeline_class_name, component_name, original_config, checkpoint, local_files_only, load_safety_checker, model_type, image_size, **kwargs)
     56 if component_name == "unet":
     57     num_in_channels = kwargs.pop("num_in_channels", None)
---> 58     unet_components = create_diffusers_unet_model_from_ldm(
     59         pipeline_class_name, original_config, checkpoint, num_in_channels=num_in_channels, image_size=image_size
     60     )
     61     return unet_components
     63 if component_name == "vae":

File ~/.conda/envs/cog/lib/python3.11/site-packages/diffusers/loaders/single_file_utils.py:1173, in create_diffusers_unet_model_from_ldm(pipeline_class_name, original_config, checkpoint, num_in_channels, upcast_attention, extract_ema, image_size)
   1171 if is_accelerate_available():
   1172     for param_name, param in diffusers_format_unet_checkpoint.items():
-> 1173         set_module_tensor_to_device(unet, param_name, "cpu", value=param)
   1174 else:
   1175     unet.load_state_dict(diffusers_format_unet_checkpoint)

File ~/.conda/envs/cog/lib/python3.11/site-packages/accelerate/utils/modeling.py:345, in set_module_tensor_to_device(module, tensor_name, device, value, dtype, fp16_statistics, tied_params_map)
    343 if value is not None:
    344     if old_value.shape != value.shape:
--> 345         raise ValueError(
    346             f'Trying to set a tensor of shape {value.shape} in "{tensor_name}" (which has shape {old_value.shape}), this look incorrect.'
    347         )
    349     if dtype is None:
    350         # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model
    351         value = value.to(old_value.dtype)

ValueError: Trying to set a tensor of shape torch.Size([320, 9, 3, 3]) in "weight" (which has shape torch.Size([320, 4, 3, 3])), this look incorrect.

[cryptexis]

found a work around

pipe_sip = StableDiffusionInpaintPipeline.from_single_file(
    "./models/majicmixRealistic_v7-inpainting.safetensors", 
    controlnet=controlnet, 
    torch_dtype=torch.float16,
    safety_checker = None,
    requires_safety_checker = False
)

pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting", 
    controlnet=controlnet, 
    torch_dtype=torch.float16,
    safety_checker = None,
    requires_safety_checker = False,
    vae = pipe_sip.vae,
    unet = pipe_sip.unet,
    tokenizer = pipe_sip.tokenizer,
    text_encoder = pipe_sip.text_encoder
)

but this does not solve the underlying issue.

[sayakpaul]

What do you mean? The problem still happens with your fix in #6964 (comment)?

[cryptexis]

@sayakpaul I made it work by first loading in one pipeline and then passing models to the other, but it's hardly a fix.
The problem is still there, that one Pipeline is able to load artifacts correctly and the other - not(throws an exception)

[sayakpaul]

You are right. I will take a look.

[sayakpaul]

Hi @cryptexis. This should have been fixed in #6975. Let us know.