nodes/city96-ComfyUI-GGUF
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge branch 'main' into convert_refactor_new

Browse Source
This commit is contained in:
City
2025-03-30 18:43:44 +02:00
parent bbd89ae547 bc5223b0e3
commit f3b968e137
3 changed files with 104 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
loader.py
+98 -16
View File
@@ -3,7 +3,7 @@ import torch
import gguf
from .ops import GGMLTensor
from .dequant import is_quantized
from .dequant import is_quantized, dequantize_tensor
IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "ltxv", "hyvid", "wan"}
TXT_ARCH_LIST = {"t5", "t5encoder", "llama"}
@@ -18,6 +18,31 @@ def get_orig_shape(reader, tensor_name):
raise TypeError(f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}")
return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data))
def get_field(reader, field_name, field_type):
field = reader.get_field(field_name)
if field is None:
return None
elif field_type == str:
# extra check here as this is used for checking arch string
if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING:
raise TypeError(f"Bad type for GGUF {field_name} key: expected string, got {field.types!r}")
return str(field.parts[field.data[-1]], encoding="utf-8")
elif field_type in [int, float, bool]:
return field_type(field.parts[field.data[-1]])
else:
raise TypeError(f"Unknown field type {field_type}")
def get_list_field(reader, field_name, field_type):
field = reader.get_field(field_name)
if field is None:
return None
elif field_type == str:
return tuple(str(field.parts[part_idx], encoding="utf-8") for part_idx in field.data)
elif field_type in [int, float, bool]:
return tuple(field_type(field.parts[part_idx][0]) for part_idx in field.data)
else:
raise TypeError(f"Unknown field type {field_type}")
def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=False):
"""
Read state dict as fake tensors
@@ -42,26 +67,27 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal
# detect and verify architecture
compat = None
arch_str = None
arch_field = reader.get_field("general.architecture")
if arch_field is not None:
if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING:
raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}")
arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8")
if arch_str not in IMG_ARCH_LIST and arch_str not in TXT_ARCH_LIST:
raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}")
else: # stable-diffusion.cpp
arch_str = get_field(reader, "general.architecture", str)
if arch_str is None: # stable-diffusion.cpp
# import here to avoid changes to convert.py breaking regular models
from .tools.convert import detect_arch
arch_str = detect_arch(set(val[0] for val in tensors)).arch
compat = "sd.cpp"
elif arch_str in ["pig"]:
from .tools.convert import detect_arch
arch_str = detect_arch(set(val[0] for val in tensors)).arch
compat = "pig"
elif arch_str not in IMG_ARCH_LIST and arch_str not in TXT_ARCH_LIST:
raise ValueError(f"Unexpected architecture type in GGUF file: {arch_str!r}")
if compat:
print(f"Warning: This model file is loaded in compatibility mode '{compat}' [arch:{arch_str}]")
# main loading loop
state_dict = {}
qtype_dict = {}
for sd_key, tensor in tensors:
tensor_name = tensor.name
tensor_type_str = str(tensor.tensor_type)
torch_tensor = torch.from_numpy(tensor.data) # mmap
shape = get_orig_shape(reader, tensor_name)
@@ -77,19 +103,20 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal
if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
torch_tensor = torch_tensor.view(*shape)
state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape)
# keep track of loaded tensor types
tensor_type_str = getattr(tensor.tensor_type, "name", repr(tensor.tensor_type))
qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1
# print loaded tensor type counts
print("gguf qtypes: " + ", ".join(f"{k} ({v})" for k, v in qtype_dict.items()))
# mark largest tensor for vram estimation
qsd = {k:v for k,v in state_dict.items() if is_quantized(v)}
if len(qsd) > 0:
max_key = max(qsd.keys(), key=lambda k: qsd[k].numel())
state_dict[max_key].is_largest_weight = True
# sanity check debug print
print("\nggml_sd_loader:")
for k,v in qtype_dict.items():
print(f" {k:30}{v:3}")
if return_arch:
return (state_dict, arch_str)
return state_dict
@@ -148,9 +175,64 @@ def llama_permute(raw_sd, n_head, n_head_kv):
sd[k] = v
return sd
def gguf_tokenizer_loader(path, temb_shape):
# convert gguf tokenizer to spiece
print(f"Attempting to recreate sentencepiece tokenizer from GGUF file metadata...")
try:
from sentencepiece import sentencepiece_model_pb2 as model
except ImportError:
raise ImportError("Please make sure sentencepiece and protobuf are installed.\npip install sentencepiece protobuf")
spm = model.ModelProto()
reader = gguf.GGUFReader(path)
if get_field(reader, "tokenizer.ggml.model", str) == "t5":
if temb_shape == (256384, 4096): # probably UMT5
spm.trainer_spec.model_type == 1 # Unigram (do we have a T5 w/ BPE?)
else:
raise NotImplementedError(f"Unknown model, can't set tokenizer!")
else:
raise NotImplementedError(f"Unknown model, can't set tokenizer!")
spm.normalizer_spec.add_dummy_prefix = get_field(reader, "tokenizer.ggml.add_space_prefix", bool)
spm.normalizer_spec.remove_extra_whitespaces = get_field(reader, "tokenizer.ggml.remove_extra_whitespaces", bool)
tokens = get_list_field(reader, "tokenizer.ggml.tokens", str)
scores = get_list_field(reader, "tokenizer.ggml.scores", float)
toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int)
for idx, (token, score, toktype) in enumerate(zip(tokens, scores, toktypes)):
# # These aren't present in the original?
# if toktype == 5 and idx >= temb_shape[0]%1000):
# continue
piece = spm.SentencePiece()
piece.piece = token
piece.score = score
piece.type = toktype
spm.pieces.append(piece)
# unsure if any of these are correct
spm.trainer_spec.byte_fallback = True
spm.trainer_spec.vocab_size = len(tokens) # split off unused?
spm.trainer_spec.max_sentence_length = 4096
spm.trainer_spec.eos_id = get_field(reader, "tokenizer.ggml.eos_token_id", int)
spm.trainer_spec.pad_id = get_field(reader, "tokenizer.ggml.padding_token_id", int)
print(f"Created tokenizer with vocab size of {len(spm.pieces)}")
del reader
return torch.ByteTensor(list(spm.SerializeToString()))
def gguf_clip_loader(path):
sd, arch = gguf_sd_loader(path, return_arch=True)
if arch in {"t5", "t5encoder"}:
temb_key = "token_embd.weight"
if temb_key in sd and sd[temb_key].shape == (256384, 4096):
# non-standard Comfy-Org tokenizer
sd["spiece_model"] = gguf_tokenizer_loader(path, sd[temb_key].shape)
# TODO: dequantizing token embed here is janky but otherwise we OOM due to tensor being massive.
print(f"Dequantizing {temb_key} to prevent runtime OOM.")
sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16)
sd = sd_map_replace(sd, T5_SD_MAP)
elif arch in {"llama"}:
temb_key = "token_embd.weight"
nodes.py
+2 -1
View File
@@ -186,6 +186,7 @@ CLIP_ENUM_MAP = {
"ltxv": "LTXV",
"hunyuan_video": "HUNYUAN_VIDEO",
"pixart": "PIXART",
"wan": "WAN",
}
def get_clip_type(name):
@@ -203,7 +204,7 @@ class CLIPLoaderGGUF:
return {
"required": {
"clip_name": (s.get_filename_list(),),
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart"],),
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "wan"],),
}
}
requirements.txt
+4 -1
View File
@@ -1,2 +1,5 @@
# main
gguf>=0.13.0
# optional - tokenizer
sentencepiece
protobuf