diff --git a/ming_image.py b/ming_image.py new file mode 100644 index 0000000..10b8f17 --- /dev/null +++ b/ming_image.py @@ -0,0 +1,190 @@ +# Ming-Image GGUF compatibility helpers. +# +# Ming-Image's Diffusers-format checkpoint stores attention Q/K/V separately. +# ComfyUI's generic conversion merges those tensors with ordinary torch tensor +# assignment. For quantized GGUF tensors, the physical packed byte width can +# differ from the logical tensor width, so that merge can fail. +# +# This helper performs the same Ming/Z-Image key mapping while preserving the +# GGML tensor type and logical shape of packed tensors. + +import logging +import torch +import comfy.utils + +from .ops import GGMLTensor +from .dequant import is_quantized + + +def _count_blocks(keys, prefix): + i = 0 + while any(key.startswith(prefix.format(i)) for key in keys): + i += 1 + return i + + +def _base_tensor(tensor): + if isinstance(tensor, GGMLTensor): + return torch.Tensor(tensor) + return tensor + + +def _merge_offset_parts(target, parts): + parts = sorted(parts, key=lambda item: item[0][1]) + + if len(parts) != 3: + raise RuntimeError( + f"Ming-Image GGUF expected 3 Q/K/V parts for {target}, got {len(parts)}" + ) + + axes = {int(item[0][0]) for item in parts} + if axes != {0}: + raise RuntimeError( + f"Ming-Image GGUF only supports axis-0 QKV merges, got {axes}" + ) + + cursor = 0 + for offset, _value in parts: + _axis, start, length = map(int, offset) + if start != cursor: + raise RuntimeError( + f"Ming-Image GGUF non-contiguous merge for {target}: " + f"expected start {cursor}, got {start}" + ) + cursor += length + + values = [item[1] for item in parts] + quantized = [is_quantized(value) for value in values] + + if any(quantized): + if not all(quantized): + raise RuntimeError( + f"Ming-Image GGUF mixed quantized/unquantized QKV for {target}" + ) + + qtypes = [getattr(value, "tensor_type", None) for value in values] + if not all(qtype == qtypes[0] for qtype in qtypes): + raise RuntimeError( + f"Ming-Image GGUF Q/K/V qtypes differ for {target}: {qtypes}" + ) + + logical_shapes = [ + tuple(int(x) for x in value.tensor_shape) for value in values + ] + + if any(len(shape) != 2 for shape in logical_shapes): + raise RuntimeError( + f"Ming-Image GGUF expected 2D Q/K/V tensors for {target}: " + f"{logical_shapes}" + ) + + input_dim = logical_shapes[0][1] + if any(shape[1] != input_dim for shape in logical_shapes): + raise RuntimeError( + f"Ming-Image GGUF Q/K/V input dims differ for {target}: " + f"{logical_shapes}" + ) + + packed = [_base_tensor(value) for value in values] + packed_tail = tuple(packed[0].size()[1:]) + + if any(tuple(value.size()[1:]) != packed_tail for value in packed): + raise RuntimeError( + f"Ming-Image GGUF packed Q/K/V row widths differ for {target}: " + f"{[tuple(value.size()) for value in packed]}" + ) + + merged_data = torch.cat(packed, dim=0) + merged_shape = torch.Size( + (sum(shape[0] for shape in logical_shapes), input_dim) + ) + + merged = GGMLTensor( + merged_data, + tensor_type=qtypes[0], + tensor_shape=merged_shape, + ) + + if any(getattr(value, "is_largest_weight", False) for value in values): + merged.is_largest_weight = True + + return merged + + return torch.cat([_base_tensor(value) for value in values], dim=0) + + +def convert_ming_image_diffusers_gguf(state_dict): + """Convert Ming-Image Diffusers keys without unpacking GGUF tensors.""" + + required = ( + "all_x_embedder.2-1.weight", + "all_final_layer.2-1.linear.weight", + "noise_refiner.0.attention.to_q.weight", + "noise_refiner.0.attention.to_k.weight", + "noise_refiner.0.attention.to_v.weight", + "layers.0.attention.to_q.weight", + ) + + if not all(key in state_dict for key in required): + return state_dict + + keys = list(state_dict.keys()) + n_layers = _count_blocks(keys, "layers.{}.") + + if n_layers <= 0: + raise RuntimeError( + "Ming-Image GGUF could not determine transformer layer count" + ) + + dim = int(state_dict["noise_refiner.0.attention.to_k.weight"].shape[0]) + + # Ming currently uses ComfyUI's Z-Image Diffusers key mapping internally. + state_dict_map = comfy.utils.z_image_to_diffusers( + {"n_layers": n_layers, "dim": dim}, + output_prefix="", + ) + + # Preserve Ming-specific keys not handled by the shared mapping. + for key in keys: + state_dict_map.setdefault(key, key) + + output = {} + pending = {} + merged_count = 0 + + for key in list(state_dict.keys()): + value = state_dict.pop(key) + target = state_dict_map.get(key, key) + + if isinstance(target, str): + output[target] = value + continue + + transform = target[2] if len(target) > 2 else (lambda x: x) + offset = target[1] + destination = target[0] + + if offset is None: + output[destination] = transform(value) + continue + + bucket = pending.setdefault(destination, []) + bucket.append((offset, transform(value))) + + if len(bucket) == 3: + output[destination] = _merge_offset_parts(destination, bucket) + del pending[destination] + merged_count += 1 + + if pending: + detail = {key: len(value) for key, value in list(pending.items())[:8]} + raise RuntimeError( + f"Ming-Image GGUF incomplete QKV groups remain: {detail}" + ) + + logging.info( + "ComfyUI-GGUF: Ming-Image quant-safe conversion complete " + f"({merged_count} packed QKV groups merged, {n_layers} main layers)" + ) + + return output diff --git a/nodes.py b/nodes.py index 4683514..dda0665 100644 --- a/nodes.py +++ b/nodes.py @@ -1,329 +1,341 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -import torch -import logging -import inspect -import collections - -import nodes -import comfy.sd -import comfy.lora -import comfy.float -import comfy.utils -import comfy.model_patcher -import comfy.model_management -import folder_paths - -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 - -def update_folder_names_and_paths(key, targets=[]): - # check for existing key - base = folder_paths.folder_names_and_paths.get(key, ([], {})) - base = base[0] if isinstance(base[0], (list, set, tuple)) else [] - # find base key & add w/ fallback, sanity check + warning - target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0]) - orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {})) - folder_paths.folder_names_and_paths[key] = (orig or base, {".gguf"}) - if base and base != orig: - logging.warning(f"Unknown file list already present on key {key}: {base}") - -# Add a custom keys for files ending in .gguf -update_folder_names_and_paths("unet_gguf", ["diffusion_models", "unet"]) -update_folder_names_and_paths("clip_gguf", ["text_encoders", "clip"]) - -class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): - patch_on_device = False - - def patch_weight_to_device(self, key, device_to=None, inplace_update=False): - if key not in self.patches: - return - weight = comfy.utils.get_attr(self.model, key) - - patches = self.patches[key] - 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) - # TODO: do we ever have legitimate duplicate patches? (i.e. patch on top of patched weight) - out_weight.patches = [(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: - temp_weight = weight.to(torch.float32, copy=True) - - out_weight = comfy.lora.calculate_weight(patches, temp_weight, key) - out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype) - - if inplace_update: - comfy.utils.copy_to_param(self.model, key, out_weight) - 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 = [] - # TODO: Find another way to not unload after patches - return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) - - - def pin_weight_to_device(self, key): - op_key = key.rsplit('.', 1)[0] - if not self.mmap_released and op_key in self.named_modules_to_munmap: - # TODO: possible to OOM, find better way to detach - self.named_modules_to_munmap[op_key].to(self.load_device).to(self.offload_device) - del self.named_modules_to_munmap[op_key] - super().pin_weight_to_device(key) - - mmap_released = False - named_modules_to_munmap = {} - - def load(self, *args, force_patch_weights=False, **kwargs): - if not self.mmap_released: - self.named_modules_to_munmap = dict(self.model.named_modules()) - - # always call `patch_weight_to_device` even for lowvram - super().load(*args, force_patch_weights=True, **kwargs) - - # make sure nothing stays linked to mmap after first load - if not self.mmap_released: - linked = [] - if kwargs.get("lowvram_model_memory", 0) > 0: - for n, m in self.named_modules_to_munmap.items(): - if hasattr(m, "weight"): - device = getattr(m.weight, "device", None) - if device == self.offload_device: - linked.append((n, m)) - continue - if hasattr(m, "bias"): - device = getattr(m.bias, "device", None) - if device == self.offload_device: - linked.append((n, m)) - continue - if linked and self.load_device != self.offload_device: - logging.info(f"Attempting to release mmap ({len(linked)})") - for n, m in linked: - # TODO: possible to OOM, find better way to detach - m.to(self.load_device).to(self.offload_device) - self.mmap_released = True - self.named_modules_to_munmap = {} - - def clone(self, *args, **kwargs): - src_cls = self.__class__ - self.__class__ = GGUFModelPatcher - n = super().clone(*args, **kwargs) - n.__class__ = GGUFModelPatcher - self.__class__ = src_cls - # GGUF specific clone values below - n.patch_on_device = getattr(self, "patch_on_device", False) - n.mmap_released = getattr(self, "mmap_released", False) - if src_cls != GGUFModelPatcher: - n.size = 0 # force recalc - return n - -class UnetLoaderGGUF: - @classmethod - def INPUT_TYPES(s): - unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] - return { - "required": { - "unet_name": (unet_names,), - } - } - - RETURN_TYPES = ("MODEL",) - FUNCTION = "load_unet" - CATEGORY = "bootleg" - TITLE = "Unet Loader (GGUF)" - - def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): - ops = GGMLOps() - - 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): - ops.Linear.patch_dtype = None - elif patch_dtype in ["target"]: - ops.Linear.patch_dtype = patch_dtype - else: - ops.Linear.patch_dtype = getattr(torch, patch_dtype) - - # init model - unet_path = folder_paths.get_full_path("unet", unet_name) - sd, extra = gguf_sd_loader(unet_path) - - kwargs = {} - valid_params = inspect.signature(comfy.sd.load_diffusion_model_state_dict).parameters - if "metadata" in valid_params: - kwargs["metadata"] = extra.get("metadata", {}) - - model = comfy.sd.load_diffusion_model_state_dict( - sd, model_options={"custom_operations": ops}, **kwargs, - ) - if model is None: - logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) - raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) - model = GGUFModelPatcher.clone(model) - model.patch_on_device = patch_on_device - return (model,) - -class UnetLoaderGGUFAdvanced(UnetLoaderGGUF): - @classmethod - def INPUT_TYPES(s): - unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] - return { - "required": { - "unet_name": (unet_names,), - "dequant_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), - "patch_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), - "patch_on_device": ("BOOLEAN", {"default": False}), - } - } - TITLE = "Unet Loader (GGUF/Advanced)" - -class CLIPLoaderGGUF: - @classmethod - def INPUT_TYPES(s): - base = nodes.CLIPLoader.INPUT_TYPES() - return { - "required": { - "clip_name": (s.get_filename_list(),), - "type": base["required"]["type"], - } - } - - RETURN_TYPES = ("CLIP",) - FUNCTION = "load_clip" - CATEGORY = "bootleg" - TITLE = "CLIPLoader (GGUF)" - - @classmethod - def get_filename_list(s): - files = [] - files += folder_paths.get_filename_list("clip") - files += folder_paths.get_filename_list("clip_gguf") - return sorted(files) - - def load_data(self, ckpt_paths): - clip_data = [] - for p in ckpt_paths: - if p.endswith(".gguf"): - sd = gguf_clip_loader(p) - else: - sd = comfy.utils.load_torch_file(p, safe_load=True) - if "scaled_fp8" in sd: # NOTE: Scaled FP8 would require different custom ops, but only one can be active - raise NotImplementedError(f"Mixing scaled FP8 with GGUF is not supported! Use regular CLIP loader or switch model(s)\n({p})") - clip_data.append(sd) - return clip_data - - def load_patcher(self, clip_paths, clip_type, clip_data): - clip = comfy.sd.load_text_encoder_state_dicts( - clip_type = clip_type, - state_dicts = clip_data, - model_options = { - "custom_operations": GGMLOps, - "initial_device": comfy.model_management.text_encoder_offload_device() - }, - embedding_directory = folder_paths.get_folder_paths("embeddings"), - ) - clip.patcher = GGUFModelPatcher.clone(clip.patcher) - return clip - - def load_clip(self, clip_name, type="stable_diffusion"): - clip_path = folder_paths.get_full_path("clip", clip_name) - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),) - -class DualCLIPLoaderGGUF(CLIPLoaderGGUF): - @classmethod - def INPUT_TYPES(s): - base = nodes.DualCLIPLoader.INPUT_TYPES() - file_options = (s.get_filename_list(), ) - return { - "required": { - "clip_name1": file_options, - "clip_name2": file_options, - "type": base["required"]["type"], - } - } - - TITLE = "DualCLIPLoader (GGUF)" - - 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_type = getattr(comfy.sd.CLIPType, type.upper(), 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": file_options, - "clip_name2": file_options, - "clip_name3": file_options, - } - } - - TITLE = "TripleCLIPLoader (GGUF)" - - def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"): - 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_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) - -class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF): - @classmethod - def INPUT_TYPES(s): - file_options = (s.get_filename_list(), ) - return { - "required": { - "clip_name1": file_options, - "clip_name2": file_options, - "clip_name3": file_options, - "clip_name4": file_options, - } - } - - TITLE = "QuadrupleCLIPLoader (GGUF)" - - def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion"): - 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_path4 = folder_paths.get_full_path("clip", clip_name4) - clip_paths = (clip_path1, clip_path2, clip_path3, clip_path4) - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) - -NODE_CLASS_MAPPINGS = { - "UnetLoaderGGUF": UnetLoaderGGUF, - "CLIPLoaderGGUF": CLIPLoaderGGUF, - "DualCLIPLoaderGGUF": DualCLIPLoaderGGUF, - "TripleCLIPLoaderGGUF": TripleCLIPLoaderGGUF, - "QuadrupleCLIPLoaderGGUF": QuadrupleCLIPLoaderGGUF, - "UnetLoaderGGUFAdvanced": UnetLoaderGGUFAdvanced, -} - +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import torch +import logging +import inspect +import collections + +import nodes +import comfy.sd +import comfy.lora +import comfy.float +import comfy.utils +import comfy.model_patcher +import comfy.model_management +import folder_paths + +from .ops import GGMLOps, move_patch_to_device +from .loader import gguf_sd_loader, gguf_clip_loader +from .ming_image import convert_ming_image_diffusers_gguf +from .dequant import is_quantized, is_torch_compatible + +def update_folder_names_and_paths(key, targets=[]): + # check for existing key + base = folder_paths.folder_names_and_paths.get(key, ([], {})) + base = base[0] if isinstance(base[0], (list, set, tuple)) else [] + # find base key & add w/ fallback, sanity check + warning + target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0]) + orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {})) + folder_paths.folder_names_and_paths[key] = (orig or base, {".gguf"}) + if base and base != orig: + logging.warning(f"Unknown file list already present on key {key}: {base}") + +# Add a custom keys for files ending in .gguf +update_folder_names_and_paths("unet_gguf", ["diffusion_models", "unet"]) +update_folder_names_and_paths("clip_gguf", ["text_encoders", "clip"]) + +class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): + patch_on_device = False + + def patch_weight_to_device(self, key, device_to=None, inplace_update=False): + if key not in self.patches: + return + weight = comfy.utils.get_attr(self.model, key) + + patches = self.patches[key] + 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) + # TODO: do we ever have legitimate duplicate patches? (i.e. patch on top of patched weight) + out_weight.patches = [(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: + temp_weight = weight.to(torch.float32, copy=True) + + out_weight = comfy.lora.calculate_weight(patches, temp_weight, key) + out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype) + + if inplace_update: + comfy.utils.copy_to_param(self.model, key, out_weight) + 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 = [] + # TODO: Find another way to not unload after patches + return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) + + + def pin_weight_to_device(self, key): + op_key = key.rsplit('.', 1)[0] + if not self.mmap_released and op_key in self.named_modules_to_munmap: + # TODO: possible to OOM, find better way to detach + self.named_modules_to_munmap[op_key].to(self.load_device).to(self.offload_device) + del self.named_modules_to_munmap[op_key] + super().pin_weight_to_device(key) + + mmap_released = False + named_modules_to_munmap = {} + + def load(self, *args, force_patch_weights=False, **kwargs): + if not self.mmap_released: + self.named_modules_to_munmap = dict(self.model.named_modules()) + + # always call `patch_weight_to_device` even for lowvram + super().load(*args, force_patch_weights=True, **kwargs) + + # make sure nothing stays linked to mmap after first load + if not self.mmap_released: + linked = [] + if kwargs.get("lowvram_model_memory", 0) > 0: + for n, m in self.named_modules_to_munmap.items(): + if hasattr(m, "weight"): + device = getattr(m.weight, "device", None) + if device == self.offload_device: + linked.append((n, m)) + continue + if hasattr(m, "bias"): + device = getattr(m.bias, "device", None) + if device == self.offload_device: + linked.append((n, m)) + continue + if linked and self.load_device != self.offload_device: + logging.info(f"Attempting to release mmap ({len(linked)})") + for n, m in linked: + # TODO: possible to OOM, find better way to detach + m.to(self.load_device).to(self.offload_device) + self.mmap_released = True + self.named_modules_to_munmap = {} + + def clone(self, *args, **kwargs): + src_cls = self.__class__ + self.__class__ = GGUFModelPatcher + n = super().clone(*args, **kwargs) + n.__class__ = GGUFModelPatcher + self.__class__ = src_cls + # GGUF specific clone values below + n.patch_on_device = getattr(self, "patch_on_device", False) + n.mmap_released = getattr(self, "mmap_released", False) + if src_cls != GGUFModelPatcher: + n.size = 0 # force recalc + return n + +class UnetLoaderGGUF: + @classmethod + def INPUT_TYPES(s): + unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] + return { + "required": { + "unet_name": (unet_names,), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_unet" + CATEGORY = "bootleg" + TITLE = "Unet Loader (GGUF)" + + def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): + ops = GGMLOps() + + 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): + ops.Linear.patch_dtype = None + elif patch_dtype in ["target"]: + ops.Linear.patch_dtype = patch_dtype + else: + ops.Linear.patch_dtype = getattr(torch, patch_dtype) + + # init model + unet_path = folder_paths.get_full_path("unet", unet_name) + sd, extra = gguf_sd_loader(unet_path) + + # Ming-Image stores split Q/K/V tensors. For quantized GGUF tensors, + # merge them while preserving packed storage and logical tensor shapes. + if ( + "all_x_embedder.2-1.weight" in sd + and "all_final_layer.2-1.linear.weight" in sd + and "noise_refiner.0.attention.to_q.weight" in sd + and "layers.0.attention.to_q.weight" in sd + and any(is_quantized(value) for value in sd.values()) + ): + sd = convert_ming_image_diffusers_gguf(sd) + + kwargs = {} + valid_params = inspect.signature(comfy.sd.load_diffusion_model_state_dict).parameters + if "metadata" in valid_params: + kwargs["metadata"] = extra.get("metadata", {}) + + model = comfy.sd.load_diffusion_model_state_dict( + sd, model_options={"custom_operations": ops}, **kwargs, + ) + if model is None: + logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) + raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) + model = GGUFModelPatcher.clone(model) + model.patch_on_device = patch_on_device + return (model,) + +class UnetLoaderGGUFAdvanced(UnetLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] + return { + "required": { + "unet_name": (unet_names,), + "dequant_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), + "patch_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), + "patch_on_device": ("BOOLEAN", {"default": False}), + } + } + TITLE = "Unet Loader (GGUF/Advanced)" + +class CLIPLoaderGGUF: + @classmethod + def INPUT_TYPES(s): + base = nodes.CLIPLoader.INPUT_TYPES() + return { + "required": { + "clip_name": (s.get_filename_list(),), + "type": base["required"]["type"], + } + } + + RETURN_TYPES = ("CLIP",) + FUNCTION = "load_clip" + CATEGORY = "bootleg" + TITLE = "CLIPLoader (GGUF)" + + @classmethod + def get_filename_list(s): + files = [] + files += folder_paths.get_filename_list("clip") + files += folder_paths.get_filename_list("clip_gguf") + return sorted(files) + + def load_data(self, ckpt_paths): + clip_data = [] + for p in ckpt_paths: + if p.endswith(".gguf"): + sd = gguf_clip_loader(p) + else: + sd = comfy.utils.load_torch_file(p, safe_load=True) + if "scaled_fp8" in sd: # NOTE: Scaled FP8 would require different custom ops, but only one can be active + raise NotImplementedError(f"Mixing scaled FP8 with GGUF is not supported! Use regular CLIP loader or switch model(s)\n({p})") + clip_data.append(sd) + return clip_data + + def load_patcher(self, clip_paths, clip_type, clip_data): + clip = comfy.sd.load_text_encoder_state_dicts( + clip_type = clip_type, + state_dicts = clip_data, + model_options = { + "custom_operations": GGMLOps, + "initial_device": comfy.model_management.text_encoder_offload_device() + }, + embedding_directory = folder_paths.get_folder_paths("embeddings"), + ) + clip.patcher = GGUFModelPatcher.clone(clip.patcher) + return clip + + def load_clip(self, clip_name, type="stable_diffusion"): + clip_path = folder_paths.get_full_path("clip", clip_name) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),) + +class DualCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + base = nodes.DualCLIPLoader.INPUT_TYPES() + file_options = (s.get_filename_list(), ) + return { + "required": { + "clip_name1": file_options, + "clip_name2": file_options, + "type": base["required"]["type"], + } + } + + TITLE = "DualCLIPLoader (GGUF)" + + 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_type = getattr(comfy.sd.CLIPType, type.upper(), 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": file_options, + "clip_name2": file_options, + "clip_name3": file_options, + } + } + + TITLE = "TripleCLIPLoader (GGUF)" + + def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"): + 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_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + +class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + file_options = (s.get_filename_list(), ) + return { + "required": { + "clip_name1": file_options, + "clip_name2": file_options, + "clip_name3": file_options, + "clip_name4": file_options, + } + } + + TITLE = "QuadrupleCLIPLoader (GGUF)" + + def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion"): + 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_path4 = folder_paths.get_full_path("clip", clip_name4) + clip_paths = (clip_path1, clip_path2, clip_path3, clip_path4) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + +NODE_CLASS_MAPPINGS = { + "UnetLoaderGGUF": UnetLoaderGGUF, + "CLIPLoaderGGUF": CLIPLoaderGGUF, + "DualCLIPLoaderGGUF": DualCLIPLoaderGGUF, + "TripleCLIPLoaderGGUF": TripleCLIPLoaderGGUF, + "QuadrupleCLIPLoaderGGUF": QuadrupleCLIPLoaderGGUF, + "UnetLoaderGGUFAdvanced": UnetLoaderGGUFAdvanced, +} +