From 51dded3398fa3a143ea7e08bdd0f649bb9a967b1 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 14 Dec 2024 22:35:32 +0100 Subject: [PATCH] Basic PixArt LoRA support --- comfy/lora.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/comfy/lora.py b/comfy/lora.py index b6d9a8d04..9e0ede308 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -344,7 +344,6 @@ def model_lora_keys_unet(model, key_map={}): key_lora = "lycoris_{}".format(k[:-len(".weight")].replace(".", "_")) #simpletuner lycoris format key_map[key_lora] = to - if isinstance(model, comfy.model_base.AuraFlow): #Diffusers lora AuraFlow diffusers_keys = comfy.utils.auraflow_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.") for k in diffusers_keys: @@ -353,6 +352,20 @@ def model_lora_keys_unet(model, key_map={}): key_lora = "transformer.{}".format(k[:-len(".weight")]) #simpletrainer and probably regular diffusers lora format key_map[key_lora] = to + if isinstance(model, comfy.model_base.PixArt): + diffusers_keys = comfy.utils.pixart_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.") + for k in diffusers_keys: + if k.endswith(".weight"): + to = diffusers_keys[k] + key_lora = "transformer.{}".format(k[:-len(".weight")]) #default format + key_map[key_lora] = to + + key_lora = "base_model.model.{}".format(k[:-len(".weight")]) #diffusers training script + key_map[key_lora] = to + + key_lora = "unet.base_model.model.{}".format(k[:-len(".weight")]) #old reference peft script + key_map[key_lora] = to + if isinstance(model, comfy.model_base.HunyuanDiT): for k in sdk: if k.startswith("diffusion_model.") and k.endswith(".weight"):