From 01f6e3ff6333d521e3e0aa5b0c5e8ce5e8045995 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Wed, 5 Mar 2025 04:44:53 +0100 Subject: [PATCH 1/7] spiece tokenizer from gguf metadata --- loader.py | 97 +++++++++++++++++++++++++++++++++++++++++++++++++------ nodes.py | 3 +- 2 files changed, 90 insertions(+), 10 deletions(-) diff --git a/loader.py b/loader.py index 894191e..ea3a0e1 100644 --- a/loader.py +++ b/loader.py @@ -18,6 +18,41 @@ 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_string_field(reader, field_name): + field = reader.get_field(field_name) + if field is None: + return None + if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING: + raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {field.types!r}") + return str(field.parts[field.data[-1]], encoding="utf-8") + +def get_field(reader, field_name, field_type): + # TODO: merge with function above and add checking + field = reader.get_field(field_name) + if field is None: + return None + if field_type == int: + return int(field.parts[field.data[-1]]) + elif field_type == float: + return float(field.parts[field.data[-1]]) + elif field_type == bool: + return bool(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 + if field_type == int: + return tuple(int(field.parts[part_idx][0]) for part_idx in field.data) + elif field_type == float: + return tuple(float(field.parts[part_idx][0]) for part_idx in field.data) + elif field_type == str: + return tuple(str(field.parts[part_idx], encoding="utf-8") 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,19 +77,14 @@ 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_string_field(reader, "general.architecture") + 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 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}") # main loading loop state_dict = {} @@ -148,9 +178,58 @@ 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...") + from sentencepiece import sentencepiece_model_pb2 as model + spm = model.ModelProto() + + reader = gguf.GGUFReader(path) + + if get_string_field(reader, "tokenizer.ggml.model") == "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) sd = sd_map_replace(sd, T5_SD_MAP) elif arch in {"llama"}: temb_key = "token_embd.weight" diff --git a/nodes.py b/nodes.py index 8185e9c..7f5ad83 100644 --- a/nodes.py +++ b/nodes.py @@ -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"],), } } From d99b2d931fabe97c963e091c5ac24663d9961900 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 14 Mar 2025 16:56:48 +0100 Subject: [PATCH 2/7] force dequant on umT5 to prevent OOM not the best solution but that tensor alone needs ~3GB of runtime memory to dequant --- loader.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/loader.py b/loader.py index ea3a0e1..18c1af1 100644 --- a/loader.py +++ b/loader.py @@ -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"} @@ -230,6 +230,9 @@ def gguf_clip_loader(path): 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" From 296b0c14197b1b8a58422029343d1c3129b9398b Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 14 Mar 2025 17:11:47 +0100 Subject: [PATCH 3/7] Clean up field logic --- loader.py | 32 +++++++++++--------------------- 1 file changed, 11 insertions(+), 21 deletions(-) diff --git a/loader.py b/loader.py index 18c1af1..2f690ed 100644 --- a/loader.py +++ b/loader.py @@ -18,25 +18,17 @@ 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_string_field(reader, field_name): - field = reader.get_field(field_name) - if field is None: - return None - if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING: - raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {field.types!r}") - return str(field.parts[field.data[-1]], encoding="utf-8") - def get_field(reader, field_name, field_type): - # TODO: merge with function above and add checking field = reader.get_field(field_name) if field is None: return None - if field_type == int: - return int(field.parts[field.data[-1]]) - elif field_type == float: - return float(field.parts[field.data[-1]]) - elif field_type == bool: - return bool(field.parts[field.data[-1]]) + 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}") @@ -44,12 +36,10 @@ def get_list_field(reader, field_name, field_type): field = reader.get_field(field_name) if field is None: return None - if field_type == int: - return tuple(int(field.parts[part_idx][0]) for part_idx in field.data) - elif field_type == float: - return tuple(float(field.parts[part_idx][0]) for part_idx in field.data) 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}") @@ -77,7 +67,7 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal # detect and verify architecture compat = None - arch_str = get_string_field(reader, "general.architecture") + 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 @@ -186,7 +176,7 @@ def gguf_tokenizer_loader(path, temb_shape): reader = gguf.GGUFReader(path) - if get_string_field(reader, "tokenizer.ggml.model") == "t5": + 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: From 31dffa061b76d388edce2f734dee802d43e27132 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 14 Mar 2025 17:37:14 +0100 Subject: [PATCH 4/7] Add error message for missing packages --- loader.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/loader.py b/loader.py index 2f690ed..5253092 100644 --- a/loader.py +++ b/loader.py @@ -171,7 +171,10 @@ def llama_permute(raw_sd, n_head, n_head_kv): def gguf_tokenizer_loader(path, temb_shape): # convert gguf tokenizer to spiece print(f"Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") - from sentencepiece import sentencepiece_model_pb2 as model + 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) From ad69decc8d09d8d0f6362a61ec68a4af3cd25927 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 14 Mar 2025 17:43:05 +0100 Subject: [PATCH 5/7] Update requirements.txt --- requirements.txt | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index ddb8bbe..f49905c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,5 @@ +# main gguf>=0.13.0 - +# optional - tokenizer +sentencepiece +protobuf From 4aebf1bf2ca54b9448aeabf3e28f983d108df068 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sun, 16 Mar 2025 19:38:40 +0100 Subject: [PATCH 6/7] Add compat mode pig closes #230 --- loader.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/loader.py b/loader.py index 5253092..948c50a 100644 --- a/loader.py +++ b/loader.py @@ -73,9 +73,16 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal 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 = {} From bc5223b0e37e053dbec2ea5e5f52c2fd4b8f712a Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sun, 16 Mar 2025 20:14:33 +0100 Subject: [PATCH 7/7] Fix qtype logging on newer gguf-py --- loader.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/loader.py b/loader.py index 948c50a..ff1ba3f 100644 --- a/loader.py +++ b/loader.py @@ -88,7 +88,6 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal 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) @@ -104,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