diff --git a/loader.py b/loader.py index 7cefb11..191ce38 100644 --- a/loader.py +++ b/loader.py @@ -347,7 +347,7 @@ def gguf_tokenizer_loader(path, temb_shape): if get_field(reader, "tokenizer.ggml.model", str) == "t5": if temb_shape == (256384, 4096): # probably UMT5 - spm.trainer_spec.model_type == 1 # Unigram (do we have a T5 w/ BPE?) + spm.trainer_spec.model_type = 1 # Unigram (do we have a T5 w/ BPE?) else: raise NotImplementedError("Unknown model, can't set tokenizer!") else: diff --git a/nodes.py b/nodes.py index 4683514..3a40591 100644 --- a/nodes.py +++ b/nodes.py @@ -17,20 +17,55 @@ 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=[]): + +def update_folder_names_and_paths(key, targets=None): + if targets is None: + 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]) + 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"]) + +def resolve_full_path(name, keys): + for key in keys: + path = folder_paths.get_full_path(key, name) + if path: + return path + return None + + +# ComfyUI frontend expects a config dict for dropdowns in newer builds. +def _dropdown(values, default=None): + values = list(values) + if default is None: + default = values[0] if values else "" + return (values, {"default": default}) + + +def _ensure_dropdown(spec, default=None): + if isinstance(spec, tuple): + if len(spec) >= 2 and isinstance(spec[1], dict): + return spec + if len(spec) == 1 and isinstance(spec[0], (list, tuple)): + return _dropdown(spec[0], default) + return spec + if isinstance(spec, list): + return _dropdown(spec, default) + return spec + + +# 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 @@ -43,18 +78,26 @@ def patch_weight_to_device(self, key, device_to=None, inplace_update=False): 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) + 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 + 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) + temp_weight = comfy.model_management.cast_to_device( + weight, device_to, torch.float32, copy=True + ) else: temp_weight = weight.to(torch.float32, copy=True) @@ -75,14 +118,17 @@ def unpatch_model(self, device_to=None, unpatch_weights=True): 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) - + 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] + 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) + 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) @@ -129,25 +175,28 @@ def clone(self, *args, **kwargs): 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 + 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,), - } - } + +class UnetLoaderGGUF: + @classmethod + def INPUT_TYPES(s): + unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] + return { + "required": { + "unet_name": _dropdown(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): + def load_unet( + self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None + ): ops = GGMLOps() if dequant_dtype in ("default", None): @@ -164,55 +213,76 @@ def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_de 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", {}) - + # init model + unet_path = resolve_full_path(unet_name, ("unet_gguf", "unet")) + if unet_path is None: + raise FileNotFoundError(f"Unable to find UNet file: {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)) + 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"}), + +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": _dropdown(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"], - } - } + +class CLIPLoaderGGUF: + @classmethod + def INPUT_TYPES(s): + base = nodes.CLIPLoader.INPUT_TYPES() + return { + "required": { + "clip_name": _dropdown(s.get_filename_list()), + "type": _ensure_dropdown(base["required"]["type"]), + } + } RETURN_TYPES = ("CLIP",) FUNCTION = "load_clip" CATEGORY = "bootleg" TITLE = "CLIPLoader (GGUF)" + @staticmethod + def resolve_clip_path(name): + keys = ( + ("clip_gguf", "clip") + if name.lower().endswith(".gguf") + else ("clip", "clip_gguf") + ) + return resolve_full_path(name, keys) + @classmethod def get_filename_list(s): files = [] @@ -227,97 +297,135 @@ def load_data(self, ckpt_paths): 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})") + 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 = { + clip_type=clip_type, + state_dicts=clip_data, + model_options={ "custom_operations": GGMLOps, - "initial_device": comfy.model_management.text_encoder_offload_device() + "initial_device": comfy.model_management.text_encoder_offload_device(), }, - embedding_directory = folder_paths.get_folder_paths("embeddings"), + 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) + clip_path = self.resolve_clip_path(clip_name) + if clip_path is None: + raise FileNotFoundError(f"Unable to find CLIP file: {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"], - } - } + +class DualCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + base = nodes.DualCLIPLoader.INPUT_TYPES() + file_options = _dropdown(s.get_filename_list()) + return { + "required": { + "clip_name1": file_options, + "clip_name2": file_options, + "type": _ensure_dropdown(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_path1 = self.resolve_clip_path(clip_name1) + if clip_path1 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {clip_name1}") + clip_path2 = self.resolve_clip_path(clip_name2) + if clip_path2 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {clip_name2}") clip_paths = (clip_path1, clip_path2) - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + 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, - } - } + +class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + file_options = _dropdown(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_path1 = self.resolve_clip_path(clip_name1) + if clip_path1 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {clip_name1}") + clip_path2 = self.resolve_clip_path(clip_name2) + if clip_path2 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {clip_name2}") + clip_path3 = self.resolve_clip_path(clip_name3) + if clip_path3 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {clip_name3}") clip_paths = (clip_path1, clip_path2, clip_path3) - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + 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, - } - } + +class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + file_options = _dropdown(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) + def load_clip( + self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion" + ): + clip_path1 = self.resolve_clip_path(clip_name1) + if clip_path1 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {clip_name1}") + clip_path2 = self.resolve_clip_path(clip_name2) + if clip_path2 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {clip_name2}") + clip_path3 = self.resolve_clip_path(clip_name3) + if clip_path3 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {clip_name3}") + clip_path4 = self.resolve_clip_path(clip_name4) + if clip_path4 is None: + raise FileNotFoundError(f"Unable to find CLIP file: {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) + 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, diff --git a/ops.py b/ops.py index 88a352e..2ab4df8 100644 --- a/ops.py +++ b/ops.py @@ -253,7 +253,7 @@ 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) + weight, _bias = self.cast_bias_weight(input, 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)