From 6ae0c733b0edb283bc5a80ee7d679254eb230945 Mon Sep 17 00:00:00 2001 From: blepping Date: Wed, 28 Aug 2024 10:34:58 -0600 Subject: [PATCH] Support loading u-net from stable-diffusion.cpp GGUF files Small cleanups --- nodes.py | 57 ++++++++++++++++++++++++++++++++++++++++---------------- ops.py | 36 ++++++++++++++++++++++++++++++----- 2 files changed, 72 insertions(+), 21 deletions(-) diff --git a/nodes.py b/nodes.py index 0e61662..84c5def 100644 --- a/nodes.py +++ b/nodes.py @@ -32,27 +32,50 @@ def gguf_sd_loader_get_orig_shape(reader, tensor_name): 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): +def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", require_prefix=False): """ Read state dict as fake tensors """ reader = gguf.GGUFReader(path) sd = {} dt = {} + 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"}: + raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl but got {arch_str!r}") + else: + arch_str = None + 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) + else: + has_prefix = False + if require_prefix and not has_prefix: + raise RuntimeError(f"Loader requires prefix {handle_prefix!r} which does not exist in the model") for tensor in reader.tensors: - tensor_name = str(tensor.name) + sd_key = tensor_name = tensor.name + if has_prefix: + if not tensor_name.startswith(handle_prefix): + continue + sd_key = tensor_name[prefix_len:] + 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))) - elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: torch_tensor = torch_tensor.view(*shape) - sd[tensor_name] = GGMLTensor( + sd[sd_key] = GGMLTensor( torch_tensor, tensor_type = tensor.tensor_type, - tensor_shape = shape + tensor_shape = shape, + ggml_tensor_name = tensor_name, ) - dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1 + dt[tensor_type_str] = dt.get(tensor_type_str, 0) + 1 # sanity check debug print print("\nggml_sd_loader:") @@ -172,14 +195,14 @@ class UnetLoaderGGUF: def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): ops = GGMLOps() - if dequant_dtype in ["default", None]: + if dequant_dtype in ("default", None): ops.Linear.dequant_dtype = None elif dequant_dtype in ["target"]: ops.Linear.dequant_dtype = dequant_dtype else: ops.Linear.dequant_dtype = getattr(torch, dequant_dtype) - if patch_dtype in ["default", None]: + if patch_dtype in ("default", None): ops.Linear.patch_dtype = None elif patch_dtype in ["target"]: ops.Linear.patch_dtype = patch_dtype @@ -283,11 +306,12 @@ class CLIPLoaderGGUF: class DualCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod def INPUT_TYPES(s): + file_options = (s.get_filename_list(), ) return { "required": { - "clip_name1": (s.get_filename_list(), ), - "clip_name2": (s.get_filename_list(), ), - "type": (["sdxl", "sd3", "flux"], ), + "clip_name1": file_options, + "clip_name2": file_options, + "type": (("sdxl", "sd3", "flux"), ), } } @@ -296,18 +320,19 @@ class DualCLIPLoaderGGUF(CLIPLoaderGGUF): def load_clip(self, clip_name1, clip_name2, type): 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_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)),) class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod def INPUT_TYPES(s): + file_options = (s.get_filename_list(), ) return { "required": { - "clip_name1": (s.get_filename_list(), ), - "clip_name2": (s.get_filename_list(), ), - "clip_name3": (s.get_filename_list(), ), + "clip_name1": file_options, + "clip_name2": file_options, + "clip_name3": file_options, } } @@ -317,7 +342,7 @@ class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): clip_path1 = folder_paths.get_full_path("clip", clip_name1) 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_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)),) diff --git a/ops.py b/ops.py index 88ca894..21b4d36 100644 --- a/ops.py +++ b/ops.py @@ -9,13 +9,14 @@ class GGMLTensor(torch.Tensor): """ Main tensor-like class for storing quantized weights """ - def __init__(self, *args, tensor_type, tensor_shape, patches=[], **kwargs): + def __init__(self, *args, tensor_type, tensor_shape, patches=[], ggml_tensor_name=None, **kwargs): super().__init__() self.tensor_type = tensor_type self.tensor_shape = tensor_shape self.patches = patches + self.ggml_tensor_name = ggml_tensor_name - def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): + def __new__(cls, *args, tensor_type, tensor_shape, patches=[], ggml_tensor_name=None, **kwargs): return super().__new__(cls, *args, **kwargs) def to(self, *args, **kwargs): @@ -23,6 +24,7 @@ class GGMLTensor(torch.Tensor): new.tensor_type = getattr(self, "tensor_type", None) new.tensor_shape = getattr(self, "tensor_shape", new.data.shape) new.patches = getattr(self, "patches", []).copy() + new.ggml_tensor_name = getattr(self, "ggml_tensor_name", None) return new def clone(self, *args, **kwargs): @@ -44,6 +46,7 @@ class GGMLTensor(torch.Tensor): new.tensor_type = getattr(self, "tensor_type", None) new.tensor_shape = getattr(self, "tensor_shape", new.data.shape) new.patches = getattr(self, "patches", []).copy() + new.ggml_tensor_name = getattr(self, "ggml_tensor_name", None) return new @property @@ -76,7 +79,7 @@ class GGMLLayer(torch.nn.Module): self.to(x.device) weight, bias = self.get_weights(x.dtype) - x = self._forward_operation(x, weight, bias=bias) + x = self._forward_operation(input=x, weight=weight, bias=bias) del weight, bias if device: @@ -84,10 +87,11 @@ class GGMLLayer(torch.nn.Module): return x def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + prefix_len = len(prefix) for k,v in state_dict.items(): - if k[len(prefix):] == "weight": + if k[prefix_len:] == "weight": self.weight = v - elif k[len(prefix):] == "bias": + elif k[prefix_len:] == "bias": self.bias = v else: missing_keys.append(k) @@ -164,6 +168,28 @@ class GGMLOps(comfy.ops.manual_cast): _ = kwargs.pop("kernel_size", None) self._forward_operation = partial(staticmethod(torch.nn.functional.conv2d), **kwargs) + class GroupNorm(GGMLLayer): + def __init__(self, *args, device=None, dtype=None, **kwargs): + super().__init__() + num_groups = kwargs.get("num_groups") + if num_groups is None: + if len(args) == 0 or not isinstance(args[0], int): + raise ValueError("Couldn't get num_groups in GroupNorm: missing or bad type") + num_groups = args[0] + eps = kwargs.get("eps", 1e-05) + self._forward_operation = partial(staticmethod(torch.nn.functional.group_norm), num_groups=num_groups, eps=eps) + + class LayerNorm(GGMLLayer): + def __init__(self, *args, device=None, dtype=None, **kwargs): + super().__init__() + normalized_shape = kwargs.get("normalized_shape") + if normalized_shape is None: + if len(args) == 0 or not all(isinstance(v, int) for v in args): + raise ValueError("Couldn't get normalized_shape in LayerNorm: missing or bad type") + normalized_shape = tuple(args) + eps = kwargs.get("eps", 1e-05) + self._forward_operation = partial(staticmethod(torch.nn.functional.layer_norm), normalized_shape=normalized_shape, eps=eps) + def move_patch_to_cuda(item, device): if isinstance(item, torch.Tensor):