From f3cf802dd6ce3c436eab5a376c283fb9b42e59a7 Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 25 Aug 2024 07:53:19 -0600 Subject: [PATCH] Cleanups, refactoring --- nodes.py | 18 +++--- ops.py | 71 ++++++++++------------ tools/convert.py | 153 ++++++++++++++++++++++++++++------------------- 3 files changed, 132 insertions(+), 110 deletions(-) diff --git a/nodes.py b/nodes.py index b4161fe..af15bcf 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 @@ -40,10 +39,13 @@ def gguf_sd_loader(path): for tensor in reader.tensors: tensor_name = str(tensor.name) shape = tensor_fixup.get(tensor_name) + torch_tensor = torch.from_numpy(tensor.data) # mmap 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.from_numpy(tensor.data), # mmap + torch_tensor, tensor_type = tensor.tensor_type, tensor_shape = shape ) @@ -106,21 +108,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) @@ -262,10 +264,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 6a8c400..1f9f666 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 @@ -14,7 +15,7 @@ class GGMLTensor(torch.Tensor): self.tensor_shape = tensor_shape self.patches = patches - def __new__(cls, *args, tensor_type, tensor_shape, tensor_data=None, patches=[], **kwargs): + def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): return super().__new__(cls, *args, **kwargs) def to(self, *args, **kwargs): @@ -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 @@ -130,53 +131,43 @@ class GGMLLayer(torch.nn.Module): bias = self.get_weight(self.bias, dtype) return (weight, bias) +class GGMLOpsLayer(GGMLLayer): + comfy_cast_weights = True + + def __init__(self, *args, device=None, dtype=None, **kwargs): + super().__init__(device=device, dtype=dtype) + + 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 + class GGMLOps(comfy.ops.manual_cast): """ Dequantize weights on the fly before doing the compute """ - class Linear(GGMLLayer): - comfy_cast_weights = True - - def __init__(self, *args, device=None, dtype=None, **kwargs): - super().__init__(device=device, dtype=dtype) - - 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 - - class Conv2d(GGMLLayer): - comfy_cast_weights = True + class Linear(GGMLOpsLayer): + _forward_operation = staticmethod(torch.nn.functional.linear) + class Conv2d(GGMLOpsLayer): def __init__(self, *args, device=None, dtype=None, **kwargs): super().__init__(device=device, dtype=dtype) _ = kwargs.pop("kernel_size", None) - self.kwargs=kwargs + 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.conv2d(x, weight, bias, **self.kwargs) - 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 df7f523..44aabb9 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -1,24 +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 = 63 + +# 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("--force-arch", type=str, choices=("flux", "sd1", "sdxl"), default=None, help="Advanced option, allows forcing the model architecture. May be necessary for non-Flux models.") + 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): @@ -27,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()]) @@ -40,27 +71,21 @@ def load_state_dict(path): return sd -def load_model(path, force_arch=None): - state_dict = load_state_dict(path) - # from ComfyUI model detection - if force_arch is not None: - arch = force_arch - elif "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!") +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) + 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) @@ -68,9 +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 - if (max_name_len - 4) >= 64: - raise ValueError("Can only handle tensor names up to 63 characters!") + 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}") tensor_fixup = {} for key, data in tqdm(state_dict.items()): old_dtype = data.dtype @@ -93,7 +126,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.", @@ -102,41 +135,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 - if data_qtype != gguf.GGMLQuantizationType.F32 and n_dims >= 2 and n_params >= 256*256 and (n_params / 256).is_integer() and data.shape[0] >= 256 and data.shape[1] >= 256: - orig_shape=data.shape - orig_n_params=n_params - data = data.reshape(256+((n_params - 256*256) // 256),256) + 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) tensor_fixup[key] = 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) @@ -144,16 +169,20 @@ 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) + + if not tensor_fixup: + return + fixup_str = "\n".join(" ".join((k,)+tuple(str(dimval) for dimval in v)) for k,v in tensor_fixup.items()) writer.add_string("comfy.gguf.tensor_fixup", fixup_str) if __name__ == "__main__": args = parse_args() path = args.src - writer, state_dict = load_model(path, force_arch=args.force_arch) + writer, state_dict = load_model(path) writer.add_quantization_version(gguf.GGML_QUANT_VERSION) if next(iter(state_dict.values())).dtype == torch.bfloat16: @@ -162,7 +191,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!")