Kolors UNet2DConditionModel from_single_file

[hlky]

What does this PR do?

Context #10207 (comment)

Fixes #10207

This partially enables from_single_file for other UNet2DConditionModel already in Diffusers format, others will need Diffusers keys adding to CHECKPOINT_KEY_NAMES and infer_diffusers_model_type updating.

diffusers/src/diffusers/loaders/single_file_utils.py

Line 60 in cef0e36

CHECKPOINT_KEY_NAMES = {

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

cc @a-r-r-o-w @DN6

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DN6]

Hmm seems to be a recurring issue. I don't think it's sustainable to maintain a large mapping of both diffusers keys and original keys in single_file_utils.

In the case of Flux the keys in CHECKPOINT_KEY_NAMES are original checkpoint keys with different prefixes (comfy uses the prefixed version). ControlNet is a bit of an exception since we published the initial ControlNets for SDXL when it was released, so the diffusers format checkpoint keys sort of stuck.

Single File loading wasn't really designed to support diffusers format checkpoints. I think it might be better to add logic that raises a helpful error message if we can't identify the model type.

Or we add a flag in single file, such as convert=False to skip the model inferring and conversion step and it works when you pass in a model repo id to the config argument. e.g

import torch
from diffusers import KolorsPipeline, EulerAncestralDiscreteScheduler

pipe = KolorsPipeline.from_single_file(
    "kolors8AcceleratedVersionOf8_accelerated.safetensors", config="Kwai-Kolors/Kolors-diffusers",  convert=False
)

[hlky]

Thanks @DN6, convert/config will work, I don't think it's very common to load Diffusers models with from_single_file anyway, the linked issue with Kolors is a special case because a 8 step/"turbo" version of the UNet has been released on civit, it's already in Diffusers format but unless I'm missing something we just don't have a nice way to load it to the pipeline other than something like this.

import safetensors.torch

weights = safetensors.torch.load_file("kolors8AcceleratedVersionOf8_accelerated.safetensors")
pipe = KolorsPipeline.from_pretrained(
    "Kwai-Kolors/Kolors-diffusers", variant="fp16", torch_dtype=torch.float16
)
pipe.unet.load_state_dict(weights)

Ideally the user could do something like

unet = UNet2DConditionModel.from_single_file("kolors8AcceleratedVersionOf8_accelerated.safetensors", config="Kwai-Kolors/Kolors-diffusers",  convert=False)
pipe = KolorsPipeline.from_pretrained(
    "Kwai-Kolors/Kolors-diffusers", unet=unet, variant="fp16", torch_dtype=torch.float16
)

[DN6]

Nice. I'll open a PR to add convert/config then 👍🏽