Merge branch 'main' into auto_convert

loader.py

@@ -45,7 +45,7 @@ def get_list_field(reader, field_name, field_type):
     else:
         raise TypeError(f"Unknown field type {field_type}")
 
-def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=False):
+def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=False, is_text_model=False):
     """
     Read state dict as fake tensors
     """
@@ -70,20 +70,23 @@ 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)
-    if arch_str is None: # stable-diffusion.cpp
+    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})")
+        compat = "sd.cpp" if arch_str is None else arch_str
         # 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:
+        try:
+            arch_str = detect_arch(set(val[0] for val in tensors)).arch
+        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}")
+    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}")
 
     if compat:
-        logging.warning(f"Warning: This model file is loaded in compatibility mode '{compat}' [arch:{arch_str}]")
+        logging.warning(f"Warning: This gguf model file is loaded in compatibility mode '{compat}' [arch:{arch_str}]")
 
     # main loading loop
     state_dict = {}
@@ -231,7 +234,7 @@ def gguf_tokenizer_loader(path, temb_shape):
     return torch.ByteTensor(list(spm.SerializeToString()))
 
 def gguf_clip_loader(path):
-    sd, arch = gguf_sd_loader(path, return_arch=True)
+    sd, arch = gguf_sd_loader(path, return_arch=True, is_text_model=True)
     if arch in {"t5", "t5encoder"}:
         temb_key = "token_embd.weight"
         if temb_key in sd and sd[temb_key].shape == (256384, 4096):
@@ -242,10 +245,12 @@ def gguf_clip_loader(path):
             sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16)
         sd = sd_map_replace(sd, T5_SD_MAP)
     elif arch in {"llama"}:
+        # TODO: pass model_options["vocab_size"] to loader somehow
         temb_key = "token_embd.weight"
-        if temb_key in sd and sd[temb_key].shape != (128320, 4096):
-            # This still works. Raise error?
-            logging.warning("Warning! token_embd shape may be incorrect for llama 3 model!")
+        if temb_key in sd and sd[temb_key].shape[0] >= (64 * 1024):
+            # See note above for T5.
+            logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.")
+            sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16)
         sd = sd_map_replace(sd, LLAMA_SD_MAP)
         sd = llama_permute(sd, 32, 8) # L3
     else:

nodes.py

@@ -3,7 +3,10 @@ import torch
 import logging
 import collections
 
+import nodes
 import comfy.sd
+import comfy.lora
+import comfy.float
 import comfy.utils
 import comfy.model_patcher
 import comfy.model_management
@@ -36,17 +39,12 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
             return
         weight = comfy.utils.get_attr(self.model, key)
 
-        try:
-            from comfy.lora import calculate_weight
-        except Exception:
-            calculate_weight = self.calculate_weight
-
         patches = self.patches[key]
         if is_quantized(weight):
             out_weight = weight.to(device_to)
             patches = move_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device)
             # TODO: do we ever have legitimate duplicate patches? (i.e. patch on top of patched weight)
-            out_weight.patches = [(calculate_weight, patches, key)]
+            out_weight.patches = [(patches, key)]
         else:
             inplace_update = self.weight_inplace_update or inplace_update
             if key not in self.backup:
@@ -59,7 +57,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
             else:
                 temp_weight = weight.to(torch.float32, copy=True)
 
-            out_weight = calculate_weight(patches, temp_weight, key)
+            out_weight = comfy.lora.calculate_weight(patches, temp_weight, key)
             out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype)
 
         if inplace_update:
@@ -98,7 +96,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
                         if device == self.offload_device:
                             linked.append((n, m))
                             continue
-            if linked:
+            if linked and self.load_device != self.offload_device:
                 logging.info(f"Attempting to release mmap ({len(linked)})")
                 for n, m in linked:
                     # TODO: possible to OOM, find better way to detach
@@ -113,6 +111,8 @@ 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)
+        if src_cls != GGUFModelPatcher:
+            n.size = 0 # force recalc
         return n
 
 class UnetLoaderGGUF:
@@ -174,37 +174,14 @@ class UnetLoaderGGUFAdvanced(UnetLoaderGGUF):
         }
     TITLE = "Unet Loader (GGUF/Advanced)"
 
