From cb1585cdcb95f081caf299d24ca8d7c19d2ff9ab Mon Sep 17 00:00:00 2001 From: molbal Date: Tue, 9 Jun 2026 20:43:23 +0200 Subject: [PATCH 1/6] Add Ideogram architecture detection in convert.py --- tools/convert.py | 739 ++++++++++++++++++++++++----------------------- 1 file changed, 374 insertions(+), 365 deletions(-) diff --git a/tools/convert.py b/tools/convert.py index 5029c87..0807e07 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -1,365 +1,374 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -import os -import gguf -import torch -import logging -import argparse -from tqdm import tqdm -from safetensors.torch import load_file, save_file - -QUANTIZATION_THRESHOLD = 1024 -REARRANGE_THRESHOLD = 512 -MAX_TENSOR_NAME_LENGTH = 127 -MAX_TENSOR_DIMS = 4 - -class ModelTemplate: - arch = "invalid" # string describing architecture - shape_fix = False # whether to reshape tensors - keys_detect = [] # list of lists to match in state dict - keys_banned = [] # list of keys that should mark model as invalid for conversion - keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason - keys_ignore = [] # list of strings to ignore keys by when found - - def handle_nd_tensor(self, key, data): - raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})") - -class ModelFlux(ModelTemplate): - arch = "flux" - keys_detect = [ - ("transformer_blocks.0.attn.norm_added_k.weight",), - ("double_blocks.0.img_attn.proj.weight",), - ] - keys_banned = ["transformer_blocks.0.attn.norm_added_k.weight",] - -class ModelSD3(ModelTemplate): - arch = "sd3" - keys_detect = [ - ("transformer_blocks.0.attn.add_q_proj.weight",), - ("joint_blocks.0.x_block.attn.qkv.weight",), - ] - keys_banned = ["transformer_blocks.0.attn.add_q_proj.weight",] - -class ModelAura(ModelTemplate): - arch = "aura" - keys_detect = [ - ("double_layers.3.modX.1.weight",), - ("joint_transformer_blocks.3.ff_context.out_projection.weight",), - ] - 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 = [ - # nn.parameter, can't load from BF16 ver - ".ff_i.gate.weight", - "img_emb.emb_pos" - ] - -class CosmosPredict2(ModelTemplate): - arch = "cosmos" - keys_detect = [ - ( - "blocks.0.mlp.layer1.weight", - "blocks.0.adaln_modulation_cross_attn.1.weight", - ) - ] - keys_hiprec = ["pos_embedder"] - keys_ignore = ["_extra_state", "accum_"] - -class ModelHyVid(ModelTemplate): - arch = "hyvid" - keys_detect = [ - ( - "double_blocks.0.img_attn_proj.weight", - "txt_in.individual_token_refiner.blocks.1.self_attn_qkv.weight", - ) - ] - - def handle_nd_tensor(self, key, data): - # hacky but don't have any better ideas - path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? - if os.path.isfile(path): - raise RuntimeError(f"5D tensor fix file already exists! {path}") - fsd = {key: torch.from_numpy(data)} - tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") - save_file(fsd, path) - -class ModelWan(ModelHyVid): - arch = "wan" - keys_detect = [ - ( - "blocks.0.self_attn.norm_q.weight", - "text_embedding.2.weight", - "head.modulation", - ) - ] - keys_hiprec = [ - ".modulation" # nn.parameter, can't load from BF16 ver - ] - -class ModelLTXV(ModelTemplate): - arch = "ltxv" - keys_detect = [ - ( - "adaln_single.emb.timestep_embedder.linear_2.weight", - "transformer_blocks.27.scale_shift_table", - "caption_projection.linear_2.weight", - ) - ] - keys_hiprec = [ - "scale_shift_table" # nn.parameter, can't load from BF16 base quant - ] - -class ModelSDXL(ModelTemplate): - arch = "sdxl" - shape_fix = True - keys_detect = [ - ("down_blocks.0.downsamplers.0.conv.weight", "add_embedding.linear_1.weight",), - ( - "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", - "output_blocks.2.2.conv.weight", "output_blocks.5.2.conv.weight", - ), # Non-diffusers - ("label_emb.0.0.weight",), - ] - -class ModelSD1(ModelTemplate): - arch = "sd1" - shape_fix = True - keys_detect = [ - ("down_blocks.0.downsamplers.0.conv.weight",), - ( - "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", "input_blocks.9.0.op.weight", - "output_blocks.2.1.conv.weight", "output_blocks.5.2.conv.weight", "output_blocks.8.2.conv.weight" - ), # Non-diffusers - ] - -class ModelLumina2(ModelTemplate): - arch = "lumina2" - keys_detect = [ - ("cap_embedder.1.weight", "context_refiner.0.attention.qkv.weight") - ] - -arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, - ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2] - -def is_model_arch(model, state_dict): - # check if model is correct - matched = False - invalid = False - for match_list in model.keys_detect: - if all(key in state_dict for key in match_list): - matched = True - invalid = any(key in state_dict for key in model.keys_banned) - break - assert not invalid, "Model architecture not allowed for conversion! (i.e. reference VS diffusers format)" - return matched - -def detect_arch(state_dict): - model_arch = None - for arch in arch_list: - if is_model_arch(arch, state_dict): - model_arch = arch() - break - assert model_arch is not None, "Unknown model architecture!" - return model_arch - -def parse_args(): - parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET") - parser.add_argument("--src", required=True, help="Source model ckpt file.") - parser.add_argument("--dst", help="Output unet gguf file.") - args = parser.parse_args() - - if not os.path.isfile(args.src): - parser.error("No input provided!") - - return args - -def strip_prefix(state_dict): - # prefix for mixed state dict - prefix = None - for pfx in ["model.diffusion_model.", "model."]: - if any([x.startswith(pfx) for x in state_dict.keys()]): - prefix = pfx - break - - # prefix for uniform state dict - if prefix is None: - for pfx in ["net."]: - if all([x.startswith(pfx) for x in state_dict.keys()]): - prefix = pfx - break - - # strip prefix if found - if prefix is not None: - logging.info(f"State dict prefix found: '{prefix}'") - sd = {} - for k, v in state_dict.items(): - if prefix not in k: - continue - k = k.replace(prefix, "") - sd[k] = v - else: - logging.debug("State dict has no prefix") - sd = state_dict - - return sd - -def load_state_dict(path): - if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]): - state_dict = torch.load(path, map_location="cpu", weights_only=True) - for subkey in ["model", "module"]: - if subkey in state_dict: - state_dict = state_dict[subkey] - break - if len(state_dict) < 20: - raise RuntimeError(f"pt subkey load failed: {state_dict.keys()}") - else: - state_dict = load_file(path) - - return strip_prefix(state_dict) - -def handle_tensors(writer, state_dict, model_arch): - name_lengths = tuple(sorted( - ((key, len(key)) for key in state_dict.keys()), - key=lambda item: item[1], - reverse=True, - )) - if not name_lengths: - return - max_name_len = name_lengths[0][1] - if max_name_len > MAX_TENSOR_NAME_LENGTH: - bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH) - raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}") - for key, data in tqdm(state_dict.items()): - old_dtype = data.dtype - - if any(x in key for x in model_arch.keys_ignore): - tqdm.write(f"Filtering ignored key: '{key}'") - continue - - 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() - - n_dims = len(data.shape) - data_shape = data.shape - if old_dtype == torch.bfloat16: - data_qtype = gguf.GGMLQuantizationType.BF16 - # elif old_dtype == torch.float32: - # data_qtype = gguf.GGMLQuantizationType.F32 - else: - data_qtype = gguf.GGMLQuantizationType.F16 - - # The max no. of dimensions that can be handled by the quantization code is 4 - if len(data.shape) > MAX_TENSOR_DIMS: - model_arch.handle_nd_tensor(key, data) - continue # needs to be added back later - - # get number of parameters (AKA elements) in this tensor - n_params = 1 - for dim_size in data_shape: - n_params *= dim_size - - if old_dtype in (torch.float32, torch.bfloat16): - if n_dims == 1: - # one-dimensional tensors should be kept in F32 - # also speeds up inference due to not dequantizing - data_qtype = gguf.GGMLQuantizationType.F32 - - elif n_params <= QUANTIZATION_THRESHOLD: - # very small tensors - data_qtype = gguf.GGMLQuantizationType.F32 - - elif any(x in key for x in model_arch.keys_hiprec): - # tensors that require max precision - data_qtype = gguf.GGMLQuantizationType.F32 - - if (model_arch.shape_fix # NEVER reshape for models such as flux - and n_dims > 1 # Skip one-dimensional tensors - and n_params >= REARRANGE_THRESHOLD # Only rearrange tensors meeting the size requirement - and (n_params / 256).is_integer() # Rearranging only makes sense if total elements is divisible by 256 - and not (data.shape[-1] / 256).is_integer() # Only need to rearrange if the last dimension is not divisible by 256 - ): - orig_shape = data.shape - data = data.reshape(n_params // 256, 256) - writer.add_array(f"comfy.gguf.orig_shape.{key}", tuple(int(dim) for dim in orig_shape)) - - try: - data = gguf.quants.quantize(data, data_qtype) - except (AttributeError, gguf.QuantError) as e: - tqdm.write(f"falling back to F16: {e}") - data_qtype = gguf.GGMLQuantizationType.F16 - data = gguf.quants.quantize(data, data_qtype) - - new_name = key # do we need to rename? - - shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" - tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") - - writer.add_tensor(new_name, data, raw_dtype=data_qtype) - -def convert_file(path, dst_path=None, interact=True, overwrite=False): - # load & run model detection logic - state_dict = load_state_dict(path) - model_arch = detect_arch(state_dict) - logging.info(f"* Architecture detected from input: {model_arch.arch}") - - # detect & set dtype for output file - dtypes = [x.dtype for x in state_dict.values()] - dtypes = {x:dtypes.count(x) for x in set(dtypes)} - main_dtype = max(dtypes, key=dtypes.get) - - if main_dtype == torch.bfloat16: - ftype_name = "BF16" - ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16 - # elif main_dtype == torch.float32: - # ftype_name = "F32" - # ftype_gguf = None - else: - ftype_name = "F16" - ftype_gguf = gguf.LlamaFileType.MOSTLY_F16 - - if dst_path is None: - dst_path = f"{os.path.splitext(path)[0]}-{ftype_name}.gguf" - elif "{ftype}" in dst_path: # lcpp logic - dst_path = dst_path.replace("{ftype}", ftype_name) - - if os.path.isfile(dst_path) and not overwrite: - if interact: - input("Output exists enter to continue or ctrl+c to abort!") - else: - raise OSError("Output exists and overwriting is disabled!") - - # handle actual file - writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) - writer.add_quantization_version(gguf.GGML_QUANT_VERSION) - if ftype_gguf is not None: - writer.add_file_type(ftype_gguf) - - handle_tensors(writer, state_dict, model_arch) - writer.write_header_to_file(path=dst_path) - writer.write_kv_data_to_file() - writer.write_tensors_to_file(progress=True) - writer.close() - - fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors" - if os.path.isfile(fix): - logging.warning(f"\n### Warning! Fix file found at '{fix}'") - logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.") - - return dst_path, model_arch - -if __name__ == "__main__": - args = parse_args() - convert_file(args.src, args.dst) - +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import os +import gguf +import torch +import logging +import argparse +from tqdm import tqdm +from safetensors.torch import load_file, save_file + +QUANTIZATION_THRESHOLD = 1024 +REARRANGE_THRESHOLD = 512 +MAX_TENSOR_NAME_LENGTH = 127 +MAX_TENSOR_DIMS = 4 + +class ModelTemplate: + arch = "invalid" # string describing architecture + shape_fix = False # whether to reshape tensors + keys_detect = [] # list of lists to match in state dict + keys_banned = [] # list of keys that should mark model as invalid for conversion + keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason + keys_ignore = [] # list of strings to ignore keys by when found + + def handle_nd_tensor(self, key, data): + raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})") + +class ModelFlux(ModelTemplate): + arch = "flux" + keys_detect = [ + ("transformer_blocks.0.attn.norm_added_k.weight",), + ("double_blocks.0.img_attn.proj.weight",), + ] + keys_banned = ["transformer_blocks.0.attn.norm_added_k.weight",] + +class ModelSD3(ModelTemplate): + arch = "sd3" + keys_detect = [ + ("transformer_blocks.0.attn.add_q_proj.weight",), + ("joint_blocks.0.x_block.attn.qkv.weight",), + ] + keys_banned = ["transformer_blocks.0.attn.add_q_proj.weight",] + +class ModelAura(ModelTemplate): + arch = "aura" + keys_detect = [ + ("double_layers.3.modX.1.weight",), + ("joint_transformer_blocks.3.ff_context.out_projection.weight",), + ] + 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 = [ + # nn.parameter, can't load from BF16 ver + ".ff_i.gate.weight", + "img_emb.emb_pos" + ] + +class CosmosPredict2(ModelTemplate): + arch = "cosmos" + keys_detect = [ + ( + "blocks.0.mlp.layer1.weight", + "blocks.0.adaln_modulation_cross_attn.1.weight", + ) + ] + keys_hiprec = ["pos_embedder"] + keys_ignore = ["_extra_state", "accum_"] + +class ModelHyVid(ModelTemplate): + arch = "hyvid" + keys_detect = [ + ( + "double_blocks.0.img_attn_proj.weight", + "txt_in.individual_token_refiner.blocks.1.self_attn_qkv.weight", + ) + ] + + def handle_nd_tensor(self, key, data): + # hacky but don't have any better ideas + path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? + if os.path.isfile(path): + raise RuntimeError(f"5D tensor fix file already exists! {path}") + fsd = {key: torch.from_numpy(data)} + tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") + save_file(fsd, path) + +class ModelWan(ModelHyVid): + arch = "wan" + keys_detect = [ + ( + "blocks.0.self_attn.norm_q.weight", + "text_embedding.2.weight", + "head.modulation", + ) + ] + keys_hiprec = [ + ".modulation" # nn.parameter, can't load from BF16 ver + ] + +class ModelLTXV(ModelTemplate): + arch = "ltxv" + keys_detect = [ + ( + "adaln_single.emb.timestep_embedder.linear_2.weight", + "transformer_blocks.27.scale_shift_table", + "caption_projection.linear_2.weight", + ) + ] + keys_hiprec = [ + "scale_shift_table" # nn.parameter, can't load from BF16 base quant + ] + +class ModelSDXL(ModelTemplate): + arch = "sdxl" + shape_fix = True + keys_detect = [ + ("down_blocks.0.downsamplers.0.conv.weight", "add_embedding.linear_1.weight",), + ( + "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", + "output_blocks.2.2.conv.weight", "output_blocks.5.2.conv.weight", + ), # Non-diffusers + ("label_emb.0.0.weight",), + ] + +class ModelSD1(ModelTemplate): + arch = "sd1" + shape_fix = True + keys_detect = [ + ("down_blocks.0.downsamplers.0.conv.weight",), + ( + "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", "input_blocks.9.0.op.weight", + "output_blocks.2.1.conv.weight", "output_blocks.5.2.conv.weight", "output_blocks.8.2.conv.weight" + ), # Non-diffusers + ] + +class ModelLumina2(ModelTemplate): + arch = "lumina2" + keys_detect = [ + ("cap_embedder.1.weight", "context_refiner.0.attention.qkv.weight") + ] + +class ModelIdeogram(ModelTemplate): + arch = "ideogram" + keys_detect = [ + ( + "t_embedding.mlp_in.weight", + "layers.0.attention.qkv.weight", + "final_layer.linear.weight", + ) + ] + +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, + ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2, ModelIdeogram] + +def is_model_arch(model, state_dict): + # check if model is correct + matched = False + invalid = False + for match_list in model.keys_detect: + if all(key in state_dict for key in match_list): + matched = True + invalid = any(key in state_dict for key in model.keys_banned) + break + assert not invalid, "Model architecture not allowed for conversion! (i.e. reference VS diffusers format)" + return matched + +def detect_arch(state_dict): + model_arch = None + for arch in arch_list: + if is_model_arch(arch, state_dict): + model_arch = arch() + break + assert model_arch is not None, "Unknown model architecture!" + return model_arch + +def parse_args(): + parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET") + parser.add_argument("--src", required=True, help="Source model ckpt file.") + parser.add_argument("--dst", help="Output unet gguf file.") + args = parser.parse_args() + + if not os.path.isfile(args.src): + parser.error("No input provided!") + + return args + +def strip_prefix(state_dict): + # prefix for mixed state dict + prefix = None + for pfx in ["model.diffusion_model.", "model."]: + if any([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + + # prefix for uniform state dict + if prefix is None: + for pfx in ["net."]: + if all([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + + # strip prefix if found + if prefix is not None: + logging.info(f"State dict prefix found: '{prefix}'") + sd = {} + for k, v in state_dict.items(): + if prefix not in k: + continue + k = k.replace(prefix, "") + sd[k] = v + else: + logging.debug("State dict has no prefix") + sd = state_dict + + return sd + +def load_state_dict(path): + if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]): + state_dict = torch.load(path, map_location="cpu", weights_only=True) + for subkey in ["model", "module"]: + if subkey in state_dict: + state_dict = state_dict[subkey] + break + if len(state_dict) < 20: + raise RuntimeError(f"pt subkey load failed: {state_dict.keys()}") + else: + state_dict = load_file(path) + + return strip_prefix(state_dict) + +def handle_tensors(writer, state_dict, model_arch): + name_lengths = tuple(sorted( + ((key, len(key)) for key in state_dict.keys()), + key=lambda item: item[1], + reverse=True, + )) + if not name_lengths: + return + max_name_len = name_lengths[0][1] + if max_name_len > MAX_TENSOR_NAME_LENGTH: + bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH) + raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}") + for key, data in tqdm(state_dict.items()): + old_dtype = data.dtype + + if any(x in key for x in model_arch.keys_ignore): + tqdm.write(f"Filtering ignored key: '{key}'") + continue + + 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() + + n_dims = len(data.shape) + data_shape = data.shape + if old_dtype == torch.bfloat16: + data_qtype = gguf.GGMLQuantizationType.BF16 + # elif old_dtype == torch.float32: + # data_qtype = gguf.GGMLQuantizationType.F32 + else: + data_qtype = gguf.GGMLQuantizationType.F16 + + # The max no. of dimensions that can be handled by the quantization code is 4 + if len(data.shape) > MAX_TENSOR_DIMS: + model_arch.handle_nd_tensor(key, data) + continue # needs to be added back later + + # get number of parameters (AKA elements) in this tensor + n_params = 1 + for dim_size in data_shape: + n_params *= dim_size + + if old_dtype in (torch.float32, torch.bfloat16): + if n_dims == 1: + # one-dimensional tensors should be kept in F32 + # also speeds up inference due to not dequantizing + data_qtype = gguf.GGMLQuantizationType.F32 + + elif n_params <= QUANTIZATION_THRESHOLD: + # very small tensors + data_qtype = gguf.GGMLQuantizationType.F32 + + elif any(x in key for x in model_arch.keys_hiprec): + # tensors that require max precision + data_qtype = gguf.GGMLQuantizationType.F32 + + if (model_arch.shape_fix # NEVER reshape for models such as flux + and n_dims > 1 # Skip one-dimensional tensors + and n_params >= REARRANGE_THRESHOLD # Only rearrange tensors meeting the size requirement + and (n_params / 256).is_integer() # Rearranging only makes sense if total elements is divisible by 256 + and not (data.shape[-1] / 256).is_integer() # Only need to rearrange if the last dimension is not divisible by 256 + ): + orig_shape = data.shape + data = data.reshape(n_params // 256, 256) + writer.add_array(f"comfy.gguf.orig_shape.{key}", tuple(int(dim) for dim in orig_shape)) + + try: + data = gguf.quants.quantize(data, data_qtype) + except (AttributeError, gguf.QuantError) as e: + tqdm.write(f"falling back to F16: {e}") + data_qtype = gguf.GGMLQuantizationType.F16 + data = gguf.quants.quantize(data, data_qtype) + + new_name = key # do we need to rename? + + shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" + tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") + + writer.add_tensor(new_name, data, raw_dtype=data_qtype) + +def convert_file(path, dst_path=None, interact=True, overwrite=False): + # load & run model detection logic + state_dict = load_state_dict(path) + model_arch = detect_arch(state_dict) + logging.info(f"* Architecture detected from input: {model_arch.arch}") + + # detect & set dtype for output file + dtypes = [x.dtype for x in state_dict.values()] + dtypes = {x:dtypes.count(x) for x in set(dtypes)} + main_dtype = max(dtypes, key=dtypes.get) + + if main_dtype == torch.bfloat16: + ftype_name = "BF16" + ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16 + # elif main_dtype == torch.float32: + # ftype_name = "F32" + # ftype_gguf = None + else: + ftype_name = "F16" + ftype_gguf = gguf.LlamaFileType.MOSTLY_F16 + + if dst_path is None: + dst_path = f"{os.path.splitext(path)[0]}-{ftype_name}.gguf" + elif "{ftype}" in dst_path: # lcpp logic + dst_path = dst_path.replace("{ftype}", ftype_name) + + if os.path.isfile(dst_path) and not overwrite: + if interact: + input("Output exists enter to continue or ctrl+c to abort!") + else: + raise OSError("Output exists and overwriting is disabled!") + + # handle actual file + writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) + writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + if ftype_gguf is not None: + writer.add_file_type(ftype_gguf) + + handle_tensors(writer, state_dict, model_arch) + writer.write_header_to_file(path=dst_path) + writer.write_kv_data_to_file() + writer.write_tensors_to_file(progress=True) + writer.close() + + fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors" + if os.path.isfile(fix): + logging.warning(f"\n### Warning! Fix file found at '{fix}'") + logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.") + + return dst_path, model_arch + +if __name__ == "__main__": + args = parse_args() + convert_file(args.src, args.dst) From db3b45ea23408884baee2cf1a0d99824578ebcb3 Mon Sep 17 00:00:00 2001 From: molbal Date: Tue, 9 Jun 2026 20:44:34 +0200 Subject: [PATCH 2/6] Add Ideogram to IMG_ARCH_LIST in loader.py --- loader.py | 1012 ++++++++++++++++++++++++++--------------------------- 1 file changed, 506 insertions(+), 506 deletions(-) diff --git a/loader.py b/loader.py index 7cefb11..8d2ae09 100644 --- a/loader.py +++ b/loader.py @@ -1,506 +1,506 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -import warnings -import logging -import torch -import gguf -import re -import os - -from .ops import GGMLTensor -from .dequant import is_quantized, dequantize_tensor - -IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "cosmos", "ltxv", "hyvid", "wan", "lumina2", "qwen_image"} -TXT_ARCH_LIST = {"t5", "t5encoder", "llama", "qwen2vl", "qwen3", "qwen3vl", "gemma3"} -VIS_TYPE_LIST = {"clip-vision", "mmproj"} - -def get_orig_shape(reader, tensor_name): - field_key = f"comfy.gguf.orig_shape.{tensor_name}" - field = reader.get_field(field_key) - if field is None: - return None - # Has original shape metadata, so we try to decode it. - if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32: - 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 get_field(reader, field_name, field_type): - field = reader.get_field(field_name) - if field is None: - return None - elif field_type == str: - # extra check here as this is used for checking arch string - if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING: - raise TypeError(f"Bad type for GGUF {field_name} key: expected string, got {field.types!r}") - return str(field.parts[field.data[-1]], encoding="utf-8") - elif field_type in [int, float, bool]: - return field_type(field.parts[field.data[-1]].item()) - else: - raise TypeError(f"Unknown field type {field_type}") - -def get_list_field(reader, field_name, field_type): - field = reader.get_field(field_name) - if field is None: - return None - elif field_type == str: - return tuple(str(field.parts[part_idx], encoding="utf-8") for part_idx in field.data) - elif field_type in [int, float, bool]: - return tuple(field_type(field.parts[part_idx][0]) for part_idx in field.data) - else: - raise TypeError(f"Unknown field type {field_type}") - -def get_gguf_metadata(reader): - """Extract all simple metadata fields like safetensors""" - metadata = {} - for field_name in reader.fields: - try: - field = reader.get_field(field_name) - if len(field.types) == 1: # Simple scalar fields only - if field.types[0] == gguf.GGUFValueType.STRING: - metadata[field_name] = str(field.parts[field.data[-1]], "utf-8") - elif field.types[0] == gguf.GGUFValueType.INT32: - metadata[field_name] = int(field.parts[field.data[-1]]) - elif field.types[0] == gguf.GGUFValueType.F32: - metadata[field_name] = float(field.parts[field.data[-1]]) - elif field.types[0] == gguf.GGUFValueType.BOOL: - metadata[field_name] = bool(field.parts[field.data[-1]]) - except: - continue - return metadata - -def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=False): - """ - Read state dict as fake tensors - """ - reader = gguf.GGUFReader(path) - - # 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: - 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 = get_field(reader, "general.architecture", str) - type_str = get_field(reader, "general.type", str) - if arch_str in [None, "pig", "cow"]: - if is_text_model: - raise ValueError(f"This gguf file is incompatible with llama.cpp!\nConsider using safetensors or a compatible gguf file\n({path})") - compat = "sd.cpp" if arch_str is None else arch_str - # import here to avoid changes to convert.py breaking regular models - from .tools.convert import detect_arch - try: - arch_str = detect_arch(set(val[0] for val in tensors)).arch - except Exception as e: - raise ValueError(f"This model is not currently supported - ({e})") - elif arch_str not in TXT_ARCH_LIST and is_text_model: - if type_str not in VIS_TYPE_LIST: - raise ValueError(f"Unexpected text model architecture type in GGUF file: {arch_str!r}") - elif arch_str not in IMG_ARCH_LIST and not is_text_model: - raise ValueError(f"Unexpected architecture type in GGUF file: {arch_str!r}") - - if compat: - logging.warning(f"Warning: This gguf 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 - - # 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: - shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) - # 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) - state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) - - # 1D tensors shouldn't be quantized, this is a fix for BF16 - if len(shape) <= 1 and tensor.tensor_type == gguf.GGMLQuantizationType.BF16: - state_dict[sd_key] = dequantize_tensor(state_dict[sd_key], dtype=torch.float32) - - # keep track of loaded tensor types - tensor_type_str = getattr(tensor.tensor_type, "name", repr(tensor.tensor_type)) - qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1 - - # print loaded tensor type counts - 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)} - if len(qsd) > 0: - max_key = max(qsd.keys(), key=lambda k: qsd[k].numel()) - state_dict[max_key].is_largest_weight = True - - # extra info to return - extra = { - "arch_str": arch_str, - "metadata": get_gguf_metadata(reader) - } - return (state_dict, extra) - -# for remapping llama.cpp -> original key names -T5_SD_MAP = { - "enc.": "encoder.", - ".blk.": ".block.", - "token_embd": "shared", - "output_norm": "final_layer_norm", - "attn_q": "layer.0.SelfAttention.q", - "attn_k": "layer.0.SelfAttention.k", - "attn_v": "layer.0.SelfAttention.v", - "attn_o": "layer.0.SelfAttention.o", - "attn_norm": "layer.0.layer_norm", - "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias", - "ffn_up": "layer.1.DenseReluDense.wi_1", - "ffn_down": "layer.1.DenseReluDense.wo", - "ffn_gate": "layer.1.DenseReluDense.wi_0", - "ffn_norm": "layer.1.layer_norm", -} - -LLAMA_SD_MAP = { - "blk.": "model.layers.", - "attn_norm": "input_layernorm", - "attn_q_norm.": "self_attn.q_norm.", - "attn_k_norm.": "self_attn.k_norm.", - "attn_v_norm.": "self_attn.v_norm.", - "attn_q": "self_attn.q_proj", - "attn_k": "self_attn.k_proj", - "attn_v": "self_attn.v_proj", - "attn_output": "self_attn.o_proj", - "ffn_up": "mlp.up_proj", - "ffn_down": "mlp.down_proj", - "ffn_gate": "mlp.gate_proj", - "ffn_norm": "post_attention_layernorm", - "token_embd": "model.embed_tokens", - "output_norm": "model.norm", - "output.weight": "lm_head.weight", -} - -GEMMA3_SD_MAP = LLAMA_SD_MAP.copy() -GEMMA3_SD_MAP.update({ - "ffn_norm": "pre_feedforward_layernorm", - "post_ffw_norm": "post_feedforward_layernorm", - "post_attention_norm": "post_attention_layernorm", -}) - -CLIP_VISION_SD_MAP = { - "mm.": "visual.merger.mlp.", - "v.post_ln.": "visual.merger.ln_q.", - "v.patch_embd": "visual.patch_embed.proj", - "v.blk.": "visual.blocks.", - "ffn_up": "mlp.up_proj", - "ffn_down": "mlp.down_proj", - "ffn_gate": "mlp.gate_proj", - "attn_out.": "attn.proj.", - "ln1.": "norm1.", - "ln2.": "norm2.", -} - -def sd_map_replace(raw_sd, key_map): - sd = {} - for k,v in raw_sd.items(): - for s,d in key_map.items(): - k = k.replace(s,d) - sd[k] = v - return sd - -def llama_permute(raw_sd, n_head, n_head_kv): - # Reverse version of LlamaModel.permute in llama.cpp convert script - sd = {} - permute = lambda x,h: x.reshape(h, x.shape[0] // h // 2, 2, *x.shape[1:]).swapaxes(1, 2).reshape(x.shape) - for k,v in raw_sd.items(): - if k.endswith(("q_proj.weight", "q_proj.bias")): - v.data = permute(v.data, n_head) - if k.endswith(("k_proj.weight", "k_proj.bias")): - v.data = permute(v.data, n_head_kv) - sd[k] = v - return sd - -def gemma3_norm_corrections(sd): - # Reverse change from Gemma3Model modify_tensors in llama.cpp convert script - norm_patterns = [ - "input_layernorm.weight", - "post_attention_layernorm.weight", - "pre_feedforward_layernorm.weight", - "post_feedforward_layernorm.weight", - "self_attn.q_norm.weight", - "self_attn.k_norm.weight", - "model.norm.weight" - ] - corrected = 0 - for key in list(sd.keys()): - if any(p in key for p in norm_patterns): - if is_quantized(sd[key]): - sd[key] = dequantize_tensor(sd[key], dtype=torch.float32) - 1.0 - else: - sd[key] = sd[key].float() - 1.0 - corrected += 1 - #logging.info(f"Gemma3: Applied -1 norm correction to {corrected} tensors") - return sd - -def strip_quant_suffix(name): - pattern = r"[-_]?(?:ud-)?i?q[0-9]_[a-z0-9_\-]{1,8}$" - match = re.search(pattern, name, re.IGNORECASE) - if match: - name = name[:match.start()] - return name - -def gguf_mmproj_loader(path): - # Reverse version of Qwen2VLVisionModel.modify_tensors - logging.info("Attenpting to find mmproj file for text encoder...") - - # get name to match w/o quant suffix - tenc_fname = os.path.basename(path) - tenc = os.path.splitext(tenc_fname)[0].lower() - tenc = strip_quant_suffix(tenc) - - # try and find matching mmproj - target = [] - root = os.path.dirname(path) - for fname in os.listdir(root): - name, ext = os.path.splitext(fname) - if ext.lower() != ".gguf": - continue - if "mmproj" not in name.lower(): - continue - if tenc in name.lower(): - target.append(fname) - - if len(target) == 0: - logging.error(f"Error: Can't find mmproj file for '{tenc_fname}' (matching:'{tenc}')! Qwen-Image-Edit will be broken!") - return {} - if len(target) > 1: - logging.error(f"Ambiguous mmproj for text encoder '{tenc_fname}', will use first match.") - - logging.info(f"Using mmproj '{target[0]}' for text encoder '{tenc_fname}'.") - target = os.path.join(root, target[0]) - vsd, _ = gguf_sd_loader(target, is_text_model=True) - - # concat 4D to 5D - if "v.patch_embd.weight.1" in vsd: - w1 = dequantize_tensor(vsd.pop("v.patch_embd.weight"), dtype=torch.float32) - w2 = dequantize_tensor(vsd.pop("v.patch_embd.weight.1"), dtype=torch.float32) - vsd["v.patch_embd.weight"] = torch.stack([w1, w2], dim=2) - - # run main replacement - vsd = sd_map_replace(vsd, CLIP_VISION_SD_MAP) - - # handle split Q/K/V - if "visual.blocks.0.attn_q.weight" in vsd: - attns = {} - # filter out attentions + group - for k,v in vsd.items(): - if any(x in k for x in ["attn_q", "attn_k", "attn_v"]): - k_attn, k_name = k.rsplit(".attn_", 1) - k_attn += ".attn.qkv." + k_name.split(".")[-1] - if k_attn not in attns: - attns[k_attn] = {} - attns[k_attn][k_name] = dequantize_tensor( - v, dtype=(torch.bfloat16 if is_quantized(v) else torch.float16) - ) - - # recombine - for k,v in attns.items(): - suffix = k.split(".")[-1] - vsd[k] = torch.cat([ - v[f"q.{suffix}"], - v[f"k.{suffix}"], - v[f"v.{suffix}"], - ], dim=0) - del attns - - return vsd - -def gguf_tokenizer_loader(path, temb_shape): - # convert gguf tokenizer to spiece - logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") - try: - from sentencepiece import sentencepiece_model_pb2 as model - except ImportError: - raise ImportError("Please make sure sentencepiece and protobuf are installed.\npip install sentencepiece protobuf") - spm = model.ModelProto() - - reader = gguf.GGUFReader(path) - - 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?) - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - else: - 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) - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - scores = get_list_field(reader, "tokenizer.ggml.scores", float) - toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) - - for idx, (token, score, toktype) in enumerate(zip(tokens, scores, toktypes)): - # # These aren't present in the original? - # if toktype == 5 and idx >= temb_shape[0]%1000): - # continue - - piece = spm.SentencePiece() - piece.piece = token - piece.score = score - piece.type = toktype - spm.pieces.append(piece) - - # unsure if any of these are correct - spm.trainer_spec.byte_fallback = True - spm.trainer_spec.vocab_size = len(tokens) # split off unused? - spm.trainer_spec.max_sentence_length = 4096 - 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) - - logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") - del reader - return torch.ByteTensor(list(spm.SerializeToString())) - -def gguf_tekken_tokenizer_loader(path, temb_shape): - # convert ggml (hf) tokenizer metadata to tekken/comfy data - logging.info("Attempting to recreate tekken tokenizer from GGUF file metadata...") - import json - import base64 - from transformers.convert_slow_tokenizer import bytes_to_unicode - - reader = gguf.GGUFReader(path) - - model_str = get_field(reader, "tokenizer.ggml.model", str) - if model_str == "gpt2": - if temb_shape == (131072, 5120): # probably Mistral - data = { - "config": {"num_vocab_tokens": 150000, "default_vocab_size": 131072}, - "vocab": [], - "special_tokens": [], - } - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) - - decoder = {v: k for k, v in bytes_to_unicode().items()} - for idx, (token, toktype) in enumerate(zip(tokens, toktypes)): - if toktype == 3: - data["special_tokens"].append( - {'rank': idx, 'token_str': token, 'is_control': True} - ) - else: - tok = bytes([decoder[char] for char in token]) - data["vocab"].append({ - "rank": len(data["vocab"]), - "token_bytes": base64.b64encode(tok).decode("ascii"), - "token_str": tok.decode("utf-8", errors="replace") # ? - }) - - logging.info(f"Created tekken tokenizer with vocab size of {len(data['vocab'])} (+{len(data['special_tokens'])})") - del reader - return torch.ByteTensor(list(json.dumps(data).encode('utf-8'))) - -def gguf_gemma3_tokenizer_loader(path): - #TODO: merge into gguf_tokenizer_loader - logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") - try: - from sentencepiece import sentencepiece_model_pb2 as model - except ImportError: - raise ImportError("Please install sentencepiece and protobuf.\npip install sentencepiece protobuf") - spm = model.ModelProto() - reader = gguf.GGUFReader(path) - - spm.normalizer_spec.name = "identity" - spm.normalizer_spec.add_dummy_prefix = False - spm.trainer_spec.model_type = 2 - spm.trainer_spec.input_format = "tsv" - spm.trainer_spec.byte_fallback = True - spm.trainer_spec.max_sentence_length = 4192 - spm.trainer_spec.bos_piece = "" - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - scores = get_list_field(reader, "tokenizer.ggml.scores", float) - toktype = get_list_field(reader, "tokenizer.ggml.token_type", int) - - if not tokens or not scores or not toktype: - raise ValueError("Missing tokenizer metadata") - - for idx in range(len(tokens)): - piece = spm.SentencePiece() - piece.piece = tokens[idx] - if idx == 3: # UNK position - piece.type = 2 # UNK Token - piece.score = 0.0 # UNK Score - else: - piece.type = toktype[idx] - piece.score = scores[idx] - spm.pieces.append(piece) - - spm.trainer_spec.vocab_size = len(spm.pieces) - logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") - - del reader - return torch.ByteTensor(list(spm.SerializeToString())) - -def gguf_clip_loader(path): - sd, extra = gguf_sd_loader(path, is_text_model=True) - arch = extra.get("arch_str", None) - if arch in {"t5", "t5encoder"}: - temb_key = "token_embd.weight" - if temb_key in sd and sd[temb_key].shape == (256384, 4096): - # 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. - 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", "qwen2vl", "qwen3", "qwen3vl", "gemma3"}: - # TODO: pass model_options["vocab_size"] to loader somehow - temb_key = "token_embd.weight" - if temb_key in sd and sd[temb_key].shape[0] >= (64 * 1024): - if arch == "llama" and sd[temb_key].shape == (131072, 5120): - # non-standard Comfy-Org tokenizer - sd["tekken_model"] = gguf_tekken_tokenizer_loader(path, sd[temb_key].shape) - elif arch == "gemma3": - sd["spiece_model"] = gguf_gemma3_tokenizer_loader(path) - # See note above for T5. - logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") - sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) - if arch == "gemma3": - sd = sd_map_replace(sd, GEMMA3_SD_MAP) - sd = gemma3_norm_corrections(sd) - else: - sd = sd_map_replace(sd, LLAMA_SD_MAP) - if arch == "llama": - sd = llama_permute(sd, 32, 8) # L3 / Mistral - if arch == "qwen2vl": - vsd = gguf_mmproj_loader(path) - sd.update(vsd) - else: - pass - return sd +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import warnings +import logging +import torch +import gguf +import re +import os + +from .ops import GGMLTensor +from .dequant import is_quantized, dequantize_tensor + +IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "cosmos", "ltxv", "hyvid", "wan", "lumina2", "qwen_image", "ideogram"} +TXT_ARCH_LIST = {"t5", "t5encoder", "llama", "qwen2vl", "qwen3", "qwen3vl", "gemma3"} +VIS_TYPE_LIST = {"clip-vision", "mmproj"} + +def get_orig_shape(reader, tensor_name): + field_key = f"comfy.gguf.orig_shape.{tensor_name}" + field = reader.get_field(field_key) + if field is None: + return None + # Has original shape metadata, so we try to decode it. + if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32: + 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 get_field(reader, field_name, field_type): + field = reader.get_field(field_name) + if field is None: + return None + elif field_type == str: + # extra check here as this is used for checking arch string + if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING: + raise TypeError(f"Bad type for GGUF {field_name} key: expected string, got {field.types!r}") + return str(field.parts[field.data[-1]], encoding="utf-8") + elif field_type in [int, float, bool]: + return field_type(field.parts[field.data[-1]].item()) + else: + raise TypeError(f"Unknown field type {field_type}") + +def get_list_field(reader, field_name, field_type): + field = reader.get_field(field_name) + if field is None: + return None + elif field_type == str: + return tuple(str(field.parts[part_idx], encoding="utf-8") for part_idx in field.data) + elif field_type in [int, float, bool]: + return tuple(field_type(field.parts[part_idx][0]) for part_idx in field.data) + else: + raise TypeError(f"Unknown field type {field_type}") + +def get_gguf_metadata(reader): + """Extract all simple metadata fields like safetensors""" + metadata = {} + for field_name in reader.fields: + try: + field = reader.get_field(field_name) + if len(field.types) == 1: # Simple scalar fields only + if field.types[0] == gguf.GGUFValueType.STRING: + metadata[field_name] = str(field.parts[field.data[-1]], "utf-8") + elif field.types[0] == gguf.GGUFValueType.INT32: + metadata[field_name] = int(field.parts[field.data[-1]]) + elif field.types[0] == gguf.GGUFValueType.F32: + metadata[field_name] = float(field.parts[field.data[-1]]) + elif field.types[0] == gguf.GGUFValueType.BOOL: + metadata[field_name] = bool(field.parts[field.data[-1]]) + except: + continue + return metadata + +def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=False): + """ + Read state dict as fake tensors + """ + reader = gguf.GGUFReader(path) + + # 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: + 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 = get_field(reader, "general.architecture", str) + type_str = get_field(reader, "general.type", str) + if arch_str in [None, "pig", "cow"]: + if is_text_model: + raise ValueError(f"This gguf file is incompatible with llama.cpp!\nConsider using safetensors or a compatible gguf file\n({path})") + compat = "sd.cpp" if arch_str is None else arch_str + # import here to avoid changes to convert.py breaking regular models + from .tools.convert import detect_arch + try: + arch_str = detect_arch(set(val[0] for val in tensors)).arch + except Exception as e: + raise ValueError(f"This model is not currently supported - ({e})") + elif arch_str not in TXT_ARCH_LIST and is_text_model: + if type_str not in VIS_TYPE_LIST: + raise ValueError(f"Unexpected text model architecture type in GGUF file: {arch_str!r}") + elif arch_str not in IMG_ARCH_LIST and not is_text_model: + raise ValueError(f"Unexpected architecture type in GGUF file: {arch_str!r}") + + if compat: + logging.warning(f"Warning: This gguf 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 + + # 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: + shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + # 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) + state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) + + # 1D tensors shouldn't be quantized, this is a fix for BF16 + if len(shape) <= 1 and tensor.tensor_type == gguf.GGMLQuantizationType.BF16: + state_dict[sd_key] = dequantize_tensor(state_dict[sd_key], dtype=torch.float32) + + # keep track of loaded tensor types + tensor_type_str = getattr(tensor.tensor_type, "name", repr(tensor.tensor_type)) + qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1 + + # print loaded tensor type counts + 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)} + if len(qsd) > 0: + max_key = max(qsd.keys(), key=lambda k: qsd[k].numel()) + state_dict[max_key].is_largest_weight = True + + # extra info to return + extra = { + "arch_str": arch_str, + "metadata": get_gguf_metadata(reader) + } + return (state_dict, extra) + +# for remapping llama.cpp -> original key names +T5_SD_MAP = { + "enc.": "encoder.", + ".blk.": ".block.", + "token_embd": "shared", + "output_norm": "final_layer_norm", + "attn_q": "layer.0.SelfAttention.q", + "attn_k": "layer.0.SelfAttention.k", + "attn_v": "layer.0.SelfAttention.v", + "attn_o": "layer.0.SelfAttention.o", + "attn_norm": "layer.0.layer_norm", + "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias", + "ffn_up": "layer.1.DenseReluDense.wi_1", + "ffn_down": "layer.1.DenseReluDense.wo", + "ffn_gate": "layer.1.DenseReluDense.wi_0", + "ffn_norm": "layer.1.layer_norm", +} + +LLAMA_SD_MAP = { + "blk.": "model.layers.", + "attn_norm": "input_layernorm", + "attn_q_norm.": "self_attn.q_norm.", + "attn_k_norm.": "self_attn.k_norm.", + "attn_v_norm.": "self_attn.v_norm.", + "attn_q": "self_attn.q_proj", + "attn_k": "self_attn.k_proj", + "attn_v": "self_attn.v_proj", + "attn_output": "self_attn.o_proj", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", + "ffn_gate": "mlp.gate_proj", + "ffn_norm": "post_attention_layernorm", + "token_embd": "model.embed_tokens", + "output_norm": "model.norm", + "output.weight": "lm_head.weight", +} + +GEMMA3_SD_MAP = LLAMA_SD_MAP.copy() +GEMMA3_SD_MAP.update({ + "ffn_norm": "pre_feedforward_layernorm", + "post_ffw_norm": "post_feedforward_layernorm", + "post_attention_norm": "post_attention_layernorm", +}) + +CLIP_VISION_SD_MAP = { + "mm.": "visual.merger.mlp.", + "v.post_ln.": "visual.merger.ln_q.", + "v.patch_embd": "visual.patch_embed.proj", + "v.blk.": "visual.blocks.", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", + "ffn_gate": "mlp.gate_proj", + "attn_out.": "attn.proj.", + "ln1.": "norm1.", + "ln2.": "norm2.", +} + +def sd_map_replace(raw_sd, key_map): + sd = {} + for k,v in raw_sd.items(): + for s,d in key_map.items(): + k = k.replace(s,d) + sd[k] = v + return sd + +def llama_permute(raw_sd, n_head, n_head_kv): + # Reverse version of LlamaModel.permute in llama.cpp convert script + sd = {} + permute = lambda x,h: x.reshape(h, x.shape[0] // h // 2, 2, *x.shape[1:]).swapaxes(1, 2).reshape(x.shape) + for k,v in raw_sd.items(): + if k.endswith(("q_proj.weight", "q_proj.bias")): + v.data = permute(v.data, n_head) + if k.endswith(("k_proj.weight", "k_proj.bias")): + v.data = permute(v.data, n_head_kv) + sd[k] = v + return sd + +def gemma3_norm_corrections(sd): + # Reverse change from Gemma3Model modify_tensors in llama.cpp convert script + norm_patterns = [ + "input_layernorm.weight", + "post_attention_layernorm.weight", + "pre_feedforward_layernorm.weight", + "post_feedforward_layernorm.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + "model.norm.weight" + ] + corrected = 0 + for key in list(sd.keys()): + if any(p in key for p in norm_patterns): + if is_quantized(sd[key]): + sd[key] = dequantize_tensor(sd[key], dtype=torch.float32) - 1.0 + else: + sd[key] = sd[key].float() - 1.0 + corrected += 1 + #logging.info(f"Gemma3: Applied -1 norm correction to {corrected} tensors") + return sd + +def strip_quant_suffix(name): + pattern = r"[-_]?(?:ud-)?i?q[0-9]_[a-z0-9_\-]{1,8}$" + match = re.search(pattern, name, re.IGNORECASE) + if match: + name = name[:match.start()] + return name + +def gguf_mmproj_loader(path): + # Reverse version of Qwen2VLVisionModel.modify_tensors + logging.info("Attenpting to find mmproj file for text encoder...") + + # get name to match w/o quant suffix + tenc_fname = os.path.basename(path) + tenc = os.path.splitext(tenc_fname)[0].lower() + tenc = strip_quant_suffix(tenc) + + # try and find matching mmproj + target = [] + root = os.path.dirname(path) + for fname in os.listdir(root): + name, ext = os.path.splitext(fname) + if ext.lower() != ".gguf": + continue + if "mmproj" not in name.lower(): + continue + if tenc in name.lower(): + target.append(fname) + + if len(target) == 0: + logging.error(f"Error: Can't find mmproj file for '{tenc_fname}' (matching:'{tenc}')! Qwen-Image-Edit will be broken!") + return {} + if len(target) > 1: + logging.error(f"Ambiguous mmproj for text encoder '{tenc_fname}', will use first match.") + + logging.info(f"Using mmproj '{target[0]}' for text encoder '{tenc_fname}'.") + target = os.path.join(root, target[0]) + vsd, _ = gguf_sd_loader(target, is_text_model=True) + + # concat 4D to 5D + if "v.patch_embd.weight.1" in vsd: + w1 = dequantize_tensor(vsd.pop("v.patch_embd.weight"), dtype=torch.float32) + w2 = dequantize_tensor(vsd.pop("v.patch_embd.weight.1"), dtype=torch.float32) + vsd["v.patch_embd.weight"] = torch.stack([w1, w2], dim=2) + + # run main replacement + vsd = sd_map_replace(vsd, CLIP_VISION_SD_MAP) + + # handle split Q/K/V + if "visual.blocks.0.attn_q.weight" in vsd: + attns = {} + # filter out attentions + group + for k,v in vsd.items(): + if any(x in k for x in ["attn_q", "attn_k", "attn_v"]): + k_attn, k_name = k.rsplit(".attn_", 1) + k_attn += ".attn.qkv." + k_name.split(".")[-1] + if k_attn not in attns: + attns[k_attn] = {} + attns[k_attn][k_name] = dequantize_tensor( + v, dtype=(torch.bfloat16 if is_quantized(v) else torch.float16) + ) + + # recombine + for k,v in attns.items(): + suffix = k.split(".")[-1] + vsd[k] = torch.cat([ + v[f"q.{suffix}"], + v[f"k.{suffix}"], + v[f"v.{suffix}"], + ], dim=0) + del attns + + return vsd + +def gguf_tokenizer_loader(path, temb_shape): + # convert gguf tokenizer to spiece + logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") + try: + from sentencepiece import sentencepiece_model_pb2 as model + except ImportError: + raise ImportError("Please make sure sentencepiece and protobuf are installed.\npip install sentencepiece protobuf") + spm = model.ModelProto() + + reader = gguf.GGUFReader(path) + + 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?) + else: + raise NotImplementedError("Unknown model, can't set tokenizer!") + else: + 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) + + tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) + scores = get_list_field(reader, "tokenizer.ggml.scores", float) + toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) + + for idx, (token, score, toktype) in enumerate(zip(tokens, scores, toktypes)): + # # These aren't present in the original? + # if toktype == 5 and idx >= temb_shape[0]%1000): + # continue + + piece = spm.SentencePiece() + piece.piece = token + piece.score = score + piece.type = toktype + spm.pieces.append(piece) + + # unsure if any of these are correct + spm.trainer_spec.byte_fallback = True + spm.trainer_spec.vocab_size = len(tokens) # split off unused? + spm.trainer_spec.max_sentence_length = 4096 + 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) + + logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") + del reader + return torch.ByteTensor(list(spm.SerializeToString())) + +def gguf_tekken_tokenizer_loader(path, temb_shape): + # convert ggml (hf) tokenizer metadata to tekken/comfy data + logging.info("Attempting to recreate tekken tokenizer from GGUF file metadata...") + import json + import base64 + from transformers.convert_slow_tokenizer import bytes_to_unicode + + reader = gguf.GGUFReader(path) + + model_str = get_field(reader, "tokenizer.ggml.model", str) + if model_str == "gpt2": + if temb_shape == (131072, 5120): # probably Mistral + data = { + "config": {"num_vocab_tokens": 150000, "default_vocab_size": 131072}, + "vocab": [], + "special_tokens": [], + } + else: + raise NotImplementedError("Unknown model, can't set tokenizer!") + else: + raise NotImplementedError("Unknown model, can't set tokenizer!") + + tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) + toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) + + decoder = {v: k for k, v in bytes_to_unicode().items()} + for idx, (token, toktype) in enumerate(zip(tokens, toktypes)): + if toktype == 3: + data["special_tokens"].append( + {'rank': idx, 'token_str': token, 'is_control': True} + ) + else: + tok = bytes([decoder[char] for char in token]) + data["vocab"].append({ + "rank": len(data["vocab"]), + "token_bytes": base64.b64encode(tok).decode("ascii"), + "token_str": tok.decode("utf-8", errors="replace") # ? + }) + + logging.info(f"Created tekken tokenizer with vocab size of {len(data['vocab'])} (+{len(data['special_tokens'])})") + del reader + return torch.ByteTensor(list(json.dumps(data).encode('utf-8'))) + +def gguf_gemma3_tokenizer_loader(path): + #TODO: merge into gguf_tokenizer_loader + logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") + try: + from sentencepiece import sentencepiece_model_pb2 as model + except ImportError: + raise ImportError("Please install sentencepiece and protobuf.\npip install sentencepiece protobuf") + spm = model.ModelProto() + reader = gguf.GGUFReader(path) + + spm.normalizer_spec.name = "identity" + spm.normalizer_spec.add_dummy_prefix = False + spm.trainer_spec.model_type = 2 + spm.trainer_spec.input_format = "tsv" + spm.trainer_spec.byte_fallback = True + spm.trainer_spec.max_sentence_length = 4192 + spm.trainer_spec.bos_piece = "" + + tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) + scores = get_list_field(reader, "tokenizer.ggml.scores", float) + toktype = get_list_field(reader, "tokenizer.ggml.token_type", int) + + if not tokens or not scores or not toktype: + raise ValueError("Missing tokenizer metadata") + + for idx in range(len(tokens)): + piece = spm.SentencePiece() + piece.piece = tokens[idx] + if idx == 3: # UNK position + piece.type = 2 # UNK Token + piece.score = 0.0 # UNK Score + else: + piece.type = toktype[idx] + piece.score = scores[idx] + spm.pieces.append(piece) + + spm.trainer_spec.vocab_size = len(spm.pieces) + logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") + + del reader + return torch.ByteTensor(list(spm.SerializeToString())) + +def gguf_clip_loader(path): + sd, extra = gguf_sd_loader(path, is_text_model=True) + arch = extra.get("arch_str", None) + if arch in {"t5", "t5encoder"}: + temb_key = "token_embd.weight" + if temb_key in sd and sd[temb_key].shape == (256384, 4096): + # 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. + 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", "qwen2vl", "qwen3", "qwen3vl", "gemma3"}: + # TODO: pass model_options["vocab_size"] to loader somehow + temb_key = "token_embd.weight" + if temb_key in sd and sd[temb_key].shape[0] >= (64 * 1024): + if arch == "llama" and sd[temb_key].shape == (131072, 5120): + # non-standard Comfy-Org tokenizer + sd["tekken_model"] = gguf_tekken_tokenizer_loader(path, sd[temb_key].shape) + elif arch == "gemma3": + sd["spiece_model"] = gguf_gemma3_tokenizer_loader(path) + # See note above for T5. + logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") + sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) + if arch == "gemma3": + sd = sd_map_replace(sd, GEMMA3_SD_MAP) + sd = gemma3_norm_corrections(sd) + else: + sd = sd_map_replace(sd, LLAMA_SD_MAP) + if arch == "llama": + sd = llama_permute(sd, 32, 8) # L3 / Mistral + if arch == "qwen2vl": + vsd = gguf_mmproj_loader(path) + sd.update(vsd) + else: + pass + return sd From 61dde54549d4b860125539a81d1e8fc96f7254fb Mon Sep 17 00:00:00 2001 From: molbal Date: Tue, 9 Jun 2026 20:52:45 +0200 Subject: [PATCH 3/6] Add dequantization handling for Ideogram inference --- loader.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/loader.py b/loader.py index 8d2ae09..a2e9366 100644 --- a/loader.py +++ b/loader.py @@ -501,6 +501,16 @@ def gguf_clip_loader(path): if arch == "qwen2vl": vsd = gguf_mmproj_loader(path) sd.update(vsd) + elif arch == "ideogram": + # Dequantize Ideogram model for inference + logging.info("Dequantizing Ideogram model for inference...") + dequantized_count = 0 + for key in list(sd.keys()): + if is_quantized(sd[key]): + # Dequantize to BF16 to save VRAM while maintaining quality + sd[key] = dequantize_tensor(sd[key], dtype=torch.bfloat16) + dequantized_count += 1 + logging.info(f"Dequantized {dequantized_count} tensors for Ideogram model") else: pass return sd From fd65e64030119b610f3055f407b3cd740fd8d6c7 Mon Sep 17 00:00:00 2001 From: molbal Date: Tue, 9 Jun 2026 22:28:05 +0200 Subject: [PATCH 4/6] Fix BF16 tensor dequantization logic Adjust handling of BF16 tensors to ensure proper dequantization based on tensor shape. --- loader.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/loader.py b/loader.py index a2e9366..6710483 100644 --- a/loader.py +++ b/loader.py @@ -138,9 +138,10 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=F torch_tensor = torch_tensor.view(*shape) state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) - # 1D tensors shouldn't be quantized, this is a fix for BF16 - if len(shape) <= 1 and tensor.tensor_type == gguf.GGMLQuantizationType.BF16: - state_dict[sd_key] = dequantize_tensor(state_dict[sd_key], dtype=torch.float32) + # BF16 GGUF tensors are full-precision storage, not compressed quants. + if tensor.tensor_type == gguf.GGMLQuantizationType.BF16: + dtype = torch.float32 if len(shape) <= 1 else torch.bfloat16 + state_dict[sd_key] = dequantize_tensor(state_dict[sd_key], dtype=dtype) # keep track of loaded tensor types tensor_type_str = getattr(tensor.tensor_type, "name", repr(tensor.tensor_type)) From e4c7598e98136e83d6dad25c53bfd8defb9d21cf Mon Sep 17 00:00:00 2001 From: molbal Date: Tue, 9 Jun 2026 22:28:29 +0200 Subject: [PATCH 5/6] Refactor dequantization functions and imports --- dequant.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dequant.py b/dequant.py index 78f5f26..809db16 100644 --- a/dequant.py +++ b/dequant.py @@ -18,6 +18,9 @@ def dequantize_tensor(tensor, dtype=None, dequant_dtype=None): if qtype in TORCH_COMPATIBLE_QTYPES: return tensor.to(dtype) + elif qtype == gguf.GGMLQuantizationType.BF16: + tensor = torch.Tensor(tensor.data.view(torch.bfloat16).reshape(oshape)) + return tensor if dtype is None or dtype == torch.bfloat16 else tensor.to(dtype) elif qtype in dequantize_functions: dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype return dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype).to(dtype) From b111cdbd30fb34cd4a26c7ef1808830fbd4aeb28 Mon Sep 17 00:00:00 2001 From: molbal Date: Tue, 9 Jun 2026 22:29:04 +0200 Subject: [PATCH 6/6] Add newline at end of convert.py Fix missing newline at end of file in convert.py