Merge pull request #63 from blepping/feat_sd1_sdxl

Add support for SD1x and SDXL models

dequant.py

@@ -6,9 +6,9 @@ from tqdm import tqdm
 def dequantize_tensor(tensor, dtype=None, dequant_dtype=None):
     data = torch.tensor(tensor.data)
     qtype = getattr(tensor, "tensor_type", None)
-    oshape = getattr(tensor, "tensor_shape", tensor.data.shape)
+    oshape = getattr(tensor, "tensor_shape", data.shape)
 
-    if qtype in [None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16]:
+    if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
         return data.to(dtype)
     elif qtype in dequantize_functions:
         dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype
@@ -37,7 +37,7 @@ def dequantize(data, qtype, oshape, dtype=None):
 
 def to_uint32(x):
     # no uint32 :(
-    x = x.view(torch.uint8).to(torch.int32) 
+    x = x.view(torch.uint8).to(torch.int32)
     return (x[:, 0] | x[:, 1] << 8 | x[:, 2] << 16 | x[:, 3] << 24).unsqueeze(1)
 
 def split_block_dims(blocks, *args):

nodes.py

@@ -3,7 +3,6 @@ import torch
 import gguf
 import copy
 import logging
-import numpy as np
 
 import comfy.sd
 import comfy.utils
@@ -23,6 +22,16 @@ if "clip_gguf" not in folder_paths.folder_names_and_paths:
     orig = folder_paths.folder_names_and_paths.get("clip", [[], set()])
     folder_paths.folder_names_and_paths["clip_gguf"] = (orig[0], {".gguf"})
 
