From cea1b47ab350ec537a94cf47efa33a6a13e5081c Mon Sep 17 00:00:00 2001 From: blepping Date: Fri, 23 Aug 2024 11:59:21 -0600 Subject: [PATCH 1/7] Add support for SD1x and SDXL models --- nodes.py | 4 +--- tools/convert.py | 12 +++++++----- tools/lcpp.patch | 32 ++++++++++++++++++-------------- 3 files changed, 26 insertions(+), 22 deletions(-) diff --git a/nodes.py b/nodes.py index 04e69ea..628f66c 100644 --- a/nodes.py +++ b/nodes.py @@ -34,9 +34,7 @@ def gguf_sd_loader(path): sd[str(tensor.name)] = GGMLTensor( torch.from_numpy(tensor.data), # mmap tensor_type = tensor.tensor_type, - tensor_shape = torch.Size( - np.flip(list(tensor.shape)) - ) + tensor_shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) ) dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1 diff --git a/tools/convert.py b/tools/convert.py index b6b5828..178b664 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -13,6 +13,7 @@ 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.") args = parser.parse_args() if not os.path.isfile(args.src): @@ -39,11 +40,12 @@ def load_state_dict(path): return sd -def load_model(path): +def load_model(path, force_arch=None): state_dict = load_state_dict(path) - # from ComfyUI model detection - if "transformer_blocks.0.attn.norm_added_k.weight" in state_dict: + 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: @@ -54,7 +56,7 @@ def load_model(path): if "add_embedding.linear_1.weight" in state_dict: arch = "sdxl" else: - arch = "sd1" + arch = "sd1" else: breakpoint() raise ValueError(f"Unknown model architecture!") @@ -141,7 +143,7 @@ def handle_tensors(args, writer, state_dict): if __name__ == "__main__": args = parse_args() path = args.src - writer, state_dict = load_model(path) + writer, state_dict = load_model(path, force_arch=args.force_arch) writer.add_quantization_version(gguf.GGML_QUANT_VERSION) if next(iter(state_dict.values())).dtype == torch.bfloat16: diff --git a/tools/lcpp.patch b/tools/lcpp.patch index 2c12ffe..ac1d59c 100644 --- a/tools/lcpp.patch +++ b/tools/lcpp.patch @@ -1,40 +1,43 @@ diff --git a/src/llama.cpp b/src/llama.cpp -index 5ab65ea9..44fe3fe6 100644 +index 5ab65ea9..ef63e28f 100644 --- a/src/llama.cpp +++ b/src/llama.cpp -@@ -212,6 +212,8 @@ enum llm_arch { +@@ -212,6 +212,9 @@ enum llm_arch { LLM_ARCH_JAIS, LLM_ARCH_NEMOTRON, LLM_ARCH_EXAONE, + LLM_ARCH_FLUX, + LLM_ARCH_SD1, ++ LLM_ARCH_SDXL, LLM_ARCH_UNKNOWN, }; -@@ -259,6 +261,8 @@ static const std::map LLM_ARCH_NAMES = { +@@ -259,6 +262,9 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_JAIS, "jais" }, { LLM_ARCH_NEMOTRON, "nemotron" }, { LLM_ARCH_EXAONE, "exaone" }, + { LLM_ARCH_FLUX, "flux" }, + { LLM_ARCH_SD1, "sd1" }, ++ { LLM_ARCH_SDXL, "sdxl" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; -@@ -1337,6 +1341,8 @@ static const std::map> LLM_TENSOR_NA +@@ -1337,6 +1343,9 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" }, }, }, + { LLM_ARCH_FLUX, { {} }, }, + { LLM_ARCH_SD1, { {} }, }, ++ { LLM_ARCH_SDXL, { {} }, }, { LLM_ARCH_UNKNOWN, { -@@ -4629,6 +4635,13 @@ static void llm_load_hparams( +@@ -4629,6 +4638,13 @@ 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) { ++ if (model.arch == LLM_ARCH_FLUX || model.arch == LLM_ARCH_SD1 || model.arch == LLM_ARCH_SDXL) { + model.ftype = ml.ftype; + hparams.rope_type = llama_rope_type(&model); + return; @@ -43,7 +46,7 @@ index 5ab65ea9..44fe3fe6 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); -@@ -15854,26 +15867,40 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n +@@ -15854,26 +15870,40 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n return std::make_pair(i_layer, n_layer); }; @@ -102,7 +105,7 @@ index 5ab65ea9..44fe3fe6 100644 if (qs.params->token_embedding_type < GGML_TYPE_COUNT) { new_type = qs.params->token_embedding_type; } else { -@@ -15891,10 +15918,13 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n +@@ -15891,10 +15921,13 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n new_type == GGML_TYPE_Q4_0_8_8) { new_type = GGML_TYPE_Q4_0; } @@ -117,7 +120,7 @@ index 5ab65ea9..44fe3fe6 100644 if (qs.model.hparams.n_gqa() >= 4 || qs.model.hparams.n_expert >= 4) new_type = GGML_TYPE_Q4_K; else new_type = ftype == LLAMA_FTYPE_MOSTLY_IQ2_S || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M ? GGML_TYPE_IQ3_S : GGML_TYPE_Q2_K; ++qs.i_attention_wv; -@@ -15916,7 +15946,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n +@@ -15916,7 +15949,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n else if (ftype == LLAMA_FTYPE_MOSTLY_IQ2_S || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M) new_type = GGML_TYPE_IQ3_S; } } @@ -126,7 +129,7 @@ index 5ab65ea9..44fe3fe6 100644 if (ftype == LLAMA_FTYPE_MOSTLY_Q2_K) { new_type = qs.model.hparams.n_gqa() >= 4 ? GGML_TYPE_Q4_K : GGML_TYPE_Q3_K; } -@@ -15954,7 +15984,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n +@@ -15954,7 +15987,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n new_type = GGML_TYPE_Q8_0; } ++qs.i_attention_wv; @@ -135,7 +138,7 @@ index 5ab65ea9..44fe3fe6 100644 if (qs.model.hparams.n_expert == 8) { // for the 8-expert model, bumping this to Q8_0 trades just ~128MB // TODO: explore better strategies -@@ -15966,7 +15996,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n +@@ -15966,7 +15999,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS) { new_type = GGML_TYPE_IQ2_S; } @@ -144,7 +147,7 @@ index 5ab65ea9..44fe3fe6 100644 if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XS) { new_type = GGML_TYPE_IQ3_XXS; } -@@ -16038,7 +16068,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n +@@ -16038,7 +16071,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L) new_type = GGML_TYPE_Q4_K; } } @@ -153,7 +156,7 @@ index 5ab65ea9..44fe3fe6 100644 if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M || ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L || ftype == LLAMA_FTYPE_MOSTLY_IQ3_M) { new_type = GGML_TYPE_Q4_K; } -@@ -16107,6 +16137,8 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n +@@ -16107,6 +16140,8 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n } LLAMA_LOG_WARN(" - using fallback quantization %s\n", ggml_type_name(new_type)); ++qs.n_fallback; @@ -162,12 +165,13 @@ index 5ab65ea9..44fe3fe6 100644 } return new_type; -@@ -17432,6 +17464,8 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) { +@@ -17432,6 +17467,9 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) { case LLM_ARCH_T5: case LLM_ARCH_T5ENCODER: case LLM_ARCH_JAIS: + case LLM_ARCH_FLUX: + case LLM_ARCH_SD1: ++ case LLM_ARCH_SDXL: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values From 489aaecf7a3b8f25c806af341b3911bd3a3e290f Mon Sep 17 00:00:00 2001 From: blepping Date: Fri, 23 Aug 2024 12:37:31 -0600 Subject: [PATCH 2/7] Show the correct type when using fallback quantization --- tools/lcpp.patch | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tools/lcpp.patch b/tools/lcpp.patch index ac1d59c..01fe0b8 100644 --- a/tools/lcpp.patch +++ b/tools/lcpp.patch @@ -1,5 +1,5 @@ diff --git a/src/llama.cpp b/src/llama.cpp -index 5ab65ea9..ef63e28f 100644 +index 5ab65ea9..6d684fef 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -212,6 +212,9 @@ enum llm_arch { @@ -156,15 +156,15 @@ index 5ab65ea9..ef63e28f 100644 if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M || ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L || ftype == LLAMA_FTYPE_MOSTLY_IQ3_M) { new_type = GGML_TYPE_Q4_K; } -@@ -16107,6 +16140,8 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n +@@ -16105,6 +16138,8 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n + case GGML_TYPE_Q6_K: new_type = GGML_TYPE_Q8_0; break; + default: throw std::runtime_error("\nUnsupported tensor size encountered\n"); } - LLAMA_LOG_WARN(" - using fallback quantization %s\n", ggml_type_name(new_type)); - ++qs.n_fallback; + // Force FP16 fallback - needed due to Conv2D + new_type = GGML_TYPE_F16; + LLAMA_LOG_WARN(" - using fallback quantization %s\n", ggml_type_name(new_type)); + ++qs.n_fallback; } - - return new_type; @@ -17432,6 +17467,9 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) { case LLM_ARCH_T5: case LLM_ARCH_T5ENCODER: From c93e622a19d4554f0d81bbf74ed0c0c9016c2ec7 Mon Sep 17 00:00:00 2001 From: blepping Date: Fri, 23 Aug 2024 20:19:48 -0600 Subject: [PATCH 3/7] SD1/SDXL fixes --- dequant.py | 12 ++++-------- nodes.py | 15 +++++++++++++-- ops.py | 25 ++++++++++++++++++++++++- tools/convert.py | 10 ++++++++++ 4 files changed, 51 insertions(+), 11 deletions(-) diff --git a/dequant.py b/dequant.py index f23aca6..949e5d6 100644 --- a/dequant.py +++ b/dequant.py @@ -7,15 +7,11 @@ from tqdm import tqdm def dequantize_tensor(tensor, dtype=torch.float16): 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 == None: - return data.to(dtype) - elif qtype == gguf.GGMLQuantizationType.F32: - return data.to(dtype) - elif qtype == gguf.GGMLQuantizationType.F16: - return data.to(dtype) - elif qtype in dequantize_functions: + if qtype in {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + return data.to(dtype).view(oshape) + if qtype in dequantize_functions: # dequantize in fp16 then convert instead of keeping FP32 out = dequantize(data, qtype, oshape, dtype=None) return out.to(dtype) if out.dtype != dtype else out # why is .to() not a no-op? diff --git a/nodes.py b/nodes.py index 628f66c..97686bd 100644 --- a/nodes.py +++ b/nodes.py @@ -28,13 +28,24 @@ def gguf_sd_loader(path): Read state dict as fake tensors """ reader = gguf.GGUFReader(path) + fixup_field = reader.get_field("comfy.gguf.tensor_fixup") + tensor_fixup = {} + if fixup_field is not None: + tensor_fixup_str=str(fixup_field.parts[fixup_field.data[-1]], encoding="utf-8") + for line in str(tensor_fixup_str).split("\n"): + key, *fixup_items = line.split(" ") + tensor_fixup[key] = torch.Size(tuple(int(i) for i in fixup_items)) sd = {} dt = {} for tensor in reader.tensors: - sd[str(tensor.name)] = GGMLTensor( + tensor_name = str(tensor.name) + shape = tensor_fixup.get(tensor_name) + if shape is None: + shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + sd[tensor_name] = GGMLTensor( torch.from_numpy(tensor.data), # mmap tensor_type = tensor.tensor_type, - tensor_shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + tensor_shape = shape ) dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1 diff --git a/ops.py b/ops.py index 6bc1e2a..fac78d1 100644 --- a/ops.py +++ b/ops.py @@ -16,7 +16,7 @@ class GGMLTensor(torch.Tensor): self.tensor_shape = tensor_shape self.patches = patches - def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): + def __new__(cls, *args, tensor_type, tensor_shape, tensor_data=None, patches=[], **kwargs): return super().__new__(cls, *args, **kwargs) def to(self, *args, **kwargs): @@ -150,3 +150,26 @@ class GGMLOps(comfy.ops.manual_cast): if device: self.to(device) return x + + class Conv2d(GGMLLayer): + comfy_cast_weights = True + + def __init__(self, *args, device=None, dtype=None, **kwargs): + super().__init__(device=device, dtype=dtype) + _ = kwargs.pop("kernel_size", None) + self.kwargs=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 diff --git a/tools/convert.py b/tools/convert.py index 178b664..df7f523 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -69,6 +69,9 @@ def handle_tensors(args, writer, state_dict): # - 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!") + tensor_fixup = {} for key, data in tqdm(state_dict.items()): old_dtype = data.dtype @@ -122,6 +125,11 @@ def handle_tensors(args, writer, state_dict): #data = np.squeeze(data) # don't do this data_qtype = gguf.GGMLQuantizationType.F16 + 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) + tensor_fixup[key] = orig_shape try: data = gguf.quants.quantize(data, data_qtype) except gguf.QuantError as e: @@ -139,6 +147,8 @@ def handle_tensors(args, writer, state_dict): tqdm.write(f"{f'%-{max_name_len}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") writer.add_tensor(new_name, data, raw_dtype=data_qtype) + 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() From f3cf802dd6ce3c436eab5a376c283fb9b42e59a7 Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 25 Aug 2024 07:53:19 -0600 Subject: [PATCH 4/7] 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!") From 7413de86e6ea4b63908625143525f10ca1a468dc Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 25 Aug 2024 10:35:39 -0600 Subject: [PATCH 5/7] Revert tensor construction change in dequantize_tensor Eliminate GGMLOpsLayer, include functionality in GGMLLayer --- dequant.py | 4 ++-- ops.py | 48 ++++++++++++++++++++++-------------------------- 2 files changed, 24 insertions(+), 28 deletions(-) diff --git a/dequant.py b/dequant.py index cd9a8b7..1750cd0 100644 --- a/dequant.py +++ b/dequant.py @@ -4,11 +4,11 @@ import torch from tqdm import tqdm def dequantize_tensor(tensor, dtype=None, dequant_dtype=None): - data = tensor.data.clone().detach() + data = torch.tensor(tensor.data) qtype = getattr(tensor, "tensor_type", None) 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 diff --git a/ops.py b/ops.py index 1f9f666..2d67e23 100644 --- a/ops.py +++ b/ops.py @@ -56,6 +56,7 @@ 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 @@ -64,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": @@ -131,40 +150,17 @@ 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(GGMLOpsLayer): + class Linear(GGMLLayer): _forward_operation = staticmethod(torch.nn.functional.linear) - class Conv2d(GGMLOpsLayer): + 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) From 0aec56ec4c5c7a32f8fb465503d26d3c78bed685 Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 25 Aug 2024 12:00:28 -0600 Subject: [PATCH 6/7] Split orig shape metadata into individual fields --- nodes.py | 22 ++++++++++++++-------- tools/convert.py | 14 +++++--------- tools/lcpp.patch | 13 +++++++++++++ 3 files changed, 32 insertions(+), 17 deletions(-) diff --git a/nodes.py b/nodes.py index af15bcf..a31f778 100644 --- a/nodes.py +++ b/nodes.py @@ -22,24 +22,30 @@ 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. + part_idx = field.data[-1] + part_type = field.types[0] if len(field.types) == 1 else "INVALID" + if part_type != gguf.GGUFValueType.STRING: + raise TypeError(f"Bad original shape metadata for {field_key}: Expected STRING, got {part_type}") + shape_str = str(field.parts[part_idx], encoding="utf-8") + return torch.Size(tuple(int(dim) for dim in shape_str.split(","))) + def gguf_sd_loader(path): """ Read state dict as fake tensors """ reader = gguf.GGUFReader(path) - fixup_field = reader.get_field("comfy.gguf.tensor_fixup") - tensor_fixup = {} - if fixup_field is not None: - tensor_fixup_str=str(fixup_field.parts[fixup_field.data[-1]], encoding="utf-8") - for line in str(tensor_fixup_str).split("\n"): - key, *fixup_items = line.split(" ") - tensor_fixup[key] = torch.Size(tuple(int(i) for i in fixup_items)) sd = {} dt = {} for tensor in reader.tensors: tensor_name = str(tensor.name) - shape = tensor_fixup.get(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}: diff --git a/tools/convert.py b/tools/convert.py index 44aabb9..99f4989 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -9,7 +9,7 @@ from safetensors.torch import load_file QUANTIZATION_THRESHOLD = 1024 REARRANGE_THRESHOLD = 512 -MAX_TENSOR_NAME_LENGTH = 63 +MAX_TENSOR_NAME_LENGTH = 127 # Tuple of arch_name, match_lists. # Each item in match_lists is a tuple of keys that must match. @@ -104,7 +104,6 @@ def handle_tensors(args, writer, state_dict): 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 @@ -157,7 +156,10 @@ def handle_tensors(args, writer, state_dict): ): orig_shape = data.shape data = data.reshape(n_params // 256, 256) - tensor_fixup[key] = orig_shape + writer.add_string( + f"comfy.gguf.orig_shape.{key}", + ", ".join(str(dim) for dim in orig_shape), + ) try: data = gguf.quants.quantize(data, data_qtype) @@ -173,12 +175,6 @@ def handle_tensors(args, writer, state_dict): 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 diff --git a/tools/lcpp.patch b/tools/lcpp.patch index 01fe0b8..6355cb8 100644 --- a/tools/lcpp.patch +++ b/tools/lcpp.patch @@ -1,3 +1,16 @@ +diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h +index 1d2a3540..7d354dc6 100644 +--- a/ggml/include/ggml.h ++++ b/ggml/include/ggml.h +@@ -230,7 +230,7 @@ + #define GGML_MAX_CONTEXTS 64 + #define GGML_MAX_SRC 10 + #ifndef GGML_MAX_NAME +-#define GGML_MAX_NAME 64 ++#define GGML_MAX_NAME 128 + #endif + #define GGML_MAX_OP_PARAMS 64 + #define GGML_DEFAULT_N_THREADS 4 diff --git a/src/llama.cpp b/src/llama.cpp index 5ab65ea9..6d684fef 100644 --- a/src/llama.cpp From c1693d0f5551f7f0b79ee9ad930c5b14869eb934 Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 25 Aug 2024 13:18:35 -0600 Subject: [PATCH 7/7] Use a GGUF array for orig shape metadata --- nodes.py | 9 +++------ tools/convert.py | 5 +---- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/nodes.py b/nodes.py index a31f778..0e61662 100644 --- a/nodes.py +++ b/nodes.py @@ -28,12 +28,9 @@ def gguf_sd_loader_get_orig_shape(reader, tensor_name): if field is None: return None # Has original shape metadata, so we try to decode it. - part_idx = field.data[-1] - part_type = field.types[0] if len(field.types) == 1 else "INVALID" - if part_type != gguf.GGUFValueType.STRING: - raise TypeError(f"Bad original shape metadata for {field_key}: Expected STRING, got {part_type}") - shape_str = str(field.parts[part_idx], encoding="utf-8") - return torch.Size(tuple(int(dim) for dim in shape_str.split(","))) + 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): """ diff --git a/tools/convert.py b/tools/convert.py index 99f4989..a49ef1c 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -156,10 +156,7 @@ def handle_tensors(args, writer, state_dict): ): orig_shape = data.shape data = data.reshape(n_params // 256, 256) - writer.add_string( - f"comfy.gguf.orig_shape.{key}", - ", ".join(str(dim) for dim in orig_shape), - ) + 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)