some extra lora support, inc. new glora (#2715)

* support new glora (via ComfyUI)
* support BFL FluxTools loras (mostly via ComfyUI)
* also support using loras (like Hyper, Turbo) with FluxTools models
This commit is contained in:
DenOfEquity
2025-03-04 00:26:43 +00:00
committed by GitHub
parent 4f825bc070
commit 5e1dcd35a8
2 changed files with 91 additions and 28 deletions

View File

@@ -30,6 +30,19 @@ LORA_CLIP_MAP = {
def load_lora(lora, to_load):
# BFL loras for Flux; from ComfyUI: comfy/lora_convert.py
def convert_lora_bfl_control(sd):
import torch
sd_out = {}
for k in sd:
k_to = "diffusion_model.{}".format(k.replace(".lora_B.bias", ".diff_b").replace("_norm.scale", "_norm.scale.set_weight"))
sd_out[k_to] = sd[k]
return sd_out
if "img_in.lora_A.weight" in lora and "single_blocks.0.norm.key_norm.scale" in lora:
lora = convert_lora_bfl_control(lora)
patch_dict = {}
loaded_keys = set()
for x in to_load:
@@ -189,6 +202,12 @@ def load_lora(lora, to_load):
patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (diff_bias,))
loaded_keys.add(diff_bias_name)
set_weight_name = "{}.set_weight".format(x)
set_weight = lora.get(set_weight_name, None)
if set_weight is not None:
patch_dict[to_load[x]] = ("set", (set_weight,))
loaded_keys.add(set_weight_name)
remaining_dict = {x: y for x, y in lora.items() if x not in loaded_keys}
return patch_dict, remaining_dict
@@ -269,11 +288,13 @@ def model_lora_keys_unet(model, key_map={}):
sdk = sd.keys()
for k in sdk:
if k.startswith("diffusion_model.") and k.endswith(".weight"):
key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
key_map["lora_unet_{}".format(key_lora)] = k
key_map["lora_prior_unet_{}".format(key_lora)] = k #cascade lora: TODO put lora key prefix in the model config
key_map["{}".format(k[:-len(".weight")])] = k #generic lora format without any weird key names
if k.startswith("diffusion_model."):
if k.endswith(".weight"):
key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
key_map["lora_unet_{}".format(key_lora)] = k
key_map["{}".format(k[:-len(".weight")])] = k #generic lora format without any weird key names
else:
key_map["{}".format(k)] = k #generic lora format for not .weight without any weird key names
diffusers_keys = utils.unet_to_diffusers(model.diffusion_model.config)
for k in diffusers_keys:
@@ -281,7 +302,8 @@ def model_lora_keys_unet(model, key_map={}):
unet_key = "diffusion_model.{}".format(diffusers_keys[k])
key_lora = k[:-len(".weight")].replace(".", "_")
key_map["lora_unet_{}".format(key_lora)] = unet_key
key_map["lycoris_{}".format(key_lora)] = unet_key #simpletuner lycoris format
diffusers_lora_prefix = ["", "unet."]
for p in diffusers_lora_prefix:
diffusers_lora_key = "{}{}".format(p, k[:-len(".weight")].replace(".to_", ".processor.to_"))
@@ -289,19 +311,19 @@ def model_lora_keys_unet(model, key_map={}):
diffusers_lora_key = diffusers_lora_key[:-2]
key_map[diffusers_lora_key] = unet_key
# if isinstance(model, comfy.model_base.SD3): #Diffusers lora SD3
# diffusers_keys = utils.mmdit_to_diffusers(model.diffusion_model.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")]) #regular diffusers sd3 lora format
# key_map[key_lora] = to
#
# key_lora = "base_model.model.{}".format(k[:-len(".weight")]) #format for flash-sd3 lora and others?
# key_map[key_lora] = to
#
# key_lora = "lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_")) #OneTrainer lora
# key_map[key_lora] = to
# if 'stable-diffusion-3' in model.config.huggingface_repo.lower(): #Diffusers lora SD3
# diffusers_keys = utils.mmdit_to_diffusers(model.diffusion_model.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")]) #regular diffusers sd3 lora format
# key_map[key_lora] = to
# key_lora = "base_model.model.{}".format(k[:-len(".weight")]) #format for flash-sd3 lora and others?
# key_map[key_lora] = to
# key_lora = "lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_")) #OneTrainer lora
# key_map[key_lora] = to
#
# if isinstance(model, comfy.model_base.AuraFlow): #Diffusers lora AuraFlow
# diffusers_keys = utils.auraflow_to_diffusers(model.diffusion_model.config, output_prefix="diffusion_model.")