diff --git a/dequant.py b/dequant.py index 1750cd0..8232500 100644 --- a/dequant.py +++ b/dequant.py @@ -3,21 +3,29 @@ import gguf import torch from tqdm import tqdm -def dequantize_tensor(tensor, dtype=None, dequant_dtype=None): - 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): - return data.to(dtype) +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 TORCH_COMPATIBLE_QTYPES: + return tensor.to(dtype) elif qtype in dequantize_functions: dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype - return dequantize(data, 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}") - new = gguf.quants.dequantize(data.cpu().numpy(), qtype) - return torch.from_numpy(new).to(data.device, dtype=dtype) + new = gguf.quants.dequantize(tensor.cpu().numpy(), qtype) + return torch.from_numpy(new).to(tensor.device, dtype=dtype) def dequantize(data, qtype, oshape, dtype=None): """ @@ -45,6 +53,10 @@ def split_block_dims(blocks, *args): dims = list(args) + [n_max - sum(args)] return torch.split(blocks, dims, dim=1) +# Full weights # +def dequantize_blocks_BF16(blocks, block_size, type_size, dtype=None): + return (blocks.view(torch.int16).to(torch.int32) << 16).view(torch.float32) + # Legacy Quants # def dequantize_blocks_Q8_0(blocks, block_size, type_size, dtype=None): d, x = split_block_dims(blocks, 2) @@ -222,6 +234,7 @@ def dequantize_blocks_Q2_K(blocks, block_size, type_size, dtype=None): return qs.reshape((n_blocks, -1)) dequantize_functions = { + gguf.GGMLQuantizationType.BF16: dequantize_blocks_BF16, gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0, gguf.GGMLQuantizationType.Q5_1: dequantize_blocks_Q5_1, gguf.GGMLQuantizationType.Q5_0: dequantize_blocks_Q5_0, diff --git a/nodes.py b/nodes.py index 70823f4..e129988 100644 --- a/nodes.py +++ b/nodes.py @@ -10,8 +10,8 @@ import comfy.model_management import comfy.model_patcher import folder_paths -from .ops import GGMLTensor, GGMLOps, move_patch_to_cuda -from .dequant import dequantize_tensor +from .ops import GGMLTensor, GGMLOps, move_patch_to_device +from .dequant import is_quantized, is_torch_compatible # Add a custom keys for files ending in .gguf if "unet_gguf" not in folder_paths.folder_names_and_paths: @@ -32,34 +32,73 @@ 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."): """ Read state dict as fake tensors """ reader = gguf.GGUFReader(path) - sd = {} - dt = {} + + # 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: - 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:] + 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", "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)) + 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))) - elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + # 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) - sd[tensor_name] = GGMLTensor( - torch_tensor, - tensor_type = tensor.tensor_type, - tensor_shape = shape - ) - dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1 + 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 dt.items(): + for k,v in qtype_dict.items(): print(f" {k:30}{v:3}") - print("\n") - return sd + + return state_dict # for remapping llama.cpp -> original key names clip_sd_map = { @@ -78,10 +117,6 @@ clip_sd_map = { "ffn_gate": "layer.1.DenseReluDense.wi_0", "ffn_norm": "layer.1.layer_norm", } -# weights that should be dequantized on load -clip_sd_dequant = { - "shared.weight", -} def gguf_clip_loader(path): raw_sd = gguf_sd_loader(path) @@ -90,9 +125,6 @@ def gguf_clip_loader(path): for k,v in raw_sd.items(): for s,d in clip_sd_map.items(): k = k.replace(s,d) - if k in clip_sd_dequant: - v = dequantize_tensor(v, torch.float32).to(torch.float16) - v = GGMLTensor(v, tensor_type=gguf.GGMLQuantizationType.F16, tensor_shape=v.shape) sd[k] = v return sd @@ -105,9 +137,6 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): if key not in self.patches: return weight = comfy.utils.get_attr(self.model, key) - inplace_update = self.weight_inplace_update or inplace_update - if key not in self.backup: - self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update) try: from comfy.lora import calculate_weight @@ -115,18 +144,17 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): 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 device_to is not None: - out_weight = weight.to(device_to, copy=True) - else: - out_weight = weight.to(weight.device, copy=True) # make sure patches are copied - - if self.patch_on_device: - patches = move_patch_to_cuda(patches, self.load_device) + if is_quantized(weight): + out_weight = weight.to(device_to) + patches = move_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device) out_weight.patches.append((calculate_weight, patches, key)) - else: + inplace_update = self.weight_inplace_update or inplace_update + if key not in self.backup: + self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])( + weight.to(device=self.offload_device, copy=inplace_update), inplace_update + ) + if device_to is not None: temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True) else: @@ -140,6 +168,16 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): else: comfy.utils.set_attr_param(self.model, key, out_weight) + def unpatch_model(self, device_to=None, unpatch_weights=True): + if unpatch_weights: + for p in self.model.parameters(): + if is_torch_compatible(p): + continue + patches = getattr(p, "patches", []) + if len(patches) > 0: + p.patches = [] + return super().unpatch_model(device_to=None, unpatch_weights=unpatch_weights) + def load(self, *args, force_patch_weights=False, **kwargs): # always call `patch_weight_to_device` even for lowvram return super().load(*args, force_patch_weights=True, **kwargs) @@ -176,14 +214,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 @@ -287,11 +325,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"), ), } } @@ -300,18 +339,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, } } @@ -321,7 +361,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..62e797c 100644 --- a/ops.py +++ b/ops.py @@ -1,9 +1,10 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import gguf import torch -from functools import partial import comfy.ops -from .dequant import dequantize_tensor +import comfy.model_management +from .dequant import dequantize_tensor, is_quantized class GGMLTensor(torch.Tensor): """ @@ -59,67 +60,49 @@ class GGMLLayer(torch.nn.Module): comfy_cast_weights = True dequant_dtype = None patch_dtype = None + torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} - def __init__(self, *args, **kwargs): - super().__init__() - self.weight = GGMLTensor(1, tensor_type=None, tensor_shape=None) - self.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 + return is_quantized(weight) or is_quantized(bias) - def _forward_operation(self, x, weight, bias): - raise NotImplementedError + 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.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 forward(self, x): - # lowvram hack - device = None - if self.weight.device != x.device: - device = self.weight.device - self.to(x.device) + def _save_to_state_dict(self, *args, **kwargs): + if self.is_ggml_quantized(): + return self.ggml_save_to_state_dict(*args, **kwargs) + return super()._save_to_state_dict(*args, **kwargs) - weight, bias = self.get_weights(x.dtype) - x = self._forward_operation(x, weight, bias=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): + 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[len(prefix):] == "weight": - self.weight = v - elif k[len(prefix):] == "bias": - self.bias = v + if k[prefix_len:] == "weight": + self.weight = torch.nn.Parameter(v, requires_grad=False) + 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 - if self.weight is not None: - weight = torch.zeros_like(self.weight, device=torch.device("meta")) - destination[f"{prefix}weight"] = weight + weight = torch.zeros_like(self.weight, device=torch.device("meta")) + destination[prefix + "weight"] = weight if self.bias is not None: bias = torch.zeros_like(self.bias, device=torch.device("meta")) - destination[f"{prefix}bias"] = bias + destination[prefix + "bias"] = bias return # This would return the actual state dict - weight, bias = self.get_weights() - if weight is not None: - destination[f"{prefix}weight"] = weight + destination[prefix + "weight"] = self.get_weight(self.weight) if bias is not None: - destination[f"{prefix}bias"] = weight - - def _apply(self, fn): - if self.weight is not None: - try: - self.weight = fn(self.weight) - except TypeError: - # TODO: Figure out why this happens - pass - if self.bias is not None: - self.bias = fn(self.bias) - super()._apply(fn) - return self + destination[prefix + "bias"] = self.get_weight(self.bias) def get_weight(self, tensor, dtype): if tensor is None: @@ -128,9 +111,8 @@ class GGMLLayer(torch.nn.Module): # consolidate and load patches to GPU in async patch_list = [] device = tensor.device - t_move = lambda x: x.to(device) if torch.is_tensor(x) else x for function, patches, key in getattr(tensor, "patches", []): - patch_list += move_patch_to_cuda(patches, device) + patch_list += move_patch_to_device(patches, device) # dequantize tensor while patches load weight = dequantize_tensor(tensor, dtype, self.dequant_dtype) @@ -145,32 +127,75 @@ class GGMLLayer(torch.nn.Module): weight = function(patch_list, weight, key, patch_dtype) return weight - def get_weights(self, dtype=torch.float16): - weight = self.get_weight(self.weight, dtype) - bias = self.get_weight(self.bias, dtype) - return (weight, bias) + def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None): + if input is not None: + if dtype is None: + dtype = getattr(input, "dtype", torch.float32) + if bias_dtype is None: + bias_dtype = dtype + if device is None: + device = input.device + bias = None + non_blocking = comfy.model_management.device_supports_non_blocking(device) + if s.bias is not None: + bias = s.get_weight(s.bias.to(device), dtype) + bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False) + + weight = s.get_weight(s.weight.to(device), dtype) + 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.is_ggml_quantized(): + return self.forward_ggml_cast_weights(input, *args, **kwargs) + return super().forward_comfy_cast_weights(input, *args, **kwargs) + + def forward_ggml_cast_weights(self, input): + raise NotImplementedError class GGMLOps(comfy.ops.manual_cast): """ Dequantize weights on the fly before doing the compute """ - class Linear(GGMLLayer): - _forward_operation = staticmethod(torch.nn.functional.linear) + class Linear(GGMLLayer, comfy.ops.manual_cast.Linear): + def forward_ggml_cast_weights(self, input): + weight, bias = self.cast_bias_weight(input) + return torch.nn.functional.linear(input, weight, bias) - class Conv2d(GGMLLayer): - def __init__(self, *args, device=None, dtype=None, **kwargs): - super().__init__() - _ = kwargs.pop("kernel_size", None) - self._forward_operation = partial(staticmethod(torch.nn.functional.conv2d), **kwargs) + class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d): + def forward_ggml_cast_weights(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_ggml_cast_weights(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) + 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) -def move_patch_to_cuda(item, device): + class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm): + def forward_ggml_cast_weights(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 forward_ggml_cast_weights(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_device(item, device): if isinstance(item, torch.Tensor): return item.to(device, non_blocking=True) elif isinstance(item, tuple): - return tuple(move_patch_to_cuda(x, device) for x in item) + return tuple(move_patch_to_device(x, device) for x in item) elif isinstance(item, list): - return [move_patch_to_cuda(x, device) for x in item] + return [move_patch_to_device(x, device) for x in item] else: return item