Use multiple controlnets tensor size error

Describe the bug

I get a wrong tensor size error when trying to run multiple controlnets

Reproduction

import argparse
import os
import random

import cv2
import numpy as np
import torch
from PIL import Image
from diffusers import (
    AutoencoderKL,
    EulerAncestralDiscreteScheduler, )
from diffusers.utils import load_image
from torchvision import transforms

from replace_bg.model.controlnet import ControlNetModel
from replace_bg.model.pipeline_controlnet_sd_xl import StableDiffusionXLControlNetPipeline
from replace_bg.utilities import resize_image, remove_bg_from_image, paste_fg_over_image, get_control_image_tensor

import torch

from transformers import DPTFeatureExtractor, DPTForDepthEstimation, DPTImageProcessor

depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
def get_depth_map(image):
    image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
    with torch.no_grad(), torch.autocast("cuda"):
        depth_map = depth_estimator(image).predicted_depth
    depth_map = torch.nn.functional.interpolate(
        depth_map.unsqueeze(1),
        size=(1024, 1024),
        mode="bicubic",
        align_corners=False,
    )
    depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
    depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
    depth_map = (depth_map - depth_min) / (depth_max - depth_min)
    image = torch.cat([depth_map] * 3, dim=1)
    image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
    image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8)).resize((1024,1024))
    return image




def run_sdxl_multicontrolnet(args, prompt, negative_prompt):
    cd = ControlNetModel.from_pretrained(
        "diffusers/controlnet-depth-sdxl-1.0",
        torch_dtype=torch.float16, variant="fp16")
    cc = ControlNetModel.from_pretrained("diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16,variant="fp16")
    controlnets = [cd,cc]
    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
        'stabilityai/stable-diffusion-xl-base-1.0',
        controlnet=controlnets, torch_dtype=torch.float16, variant="fp16",vae=vae)




    pipe.to("cuda")
    pipe.enable_model_cpu_offload()


    image = load_image(args.image_path)
    image = resize_image(image)
    image.save('./tmp.jpg')
    depth_image = get_depth_map(image)


    low_threshold, high_threshold = 100, 200
    input_image = cv2.Canny(image, low_threshold, high_threshold)
    input_image = input_image[:, :, None]
    input_image = np.concatenate([input_image, input_image, input_image], axis=2).astype(np.uint8)
    canny_image = Image.fromarray(input_image)


    assert depth_image.size==canny_image.size
    images = [depth_image.resize((1024, 1024)),canny_image.resize((1024, 1024))]
    gen_img = pipe(
        negative_prompt=negative_prompt,
        prompt=prompt ,
        controlnet_conditioning_scale=[1.0,1.0],
        num_inference_steps=30,
        image=images,

    ).images[0]
    gen_img.save('gen.png')
    depth_image.save('depth.png')
    canny_image.save('canny.png')

Logs

0%|          | 0/30 [00:01<?, ?it/s]
Traceback (most recent call last):
  File "/home/ilias.papastratis/workdir/bria_models/run_bria_modelsv2.py", line 375, in <module>
    run_sdxl_multicontrolnet(args, args.prompt, args.negative_prompt)
  File "/home/ilias.papastratis/workdir/bria_models/run_bria_modelsv2.py", line 227, in run_sdxl_multicontrolnet
    gen_img = pipe(
  File "/home/ilias.papastratis/workdir/envs/object_env/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/ilias.papastratis/workdir/bria_models/replace_bg/model/pipeline_controlnet_sd_xl.py", line 1377, in __call__
    down_block_res_samples, mid_block_res_sample = self.controlnet(
  File "/home/ilias.papastratis/workdir/envs/object_env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ilias.papastratis/workdir/envs/object_env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ilias.papastratis/workdir/envs/object_env/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/home/ilias.papastratis/workdir/envs/object_env/lib/python3.10/site-packages/diffusers/pipelines/controlnet/multicontrolnet.py", line 48, in forward
    down_samples, mid_sample = controlnet(
  File "/home/ilias.papastratis/workdir/envs/object_env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ilias.papastratis/workdir/envs/object_env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ilias.papastratis/workdir/bria_models/replace_bg/model/controlnet.py", line 801, in forward
    sample = sample + controlnet_cond
RuntimeError: The size of tensor a (128) must match the size of tensor b (1024) at non-singleton dimension

System Info

diffusers==0.26.0
Python=3.10

Who can help?

No response

まず、diffusers 0.26.0 を使用して報告された multi-ControlNet の再現を行い、次に pipeline_controlnet_sd_xl.py の self.controlnet の呼び出し付近を調べます。失敗している sample + controlnet_cond 操作について、diffusers/pipelines/controlnet/multicontrolnet.py と replace_bg/model/controlnet.py を通してテンソル形状を追跡します。完了条件は、提供された 2 つの ControlNet の例が 128 対 1024 のサイズエラーなしで実行されることです。