-# Mapping from common name to name used in comfy.sd.CLIPType enum
-CLIP_ENUM_MAP = {
-    "stable_diffusion": "STABLE_DIFFUSION",
-    "stable_cascade":   "STABLE_CASCADE",
-    "stable_audio":     "STABLE_AUDIO",
-    "sdxl":             "STABLE_DIFFUSION",
-    "sd3":              "SD3",
-    "flux":             "FLUX",
-    "mochi":            "MOCHI",
-    "ltxv":             "LTXV",
-    "hunyuan_video":    "HUNYUAN_VIDEO",
-    "pixart":           "PIXART",
-    "wan":              "WAN",
-}
-
-def get_clip_type(name):
-    enum_name = CLIP_ENUM_MAP.get(name, None)
-    if enum_name is None:
-        raise ValueError(f"Unknown CLIP model type {name}")
-    clip_type = getattr(comfy.sd.CLIPType, CLIP_ENUM_MAP[name], None)
-    if clip_type is None:
-        raise ValueError(f"Unsupported CLIP model type {name} (Update ComfyUI)")
-    return clip_type
-
 class CLIPLoaderGGUF:
     @classmethod
     def INPUT_TYPES(s):
+        base = nodes.CLIPLoader.INPUT_TYPES()
         return {
             "required": {
                 "clip_name": (s.get_filename_list(),),
-                "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "wan"],),
+                "type": base["required"]["type"],
             }
         }
 
@@ -247,17 +224,19 @@ class CLIPLoaderGGUF:
 
     def load_clip(self, clip_name, type="stable_diffusion"):
         clip_path = folder_paths.get_full_path("clip", clip_name)
-        return (self.load_patcher([clip_path], get_clip_type(type), self.load_data([clip_path])),)
+        clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
+        return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),)
 
 class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
     @classmethod
     def INPUT_TYPES(s):
+        base = nodes.DualCLIPLoader.INPUT_TYPES()
         file_options = (s.get_filename_list(), )
         return {
             "required": {
                 "clip_name1": file_options,
                 "clip_name2": file_options,
-                "type": (("sdxl", "sd3", "flux", "hunyuan_video"),),
+                "type": base["required"]["type"],
             }
         }
 
@@ -267,7 +246,8 @@ class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
         clip_path1 = folder_paths.get_full_path("clip", clip_name1)
         clip_path2 = folder_paths.get_full_path("clip", clip_name2)
         clip_paths = (clip_path1, clip_path2)
-        return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),)
+        clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
+        return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
 
 class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
     @classmethod
@@ -288,7 +268,8 @@ class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
         clip_path2 = folder_paths.get_full_path("clip", clip_name2)
         clip_path3 = folder_paths.get_full_path("clip", clip_name3)
         clip_paths = (clip_path1, clip_path2, clip_path3)
-        return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),)
+        clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
+        return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
 
 class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF):
     @classmethod
@@ -311,7 +292,8 @@ class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF):
         clip_path3 = folder_paths.get_full_path("clip", clip_name3)
         clip_path4 = folder_paths.get_full_path("clip", clip_name4)
         clip_paths = (clip_path1, clip_path2, clip_path3, clip_path4)
-        return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),)
+        clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
+        return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
 
 NODE_CLASS_MAPPINGS = {
     "UnetLoaderGGUF": UnetLoaderGGUF,

ops.py

@@ -4,6 +4,7 @@ import torch
 import logging
 
 import comfy.ops
+import comfy.lora
 import comfy.model_management
 from .dequant import dequantize_tensor, is_quantized
 
@@ -111,6 +112,9 @@ class GGMLLayer(torch.nn.Module):
         # NOTE: using modified load for linear due to not initializing on creation, see GGMLOps todo
         if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(self, torch.nn.Linear):
             return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        # Not strictly required, but fixes embedding shape mismatch. Threshold set in loader.py
+        if isinstance(self, torch.nn.Embedding) and self.weight.shape[0] >= (64 * 1024):
+            return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
         return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
 
     def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
@@ -166,7 +170,7 @@ class GGMLLayer(torch.nn.Module):
         # consolidate and load patches to GPU in async
         patch_list = []
         device = tensor.device
-        for function, patches, key in getattr(tensor, "patches", []):
+        for patches, key in getattr(tensor, "patches", []):
             patch_list += move_patch_to_device(patches, device)
 
         # dequantize tensor while patches load
@@ -177,13 +181,13 @@ class GGMLLayer(torch.nn.Module):
             weight = torch.Tensor(weight)
 
         # apply patches
-        if patch_list:
+        if len(patch_list) > 0:
             if self.patch_dtype is None:
-                weight = function(patch_list, weight, key)
+                weight = comfy.lora.calculate_weight(patch_list, weight, key)
             else:
                 # for testing, may degrade image quality
                 patch_dtype = dtype if self.patch_dtype == "target" else self.patch_dtype
-                weight = function(patch_list, weight, key, patch_dtype)
+                weight = comfy.lora.calculate_weight(patch_list, weight, key, patch_dtype)
         return weight
 
     @torch_compiler_disable()

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "comfyui-gguf"
+description = "GGUF Quantization support for native ComfyUI models."
+version = "2.0.0" # 2.0.0 = GitHub main, 1.X.X = ComfyUI Registry
+license = { file = "LICENSE" }
+dependencies = ["gguf>=0.13.0", "sentencepiece", "protobuf"]
+
+[project.urls]
+Repository = "https://github.com/city96/ComfyUI-GGUF"
+
+[tool.comfy]
+PublisherId = "city96"
+DisplayName = "ComfyUI-GGUF"
+Icon = ""

tools/convert.py

@@ -2,6 +2,7 @@
 import os
 import gguf
 import torch
+import logging
 import argparse
 from tqdm import tqdm
 from safetensors.torch import load_file, save_file
@@ -52,7 +53,9 @@ class ModelHiDream(ModelTemplate):
         )
     ]
     keys_hiprec = [
-        ".ff_i.gate.weight" # nn.parameter, can't load from BF16 ver
+        # nn.parameter, can't load from BF16 ver
+        ".ff_i.gate.weight",
+        "img_emb.emb_pos"
     ]
 
 class ModelHyVid(ModelTemplate):
