From 594ed8ba788aaf97f17d3a97f8a6795b95a375ca Mon Sep 17 00:00:00 2001 From: hlky Date: Fri, 13 Dec 2024 14:21:13 +0000 Subject: [PATCH 1/3] Kolors UNet from_single_file --- src/diffusers/loaders/single_file_utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index 21ff2841700d..a4553936f869 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -100,6 +100,7 @@ ], "autoencoder-dc": "decoder.stages.1.op_list.0.main.conv.conv.bias", "autoencoder-dc-sana": "encoder.project_in.conv.bias", + "kolors": "down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.weight" } DIFFUSERS_DEFAULT_PIPELINE_PATHS = { @@ -151,6 +152,7 @@ "autoencoder-dc-f64c128": {"pretrained_model_name_or_path": "mit-han-lab/dc-ae-f64c128-mix-1.0-diffusers"}, "autoencoder-dc-f32c32": {"pretrained_model_name_or_path": "mit-han-lab/dc-ae-f32c32-mix-1.0-diffusers"}, "autoencoder-dc-f32c32-sana": {"pretrained_model_name_or_path": "mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers"}, + "kolors": {"pretrained_model_name_or_path": "Kwai-Kolors/Kolors-diffusers"}, } # Use to configure model sample size when original config is provided @@ -597,6 +599,9 @@ def infer_diffusers_model_type(checkpoint): else: model_type = "autoencoder-dc-f128c512" + elif CHECKPOINT_KEY_NAMES["kolors"] in checkpoint and checkpoint[CHECKPOINT_KEY_NAMES["kolors"]].shape[-1] == 2048: + model_type = "kolors" + else: model_type = "v1" From e0b71fb4153144bfa2efa593d9316063057aa8b6 Mon Sep 17 00:00:00 2001 From: hlky Date: Fri, 13 Dec 2024 14:31:25 +0000 Subject: [PATCH 2/3] convert_ldm_unet_checkpoint is_diffusers --- src/diffusers/loaders/single_file_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index a4553936f869..8737cb7f41f6 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -966,6 +966,9 @@ def convert_ldm_unet_checkpoint(checkpoint, config, extract_ema=False, **kwargs) """ Takes a state dict and a config, and returns a converted checkpoint. """ + is_diffusers = "time_embedding.linear_1.weight" in checkpoint + if is_diffusers: + return checkpoint # extract state_dict for UNet unet_state_dict = {} keys = list(checkpoint.keys()) From ecaf1db3eaa5c8d83baa2e524b543cdc0582fa75 Mon Sep 17 00:00:00 2001 From: hlky Date: Fri, 13 Dec 2024 14:33:43 +0000 Subject: [PATCH 3/3] make --- src/diffusers/loaders/single_file_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index 8737cb7f41f6..2150b70192a3 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -100,7 +100,7 @@ ], "autoencoder-dc": "decoder.stages.1.op_list.0.main.conv.conv.bias", "autoencoder-dc-sana": "encoder.project_in.conv.bias", - "kolors": "down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.weight" + "kolors": "down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.weight", } DIFFUSERS_DEFAULT_PIPELINE_PATHS = {