From a7c894e8b5475c0eeef68325b9f40db2a94de86c Mon Sep 17 00:00:00 2001 From: smthemex Date: Tue, 4 Aug 2026 20:01:35 +0800 Subject: [PATCH] init --- __init__.py | 11 +++- dif_gguf_node.py | 145 +++++++++++++++++++++++++++++++++++++++++++++- ggml_ops.py | 8 +++ tokenizer_map.py | 148 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 310 insertions(+), 2 deletions(-) create mode 100644 tokenizer_map.py diff --git a/__init__.py b/__init__.py index a5a971f..4c0a825 100644 --- a/__init__.py +++ b/__init__.py @@ -1,15 +1,24 @@ # ComfyUI_Dif_GGUF - GGUF model loader for ComfyUI # Reads GGUF files, preserves quantization, injects weights bypassing shape checks. -from .dif_gguf_node import Dif_GGUF_Loader, DifGGUFExtension +from .dif_gguf_node import ( + Dif_GGUF_Loader, + CLIP_GGUF_Loader, + DualCLIP_GGUF_Loader, + DifGGUFExtension, +) from .ggml_ops import GGMLTensor, GGMLOps, GGMLLayer, GGUFModelPatcher NODE_CLASS_MAPPINGS = { "Dif_GGUF_Loader": Dif_GGUF_Loader, + "CLIP_GGUF_Loader": CLIP_GGUF_Loader, + "DualCLIP_GGUF_Loader": DualCLIP_GGUF_Loader, } __all__ = [ "Dif_GGUF_Loader", + "CLIP_GGUF_Loader", + "DualCLIP_GGUF_Loader", "DifGGUFExtension", "GGMLTensor", "GGMLOps", diff --git a/dif_gguf_node.py b/dif_gguf_node.py index b09d5ec..a7a12ca 100644 --- a/dif_gguf_node.py +++ b/dif_gguf_node.py @@ -9,6 +9,7 @@ from comfy_api.latest import IO, ComfyExtension import comfy import comfy.utils +import comfy.sd import comfy.model_management import comfy.model_detection import comfy.model_patcher @@ -27,6 +28,7 @@ def update_folder_names_and_paths(key, targets=[]): # Add a custom keys for files ending in .gguf update_folder_names_and_paths("unet_gguf", ["diffusion_models", "unet"]) +update_folder_names_and_paths("clip_gguf", ["text_encoders", "clip"]) # add gguf folder to comfy optional weigths_gguf_current_path = os.path.join(folder_paths.models_dir, "gguf") @@ -38,6 +40,7 @@ folder_paths.add_model_folder_path("gguf", weigths_gguf_current_path) # Some codes from @city96 https://github.com/city96/ComfyUI-GGUF and used for diffusion_gguf model import gguf as gguf_lib from .ggml_ops import (GGMLTensor, GGMLOps,is_quantized, GGUFModelPatcher) +from .tokenizer_map import convert_text_encoder_state_dict def get_field(reader, field_name, field_type): @@ -282,10 +285,150 @@ class Dif_GGUF_Loader(IO.ComfyNode): return IO.NodeOutput(model) +# ── Text encoder (CLIP) GGUF ──────────────────────────────────────────── +# Mirrors the unet path above, but hands the state dict to ComfyUI's +# load_text_encoder_state_dicts() with custom_operations=GGMLOps so the +# quantized tensors stay quantized and are dequantized lazily at forward time. +# No ComfyUI source file is modified: model_options is a public entry point. + +CLIP_GGUF_TYPES = [ + "stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", + "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", + "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", + "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage", + "mage", "minimax", +] + +CLIP_GGUF_DUAL_TYPES = [ + "sdxl", "sd3", "flux", "hunyuan_video", "hidream", "hunyuan_image", + "hunyuan_video_15", "kandinsky5", "kandinsky5_image", "ltxv", "newbie", "ace", +] + + +def _clip_gguf_full_path(name): + """Resolve a text encoder file from either the gguf or text_encoders folder.""" + for folder in ("clip_gguf", "gguf", "text_encoders"): + try: + path = folder_paths.get_full_path(folder, name) + except Exception: + path = None + if path is not None: + return path + return None + + +def read_text_encoder_gguf(path): + """Read a text encoder GGUF and normalize its key names.""" + sd = read_gguf(path) + sd, arch = convert_text_encoder_state_dict(sd) + logging.info(f"GGUF TE: {len(sd)} tensors (arch={arch})") + return sd + + +def load_text_encoder_gguf(paths, clip_type_name, device="default"): + """Load one or more text encoders, any of which may be a .gguf file.""" + clip_type = getattr(comfy.sd.CLIPType, clip_type_name.upper(), + comfy.sd.CLIPType.STABLE_DIFFUSION) + + state_dicts = [] + has_quant = False + for path in paths: + if path.endswith(".gguf"): + sd = read_text_encoder_gguf(path) + if any(is_quantized(v) for v in sd.values()): + has_quant = True + else: + sd = comfy.utils.load_torch_file(path, safe_load=True) + state_dicts.append(sd) + + model_options = {} + if device == "cpu": + model_options["load_device"] = model_options["offload_device"] = torch.device("cpu") + if has_quant: + # Public hook: sd1_clip reads model_options["custom_operations"]. + model_options["custom_operations"] = GGMLOps + # GGUF norm tensors are F32, so ComfyUI's t5xxl_detect/llama_detect would + # pick float32 as the compute dtype. Force a sane one instead. + compute_dtype = comfy.model_management.text_encoder_dtype( + comfy.model_management.text_encoder_device() + ) + model_options["dtype_t5"] = compute_dtype + model_options["dtype_llama"] = compute_dtype + + clip = comfy.sd.load_text_encoder_state_dicts( + state_dicts, + embedding_directory=folder_paths.get_folder_paths("embeddings"), + clip_type=clip_type, + model_options=model_options, + ) + + if has_quant: + clip.patcher = GGUFModelPatcher.clone(clip.patcher) + clip.patcher.size = 0 + return clip + + +class CLIP_GGUF_Loader(IO.ComfyNode): + @classmethod + def define_schema(cls): + options = ["none"] + folder_paths.get_filename_list("clip_gguf") \ + + folder_paths.get_filename_list("gguf") + return IO.Schema( + node_id="CLIP_GGUF_Loader", + display_name="CLIP_GGUF_Loader", + description="Load a text encoder from GGUF (quantized lazy dequant)", + category="model/loaders", + inputs=[ + IO.Combo.Input("gguf", options=options), + IO.Combo.Input("type", options=CLIP_GGUF_TYPES), + IO.Combo.Input("device", options=["default", "cpu"], optional=True), + ], + outputs=[IO.Clip.Output(display_name="CLIP")], + ) + + @classmethod + def execute(cls, gguf, type="stable_diffusion", device="default") -> IO.NodeOutput: + path = _clip_gguf_full_path(gguf) if gguf != "none" else None + assert path is not None, f"Text encoder not found: {gguf}" + clip = load_text_encoder_gguf([path], type, device) + logging.info(f"Loaded text encoder: {type}") + return IO.NodeOutput(clip) + + +class DualCLIP_GGUF_Loader(IO.ComfyNode): + @classmethod + def define_schema(cls): + options = ["none"] + folder_paths.get_filename_list("clip_gguf") \ + + folder_paths.get_filename_list("gguf") + return IO.Schema( + node_id="DualCLIP_GGUF_Loader", + display_name="DualCLIP_GGUF_Loader", + description="Load two text encoders, either of which may be GGUF", + category="model/loaders", + inputs=[ + IO.Combo.Input("gguf1", options=options), + IO.Combo.Input("gguf2", options=options), + IO.Combo.Input("type", options=CLIP_GGUF_DUAL_TYPES), + IO.Combo.Input("device", options=["default", "cpu"], optional=True), + ], + outputs=[IO.Clip.Output(display_name="CLIP")], + ) + + @classmethod + def execute(cls, gguf1, gguf2, type="flux", device="default") -> IO.NodeOutput: + path1 = _clip_gguf_full_path(gguf1) if gguf1 != "none" else None + path2 = _clip_gguf_full_path(gguf2) if gguf2 != "none" else None + assert path1 is not None, f"Text encoder not found: {gguf1}" + assert path2 is not None, f"Text encoder not found: {gguf2}" + clip = load_text_encoder_gguf([path1, path2], type, device) + logging.info(f"Loaded dual text encoder: {type}") + return IO.NodeOutput(clip) + + class DifGGUFExtension(ComfyExtension): @override async def get_node_list(self) -> list[type[IO.ComfyNode]]: - return [Dif_GGUF_Loader,] + return [Dif_GGUF_Loader, CLIP_GGUF_Loader, DualCLIP_GGUF_Loader] async def comfy_entrypoint() -> DifGGUFExtension: return DifGGUFExtension() \ No newline at end of file diff --git a/ggml_ops.py b/ggml_ops.py index 8911b6f..831b097 100644 --- a/ggml_ops.py +++ b/ggml_ops.py @@ -414,6 +414,14 @@ class GGMLOps(comfy.ops.manual_cast): weight, bias = self.cast_bias_weight(input) return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps) + class RMSNorm(GGMLLayer, comfy.ops.manual_cast.RMSNorm): + # Required by T5 / LLaMA / Gemma text encoders. + def forward_ggml_cast_weights(self, input): + if self.weight is None: + return super().forward_comfy_cast_weights(input) + weight, _bias = self.cast_bias_weight(input) + return torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps) + class GroupNorm(GGMLLayer, comfy.ops.manual_cast.GroupNorm): def forward_ggml_cast_weights(self, input): weight, bias = self.cast_bias_weight(input) diff --git a/tokenizer_map.py b/tokenizer_map.py new file mode 100644 index 0000000..53944b8 --- /dev/null +++ b/tokenizer_map.py @@ -0,0 +1,148 @@ +# Key remapping for text-encoder GGUF files. +# +# Unet GGUF files (city96 style) keep the original ComfyUI tensor names, so the +# unet loader can consume them directly. Text-encoder GGUF files produced by +# llama.cpp's convert scripts instead use llama.cpp naming ("blk.0.attn_q.weight", +# "token_embd.weight", ...). This module converts those back to the names that +# ComfyUI's T5 / llama / gemma implementations expect. +# +# Nothing here touches ComfyUI source; it only rewrites a state dict in memory. + +import logging + + +def sd_prefix_from_state_dict(sd): + """Guess the architecture of a text encoder state dict.""" + keys = sd.keys() + if any(k.startswith("enc.") or k.startswith("dec.") for k in keys): + return "t5" + if "token_embd.weight" in sd and any(k.endswith("attn_norm.weight") for k in keys): + return "llama" + if any(k.startswith("text_model.") for k in keys): + return "clip" + if any(k.startswith("encoder.block.") for k in keys): + return "t5_native" + return None + + +def _t5_block_map(n_block, prefix="encoder"): + """llama.cpp T5 block names -> ComfyUI T5 block names.""" + src = f"enc.blk.{n_block}" + dst = f"{prefix}.block.{n_block}" + return { + f"{src}.attn_q.weight": f"{dst}.layer.0.SelfAttention.q.weight", + f"{src}.attn_k.weight": f"{dst}.layer.0.SelfAttention.k.weight", + f"{src}.attn_v.weight": f"{dst}.layer.0.SelfAttention.v.weight", + f"{src}.attn_o.weight": f"{dst}.layer.0.SelfAttention.o.weight", + f"{src}.attn_norm.weight": f"{dst}.layer.0.layer_norm.weight", + f"{src}.attn_rel_b.weight": f"{dst}.layer.0.SelfAttention.relative_attention_bias.weight", + f"{src}.ffn_up.weight": f"{dst}.layer.1.DenseReluDense.wi_1.weight", + f"{src}.ffn_gate.weight": f"{dst}.layer.1.DenseReluDense.wi_0.weight", + f"{src}.ffn_down.weight": f"{dst}.layer.1.DenseReluDense.wo.weight", + f"{src}.ffn_norm.weight": f"{dst}.layer.1.layer_norm.weight", + } + + +def _llama_block_map(n_block): + """llama.cpp block names -> ComfyUI llama/gemma block names.""" + src = f"blk.{n_block}" + dst = f"model.layers.{n_block}" + return { + f"{src}.attn_q.weight": f"{dst}.self_attn.q_proj.weight", + f"{src}.attn_k.weight": f"{dst}.self_attn.k_proj.weight", + f"{src}.attn_v.weight": f"{dst}.self_attn.v_proj.weight", + f"{src}.attn_output.weight": f"{dst}.self_attn.o_proj.weight", + f"{src}.attn_norm.weight": f"{dst}.input_layernorm.weight", + f"{src}.ffn_up.weight": f"{dst}.mlp.up_proj.weight", + f"{src}.ffn_gate.weight": f"{dst}.mlp.gate_proj.weight", + f"{src}.ffn_down.weight": f"{dst}.mlp.down_proj.weight", + f"{src}.ffn_norm.weight": f"{dst}.post_attention_layernorm.weight", + } + + +def _count_blocks(sd, pattern): + n = 0 + while True: + if not any(k.startswith(pattern.format(n)) for k in sd): + break + n += 1 + return n + + +def convert_t5_state_dict(sd): + """llama.cpp T5 GGUF -> ComfyUI T5 naming.""" + n_blocks = _count_blocks(sd, "enc.blk.{}.") + key_map = {} + for i in range(n_blocks): + key_map.update(_t5_block_map(i)) + key_map.update({ + "token_embd.weight": "shared.weight", + "enc.output_norm.weight": "encoder.final_layer_norm.weight", + "output.weight": "lm_head.weight", + }) + + new_sd = {} + for k, v in sd.items(): + # Drop the decoder: ComfyUI only runs the T5 encoder. + if k.startswith("dec."): + continue + dst = key_map.get(k) + if dst is None: + if k.startswith("enc.") or k == "output.weight": + logging.debug(f"GGUF TE: dropping unmapped key {k}") + continue + dst = k + new_sd[dst] = v + + # ComfyUI's T5 keeps the relative attention bias only on block 0. + rel_b = "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight" + for i in range(1, n_blocks): + new_sd.pop( + f"encoder.block.{i}.layer.0.SelfAttention.relative_attention_bias.weight", + None, + ) + if rel_b not in new_sd: + logging.warning("GGUF TE: T5 relative_attention_bias missing") + return new_sd + + +def convert_llama_state_dict(sd): + """llama.cpp LLaMA/Gemma GGUF -> ComfyUI naming.""" + n_blocks = _count_blocks(sd, "blk.{}.") + key_map = {} + for i in range(n_blocks): + key_map.update(_llama_block_map(i)) + key_map.update({ + "token_embd.weight": "model.embed_tokens.weight", + "output_norm.weight": "model.norm.weight", + "output.weight": "lm_head.weight", + }) + + new_sd = {} + for k, v in sd.items(): + dst = key_map.get(k) + if dst is None: + if k.startswith("blk.") or k.startswith("output"): + logging.debug(f"GGUF TE: dropping unmapped key {k}") + continue + dst = k + new_sd[dst] = v + return new_sd + + +def convert_text_encoder_state_dict(sd): + """ + Detect the text encoder architecture and normalize its key names. + + Returns (state_dict, arch). If the file already uses ComfyUI naming it is + returned unchanged. + """ + arch = sd_prefix_from_state_dict(sd) + if arch == "t5": + logging.info("GGUF TE: detected T5 (llama.cpp naming), converting keys") + return convert_t5_state_dict(sd), "t5" + if arch == "llama": + logging.info("GGUF TE: detected LLaMA/Gemma (llama.cpp naming), converting keys") + return convert_llama_state_dict(sd), "llama" + logging.info(f"GGUF TE: using keys as-is (arch={arch})") + return sd, arch