@@ -283,14 +286,14 @@ def handle_tensors(writer, state_dict, model_arch, allow_fp32=False):
         tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{key}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
 
         writer.add_tensor(key, data, raw_dtype=data_qtype)
-    
+
     return quantized_tensors, invalid_tensors
 
 def convert_file(path, dst_path=None, interact=True, overwrite=False, allow_fp32=False):
     # load & run model detection logic
     state_dict = load_state_dict(path)
     model_arch = detect_arch(state_dict)
-    print(f"* Architecture detected from input: {model_arch.arch}")
+    logging.info(f"* Architecture detected from input: {model_arch.arch}")
 
     ftype_name, ftype_gguf = find_main_dtype(state_dict, allow_fp32=allow_fp32)
 
@@ -325,8 +328,8 @@ def convert_file(path, dst_path=None, interact=True, overwrite=False, allow_fp32
 
         invalid_tensors = {k:torch.from_numpy(v.copy()) for k,v in invalid_tensors.items()}
         save_file(invalid_tensors, fix_path)
-        print(f"\n### Warning! Fix file saved to '{fix_path}'")
-        print(f" you most likely need to run 'fix_5d_tensors.py' after quantization.")
+        logging.warning(f"\n### Warning! Fix file found at '{fix_path}'")
+        logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.")
     else:
         fix_path = None
 

tools/lcpp.patch

@@ -39,7 +39,7 @@ index b16c462f..6d1568f1 100644
      const int idx = gguf_find_tensor(ctx, name);
      if (idx < 0) {
 diff --git a/src/llama.cpp b/src/llama.cpp
-index 24e1f1f0..5e944837 100644
+index 24e1f1f0..39045ca5 100644
 --- a/src/llama.cpp
 +++ b/src/llama.cpp
 @@ -205,6 +205,15 @@ enum llm_arch {
@@ -255,7 +255,7 @@ index 24e1f1f0..5e944837 100644
      }
  
      size_t total_size_org = 0;
-@@ -18547,6 +18716,45 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
+@@ -18547,6 +18716,51 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
              ctx_outs[i_split] = gguf_init_empty();
          }
          gguf_add_tensor(ctx_outs[i_split], tensor);
@@ -297,11 +297,17 @@ index 24e1f1f0..5e944837 100644
 +                gguf_set_tensor_ndim(ctx_outs[i_split], tensor->name, n_dim);
 +                LLAMA_LOG_INFO("\n%s: Correcting shape for Wan: [key:%s]\n", __func__, tensor->name);
 +            }
++            // FLF2V model only
++            if (name == "img_emb.emb_pos") {
++                const int n_dim = 3;
++                gguf_set_tensor_ndim(ctx_outs[i_split], tensor->name, n_dim);
++                LLAMA_LOG_INFO("\n%s: Correcting shape for Wan FLF2V: [key:%s]\n", __func__, tensor->name);
++            }
 +        }
      }
  
      // Set split info if needed
-@@ -18647,6 +18855,92 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
+@@ -18647,6 +18861,92 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
          // do not quantize relative position bias (T5)
          quantize &= name.find("attn_rel_b.weight") == std::string::npos;
  
@@ -394,7 +400,7 @@ index 24e1f1f0..5e944837 100644
          enum ggml_type new_type;
          void * new_data;
          size_t new_size;
-@@ -18655,6 +18949,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
+@@ -18655,6 +18955,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
              new_type = default_type;
  
              // get more optimal quantization type based on the tensor shape, layer, etc.
@@ -404,7 +410,7 @@ index 24e1f1f0..5e944837 100644
              if (!params->pure && ggml_is_quantized(default_type)) {
                  new_type = llama_tensor_get_type(qs, new_type, tensor, ftype);
              }
-@@ -18664,6 +18961,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
+@@ -18664,6 +18967,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
              if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) {
                  new_type = params->output_tensor_type;
              }