diff --git a/loader.py b/loader.py index 56bbd57..032c739 100644 --- a/loader.py +++ b/loader.py @@ -3,12 +3,15 @@ import warnings import logging import torch import gguf +import re +import os from .ops import GGMLTensor from .dequant import is_quantized, dequantize_tensor IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "cosmos", "ltxv", "hyvid", "wan", "lumina2", "qwen_image"} TXT_ARCH_LIST = {"t5", "t5encoder", "llama", "qwen2vl"} +VIS_TYPE_LIST = {"clip-vision", "mmproj"} def get_orig_shape(reader, tensor_name): field_key = f"comfy.gguf.orig_shape.{tensor_name}" @@ -70,6 +73,7 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal # detect and verify architecture compat = None arch_str = get_field(reader, "general.architecture", str) + type_str = get_field(reader, "general.type", str) if arch_str in [None, "pig"]: if is_text_model: raise ValueError(f"This text model is incompatible with llama.cpp!\nConsider using the safetensors version\n({path})") @@ -81,7 +85,8 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal except Exception as e: raise ValueError(f"This model is not currently supported - ({e})") elif arch_str not in TXT_ARCH_LIST and is_text_model: - raise ValueError(f"Unexpected text model architecture type in GGUF file: {arch_str!r}") + if type_str not in VIS_TYPE_LIST: + raise ValueError(f"Unexpected text model architecture type in GGUF file: {arch_str!r}") elif arch_str not in IMG_ARCH_LIST and not is_text_model: raise ValueError(f"Unexpected architecture type in GGUF file: {arch_str!r}") @@ -165,6 +170,19 @@ LLAMA_SD_MAP = { "output.weight": "lm_head.weight", } +CLIP_VISION_SD_MAP = { + "mm.": "visual.merger.mlp.", + "v.post_ln.": "visual.merger.ln_q.", + "v.patch_embd": "visual.patch_embed.proj", + "v.blk.": "visual.blocks.", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", + "ffn_gate": "mlp.gate_proj", + "attn_out.": "attn.proj.", + "ln1.": "norm1.", + "ln2.": "norm2.", +} + def sd_map_replace(raw_sd, key_map): sd = {} for k,v in raw_sd.items(): @@ -185,6 +203,79 @@ def llama_permute(raw_sd, n_head, n_head_kv): sd[k] = v return sd +def strip_quant_suffix(name): + pattern = r"[-_]?(?:ud-)?i?q[0-9]_[a-z0-9_\-]{1,8}$" + match = re.search(pattern, name, re.IGNORECASE) + if match: + name = name[:match.start()] + return name + +def gguf_mmproj_loader(path): + # Reverse version of Qwen2VLVisionModel.modify_tensors + logging.info("Attenpting to find mmproj file for text encoder...") + + # get name to match w/o quant suffix + tenc_fname = os.path.basename(path) + tenc = os.path.splitext(tenc_fname)[0].lower() + tenc = strip_quant_suffix(tenc) + + # try and find matching mmproj + target = [] + root = os.path.dirname(path) + for fname in os.listdir(root): + name, ext = os.path.splitext(fname) + if ext.lower() != ".gguf": + continue + if "mmproj" not in name.lower(): + continue + if tenc in name.lower(): + target.append(fname) + + if len(target) == 0: + logging.error(f"Error: Can't find mmproj file for '{tenc_fname}' (matching:'{tenc}')! Qwen-Image-Edit will be broken!") + return {} + if len(target) > 1: + logging.error(f"Ambiguous mmproj for text encoder '{tenc_fname}', will use first match.") + + logging.info(f"Using mmproj '{target[0]}' for text encoder '{tenc_fname}'.") + target = os.path.join(root, target[0]) + vsd = gguf_sd_loader(target, is_text_model=True) + + # concat 4D to 5D + if "v.patch_embd.weight.1" in vsd: + w1 = dequantize_tensor(vsd.pop("v.patch_embd.weight"), dtype=torch.float32) + w2 = dequantize_tensor(vsd.pop("v.patch_embd.weight.1"), dtype=torch.float32) + vsd["v.patch_embd.weight"] = torch.stack([w1, w2], dim=2) + + # run main replacement + vsd = sd_map_replace(vsd, CLIP_VISION_SD_MAP) + + # handle split Q/K/V + if "visual.blocks.0.attn_q.weight" in vsd: + attns = {} + # filter out attentions + group + for k,v in vsd.items(): + if any(x in k for x in ["attn_q", "attn_k", "attn_v"]): + k_attn, k_name = k.rsplit(".attn_", 1) + k_attn += ".attn.qkv." + k_name.split(".")[-1] + if k_attn not in attns: + attns[k_attn] = {} + attns[k_attn][k_name] = dequantize_tensor( + v, dtype=(torch.bfloat16 if is_quantized(v) else torch.float16) + ) + + # recombine + for k,v in attns.items(): + suffix = k.split(".")[-1] + vsd[k] = torch.cat([ + v[f"q.{suffix}"], + v[f"k.{suffix}"], + v[f"v.{suffix}"], + ], dim=0) + del attns + + return vsd + def gguf_tokenizer_loader(path, temb_shape): # convert gguf tokenizer to spiece logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") @@ -254,6 +345,9 @@ def gguf_clip_loader(path): sd = sd_map_replace(sd, LLAMA_SD_MAP) if arch == "llama": sd = llama_permute(sd, 32, 8) # L3 + if arch == "qwen2vl": + vsd = gguf_mmproj_loader(path) + sd.update(vsd) else: pass return sd diff --git a/nodes.py b/nodes.py index 4159142..ff5aaf0 100644 --- a/nodes.py +++ b/nodes.py @@ -76,8 +76,22 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): # TODO: Find another way to not unload after patches return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) + + def pin_weight_to_device(self, key): + op_key = key.rsplit('.', 1)[0] + if not self.mmap_released and op_key in self.named_modules_to_munmap: + # TODO: possible to OOM, find better way to detach + self.named_modules_to_munmap[op_key].to(self.load_device).to(self.offload_device) + del self.named_modules_to_munmap[op_key] + super().pin_weight_to_device(key) + mmap_released = False + named_modules_to_munmap = {} + def load(self, *args, force_patch_weights=False, **kwargs): + if not self.mmap_released: + self.named_modules_to_munmap = dict(self.model.named_modules()) + # always call `patch_weight_to_device` even for lowvram super().load(*args, force_patch_weights=True, **kwargs) @@ -85,7 +99,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): if not self.mmap_released: linked = [] if kwargs.get("lowvram_model_memory", 0) > 0: - for n, m in self.model.named_modules(): + for n, m in self.named_modules_to_munmap.items(): if hasattr(m, "weight"): device = getattr(m.weight, "device", None) if device == self.offload_device: @@ -102,6 +116,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): # TODO: possible to OOM, find better way to detach m.to(self.load_device).to(self.offload_device) self.mmap_released = True + self.named_modules_to_munmap = {} def clone(self, *args, **kwargs): src_cls = self.__class__ @@ -111,6 +126,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): self.__class__ = src_cls # GGUF specific clone values below n.patch_on_device = getattr(self, "patch_on_device", False) + n.mmap_released = getattr(self, "mmap_released", False) if src_cls != GGUFModelPatcher: n.size = 0 # force recalc return n diff --git a/ops.py b/ops.py index 41d42f5..88a352e 100644 --- a/ops.py +++ b/ops.py @@ -153,7 +153,7 @@ class GGMLLayer(torch.nn.Module): # Take into account space required for dequantizing the largest tensor if self.largest_layer: shape = getattr(self.weight, "tensor_shape", self.weight.shape) - dtype = self.dequant_dtype or torch.float16 + dtype = self.dequant_dtype if self.dequant_dtype and self.dequant_dtype != "target" else torch.float16 temp = torch.empty(*shape, device=torch.device("meta"), dtype=dtype) destination[prefix + "temp.weight"] = temp