使用diffusers中ControlNetModel.from_single_file加载模型遇KeyError问题

调用diffusers库的ControlNetModel.from_single_file加载指定路径的ControlNet模型文件时抛出KeyError: 'middle_block_out.0.weight'错误，相关信息如下：

代码 #

from diffusers import StableDiffusionControlNetPipeline ,ControlNetModel
import torch
import cv2
img = cv2.imread('/home/leo/blues.png')

controlnet_model_path = "/mnt/d/models/diffusion_pytorch_model.safetensors"
controlnet = ControlNetModel.from_single_file(controlnet_model_path, torch_dtype=torch.float16)

pipe = StableDiffusionControlNetPipeline.from_pretrained("riffusion/riffusion-model-v1", controlnet=controlnet, torch_dtype=torch.float16)

torch.manual_seed(0)
generator = torch.random.manual_seed(0)

output_path = "happy.png"

out_image = pipe(
"guitar", num_inference_steps=20, generator=generator, image=img
).images[0]

out_image.save(output_path)

错误日志 #

Traceback (most recent call last):
  File "/home/leo/PhantomHarmony/src/process/diff.py", line 9, in <module>
    controlnet = ControlNetModel.from_single_file(controlnet_model_path)
  File "/home/leo/test/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py", line 118, in _inner_fn
    return fn(*args, **kwargs)
  File "/home/leo/test/lib/python3.10/site-packages/diffusers/loaders/controlnet.py", line 120, in from_single_file
    component = create_diffusers_controlnet_model_from_ldm(
  File "/home/leo/test/lib/python3.10/site-packages/diffusers/loaders/single_file_utils.py", line 866, in create_diffusers_controlnet_model_from_ldm
    diffusers_format_controlnet_checkpoint = convert_controlnet_checkpoint(checkpoint, diffusers_config)
  File "/home/leo/test/lib/python3.10/site-packages/diffusers/loaders/single_file_utils.py", line 830, in convert_controlnet_checkpoint
    new_checkpoint["controlnet_mid_block.weight"] = controlnet_state_dict.pop("middle_block_out.0.weight")
KeyError: 'middle_block_out.0.weight'

系统环境 #

diffusers version: 0.26.0
Platform: Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Python version: 3.10.12
PyTorch version (GPU?): 2.1.2+cu121 (True)
Huggingface_hub version: 0.20.3
Transformers version: 4.37.2
Accelerate version: 0.27.0.dev0
xFormers version: not installed

这个错误核心是模型权重结构与diffusers预期格式不匹配，可按以下步骤排查解决：

• 确认模型与基础版本兼容：你使用的基础模型riffusion/riffusion-model-v1基于SD1.x架构，如果你的ControlNet是为SDXL训练的，权重结构会完全不同，必然找不到目标键。检查模型来源，确保下载的是SD1.x兼容的ControlNet。

• 检查模型格式是否正确：from_single_file仅支持加载原始LDM风格的ControlNet单文件（ckpt/safetensors），如果你的模型是从diffusers格式导出的（权重命名为diffusers风格，如包含controlnet_mid_block类键），需改用from_pretrained加载模型文件夹（需包含config.json等配置文件）：

  # 替换原加载代码
  controlnet = ControlNetModel.from_pretrained("/mnt/d/models/", torch_dtype=torch.float16)
  

• 调整diffusers版本：diffusers 0.26.0对单文件ControlNet的转换逻辑可能存在兼容性问题，尝试降级到0.25.x稳定版或升级到最新版：

  # 降级到0.25.1
  pip install diffusers==0.25.1
  # 或升级到最新版
  pip install --upgrade diffusers
  

• 手动验证权重结构：加载权重文件查看键名，确认模型是否符合预期：

  from safetensors.torch import load_file
  weights = load_file("/mnt/d/models/diffusion_pytorch_model.safetensors")
  # 查看包含middle或mid的键
  print([k for k in weights.keys() if "middle" in k.lower() or "mid" in k.lower()])
  

  如果输出中没有类似middle_block_out.0.weight的键，说明该模型确实不是SD1.x的LDM格式ControlNet，需更换对应模型。

内容的提问来源于stack exchange，提问作者leo-ke