From 6ae0c733b0edb283bc5a80ee7d679254eb230945 Mon Sep 17 00:00:00 2001 From: blepping Date: Wed, 28 Aug 2024 10:34:58 -0600 Subject: [PATCH 1/9] 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): From f7d8fbe74052e139a3b20b4a4f702bda82fdc23e Mon Sep 17 00:00:00 2001 From: blepping Date: Wed, 28 Aug 2024 11:56:51 -0600 Subject: [PATCH 2/9] Workaround for stable-diffusion.cpp SDXL GGUF files --- nodes.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/nodes.py b/nodes.py index 84c5def..9df2523 100644 --- a/nodes.py +++ b/nodes.py @@ -67,6 +67,11 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", require_prefix= 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))) + if arch_str is None and (tensor_name.endswith(".proj_in.weight") or tensor_name.endswith(".proj_out.weight")): + # Workaround for stable-diffusion.cpp SDXL detection. + # Impossible to land here for our own GGUF files as we set architecture. + while len(shape) > 2 and shape[-1] == 1: + shape = shape[:-1] if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: torch_tensor = torch_tensor.view(*shape) sd[sd_key] = GGMLTensor( From e814fc12642aa4d6ad3ae5323614c028ad1c2be1 Mon Sep 17 00:00:00 2001 From: blepping Date: Wed, 28 Aug 2024 12:34:08 -0600 Subject: [PATCH 3/9] Revert ggml_tensor_name addition --- nodes.py | 3 +-- ops.py | 7 ++----- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/nodes.py b/nodes.py index 9df2523..d702e0c 100644 --- a/nodes.py +++ b/nodes.py @@ -77,8 +77,7 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", require_prefix= sd[sd_key] = GGMLTensor( torch_tensor, tensor_type = tensor.tensor_type, - tensor_shape = shape, - ggml_tensor_name = tensor_name, + tensor_shape = shape ) dt[tensor_type_str] = dt.get(tensor_type_str, 0) + 1 diff --git a/ops.py b/ops.py index 21b4d36..8c12de7 100644 --- a/ops.py +++ b/ops.py @@ -9,14 +9,13 @@ class GGMLTensor(torch.Tensor): """ Main tensor-like class for storing quantized weights """ - def __init__(self, *args, tensor_type, tensor_shape, patches=[], ggml_tensor_name=None, **kwargs): + def __init__(self, *args, tensor_type, tensor_shape, patches=[], **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=[], ggml_tensor_name=None, **kwargs): + def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): return super().__new__(cls, *args, **kwargs) def to(self, *args, **kwargs): @@ -24,7 +23,6 @@ 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): @@ -46,7 +44,6 @@ 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 From b88cb9a344d5f8601dbb9e9af13c9d9c4b466112 Mon Sep 17 00:00:00 2001 From: blepping Date: Wed, 28 Aug 2024 12:50:39 -0600 Subject: [PATCH 4/9] Remove require_prefix in gguf_sd_loader --- nodes.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/nodes.py b/nodes.py index d702e0c..609716e 100644 --- a/nodes.py +++ b/nodes.py @@ -32,7 +32,7 @@ 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, handle_prefix="model.diffusion_model.", require_prefix=False): +def gguf_sd_loader(path, handle_prefix="model.diffusion_model."): """ Read state dict as fake tensors """ @@ -54,8 +54,6 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", require_prefix= 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: sd_key = tensor_name = tensor.name if has_prefix: From 7f3ced6c3ed2b23abf4105e0210da03ad6015b87 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sun, 1 Sep 2024 00:56:30 +0200 Subject: [PATCH 5/9] Consider FP8 as source #91 --- tools/convert.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tools/convert.py b/tools/convert.py index a49ef1c..63bded9 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -109,6 +109,9 @@ def handle_tensors(args, writer, state_dict): if data.dtype == torch.bfloat16: data = data.to(torch.float32).numpy() + # this is so we don't break torch 2.0.X + elif data.dtype in [getattr(torch, "float8_e4m3fn", "_invalid"), getattr(torch, "float8_e5m2", "_invalid")]: + data = data.to(torch.float16).numpy() else: data = data.numpy() From a0cb97ecebc0b9a1ebeaac77cd201a30948da402 Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 1 Sep 2024 05:25:32 -0600 Subject: [PATCH 6/9] Add LayerNorm and GroupNorm ops Delegate to superclass methods whenever possible --- dequant.py | 2 +- ops.py | 57 ++++++++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/dequant.py b/dequant.py index 3688d7f..5bc775b 100644 --- a/dequant.py +++ b/dequant.py @@ -11,7 +11,7 @@ def dequantize_tensor(tensor, dtype=None, dequant_dtype=None): return tensor.to(dtype) elif qtype in dequantize_functions: dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype - return dequantize(tensor, qtype, oshape, dtype=dequant_dtype).to(dtype) + return dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype).to(dtype) else: # this is incredibly slow tqdm.write(f"Falling back to numpy dequant for qtype: {qtype}") diff --git a/ops.py b/ops.py index a6f47c1..70a0111 100644 --- a/ops.py +++ b/ops.py @@ -59,18 +59,39 @@ class GGMLLayer(torch.nn.Module): comfy_cast_weights = True dequant_dtype = None patch_dtype = None + ggml_delegate_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + def ggml_check_delegate(self, *, weight=None, bias=None): + if weight is None: + weight = self.weight + if bias is None: + bias = self.bias + weight_type = getattr(weight, "tensor_type", None) + bias_type = None if bias is None else getattr(bias, "tensor_type", None) + return weight_type in self.ggml_delegate_types and bias_type in self.ggml_delegate_types + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + weight, bias = state_dict.get(f"{prefix}weight"), state_dict.get(f"{prefix}bias") + if self.ggml_check_delegate(weight=weight, bias=bias): + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def _save_to_state_dict(self, *args, **kwargs): + if self.ggml_check_delegate(): + return super()._save_to_state_dict(*args, **kwargs) + return self.ggml_save_to_state_dict(*args, **kwargs) + + def ggml_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[prefix_len] == "weight": + if k[prefix_len:] == "weight": self.weight = torch.nn.Parameter(v, requires_grad=False) - elif k[prefix_len] == "bias" and v is not None: + elif k[prefix_len:] == "bias" and v is not None: self.bias = torch.nn.Parameter(v, requires_grad=False) else: missing_keys.append(k) - def _save_to_state_dict(self, destination, prefix, keep_vars): + def ggml_save_to_state_dict(self, destination, prefix, keep_vars): # This is a fake state dict for vram estimation weight = torch.zeros_like(self.weight, device=torch.device("meta")) destination[prefix + "weight"] = weight @@ -126,30 +147,50 @@ class GGMLLayer(torch.nn.Module): weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False) return weight, bias + def forward_comfy_cast_weights(self, input, *args, **kwargs): + if self.ggml_check_delegate(): + return super().forward_comfy_cast_weights(input, *args, **kwargs) + return self.ggml_forward(input, *args, **kwargs) + + def ggml_forward(self, input): + raise NotImplementedError + class GGMLOps(comfy.ops.manual_cast): """ Dequantize weights on the fly before doing the compute """ class Linear(GGMLLayer, comfy.ops.manual_cast.Linear): - def forward_comfy_cast_weights(self, input): + def ggml_forward(self, input): weight, bias = self.cast_bias_weight(input) return torch.nn.functional.linear(input, weight, bias) class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d): - def forward_comfy_cast_weights(self, input): + def ggml_forward(self, input): weight, bias = self.cast_bias_weight(input) return self._conv_forward(input, weight, bias) class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding): - def forward_comfy_cast_weights(self, input, out_dtype=None): + def ggml_forward(self, input, out_dtype=None): output_dtype = out_dtype if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16: out_dtype = None - weight, bias = self.cast_bias_weight(self, device=input.device, dtype=out_dtype) + weight, _bias = self.cast_bias_weight(self, device=input.device, dtype=out_dtype) return torch.nn.functional.embedding( input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse ).to(dtype=output_dtype) + class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm): + def ggml_forward(self, input): + if self.weight is None: + return super().forward_comfy_cast_weights(input) + weight, bias = self.cast_bias_weight(input) + return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps) + + class GroupNorm(GGMLLayer, comfy.ops.manual_cast.GroupNorm): + def ggml_forward(self, input): + weight, bias = self.cast_bias_weight(input) + return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps) + def move_patch_to_cuda(item, device): if isinstance(item, torch.Tensor): return item.to(device, non_blocking=True) From 3bad1b4e04e839fd6ec57f8ea0b2a657e4936034 Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 1 Sep 2024 06:46:50 -0600 Subject: [PATCH 7/9] Fix loading GGUF t5xxl Can't use squeeze proj workaround for stable-diffusion.cpp SD15 --- nodes.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/nodes.py b/nodes.py index d095de2..df5678c 100644 --- a/nodes.py +++ b/nodes.py @@ -11,7 +11,7 @@ import comfy.model_patcher import folder_paths from .ops import GGMLTensor, GGMLOps, move_patch_to_cuda -from .dequant import dequantize_tensor +from .tools.convert import detect_arch # Add a custom keys for files ending in .gguf if "unet_gguf" not in folder_paths.folder_names_and_paths: @@ -44,8 +44,8 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model."): 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}") + if arch_str not in {"flux", "sd1", "sdxl", "t5encoder"}: + raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}") else: arch_str = None if handle_prefix is not None: @@ -54,18 +54,24 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model."): has_prefix = any(s.startswith(handle_prefix) for s in tensor_names) else: has_prefix = False + 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)) + detected_arch = detect_arch(set(val[0] for val in tensors)) if arch_str is None else None + 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))) - if arch_str is None and (tensor_name.endswith(".proj_in.weight") or tensor_name.endswith(".proj_out.weight")): + if arch_str is None and detected_arch == "sdxl" and \ + (tensor_name.endswith(".proj_in.weight") or tensor_name.endswith(".proj_out.weight")): # Workaround for stable-diffusion.cpp SDXL detection. # Impossible to land here for our own GGUF files as we set architecture. while len(shape) > 2 and shape[-1] == 1: From e19fe8e898231b5ec35d385e2ea3a9373f85bc6e Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 1 Sep 2024 18:04:31 -0600 Subject: [PATCH 8/9] Rename ggml_check_delegate to is_ggml_quantized, refactor accordingly --- ops.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/ops.py b/ops.py index 70a0111..930b4da 100644 --- a/ops.py +++ b/ops.py @@ -59,27 +59,27 @@ class GGMLLayer(torch.nn.Module): comfy_cast_weights = True dequant_dtype = None patch_dtype = None - ggml_delegate_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} + torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} - def ggml_check_delegate(self, *, weight=None, bias=None): + def is_ggml_quantized(self, *, weight=None, bias=None): if weight is None: weight = self.weight if bias is None: bias = self.bias weight_type = getattr(weight, "tensor_type", None) bias_type = None if bias is None else getattr(bias, "tensor_type", None) - return weight_type in self.ggml_delegate_types and bias_type in self.ggml_delegate_types + return any(t not in self.torch_compatible_tensor_types for t in (weight_type, bias_type)) def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): weight, bias = state_dict.get(f"{prefix}weight"), state_dict.get(f"{prefix}bias") - if self.ggml_check_delegate(weight=weight, bias=bias): - return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) - return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) + if self.is_ggml_quantized(weight=weight, bias=bias): + return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) def _save_to_state_dict(self, *args, **kwargs): - if self.ggml_check_delegate(): - return super()._save_to_state_dict(*args, **kwargs) - return self.ggml_save_to_state_dict(*args, **kwargs) + if self.is_ggml_quantized(): + return self.ggml_save_to_state_dict(*args, **kwargs) + return super()._save_to_state_dict(*args, **kwargs) def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): prefix_len = len(prefix) @@ -148,9 +148,9 @@ class GGMLLayer(torch.nn.Module): return weight, bias def forward_comfy_cast_weights(self, input, *args, **kwargs): - if self.ggml_check_delegate(): - return super().forward_comfy_cast_weights(input, *args, **kwargs) - return self.ggml_forward(input, *args, **kwargs) + if self.is_ggml_quantized(): + return self.ggml_forward(input, *args, **kwargs) + return super().forward_comfy_cast_weights(input, *args, **kwargs) def ggml_forward(self, input): raise NotImplementedError From 379175e7bf8b65019cdd11108bb882120a6f17df Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 1 Sep 2024 18:30:47 -0600 Subject: [PATCH 9/9] Refactor tensor type checking a bit more --- dequant.py | 11 ++++++++++- ops.py | 6 ++---- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/dequant.py b/dequant.py index 5bc775b..8232500 100644 --- a/dequant.py +++ b/dequant.py @@ -3,11 +3,20 @@ import gguf import torch from tqdm import tqdm + +TORCH_COMPATIBLE_QTYPES = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} + +def is_torch_compatible(tensor): + return tensor is None or getattr(tensor, "tensor_type", None) in TORCH_COMPATIBLE_QTYPES + +def is_quantized(tensor): + return not is_torch_compatible(tensor) + def dequantize_tensor(tensor, dtype=None, dequant_dtype=None): qtype = getattr(tensor, "tensor_type", None) oshape = getattr(tensor, "tensor_shape", tensor.shape) - if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16): + if qtype in TORCH_COMPATIBLE_QTYPES: return tensor.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 930b4da..1789b7f 100644 --- a/ops.py +++ b/ops.py @@ -3,7 +3,7 @@ import gguf import torch import comfy.ops -from .dequant import dequantize_tensor +from .dequant import dequantize_tensor, is_quantized class GGMLTensor(torch.Tensor): """ @@ -66,9 +66,7 @@ class GGMLLayer(torch.nn.Module): weight = self.weight if bias is None: bias = self.bias - weight_type = getattr(weight, "tensor_type", None) - bias_type = None if bias is None else getattr(bias, "tensor_type", None) - return any(t not in self.torch_compatible_tensor_types for t in (weight_type, bias_type)) + return is_quantized(weight) or is_quantized(bias) def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): weight, bias = state_dict.get(f"{prefix}weight"), state_dict.get(f"{prefix}bias")