diff --git a/dequant.py b/dequant.py index 05c3868..1750cd0 100644 --- a/dequant.py +++ b/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): diff --git a/nodes.py b/nodes.py index 46d5a49..0e61662 100644 --- a/nodes.py +++ b/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 diff --git a/ops.py b/ops.py index 7094304..2d67e23 100644 --- a/ops.py +++ b/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): diff --git a/tools/convert.py b/tools/convert.py index b6b5828..a49ef1c 100644 --- a/tools/convert.py +++ b/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!") diff --git a/tools/lcpp.patch b/tools/lcpp.patch index fd64ef7..e7bfbff 100644 --- a/tools/lcpp.patch +++ b/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_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_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 {