+def gguf_sd_loader_get_orig_shape(reader, tensor_name):
+    field_key = f"comfy.gguf.orig_shape.{tensor_name}"
+    field = reader.get_field(field_key)
+    if field is None:
+        return None
+    # Has original shape metadata, so we try to decode it.
+    if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32:
+        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 gguf_sd_loader(path):
     """
     Read state dict as fake tensors
@@ -31,12 +40,17 @@ def gguf_sd_loader(path):
     sd = {}
     dt = {}
     for tensor in reader.tensors:
-        sd[str(tensor.name)] = GGMLTensor(
-            torch.from_numpy(tensor.data), # mmap
+        tensor_name = str(tensor.name)
+        torch_tensor = torch.from_numpy(tensor.data) # mmap
+        shape = gguf_sd_loader_get_orig_shape(reader, tensor_name)
+        if shape is None:
+            shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
+        elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
+            torch_tensor = torch_tensor.view(*shape)
+        sd[tensor_name] = GGMLTensor(
+            torch_tensor,
             tensor_type = tensor.tensor_type,
-            tensor_shape = torch.Size(
-                np.flip(list(tensor.shape))
-            )
+            tensor_shape = shape
         )
         dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1
 
@@ -97,21 +111,21 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
 
         try:
             from comfy.lora import calculate_weight
-        except:
+        except Exception:
             calculate_weight = self.calculate_weight
 
         patches = self.patches[key]
         qtype = getattr(weight, "tensor_type", None)
-        if qtype not in [None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16]:
+        if qtype not in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
             if device_to is not None:
                 out_weight = weight.to(device_to, copy=True)
             else:
                 out_weight = weight.clone()
-            
+
             if self.patch_on_device:
                 patches = move_patch_to_cuda(patches, self.load_device)
             out_weight.patches.append((calculate_weight, patches, key))
-            
+
         else:
             if device_to is not None:
                 temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
@@ -253,10 +267,10 @@ class CLIPLoaderGGUF:
 
         # for some reason this is just missing in some SAI checkpoints
         if getattr(clip.cond_stage_model, "clip_l", None) is not None:
-            if getattr(clip.cond_stage_model.clip_l.transformer.text_projection.weight, "tensor_shape", None) == None:
+            if getattr(clip.cond_stage_model.clip_l.transformer.text_projection.weight, "tensor_shape", None) is None:
                 clip.cond_stage_model.clip_l.transformer.text_projection = comfy.ops.manual_cast.Linear(768, 768)
         if getattr(clip.cond_stage_model, "clip_g", None) is not None:
-            if getattr(clip.cond_stage_model.clip_g.transformer.text_projection.weight, "tensor_shape", None) == None:
+            if getattr(clip.cond_stage_model.clip_g.transformer.text_projection.weight, "tensor_shape", None) is None:
                 clip.cond_stage_model.clip_g.transformer.text_projection = comfy.ops.manual_cast.Linear(1280, 1280)
 
         return clip

ops.py

@@ -1,5 +1,6 @@
 # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
 import torch
+from functools import partial
 
 import comfy.ops
 from .dequant import dequantize_tensor
@@ -35,7 +36,7 @@ class GGMLTensor(torch.Tensor):
         try:
             return super().copy_(*args, **kwargs)
         except Exception as e:
-            print(f"ignoring 'copy_' on tensor")
+            print(f"ignoring 'copy_' on tensor: {e}")
 
     def __deepcopy__(self, *args, **kwargs):
         # Intel Arc fix, ref#50
@@ -48,13 +49,14 @@ class GGMLTensor(torch.Tensor):
     @property
     def shape(self):
         if not hasattr(self, "tensor_shape"):
-            self.tensor_shape = self.size() 
+            self.tensor_shape = self.size()
         return self.tensor_shape
 
 class GGMLLayer(torch.nn.Module):
     """
     This (should) be responsible for de-quantizing on the fly
     """
+    comfy_cast_weights = True
     dequant_dtype = None
     patch_dtype = None
 
@@ -63,6 +65,24 @@ class GGMLLayer(torch.nn.Module):
         self.weight = GGMLTensor(1, tensor_type=None, tensor_shape=None)
         self.bias = None
 
+    def _forward_operation(self, x, weight, bias):
+        raise NotImplementedError
+
+    def forward(self, x):
+        # lowvram hack
+        device = None
+        if self.weight.device != x.device:
+            device = self.weight.device
+            self.to(x.device)
+
+        weight, bias = self.get_weights(x.dtype)
+        x = self._forward_operation(x, weight, bias)
+        del weight, bias
+
+        if device:
+            self.to(device)
+        return x
+
     def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
         for k,v in state_dict.items():
             if k[len(prefix):] == "weight":
@@ -130,30 +150,20 @@ class GGMLLayer(torch.nn.Module):
         bias = self.get_weight(self.bias, dtype)
         return (weight, bias)
 
+
 class GGMLOps(comfy.ops.manual_cast):
     """
     Dequantize weights on the fly before doing the compute
     """
     class Linear(GGMLLayer):
-        comfy_cast_weights = True
+        _forward_operation = staticmethod(torch.nn.functional.linear)
 
+    class Conv2d(GGMLLayer):
         def __init__(self, *args, device=None, dtype=None, **kwargs):
-            super().__init__(device=device, dtype=dtype)
+            super().__init__()
+            _ = kwargs.pop("kernel_size", None)
+            self._forward_operation = partial(staticmethod(torch.nn.functional.conv2d), **kwargs)
 
-        def forward(self, x):
-            # lowvram hack
-            device = None
-            if self.weight.device != x.device:
-                device = self.weight.device
-                self.to(x.device)
-
-            weight, bias = self.get_weights(x.dtype)
-            x = torch.nn.functional.linear(x, weight, bias)
-            del weight, bias
-
-            if device:
-                self.to(device)
-            return x
 
 def move_patch_to_cuda(item, device):
     if isinstance(item, torch.Tensor):

tools/convert.py

@@ -1,23 +1,55 @@
 # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
 import os
-import sys
 import torch
-import numpy as np
 import gguf # This needs to be the llama.cpp one specifically!
 import argparse
 from tqdm import tqdm
 
 from safetensors.torch import load_file
 
+QUANTIZATION_THRESHOLD = 1024
+REARRANGE_THRESHOLD = 512
+MAX_TENSOR_NAME_LENGTH = 127
+
+# Tuple of arch_name, match_lists.
+# Each item in match_lists is a tuple of keys that must match.
+# All keys in a match_lists item must exist for the architecture to match.
+# The architectures are checked in order and the first successful match terminates the search.
+MODEL_DETECTION = (
+    ("flux", (
+        ("transformer_blocks.0.attn.norm_added_k.weight",),
+        ("double_blocks.0.img_attn.proj.weight",),
+    )),
+    ("sd3", (
+        ("transformer_blocks.0.attn.add_q_proj.weight",),
+    )),
+    ("sdxl", (
+        ("down_blocks.0.downsamplers.0.conv.weight", "add_embedding.linear_1.weight",),
+        (
+            "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight",
+            "output_blocks.2.2.conv.weight", "output_blocks.5.2.conv.weight",
+        ), # Non-diffusers
+        ("label_emb.0.0.weight",),
+    )),
+    ("sd1", (
+        ("down_blocks.0.downsamplers.0.conv.weight",),
+        (
+            "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", "input_blocks.9.0.op.weight",
+            "output_blocks.2.1.conv.weight", "output_blocks.5.2.conv.weight", "output_blocks.8.2.conv.weight"
+        ), # Non-diffusers
+    )),
+)
+
+
 def parse_args():
     parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET")
     parser.add_argument("--src", required=True, help="Source model ckpt file.")
-    parser.add_argument("--dst", help="Output  unet gguf file.")
+    parser.add_argument("--dst", help="Output unet gguf file.")
     args = parser.parse_args()
 
     if not os.path.isfile(args.src):
         parser.error("No input provided!")
-    
+
     return args
 
 def load_state_dict(path):
@@ -26,7 +58,7 @@ def load_state_dict(path):
         state_dict = state_dict.get("model", state_dict)
     else:
         state_dict = load_file(path)
-    
+
     # only keep unet with no prefix!
     sd = {}
     has_prefix = any(["model.diffusion_model." in x for x in state_dict.keys()])
@@ -39,26 +71,21 @@ def load_state_dict(path):
 
     return sd
 
+def detect_arch(state_dict):
+    for arch, match_lists in MODEL_DETECTION:
+        for match_list in match_lists:
+            if all(key in state_dict for key in match_list):
+                return arch
+    breakpoint()
+    raise ValueError("Unknown model architecture!")
+
+
 def load_model(path):
     state_dict = load_state_dict(path)
-
-    # from ComfyUI model detection
-    if "transformer_blocks.0.attn.norm_added_k.weight" in state_dict:
-        arch = "flux"
-        raise ValueError(f"The Diffusers UNET can not be used for this!")
-    elif "double_blocks.0.img_attn.proj.weight" in state_dict:
-        arch = "flux" # mmdit ...?
-    elif "transformer_blocks.0.attn.add_q_proj.weight" in state_dict:
-        arch = "sd3"
-    elif "down_blocks.0.downsamplers.0.conv.weight" in state_dict:
-        if "add_embedding.linear_1.weight" in state_dict:
-            arch = "sdxl"
-        else:
-            arch = "sd1" 
-    else:
-        breakpoint()
-        raise ValueError(f"Unknown model architecture!")
-
+    arch = detect_arch(state_dict)
+    print(f"* Architecture detected from input: {arch}")
+    if arch == "flux" and "transformer_blocks.0.attn.norm_added_k.weight" in state_dict:
+        raise ValueError("The Diffusers UNET can not be used for this!")
     writer = gguf.GGUFWriter(path=None, arch=arch)
     return (writer, state_dict)
 
@@ -66,7 +93,17 @@ def handle_tensors(args, writer, state_dict):
     # TODO list:
     # - do something about this being awful and hacky
 
-    max_name_len = max([len(s) for s in state_dict.keys()]) + 4
+    name_lengths = tuple(sorted(
+        ((key, len(key)) for key in state_dict.keys()),
+        key=lambda item: item[1],
+        reverse=True,
+    ))
+    if not name_lengths:
+        return
+    max_name_len = name_lengths[0][1]
+    if max_name_len > MAX_TENSOR_NAME_LENGTH:
+        bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH)
+        raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}")
     for key, data in tqdm(state_dict.items()):
         old_dtype = data.dtype
 
@@ -88,7 +125,7 @@ def handle_tensors(args, writer, state_dict):
             n_params *= dim_size
 
         # keys to keep as max precision
-        blacklist = [
+        blacklist = {
             "time_embedding.",
             "add_embedding.",
             "time_in.",
@@ -97,36 +134,33 @@ def handle_tensors(args, writer, state_dict):
             "img_in.",
             "guidance_in.",
             "final_layer.",
-        ]
+        }
 
-        if any([x in key for x in blacklist]) and ".weight" in key:
-            data_qtype = gguf.GGMLQuantizationType.F32
+        if old_dtype in (torch.float32, torch.bfloat16):
+            if n_dims == 1:
+                # one-dimensional tensors should be kept in F32
+                # also speeds up inference due to not dequantizing
+                data_qtype = gguf.GGMLQuantizationType.F32
 
-        if n_dims == 1: 
-            # one-dimensional tensors should be kept in F32
-            # also speeds up inference due to not dequantizing
-            data_qtype = gguf.GGMLQuantizationType.F32
-        
-        elif n_params <= 1024:
-            # very small tensors
-            data_qtype = gguf.GGMLQuantizationType.F32
-        
-        elif n_dims == 4:
-            if min(data.shape[:2]) == 4: # output tensor
-                data_qtype = gguf.GGMLQuantizationType.F16
-            elif data_shape[-1] == 3: # 3x3 kernel
-                data_qtype = gguf.GGMLQuantizationType.F16
-            elif data_shape[-1] == 1: # 1x1 kernel
-                #data = np.squeeze(data) # don't do this
-                data_qtype = gguf.GGMLQuantizationType.F16
+            elif n_params <= QUANTIZATION_THRESHOLD:
+                # very small tensors
+                data_qtype = gguf.GGMLQuantizationType.F32
+
+            elif ".weight" in key and any(x in key for x in blacklist):
+                data_qtype = gguf.GGMLQuantizationType.F32
+
+        if (    n_dims > 1                              # Skip one-dimensional tensors
+            and n_params >= REARRANGE_THRESHOLD         # Only rearrange tensors meeting the size requirement
+            and (n_params / 256).is_integer()           # Rearranging only makes sense if total elements is divisible by 256
+            and not (data.shape[-1] / 256).is_integer() # Only need to rearrange if the last dimension is not divisible by 256
+        ):
+            orig_shape = data.shape
+            data = data.reshape(n_params // 256, 256)
+            writer.add_array(f"comfy.gguf.orig_shape.{key}", tuple(int(dim) for dim in orig_shape))
 
         try:
             data = gguf.quants.quantize(data, data_qtype)
-        except gguf.QuantError as e:
-            tqdm.write(f"falling back to F16: {e}")
-            data_qtype = gguf.GGMLQuantizationType.F16
-            data = gguf.quants.quantize(data, data_qtype)
-        except AttributeError as e:
+        except (AttributeError, gguf.QuantError) as e:
             tqdm.write(f"falling back to F16: {e}")
             data_qtype = gguf.GGMLQuantizationType.F16
             data = gguf.quants.quantize(data, data_qtype)
@@ -134,7 +168,7 @@ def handle_tensors(args, writer, state_dict):
         new_name = key # do we need to rename?
 
         shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}"
-        tqdm.write(f"{f'%-{max_name_len}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
+        tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
 
         writer.add_tensor(new_name, data, raw_dtype=data_qtype)
 
@@ -150,7 +184,7 @@ if __name__ == "__main__":
     else:
         out_path = f"{os.path.splitext(path)[0]}-F16.gguf"
         writer.add_file_type(gguf.LlamaFileType.MOSTLY_F16)
-    
+
     out_path = args.dst or out_path
     if os.path.isfile(out_path):
         input("Output exists enter to continue or ctrl+c to abort!")

tools/lcpp.patch

@@ -24,7 +24,7 @@ index 5ab65ea9..35580d9d 100644
 +    LLM_ARCH_SDXL,
      LLM_ARCH_UNKNOWN,
  };
- 
+
 @@ -259,6 +262,9 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
      { LLM_ARCH_JAIS,            "jais"         },
      { LLM_ARCH_NEMOTRON,        "nemotron"     },
@@ -34,7 +34,7 @@ index 5ab65ea9..35580d9d 100644
 +    { LLM_ARCH_SDXL,            "sdxl"         },
      { LLM_ARCH_UNKNOWN,         "(unknown)"    },
  };
- 
+
 @@ -1337,6 +1343,9 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
              { LLM_TENSOR_FFN_UP,          "blk.%d.ffn_up" },
          },
@@ -48,7 +48,7 @@ index 5ab65ea9..35580d9d 100644
 @@ -4629,6 +4638,12 @@ static void llm_load_hparams(
      // get general kv
      ml.get_key(LLM_KV_GENERAL_NAME, model.name, false);
- 
+
 +    // Disable LLM metadata for image models
 +    if (model.arch == LLM_ARCH_FLUX || model.arch == LLM_ARCH_SD1 || model.arch == LLM_ARCH_SDXL) {
 +        model.ftype = ml.ftype;
@@ -57,11 +57,11 @@ index 5ab65ea9..35580d9d 100644
 +
      // get hparams kv
      ml.get_key(LLM_KV_VOCAB_SIZE, hparams.n_vocab, false) || ml.get_arr_n(LLM_KV_TOKENIZER_LIST, hparams.n_vocab);
- 
+
 @@ -15827,11 +15842,162 @@ static void llama_tensor_dequantize_internal(
      workers.clear();
  }
- 
+
 +static ggml_type img_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) {
 +    // Special function for quantizing image model tensors
 +    const std::string name = ggml_get_name(tensor);
@@ -214,10 +214,10 @@ index 5ab65ea9..35580d9d 100644
 +
  static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) {
      const std::string name = ggml_get_name(tensor);
- 
+
      // TODO: avoid hardcoded tensor names - use the TN_* constants
      const llm_arch arch = qs.model.arch;
 +    if (arch == LLM_ARCH_FLUX || arch == LLM_ARCH_SD1 || arch == LLM_ARCH_SDXL) { return img_tensor_get_type(qs, new_type, tensor, ftype); };
      const auto       tn = LLM_TN(arch);
- 
+
      auto use_more_bits = [](int i_layer, int n_layers) -> bool {