From 0a1723baae5383aac8ababf6f22f8d1e49a28857 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 20 Dec 2024 13:03:27 +0100 Subject: [PATCH] Split loader into separate file --- loader.py | 122 ++++++++++++++++++++++++++++++++++++++++++++++++ nodes.py | 136 +++++++----------------------------------------------- 2 files changed, 139 insertions(+), 119 deletions(-) create mode 100644 loader.py diff --git a/loader.py b/loader.py new file mode 100644 index 0000000..3a06ce2 --- /dev/null +++ b/loader.py @@ -0,0 +1,122 @@ +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import torch +import gguf + +from .ops import GGMLTensor + +IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "ltxv", "hyvid"} +TXT_ARCH_LIST = {"t5", "t5encoder"} + +def 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, handle_prefix="model.diffusion_model.", return_arch=False): + """ + Read state dict as fake tensors + """ + reader = gguf.GGUFReader(path) + + # filter and strip prefix + has_prefix = False + if handle_prefix is not None: + prefix_len = len(handle_prefix) + tensor_names = set(tensor.name for tensor in reader.tensors) + has_prefix = any(s.startswith(handle_prefix) for s in tensor_names) + + tensors = [] + for tensor in reader.tensors: + sd_key = tensor_name = tensor.name + if has_prefix: + if not tensor_name.startswith(handle_prefix): + continue + sd_key = tensor_name[prefix_len:] + tensors.append((sd_key, tensor)) + + # detect and verify architecture + compat = None + arch_str = None + arch_field = reader.get_field("general.architecture") + if arch_field is not None: + if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING: + raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}") + arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8") + if arch_str not in IMG_ARCH_LIST and arch_str not in TXT_ARCH_LIST: + raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}") + else: # stable-diffusion.cpp + # 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" + + # main loading loop + state_dict = {} + qtype_dict = {} + for sd_key, tensor in tensors: + tensor_name = tensor.name + tensor_type_str = str(tensor.tensor_type) + torch_tensor = torch.from_numpy(tensor.data) # mmap + + shape = get_orig_shape(reader, tensor_name) + if shape is None: + shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + # Workaround for stable-diffusion.cpp SDXL detection. + if compat == "sd.cpp" and arch_str == "sdxl": + if any([tensor_name.endswith(x) for x in (".proj_in.weight", ".proj_out.weight")]): + while len(shape) > 2 and shape[-1] == 1: + shape = shape[:-1] + + # add to state dict + if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + torch_tensor = torch_tensor.view(*shape) + state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) + qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1 + + # sanity check debug print + print("\nggml_sd_loader:") + for k,v in qtype_dict.items(): + print(f" {k:30}{v:3}") + + if return_arch: + return (state_dict, arch_str) + return state_dict + +# for remapping llama.cpp -> original key names +T5_SD_MAP = { + "enc.": "encoder.", + ".blk.": ".block.", + "token_embd": "shared", + "output_norm": "final_layer_norm", + "attn_q": "layer.0.SelfAttention.q", + "attn_k": "layer.0.SelfAttention.k", + "attn_v": "layer.0.SelfAttention.v", + "attn_o": "layer.0.SelfAttention.o", + "attn_norm": "layer.0.layer_norm", + "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias", + "ffn_up": "layer.1.DenseReluDense.wi_1", + "ffn_down": "layer.1.DenseReluDense.wo", + "ffn_gate": "layer.1.DenseReluDense.wi_0", + "ffn_norm": "layer.1.layer_norm", +} + +def sd_map_replace(raw_sd, key_map): + sd = {} + for k,v in raw_sd.items(): + for s,d in key_map.items(): + k = k.replace(s,d) + sd[k] = v + return sd + +def gguf_clip_loader(path): + sd, arch = gguf_sd_loader(path, return_arch=True) + if "enc.blk.23.ffn_up.weight" in sd or arch in {"t5", "t5encoder"}: + sd = sd_map_replace(sd, T5_SD_MAP) + else: + pass + return sd diff --git a/nodes.py b/nodes.py index 850b87d..f130f1c 100644 --- a/nodes.py +++ b/nodes.py @@ -1,16 +1,16 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) import torch -import gguf -import copy import logging +import collections import comfy.sd import comfy.utils -import comfy.model_management import comfy.model_patcher +import comfy.model_management import folder_paths -from .ops import GGMLTensor, GGMLOps, move_patch_to_device +from .ops import GGMLOps, move_patch_to_device +from .loader import gguf_sd_loader, gguf_clip_loader from .dequant import is_quantized, is_torch_compatible # Add a custom keys for files ending in .gguf @@ -22,114 +22,6 @@ if "clip_gguf" not in folder_paths.folder_names_and_paths: orig = folder_paths.folder_names_and_paths.get("text_encoders", 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, handle_prefix="model.diffusion_model."): - """ - Read state dict as fake tensors - """ - reader = gguf.GGUFReader(path) - - # filter and strip prefix - has_prefix = False - if handle_prefix is not None: - prefix_len = len(handle_prefix) - tensor_names = set(tensor.name for tensor in reader.tensors) - has_prefix = any(s.startswith(handle_prefix) for s in tensor_names) - - tensors = [] - for tensor in reader.tensors: - sd_key = tensor_name = tensor.name - if has_prefix: - if not tensor_name.startswith(handle_prefix): - continue - sd_key = tensor_name[prefix_len:] - tensors.append((sd_key, tensor)) - - # detect and verify architecture - compat = None - arch_str = None - arch_field = reader.get_field("general.architecture") - if arch_field is not None: - if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING: - raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}") - arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8") - if arch_str not in {"flux", "sd1", "sdxl", "sd3", "aura", "ltxv", "hyvid", "t5", "t5encoder"}: - raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}") - else: # stable-diffusion.cpp - # 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" - - # main loading loop - state_dict = {} - qtype_dict = {} - for sd_key, tensor in tensors: - tensor_name = tensor.name - tensor_type_str = str(tensor.tensor_type) - 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))) - # Workaround for stable-diffusion.cpp SDXL detection. - if compat == "sd.cpp" and arch_str == "sdxl": - if any([tensor_name.endswith(x) for x in (".proj_in.weight", ".proj_out.weight")]): - while len(shape) > 2 and shape[-1] == 1: - shape = shape[:-1] - - # add to state dict - if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: - torch_tensor = torch_tensor.view(*shape) - state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) - qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1 - - # sanity check debug print - print("\nggml_sd_loader:") - for k,v in qtype_dict.items(): - print(f" {k:30}{v:3}") - - return state_dict - -# for remapping llama.cpp -> original key names -clip_sd_map = { - "enc.": "encoder.", - ".blk.": ".block.", - "token_embd": "shared", - "output_norm": "final_layer_norm", - "attn_q": "layer.0.SelfAttention.q", - "attn_k": "layer.0.SelfAttention.k", - "attn_v": "layer.0.SelfAttention.v", - "attn_o": "layer.0.SelfAttention.o", - "attn_norm": "layer.0.layer_norm", - "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias", - "ffn_up": "layer.1.DenseReluDense.wi_1", - "ffn_down": "layer.1.DenseReluDense.wo", - "ffn_gate": "layer.1.DenseReluDense.wi_0", - "ffn_norm": "layer.1.layer_norm", -} - -def gguf_clip_loader(path): - raw_sd = gguf_sd_loader(path) - assert "enc.blk.23.ffn_up.weight" in raw_sd, "Invalid Text Encoder!" - sd = {} - for k,v in raw_sd.items(): - for s,d in clip_sd_map.items(): - k = k.replace(s,d) - sd[k] = v - return sd - -# TODO: Temporary fix for now -import collections class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): patch_on_device = False @@ -276,7 +168,8 @@ class UnetLoaderGGUFAdvanced(UnetLoaderGGUF): } TITLE = "Unet Loader (GGUF/Advanced)" -clip_name_dict = { +# Mapping from name to CLIPType enum with fallback +CLIP_TYPE_MAP = { "stable_diffusion": comfy.sd.CLIPType.STABLE_DIFFUSION, "stable_cascade": comfy.sd.CLIPType.STABLE_CASCADE, "stable_audio": comfy.sd.CLIPType.STABLE_AUDIO, @@ -287,6 +180,14 @@ clip_name_dict = { "ltxv": getattr(comfy.sd.CLIPType, "LTXV", None), } +def get_clip_type(type): + if type not in CLIP_TYPE_MAP: + raise ValueError(f"Unknown CLIP model type {type}") + clip_type = CLIP_TYPE_MAP[type] + if clip_type is None: + raise ValueError(f"Unsupported CLIP model type {type} (Update ComfyUI)") + return clip_type + class CLIPLoaderGGUF: @classmethod def INPUT_TYPES(s): @@ -343,8 +244,7 @@ class CLIPLoaderGGUF: def load_clip(self, clip_name, type="stable_diffusion"): clip_path = folder_paths.get_full_path("clip", clip_name) - clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),) + return (self.load_patcher([clip_path], get_clip_type(type), self.load_data([clip_path])),) class DualCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod @@ -364,8 +264,7 @@ 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) - clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),) class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod @@ -386,8 +285,7 @@ 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) - clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),) NODE_CLASS_MAPPINGS = { "UnetLoaderGGUF": UnetLoaderGGUF,