From 5a403e638c0bcffb62a4a87f26515423818bbdf7 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 7 Feb 2025 04:43:45 -0600 Subject: [PATCH] MergeFluxLoRAsQuantizeAndLoad, WIP --- __init__.py | 233 ++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 161 insertions(+), 72 deletions(-) diff --git a/__init__.py b/__init__.py index 3974ba1..8559761 100644 --- a/__init__.py +++ b/__init__.py @@ -6,11 +6,31 @@ import os from pathlib import Path import logging import folder_paths +import shutil from collections import defaultdict import hashlib +import tempfile +import subprocess +import gc +from safetensors.torch import save_file, load_file +import comfy.utils +from typing import Dict, List + + +from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS +from .nodes import ( + UnetLoaderGGUF, UnetLoaderGGUFAdvanced, + CLIPLoaderGGUF, DualCLIPLoaderGGUF, TripleCLIPLoaderGGUF, + LTXVLoader, + Florence2ModelLoader, DownloadAndLoadFlorence2Model, + CheckpointLoaderNF4, + LoadFluxControlNet, + MMAudioModelLoader, MMAudioFeatureUtilsLoader, MMAudioSampler, + PulidModelLoader, PulidInsightFaceLoader, PulidEvaClipLoader, + HyVideoModelLoader, HyVideoVAELoader, DownloadAndLoadHyVideoTextEncoder +) current_device = mm.get_torch_device() -current_offload_device = mm.get_torch_device() model_allocation_store = {} def get_torch_device_patched(): @@ -21,35 +41,7 @@ def get_torch_device_patched(): device = torch.device(current_device) return device -def text_encoder_device_patched(): - device = None - if (not torch.cuda.is_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_device).lower()): - device = torch.device("cpu") - else: - device = torch.device(current_device) - return device - -def unet_offload_device_patched(): - device = None - if (not torch.cuda.is_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_offload_device).lower()): - device = torch.device("cpu") - else: - device = torch.device(current_offload_device) - return device - -def text_encoder_offload_device_patched(): - device = None - if (not torch.cuda.is_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_offload_device).lower()): - device = torch.device("cpu") - else: - device = torch.device(current_offload_device) - return device - mm.get_torch_device = get_torch_device_patched -mm.unet_offload_device = unet_offload_device_patched -mm.text_encoder_device = text_encoder_device_patched -mm.text_encoder_offload_device = text_encoder_offload_device_patched - def create_model_hash(model, caller): @@ -217,7 +209,6 @@ def analyze_ggml_loading(model, allocations_str): return {"device_assignments": device_assignments} - def get_device_list(): import torch return ["cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())] @@ -272,6 +263,141 @@ class HunyuanVideoEmbeddingsAdapter: return ([[cond, pooled_dict]],) + +class MergeFluxLoRAsQuantizeAndLoad: + @classmethod + def INPUT_TYPES(cls): + unet_name = folder_paths.get_filename_list("diffusion_models") + loras = ["None"] + folder_paths.get_filename_list("loras") + inputs = { + "required": { + "unet_name": (unet_name,), + "switch_1": (["Off", "On"],), + "lora_name_1": (loras,), + "lora_weight_1": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), + "switch_2": (["Off", "On"],), + "lora_name_2": (loras,), + "lora_weight_2": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), + "switch_3": (["Off", "On"],), + "lora_name_3": (loras,), + "lora_weight_3": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), + "switch_4": (["Off", "On"],), + "lora_name_4": (loras,), + "lora_weight_4": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), + "quantization": (["Q2_K", "Q3_K_S", "Q4_0", "Q4_1", "Q4_K_S", "Q5_0", "Q5_1", "Q5_K_S", "Q6_K", "Q8_0", "FP16"], {"default": "Q4_K_S"}), + "delete_final_gguf": ("BOOLEAN", {"default": False}), + "new_model_name": ("STRING", {"default": "merged_model"}), + } + } + return inputs + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_and_quantize" + CATEGORY = "loaders" + + def merge_flux_loras(self, model_sd: dict, lora_paths: list, weights: list, device="cuda") -> dict: + for lora_path, weight in zip(lora_paths, weights): + logging.info(f"[DEBUG] Merging LoRA file: {lora_path} with weight: {weight}") + lora_sd = load_file(lora_path, device=device) + for key in list(lora_sd.keys()): + if "lora_down" not in key: + continue + base_name = key[: key.rfind(".lora_down")] + up_key = key.replace("lora_down", "lora_up") + module_name = base_name.replace("_", ".") + alpha_key = f"{base_name}.alpha" + if module_name not in model_sd: + logging.info(f"[DEBUG] Module {module_name} not found in model_sd; skipping key {key}") + continue + down_weight = lora_sd[key].float() + up_weight = lora_sd[up_key].float() + alpha = float(lora_sd.get(alpha_key, up_weight.shape[0])) + scale = weight * alpha / up_weight.shape[0] + logging.info(f"[DEBUG] Merging module: {module_name} with alpha: {alpha}, scale: {scale}") + target_weight = model_sd[module_name] + if len(target_weight.shape) == 2: + update = (up_weight @ down_weight) * scale + else: + if down_weight.shape[2:4] == (1, 1): + update = (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)) + update = update.unsqueeze(2).unsqueeze(3) * scale + else: + update = torch.nn.functional.conv2d( + down_weight.permute(1, 0, 2, 3), up_weight + ).permute(1, 0, 2, 3) * scale + model_sd[module_name] = target_weight + update.to(target_weight.dtype) + logging.info(f"[DEBUG] Updated module: {module_name}") + del up_weight, down_weight, update + del lora_sd + torch.cuda.empty_cache() + return model_sd + + def convert_to_gguf(self, model_path, working_dir): + base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + convert_script = os.path.join(base_path, "ComfyUI-GGUF", "tools", "convert.py") + temp_gguf = os.path.join(working_dir, "temp_converted.gguf") + logging.info("[DEBUG] Running conversion script: " + convert_script) + subprocess.run([sys.executable, convert_script, "--src", model_path, "--dst", temp_gguf], check=True) + logging.info("[DEBUG] Conversion complete.") + return temp_gguf + + def load_and_quantize(self, unet_name, quantization, delete_final_gguf, new_model_name, **kwargs): + mapping = {"FP16": "F16"} + logging.info(f"[DEBUG] Starting load_and_quantize: {new_model_name} | Quantization: {quantization}") + with tempfile.TemporaryDirectory() as merge_dir: + merged_model_path = os.path.join(merge_dir, "merged_model.safetensors") + model_path = folder_paths.get_full_path("diffusion_models", unet_name) + lora_list = [] + for i in range(1, 5): + name = kwargs.get(f"lora_name_{i}", "None") + switch = kwargs.get(f"switch_{i}", "Off") + logging.info(f"[DEBUG] Processing LoRA slot {i}: name = {name}, switch = {switch}") + if switch == "On" and name and name != "None": + lora_file_path = folder_paths.get_full_path("loras", name) + weight = kwargs.get(f"lora_weight_{i}", 1.0) + lora_list.append((lora_file_path, weight)) + logging.info(f"[DEBUG] Slot {i} active: path = {lora_file_path}, weight = {weight}") + else: + logging.info(f"[DEBUG] Slot {i} is inactive") + logging.info(f"[DEBUG] Total active LoRAs: {len(lora_list)}") + if lora_list: + model_sd = load_file(model_path, device="cuda") + model_sd = self.merge_flux_loras( + model_sd, + [lp for lp, _ in lora_list], + [w for _, w in lora_list] + ) + save_file(model_sd, merged_model_path) + del model_sd + torch.cuda.empty_cache() + else: + shutil.copy2(model_path, merged_model_path) + initial_gguf = self.convert_to_gguf(merged_model_path, merge_dir) + logging.info("[DEBUG] Initial GGUF file created.") + if quantization == "FP16": + final_gguf = os.path.join(merge_dir, f"{new_model_name}-{mapping.get(quantization, quantization)}.gguf") + shutil.copy2(initial_gguf, final_gguf) + logging.info("[DEBUG] FP16 selected; conversion skipped.") + else: + binary = os.path.join(os.path.dirname(os.path.abspath(__file__)), "binaries", "linux", "llama-quantize") + final_gguf = os.path.join(merge_dir, f"quantized_{quantization}.gguf") + subprocess.run([binary, initial_gguf, final_gguf, quantization], check=True) + logging.info("[DEBUG] Quantization completed.") + models_dir = os.path.join(folder_paths.models_dir, "unet") + os.makedirs(models_dir, exist_ok=True) + final_name = f"{new_model_name}-{mapping.get(quantization, quantization)}.gguf" + final_path = os.path.join(models_dir, final_name) + shutil.copy2(final_gguf, final_path) + logging.info("[DEBUG] Final model file copied to: " + final_path) + logging.info("[DEBUG] Loading final model.") + loader = UnetLoaderGGUF() + result = loader.load_unet(final_name) + logging.info("[DEBUG] Final model loaded.") + if delete_final_gguf: + os.unlink(final_path) + return result + + def override_class(cls): class NodeOverride(cls): @classmethod @@ -296,34 +422,6 @@ def override_class(cls): return NodeOverride -def override_class_with_offload(cls): - class NodeOverrideDiffSynth(cls): - @classmethod - def INPUT_TYPES(s): - inputs = copy.deepcopy(cls.INPUT_TYPES()) - devices = get_device_list() - default_device = devices[1] if len(devices) > 1 else devices[0] - inputs["optional"] = inputs.get("optional", {}) - inputs["optional"]["device"] = (devices, {"default": default_device}) - inputs["optional"]["offload_device"] = (devices, {"default": "cpu"}) - return inputs - - CATEGORY = "multigpu" - FUNCTION = "override" - - def override(self, *args, device=None, offload_device=None, **kwargs): - global current_device - global current_offload_device - if device is not None: - current_device = device - if offload_device is not None: - current_offload_device = offload_device - fn = getattr(super(), cls.FUNCTION) - return fn(*args, **kwargs) - - return NodeOverrideDiffSynth - - def override_class_with_distorch(cls): class NodeOverrideDisTorch(cls): @classmethod @@ -368,23 +466,14 @@ def check_module_exists(module_path): logging.info(f"MultiGPU: Found {module_path}, creating compatible MultiGPU nodes") return True -from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS -from .nodes import ( - UnetLoaderGGUF, UnetLoaderGGUFAdvanced, - CLIPLoaderGGUF, DualCLIPLoaderGGUF, TripleCLIPLoaderGGUF, - LTXVLoader, - Florence2ModelLoader, DownloadAndLoadFlorence2Model, - CheckpointLoaderNF4, - LoadFluxControlNet, - MMAudioModelLoader, MMAudioFeatureUtilsLoader, MMAudioSampler, - PulidModelLoader, PulidInsightFaceLoader, PulidEvaClipLoader, - HyVideoModelLoader, HyVideoModelLoaderDiffSynth, HyVideoVAELoader, DownloadAndLoadHyVideoTextEncoder -) - NODE_CLASS_MAPPINGS = { "DeviceSelectorMultiGPU": DeviceSelectorMultiGPU, - "HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter + "HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter, + "MergeFluxLoRAsQuantizeAndLoad": MergeFluxLoRAsQuantizeAndLoad, } + +NODE_CLASS_MAPPINGS["MergeFluxLoRAsQuantizeAndLoaddMultiGPU"] = override_class(MergeFluxLoRAsQuantizeAndLoad) + NODE_CLASS_MAPPINGS["UNETLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["UNETLoader"]) NODE_CLASS_MAPPINGS["VAELoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["VAELoader"]) NODE_CLASS_MAPPINGS["CLIPLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["CLIPLoader"])