diff --git a/comfy/lora.py b/comfy/lora.py index eb95d02ab..ad951bbaf 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -201,6 +201,9 @@ def load_lora(lora, to_load): def model_lora_keys_clip(model, key_map={}): sdk = model.state_dict().keys() + for k in sdk: + if k.endswith(".weight"): + key_map["text_encoders.{}".format(k[:-len(".weight")])] = k #generic lora format without any weird key names text_model_lora_key = "lora_te_text_model_encoder_layers_{}_{}" clip_l_present = False diff --git a/comfy_extras/nodes_lora_extract.py b/comfy_extras/nodes_lora_extract.py index dcb46f0e0..3c2f179d3 100644 --- a/comfy_extras/nodes_lora_extract.py +++ b/comfy_extras/nodes_lora_extract.py @@ -4,6 +4,7 @@ import comfy.utils import folder_paths import os import logging +from enum import Enum CLAMP_QUANTILE = 0.99 @@ -38,6 +39,38 @@ def extract_lora(diff, rank): Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1]) return (U, Vh) +class LORAType(Enum): + STANDARD = 0 + FULL_DIFF = 1 + +LORA_TYPES = {"standard": LORAType.STANDARD, + "full_diff": LORAType.FULL_DIFF} + +def calc_lora_model(model_diff, rank, prefix_model, prefix_lora, output_sd, lora_type, bias_diff=False): + comfy.model_management.load_models_gpu([model_diff], force_patch_weights=True) + sd = model_diff.model_state_dict(filter_prefix=prefix_model) + + for k in sd: + if k.endswith(".weight"): + weight_diff = sd[k] + if lora_type == LORAType.STANDARD: + if weight_diff.ndim < 2: + if bias_diff: + output_sd["{}{}.diff".format(prefix_lora, k[len(prefix_model):-7])] = weight_diff.contiguous().half().cpu() + continue + try: + out = extract_lora(weight_diff, rank) + output_sd["{}{}.lora_up.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[0].contiguous().half().cpu() + output_sd["{}{}.lora_down.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[1].contiguous().half().cpu() + except: + logging.warning("Could not generate lora weights for key {}, is the weight difference a zero?".format(k)) + elif lora_type == LORAType.FULL_DIFF: + output_sd["{}{}.diff".format(prefix_lora, k[len(prefix_model):-7])] = weight_diff.contiguous().half().cpu() + + elif bias_diff and k.endswith(".bias"): + output_sd["{}{}.diff_b".format(prefix_lora, k[len(prefix_model):-5])] = sd[k].contiguous().half().cpu() + return output_sd + class LoraSave: def __init__(self): self.output_dir = folder_paths.get_output_directory() @@ -45,9 +78,12 @@ class LoraSave: @classmethod def INPUT_TYPES(s): return {"required": {"filename_prefix": ("STRING", {"default": "loras/ComfyUI_extracted_lora"}), - "rank": ("INT", {"default": 8, "min": 1, "max": 1024, "step": 1}), + "rank": ("INT", {"default": 8, "min": 1, "max": 4096, "step": 1}), + "lora_type": (tuple(LORA_TYPES.keys()),), + "bias_diff": ("BOOLEAN", {"default": True}), }, - "optional": {"model_diff": ("MODEL",),}, + "optional": {"model_diff": ("MODEL",), + "text_encoder_diff": ("CLIP",)}, } RETURN_TYPES = () FUNCTION = "save" @@ -55,30 +91,18 @@ class LoraSave: CATEGORY = "_for_testing" - def save(self, filename_prefix, rank, model_diff=None): - if model_diff is None: + def save(self, filename_prefix, rank, lora_type, bias_diff, model_diff=None, text_encoder_diff=None): + if model_diff is None and text_encoder_diff is None: return {} + lora_type = LORA_TYPES.get(lora_type) full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) output_sd = {} - prefix_key = "diffusion_model." - stored = set() - - comfy.model_management.load_models_gpu([model_diff], force_patch_weights=True) - sd = model_diff.model_state_dict(filter_prefix=prefix_key) - - for k in sd: - if k.endswith(".weight"): - weight_diff = sd[k] - if weight_diff.ndim < 2: - continue - try: - out = extract_lora(weight_diff, rank) - output_sd["{}.lora_up.weight".format(k[:-7])] = out[0].contiguous().half().cpu() - output_sd["{}.lora_down.weight".format(k[:-7])] = out[1].contiguous().half().cpu() - except: - logging.warning("Could not generate lora weights for key {}, is the weight difference a zero?".format(k)) + if model_diff is not None: + output_sd = calc_lora_model(model_diff, rank, "diffusion_model.", "diffusion_model.", output_sd, lora_type, bias_diff=bias_diff) + if text_encoder_diff is not None: + output_sd = calc_lora_model(text_encoder_diff.patcher, rank, "", "text_encoders.", output_sd, lora_type, bias_diff=bias_diff) output_checkpoint = f"{filename}_{counter:05}_.safetensors" output_checkpoint = os.path.join(full_output_folder, output_checkpoint)