diff --git a/dequant.py b/dequant.py index 3688d7f..8232500 100644 --- a/dequant.py +++ b/dequant.py @@ -3,15 +3,24 @@ 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 - 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/nodes.py b/nodes.py index 094d8e9..ed4c242 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_device -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: @@ -32,27 +32,58 @@ 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 = {} + 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", "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: + 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 + 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)) + 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))) - elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + 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: + shape = shape[:-1] + 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 ) - 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:") @@ -133,7 +164,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): def unpatch_model(self, device_to=None, unpatch_weights=True): if unpatch_weights: for p in self.model.parameters(): - qtype = getattr(p, "tensor_type", None) + qtype = getattr(p, "tensor_type", None) patches = getattr(p, "patches", []) if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16): continue @@ -177,14 +208,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 @@ -288,11 +319,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"), ), } } @@ -301,18 +333,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, } } @@ -322,7 +355,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 82a3e8e..6e106fe 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): """ @@ -59,17 +59,37 @@ 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 _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + 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 _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 _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) + + 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": + if k[prefix_len:] == "weight": self.weight = torch.nn.Parameter(v, requires_grad=False) - elif k[len(prefix):] == "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 @@ -125,30 +145,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.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 + 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_device(item, device): if isinstance(item, torch.Tensor): return item.to(device, non_blocking=True) 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()