Merge branch 'main' into auto_convert

This commit is contained in:
City
2025-11-25 23:46:01 +01:00
3 changed files with 113 additions and 3 deletions
+95 -1
View File
@@ -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
+17 -1
View File
@@ -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
+1 -1
View File
@@ -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