Merge branch 'main' into auto_convert
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user