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/5] 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/5] 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/5] 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/5] 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/5] 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