diff --git a/loader.py b/loader.py index ff1ba3f..fbc2716 100644 --- a/loader.py +++ b/loader.py @@ -1,11 +1,13 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import warnings +import logging import torch import gguf from .ops import GGMLTensor from .dequant import is_quantized, dequantize_tensor -IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "ltxv", "hyvid", "wan"} +IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "ltxv", "hyvid", "wan"} TXT_ARCH_LIST = {"t5", "t5encoder", "llama"} def get_orig_shape(reader, tensor_name): @@ -81,14 +83,19 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal raise ValueError(f"Unexpected architecture type in GGUF file: {arch_str!r}") if compat: - print(f"Warning: This model file is loaded in compatibility mode '{compat}' [arch:{arch_str}]") + logging.warning(f"Warning: This model file is loaded in compatibility mode '{compat}' [arch:{arch_str}]") # main loading loop state_dict = {} qtype_dict = {} for sd_key, tensor in tensors: tensor_name = tensor.name - torch_tensor = torch.from_numpy(tensor.data) # mmap + # torch_tensor = torch.from_numpy(tensor.data) # mmap + + # NOTE: line above replaced with this block to avoid persistent numpy warning about mmap + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="The given NumPy array is not writable") + torch_tensor = torch.from_numpy(tensor.data) # mmap shape = get_orig_shape(reader, tensor_name) if shape is None: @@ -109,7 +116,7 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=Fal qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1 # print loaded tensor type counts - print("gguf qtypes: " + ", ".join(f"{k} ({v})" for k, v in qtype_dict.items())) + logging.info("gguf qtypes: " + ", ".join(f"{k} ({v})" for k, v in qtype_dict.items())) # mark largest tensor for vram estimation qsd = {k:v for k,v in state_dict.items() if is_quantized(v)} @@ -177,7 +184,7 @@ def llama_permute(raw_sd, n_head, n_head_kv): def gguf_tokenizer_loader(path, temb_shape): # convert gguf tokenizer to spiece - print(f"Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") + logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") try: from sentencepiece import sentencepiece_model_pb2 as model except ImportError: @@ -190,9 +197,9 @@ def gguf_tokenizer_loader(path, temb_shape): if temb_shape == (256384, 4096): # probably UMT5 spm.trainer_spec.model_type == 1 # Unigram (do we have a T5 w/ BPE?) else: - raise NotImplementedError(f"Unknown model, can't set tokenizer!") + raise NotImplementedError("Unknown model, can't set tokenizer!") else: - raise NotImplementedError(f"Unknown model, can't set tokenizer!") + raise NotImplementedError("Unknown model, can't set tokenizer!") spm.normalizer_spec.add_dummy_prefix = get_field(reader, "tokenizer.ggml.add_space_prefix", bool) spm.normalizer_spec.remove_extra_whitespaces = get_field(reader, "tokenizer.ggml.remove_extra_whitespaces", bool) @@ -219,7 +226,7 @@ def gguf_tokenizer_loader(path, temb_shape): spm.trainer_spec.eos_id = get_field(reader, "tokenizer.ggml.eos_token_id", int) spm.trainer_spec.pad_id = get_field(reader, "tokenizer.ggml.padding_token_id", int) - print(f"Created tokenizer with vocab size of {len(spm.pieces)}") + logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") del reader return torch.ByteTensor(list(spm.SerializeToString())) @@ -231,14 +238,14 @@ def gguf_clip_loader(path): # non-standard Comfy-Org tokenizer sd["spiece_model"] = gguf_tokenizer_loader(path, sd[temb_key].shape) # TODO: dequantizing token embed here is janky but otherwise we OOM due to tensor being massive. - print(f"Dequantizing {temb_key} to prevent runtime OOM.") + logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) sd = sd_map_replace(sd, T5_SD_MAP) elif arch in {"llama"}: temb_key = "token_embd.weight" if temb_key in sd and sd[temb_key].shape != (128320, 4096): # This still works. Raise error? - print("Warning! token_embd shape may be incorrect for llama 3 model!") + logging.warning("Warning! token_embd shape may be incorrect for llama 3 model!") sd = sd_map_replace(sd, LLAMA_SD_MAP) sd = llama_permute(sd, 32, 8) # L3 else: diff --git a/nodes.py b/nodes.py index 7f5ad83..f82469a 100644 --- a/nodes.py +++ b/nodes.py @@ -99,7 +99,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): linked.append((n, m)) continue if linked: - print(f"Attempting to release mmap ({len(linked)})") + 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) @@ -192,7 +192,7 @@ CLIP_ENUM_MAP = { def get_clip_type(name): enum_name = CLIP_ENUM_MAP.get(name, None) if enum_name is None: - raise ValueError(f"Unknown CLIP model type {name}") + raise ValueError(f"Unknown CLIP model type {name}") clip_type = getattr(comfy.sd.CLIPType, CLIP_ENUM_MAP[name], None) if clip_type is None: raise ValueError(f"Unsupported CLIP model type {name} (Update ComfyUI)") @@ -227,6 +227,8 @@ class CLIPLoaderGGUF: 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 @@ -288,10 +290,34 @@ class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): clip_paths = (clip_path1, clip_path2, clip_path3) return (self.load_patcher(clip_paths, get_clip_type(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) + return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),) + NODE_CLASS_MAPPINGS = { "UnetLoaderGGUF": UnetLoaderGGUF, "CLIPLoaderGGUF": CLIPLoaderGGUF, "DualCLIPLoaderGGUF": DualCLIPLoaderGGUF, "TripleCLIPLoaderGGUF": TripleCLIPLoaderGGUF, + "QuadrupleCLIPLoaderGGUF": QuadrupleCLIPLoaderGGUF, "UnetLoaderGGUFAdvanced": UnetLoaderGGUFAdvanced, } diff --git a/ops.py b/ops.py index c74296e..0a04164 100644 --- a/ops.py +++ b/ops.py @@ -1,6 +1,7 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) import gguf import torch +import logging import comfy.ops import comfy.model_management @@ -25,12 +26,16 @@ def get_torch_compiler_disable_decorator(): from packaging import version if not chained_hasattr(torch, "compiler.disable"): + logging.info("ComfyUI-GGUF: Torch too old for torch.compile - bypassing") return dummy_decorator # torch too old elif version.parse(torch.__version__) >= version.parse("2.8"): + logging.info("ComfyUI-GGUF: Allowing full torch compile") return dummy_decorator # torch compile works if chained_hasattr(torch, "_dynamo.config.nontraceable_tensor_subclasses"): + logging.info("ComfyUI-GGUF: Allowing full torch compile (nightly)") return dummy_decorator # torch compile works, nightly before 2.8 release else: + logging.info("ComfyUI-GGUF: Partial torch compile only, consider updating pytorch") return torch.compiler.disable torch_compiler_disable = get_torch_compiler_disable_decorator() @@ -66,7 +71,7 @@ class GGMLTensor(torch.Tensor): try: return super().copy_(*args, **kwargs) except Exception as e: - print(f"ignoring 'copy_' on tensor: {e}") + logging.warning(f"ignoring 'copy_' on tensor: {e}") def new_empty(self, size, *args, **kwargs): # Intel Arc fix, ref#50 @@ -103,7 +108,7 @@ class GGMLLayer(torch.nn.Module): 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") - # NOTE: using modified load for linear due to not initializing on creation, see GGMLOps todo + # NOTE: using modified load for linear due to not initializing on creation, see GGMLOps todo if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(self, torch.nn.Linear): return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) diff --git a/tools/convert.py b/tools/convert.py index c7dd146..9212646 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -43,6 +43,18 @@ class ModelAura(ModelTemplate): ] keys_banned = ["joint_transformer_blocks.3.ff_context.out_projection.weight",] +class ModelHiDream(ModelTemplate): + arch = "hidream" + keys_detect = [ + ( + "caption_projection.0.linear.weight", + "double_stream_blocks.0.block.ff_i.shared_experts.w3.weight" + ) + ] + keys_hiprec = [ + ".ff_i.gate.weight" # nn.parameter, can't load from BF16 ver + ] + class ModelHyVid(ModelTemplate): arch = "hyvid" ndims_fix = True @@ -104,7 +116,7 @@ class ModelSD1(ModelTemplate): ] # The architectures are checked in order and the first successful match terminates the search. -arch_list = [ModelFlux, ModelSD3, ModelAura, ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1] +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1] def is_model_arch(model, state_dict): # check if model is correct diff --git a/tools/lcpp.patch b/tools/lcpp.patch index 9c522ce..17cba8d 100644 --- a/tools/lcpp.patch +++ b/tools/lcpp.patch @@ -39,10 +39,10 @@ index b16c462f..6d1568f1 100644 const int idx = gguf_find_tensor(ctx, name); if (idx < 0) { diff --git a/src/llama.cpp b/src/llama.cpp -index 24e1f1f0..d3ca4756 100644 +index 24e1f1f0..5e944837 100644 --- a/src/llama.cpp +++ b/src/llama.cpp -@@ -205,6 +205,14 @@ enum llm_arch { +@@ -205,6 +205,15 @@ enum llm_arch { LLM_ARCH_GRANITE, LLM_ARCH_GRANITE_MOE, LLM_ARCH_CHAMELEON, @@ -54,10 +54,11 @@ index 24e1f1f0..d3ca4756 100644 + LLM_ARCH_LTXV, + LLM_ARCH_HYVID, + LLM_ARCH_WAN, ++ LLM_ARCH_HIDREAM, LLM_ARCH_UNKNOWN, }; -@@ -258,6 +266,14 @@ static const std::map LLM_ARCH_NAMES = { +@@ -258,6 +267,15 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE, "granite" }, { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_CHAMELEON, "chameleon" }, @@ -69,25 +70,27 @@ index 24e1f1f0..d3ca4756 100644 + { LLM_ARCH_LTXV, "ltxv" }, + { LLM_ARCH_HYVID, "hyvid" }, + { LLM_ARCH_WAN, "wan" }, ++ { LLM_ARCH_HIDREAM, "hidream" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; -@@ -1531,6 +1547,14 @@ static const std::map> LLM_TENSOR_N +@@ -1531,6 +1549,15 @@ static const std::map> LLM_TENSOR_N { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, }, }, -+ { LLM_ARCH_FLUX, {}}, -+ { LLM_ARCH_SD1, {}}, -+ { LLM_ARCH_SDXL, {}}, -+ { LLM_ARCH_SD3, {}}, -+ { LLM_ARCH_AURA, {}}, -+ { LLM_ARCH_LTXV, {}}, -+ { LLM_ARCH_HYVID,{}}, -+ { LLM_ARCH_WAN,{}}, ++ { LLM_ARCH_FLUX, {}}, ++ { LLM_ARCH_SD1, {}}, ++ { LLM_ARCH_SDXL, {}}, ++ { LLM_ARCH_SD3, {}}, ++ { LLM_ARCH_AURA, {}}, ++ { LLM_ARCH_LTXV, {}}, ++ { LLM_ARCH_HYVID, {}}, ++ { LLM_ARCH_WAN, {}}, ++ { LLM_ARCH_HIDREAM, {}}, { LLM_ARCH_UNKNOWN, { -@@ -5403,6 +5427,22 @@ static void llm_load_hparams( +@@ -5403,6 +5430,23 @@ static void llm_load_hparams( // get general kv ml.get_key(LLM_KV_GENERAL_NAME, model.name, false); @@ -101,6 +104,7 @@ index 24e1f1f0..d3ca4756 100644 + case LLM_ARCH_LTXV: + case LLM_ARCH_HYVID: + case LLM_ARCH_WAN: ++ case LLM_ARCH_HIDREAM: + model.ftype = ml.ftype; + return; + default: @@ -110,7 +114,7 @@ index 24e1f1f0..d3ca4756 100644 // get hparams kv ml.get_key(LLM_KV_VOCAB_SIZE, hparams.n_vocab, false) || ml.get_arr_n(LLM_KV_TOKENIZER_LIST, hparams.n_vocab); -@@ -18016,6 +18056,128 @@ static void llama_tensor_dequantize_internal( +@@ -18016,6 +18060,129 @@ static void llama_tensor_dequantize_internal( workers.clear(); } @@ -178,6 +182,7 @@ index 24e1f1f0..d3ca4756 100644 + } + } else if ( // Rules for ffn + (name.find("ffn_down") != std::string::npos) || ++ ((name.find("experts.") != std::string::npos) && (name.find(".w2.weight") != std::string::npos)) || + (name.find(".ffn.2.weight") != std::string::npos) || // is this even the right way around? + (name.find(".ff.net.2.weight") != std::string::npos) + ) { @@ -239,7 +244,7 @@ index 24e1f1f0..d3ca4756 100644 static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) { const std::string name = ggml_get_name(tensor); -@@ -18513,7 +18675,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18513,7 +18680,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s if (llama_model_has_encoder(&model)) { n_attn_layer *= 3; } @@ -250,7 +255,7 @@ index 24e1f1f0..d3ca4756 100644 } size_t total_size_org = 0; -@@ -18547,6 +18711,45 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18547,6 +18716,45 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s ctx_outs[i_split] = gguf_init_empty(); } gguf_add_tensor(ctx_outs[i_split], tensor); @@ -296,7 +301,7 @@ index 24e1f1f0..d3ca4756 100644 } // Set split info if needed -@@ -18647,6 +18850,83 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18647,6 +18855,92 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s // do not quantize relative position bias (T5) quantize &= name.find("attn_rel_b.weight") == std::string::npos; @@ -372,6 +377,15 @@ index 24e1f1f0..d3ca4756 100644 + quantize &= name.find("img_emb.") == std::string::npos; + quantize &= name.find("head.") == std::string::npos; + } ++ if (model.arch == LLM_ARCH_HIDREAM) { ++ image_model = true; ++ quantize &= name.find("p_embedder.") == std::string::npos; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("x_embedder.") == std::string::npos; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ quantize &= name.find(".ff_i.gate.weight") == std::string::npos; ++ quantize &= name.find("caption_projection.") == std::string::npos; ++ } + // ignore 3D/4D tensors for image models as the code was never meant to handle these + if (image_model) { + quantize &= ggml_n_dims(tensor) == 2; @@ -380,7 +394,7 @@ index 24e1f1f0..d3ca4756 100644 enum ggml_type new_type; void * new_data; size_t new_size; -@@ -18655,6 +18935,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18655,6 +18949,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s new_type = default_type; // get more optimal quantization type based on the tensor shape, layer, etc. @@ -390,7 +404,7 @@ index 24e1f1f0..d3ca4756 100644 if (!params->pure && ggml_is_quantized(default_type)) { new_type = llama_tensor_get_type(qs, new_type, tensor, ftype); } -@@ -18664,6 +18947,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18664,6 +18961,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) { new_type = params->output_tensor_type; }