diff --git a/dequant.py b/dequant.py index 78f5f26..025c682 100644 --- a/dequant.py +++ b/dequant.py @@ -1,5 +1,6 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) import gguf +import numpy as np import torch from tqdm import tqdm @@ -240,6 +241,24 @@ def dequantize_blocks_Q2_K(blocks, block_size, type_size, dtype=None): # IQ quants KVALUES = torch.tensor([-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113], dtype=torch.int8) +def _get_iq_grid(iq_cls): + iq_cls.init_grid() + return torch.from_numpy(np.array(iq_cls.grid).squeeze().copy()) + +def _get_iq_ksigns(iq_cls): + iq_cls.init_grid() + return torch.from_numpy(np.frombuffer(iq_cls.ksigns, dtype=np.uint8).copy()) + +from gguf.quants import IQ1_M as _IQ1_M, IQ1_S as _IQ1_S, IQ2_S as _IQ2_S, IQ2_XXS as _IQ2_XXS, IQ3_S as _IQ3_S, IQ3_XXS as _IQ3_XXS + +GRID_IQ3_S = _get_iq_grid(_IQ3_S) +GRID_IQ3_XXS = _get_iq_grid(_IQ3_XXS) +GRID_IQ2_S = _get_iq_grid(_IQ2_S) +GRID_IQ2_XXS = _get_iq_grid(_IQ2_XXS) +GRID_IQ1_S = _get_iq_grid(_IQ1_S) +_get_iq_grid(_IQ1_M) # IQ1_M uses the same grid as IQ1_S internally +KSIGNS_IQ2_XXS = _get_iq_ksigns(_IQ2_XXS) + def dequantize_blocks_IQ4_NL(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] @@ -284,6 +303,205 @@ def dequantize_blocks_IQ4_XS(blocks, block_size, type_size, dtype=None): return (dl * qs).reshape((n_blocks, -1)) +def dequantize_blocks_IQ3_S(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs, qh, signs, scales = split_block_dims(blocks, 2, 64, 8, 32) + d = d.view(torch.float16).to(dtype) + + scales = scales.view(torch.uint8) + scales = torch.stack([scales & 0xF, scales >> 4], dim=-1).reshape((n_blocks, 8)) + db = d * (1 + 2 * scales.to(dtype)) + db = db.reshape((n_blocks, 8, 1, 1)) + + shifts = torch.arange(8, device=d.device, dtype=torch.uint8).reshape((1, 1, 8)) + signs = (signs.unsqueeze(-1) >> shifts) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape((n_blocks, 8, 8, 4)) + + qh_bits = (qh.unsqueeze(-1) >> shifts) & 1 + qh_bits = qh_bits.reshape((n_blocks, 64)) + qs = qs.to(torch.int16) | (qh_bits.to(torch.int16) << 8) + + grid = GRID_IQ3_S.to(dtype=dtype, device=d.device) + grid_val = grid[qs.to(torch.long)].reshape((n_blocks, 8, 8, 4)) + + return (db * grid_val * signs).reshape((n_blocks, QK_K)) + +def dequantize_blocks_IQ3_XXS(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs, scales, _ = split_block_dims(blocks, 2, 64, 32) + d = d.view(torch.float16).to(dtype) + + scales = scales.reshape((n_blocks, 8, 4)).to(torch.int32) + scales = scales[:, :, 0] | scales[:, :, 1] << 8 | scales[:, :, 2] << 16 | scales[:, :, 3] << 24 + + db = d * (0.5 + ((scales >> 28) & 0xF).to(dtype)) * 0.5 + db = db.reshape((n_blocks, 8, 1, 1)) + + shifts = torch.tensor([0, 7, 14, 21], device=d.device, dtype=torch.int32).reshape((1, 1, 4)) + sign_indices = (scales.reshape((n_blocks, 8, 1)) >> shifts) & 0x7F + + ksigns = KSIGNS_IQ2_XXS.to(d.device) + sign_bytes = ksigns[sign_indices.to(torch.long)] + + shifts_bits = torch.arange(8, device=d.device, dtype=torch.uint8).reshape((1, 1, 1, 8)) + signs = (sign_bytes.unsqueeze(-1) >> shifts_bits) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape((n_blocks, 8, 4, 8)) + + grid = GRID_IQ3_XXS.to(dtype=dtype, device=d.device) + grid_val = grid[qs.to(torch.long)].reshape((n_blocks, 8, 4, 8)) + + return (db * grid_val * signs).reshape((n_blocks, QK_K)) + +def dequantize_blocks_IQ2_S(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs, signs, qh, scales = split_block_dims(blocks, 2, 32, 32, 8) + d = d.view(torch.float16).to(dtype) + + scales = scales.view(torch.uint8) + scales = torch.stack([scales & 0xF, scales >> 4], dim=-1).reshape((n_blocks, 16)) + db = d * (0.5 + scales.to(dtype)) * 0.25 + db = db.reshape((n_blocks, 16, 1, 1)) + + shifts = torch.arange(8, device=d.device, dtype=torch.uint8).reshape((1, 1, 8)) + signs = (signs.unsqueeze(-1) >> shifts) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape((n_blocks, 16, 2, 8)) + + qh_shifts = torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4)) + qh_bits = (qh.view(torch.uint8).reshape((n_blocks, 8, 1)) >> qh_shifts) & 3 + qh_bits = qh_bits.reshape((n_blocks, 32)) + + qs = qs.view(torch.uint8).to(torch.int32) + indices = qs | (qh_bits.to(torch.int32) << 8) + + grid = GRID_IQ2_S.to(dtype=dtype, device=d.device) + grid_val = grid[indices.to(torch.long)].reshape((n_blocks, 16, 2, 8)) + + return (db * grid_val * signs).reshape((n_blocks, QK_K)) + +def dequantize_blocks_IQ2_XXS(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs = split_block_dims(blocks, 2) + d = d.view(torch.float16).to(dtype) + + u32 = qs.reshape((n_blocks, 16, 4)).to(torch.int32) + u32 = u32[:, :, 0] | (u32[:, :, 1] << 8) | (u32[:, :, 2] << 16) | (u32[:, :, 3] << 24) + u32 = u32.reshape((n_blocks, 8, 2)) + + q0 = u32[:, :, 0] # grid indices + q1 = u32[:, :, 1] # scales and signs + + db = d * (0.5 + ((q1 >> 28) & 0xF).to(dtype)) * 0.25 + db = db.reshape((n_blocks, 8, 1, 1)) + + shifts = torch.tensor([0, 7, 14, 21], device=d.device, dtype=torch.int32).reshape((1, 1, 4)) + sign_indices = (q1.unsqueeze(-1) >> shifts) & 0x7F + + ksigns = KSIGNS_IQ2_XXS.to(d.device) + sign_bytes = ksigns[sign_indices.to(torch.long)] + + shifts_bits = torch.arange(8, device=d.device, dtype=torch.uint8).reshape((1, 1, 1, 8)) + signs = (sign_bytes.unsqueeze(-1) >> shifts_bits) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape((n_blocks, 8, 4, 8)) + + indices = q0.contiguous().view(torch.uint8) + grid = GRID_IQ2_XXS.to(dtype=dtype, device=d.device) + grid_val = grid[indices.to(torch.long)].reshape((n_blocks, 8, 4, 8)) + + return (db * grid_val * signs).reshape((n_blocks, QK_K)) + +def dequantize_blocks_IQ1_M(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + qs, qh, scales = split_block_dims(blocks, 32, 16) + + scales_u16 = scales.reshape((n_blocks, 4, 2)).to(torch.int32) + scales_u16 = scales_u16[:, :, 0] | (scales_u16[:, :, 1] << 8) + + d_bits = ( + ((scales_u16[:, 0] & 0xF000) >> 12) + | ((scales_u16[:, 1] & 0xF000) >> 8) + | ((scales_u16[:, 2] & 0xF000) >> 4) + | (scales_u16[:, 3] & 0xF000) + ) + d = d_bits.to(torch.int16).view(torch.float16).to(dtype).reshape((n_blocks, 1)) + + sub_shifts = torch.tensor([0, 3, 6, 9], device=d.device, dtype=torch.int32).reshape((1, 1, 4)) + sub_scales = (scales_u16.reshape((n_blocks, 4, 1)) >> sub_shifts) & 7 + dl = d.reshape((n_blocks, 1, 1)) * (2 * sub_scales.to(dtype) + 1) + dl = dl.reshape((n_blocks, 8, 2, 1, 1)) + + qh_bytes = qh.to(torch.int32) + qh_shifts = torch.tensor([0, 4], device=d.device, dtype=torch.int32).reshape((1, 1, 2)) + qh_unpacked = (qh_bytes.reshape((n_blocks, 16, 1)) >> qh_shifts).reshape((n_blocks, 32)) + + delta = torch.where( + (qh_unpacked & 8) == 0, + torch.full((1,), 0.125, dtype=dtype, device=d.device), + torch.full((1,), -0.125, dtype=dtype, device=d.device), + ).reshape((n_blocks, 8, 2, 2, 1)) + + qh_bits = qh_unpacked & 7 + qs = qs.to(torch.int32) + indices = qs | (qh_bits << 8) + + grid = GRID_IQ1_S.to(dtype=dtype, device=d.device) + grid_val = grid[indices.to(torch.long)].reshape((n_blocks, 8, 2, 2, 8)) + + return (dl * (grid_val + delta)).reshape((n_blocks, QK_K)) + +def dequantize_blocks_IQ1_S(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs, qh = split_block_dims(blocks, 2, 32) + d = d.view(torch.float16).to(dtype) + + qh = qh.view(torch.int16).to(torch.int32) & 0xFFFF + + dl = d * (2 * ((qh >> 12) & 7).to(dtype) + 1) + delta = torch.where( + (qh & 0x8000) == 0, + torch.full((1,), 0.125, dtype=dtype, device=d.device), + torch.full((1,), -0.125, dtype=dtype, device=d.device), + ) + + shifts = torch.tensor([0, 3, 6, 9], device=d.device, dtype=torch.int32).reshape((1, 1, 4)) + qh_bits = (qh.reshape((n_blocks, 8, 1)) >> shifts) & 7 + + qs = qs.view(torch.uint8).to(torch.int32).reshape((n_blocks, 8, 4)) + indices = qs | (qh_bits << 8) + + grid = GRID_IQ1_S.to(dtype=dtype, device=d.device) + grid_val = grid[indices.to(torch.long)].reshape((n_blocks, 8, 4, 8)) + + dl = dl.reshape((n_blocks, 8, 1, 1)) + delta = delta.reshape((n_blocks, 8, 1, 1)) + + return (dl * (grid_val + delta)).reshape((n_blocks, QK_K)) + dequantize_functions = { gguf.GGMLQuantizationType.BF16: dequantize_blocks_BF16, gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0, @@ -298,4 +516,10 @@ def dequantize_blocks_IQ4_XS(blocks, block_size, type_size, dtype=None): gguf.GGMLQuantizationType.Q2_K: dequantize_blocks_Q2_K, gguf.GGMLQuantizationType.IQ4_NL: dequantize_blocks_IQ4_NL, gguf.GGMLQuantizationType.IQ4_XS: dequantize_blocks_IQ4_XS, + gguf.GGMLQuantizationType.IQ3_S: dequantize_blocks_IQ3_S, + gguf.GGMLQuantizationType.IQ3_XXS: dequantize_blocks_IQ3_XXS, + gguf.GGMLQuantizationType.IQ2_S: dequantize_blocks_IQ2_S, + gguf.GGMLQuantizationType.IQ2_XXS: dequantize_blocks_IQ2_XXS, + gguf.GGMLQuantizationType.IQ1_M: dequantize_blocks_IQ1_M, + gguf.GGMLQuantizationType.IQ1_S: dequantize_blocks_IQ1_S, } diff --git a/loader.py b/loader.py index 7cefb11..582bcde 100644 --- a/loader.py +++ b/loader.py @@ -5,14 +5,72 @@ import gguf import re import os +import numpy as np +from gguf import GGUFReader, GGUFValueType from .ops import GGMLTensor from .dequant import is_quantized, dequantize_tensor +from .quant_ops import make_quantized -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"} +IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "cosmos", "ltxv", "hyvid", "wan", "lumina2", "qwen_image", "ideogram4", "krea2"} +TXT_ARCH_LIST = {"t5", "t5encoder", "llama", "qwen2vl", "qwen3", "qwen3vl", "gemma3", "gemma4"} VIS_TYPE_LIST = {"clip-vision", "mmproj"} +class LazyGGUFReader(GGUFReader): + def _get_field_parts(self, orig_offs: int, raw_type: int): + gtype = GGUFValueType(raw_type) + + if gtype == GGUFValueType.ARRAY: + raw_itype = self._get(orig_offs, np.uint32) + offs = orig_offs + int(raw_itype.nbytes) + alen = self._get(offs, np.uint64) + array_len = alen[0] + + if array_len > 1000: + offs += int(alen.nbytes) + sub_type = raw_itype[0] + types = [gtype, GGUFValueType(sub_type)] + aparts = [raw_itype, alen] + data_idxs = [] + data_view = self.data + is_swapped = (self.byte_order == 'S') + + if sub_type == 8: + for _ in range(array_len): + slen_arr = data_view[offs : offs + 8] + slen = slen_arr.view(dtype=np.uint64)[0] + if is_swapped: + slen = slen.newbyteorder('S') + + str_total_bytes = 8 + int(slen) + sdata_arr = data_view[offs + 8 : offs + str_total_bytes] + + idxs_offs = len(aparts) + aparts.append(slen_arr) + aparts.append(sdata_arr) + data_idxs.append(idxs_offs + 1) + offs += str_total_bytes + + return offs - orig_offs, aparts, data_idxs, types + else: + nptype = self.gguf_scalar_to_np.get(GGUFValueType(sub_type)) + if nptype is not None: + item_size = np.dtype(nptype).itemsize + total_bytes = array_len * item_size + total_data = data_view[offs : offs + total_bytes].view(dtype=nptype) + if is_swapped: + total_data = total_data.newbyteorder('S') + + idxs_offs = len(aparts) + aparts.extend(total_data[i : i + 1] for i in range(array_len)) + data_idxs = list(range(idxs_offs, idxs_offs + array_len)) + offs += total_bytes + + return offs - orig_offs, aparts, data_idxs, types + + return super()._get_field_parts(orig_offs, raw_type) + + def get_orig_shape(reader, tensor_name): field_key = f"comfy.gguf.orig_shape.{tensor_name}" field = reader.get_field(field_key) @@ -67,11 +125,11 @@ def get_gguf_metadata(reader): continue return metadata -def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=False): +def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=False, dynamic=False): """ Read state dict as fake tensors """ - reader = gguf.GGUFReader(path) + reader = LazyGGUFReader(path) # filter and strip prefix has_prefix = False @@ -134,13 +192,14 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=F 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) + if tensor.tensor_type == gguf.GGMLQuantizationType.BF16: + state_dict[sd_key] = torch_tensor.view(torch.bfloat16).reshape(*shape) + elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + state_dict[sd_key] = torch_tensor.view(*shape) + elif dynamic: + state_dict[sd_key] = make_quantized(torch_tensor, tensor.tensor_type, shape) + else: + state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) # keep track of loaded tensor types tensor_type_str = getattr(tensor.tensor_type, "name", repr(tensor.tensor_type)) @@ -160,6 +219,8 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=F "arch_str": arch_str, "metadata": get_gguf_metadata(reader) } + if is_text_model: + extra["reader"] = reader return (state_dict, extra) # for remapping llama.cpp -> original key names @@ -206,6 +267,19 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=F "post_attention_norm": "post_attention_layernorm", }) +GEMMA4_SD_MAP = { + "per_layer_token_embd": "model.embed_tokens_per_layer", + "per_layer_model_proj": "model.per_layer_model_projection", + "proj.weight": "per_layer_projection.weight", +} +GEMMA4_SD_MAP.update(GEMMA3_SD_MAP) +GEMMA4_SD_MAP.update({ + "layer_output_scale.weight": "layer_scalar", + "inp_gate.weight": "per_layer_input_gate.weight", + "post_norm.weight": "post_per_layer_input_norm.weight", + "per_layer_proj_norm": "model.per_layer_projection_norm", +}) + CLIP_VISION_SD_MAP = { "mm.": "visual.merger.mlp.", "v.post_ln.": "visual.merger.ln_q.", @@ -219,6 +293,53 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=F "ln2.": "norm2.", } +CLIP_VISION_QWEN3_MAP = { + "v.blk": "model.visual.blocks", + ".fc": ".linear_fc", + "ck.8.": "st.0.", + "ck.16.": "st.1.", + "ck.24.": "st.2.", + "ck.5.": "st.0.", + "ck.11.": "st.1.", + "ck.17.": "st.2.", + "attn_out": "attn.proj", + "ln1": "norm1", + "ln2": "norm2", + "attn_qkv": "attn.qkv", + "ffn_up": "mlp.linear_fc1", + "ffn_down": "mlp.linear_fc2", + "mm.0": "model.visual.merger.linear_fc1", + "mm.2": "model.visual.merger.linear_fc2", + "v.post_ln": "model.visual.merger.norm", + "v.patch_embd": "model.visual.patch_embed.proj", + "v.position_embd.weight": "visual.pos_embed.weight", + "v.deepstast.": "model.visual.deepstack_merger_list.", +} + +CLIP_VISION_GEMMA4_MAP = { + "v.position_embd.weight": "vision_model.patch_embedder.position_embedding_table", + "mm.input_projection": "multi_modal_projector.embedding_projection", + "mm.a.input_projection": "audio_projector.embedding_projection", + "attn_post_norm": "post_attention_layernorm", + "ffn_post_norm": "post_feedforward_layernorm", + "attn_k_norm": "self_attn.k_norm", + "attn_q_norm": "self_attn.q_norm", + "ln1.weight": "input_layernorm.weight", + "ln2.weight": "pre_feedforward_layernorm.weight", + "attn_out.": "self_attn.o_proj.", + "attn_k.": "self_attn.k_proj.", + "attn_q.": "self_attn.q_proj.", + "attn_v.": "self_attn.v_proj.", + "ffn_down.": "mlp.down_proj.", + "ffn_gate.": "mlp.gate_proj.", + "ffn_up.": "mlp.up_proj.", + "r0.weight": "r0.conv.weight", + "r1.weight": "r1.conv.weight", + "v.blk": "vision_model.encoder.layers", + "_proj.weight": "_proj.linear.weight", + "v.patch_embd.": "vision_model.patch_embedder.input_proj.", +} + def sd_map_replace(raw_sd, key_map): sd = {} for k,v in raw_sd.items(): @@ -268,9 +389,9 @@ def strip_quant_suffix(name): name = name[:match.start()] return name -def gguf_mmproj_loader(path): +def gguf_mmproj_loader(path, dynamic=False): # Reverse version of Qwen2VLVisionModel.modify_tensors - logging.info("Attenpting to find mmproj file for text encoder...") + logging.info("Attempting to find mmproj file for text encoder...") # get name to match w/o quant suffix tenc_fname = os.path.basename(path) @@ -290,14 +411,19 @@ def gguf_mmproj_loader(path): 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!") + logging.warning(f"Can't find mmproj file for '{tenc_fname}' (matching:'{tenc}'), vision function will not work!") return {} if len(target) > 1: - logging.error(f"Ambiguous mmproj for text encoder '{tenc_fname}', will use first match.") + logging.info(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) + vsd, _ = gguf_sd_loader(target, is_text_model=True, dynamic=dynamic) + + # gemma4 + if "mm.a.input_projection.weight" in vsd: + vsd["v.patch_embd.weight"] = vsd["v.patch_embd.weight"].permute(0, 2, 3, 1).flatten(start_dim=1) + return sd_map_replace(vsd, CLIP_VISION_GEMMA4_MAP) # concat 4D to 5D if "v.patch_embd.weight.1" in vsd: @@ -305,7 +431,11 @@ def gguf_mmproj_loader(path): 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 + # qwen3vl + if any("deepstack" in key for key in vsd): + return sd_map_replace(vsd, CLIP_VISION_QWEN3_MAP) + + # qwen2vl vsd = sd_map_replace(vsd, CLIP_VISION_SD_MAP) # handle split Q/K/V @@ -334,7 +464,7 @@ def gguf_mmproj_loader(path): return vsd -def gguf_tokenizer_loader(path, temb_shape): +def gguf_tokenizer_loader(reader, temb_shape): # convert gguf tokenizer to spiece logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") try: @@ -343,11 +473,9 @@ def gguf_tokenizer_loader(path, temb_shape): 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?) + spm.trainer_spec.model_type = 1 # Unigram (do we have a T5 w/ BPE?) else: raise NotImplementedError("Unknown model, can't set tokenizer!") else: @@ -380,17 +508,17 @@ def gguf_tokenizer_loader(path, temb_shape): logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") del reader - return torch.ByteTensor(list(spm.SerializeToString())) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="The given buffer is not writable") + return torch.frombuffer(spm.SerializeToString(), dtype=torch.uint8) -def gguf_tekken_tokenizer_loader(path, temb_shape): +def gguf_tekken_tokenizer_loader(reader, 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 @@ -423,9 +551,11 @@ def gguf_tekken_tokenizer_loader(path, temb_shape): 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'))) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="The given buffer is not writable") + return torch.frombuffer(json.dumps(data).encode('utf-8'), dtype=torch.uint8) -def gguf_gemma3_tokenizer_loader(path): +def gguf_gemma3_tokenizer_loader(reader): #TODO: merge into gguf_tokenizer_loader logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") try: @@ -433,7 +563,6 @@ def gguf_gemma3_tokenizer_loader(path): 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 @@ -465,42 +594,202 @@ def gguf_gemma3_tokenizer_loader(path): logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") del reader - return torch.ByteTensor(list(spm.SerializeToString())) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="The given buffer is not writable") + return torch.frombuffer(spm.SerializeToString(), dtype=torch.uint8) + +def gguf_gemma4_tokenizer_loader(reader): + # convert gguf tokenizer to spiece + logging.info("Attempting to recreate tokenizer from GGUF file metadata...") + import json + + tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) + merges = get_list_field(reader, "tokenizer.ggml.merges", str) + del reader + + if not tokens or not merges: + raise ValueError("Missing tokenizer metadata") + + vocab = {token: idx for idx, token in enumerate(tokens)} + target_special_ids = [ + 0, 1, 2, 3, 4, 46, 47, 48, 49, 50, 51, 52, 98, 100, 101, 105, 106, 255999, 256000, 258880, 258881, 258882, 258883, 258884 + ] + + added_tokens = [] + for sp_id in target_special_ids: + if sp_id < len(tokens): + added_tokens.append({ + "id": sp_id, + "content": tokens[sp_id], + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True + }) + + tokenizer_dict = { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": added_tokens, + "normalizer": { + "type": "Replace", + "pattern": {"String": " "}, + "content": "\u2581" + }, + "pre_tokenizer": { + "type": "Split", + "pattern": {"String": " "}, + "behavior": "MergedWithPrevious", + "invert": False + }, + "post_processor": { + "type": "TemplateProcessing", + "single": [{"Sequence": {"id": "A", "type_id": 0}}], + "pair": [ + {"Sequence": {"id": "A", "type_id": 0}}, + {"Sequence": {"id": "B", "type_id": 1}} + ], + "special_tokens": {} + }, + "decoder": { + "type": "Sequence", + "decoders": [ + {"type": "Replace", "pattern": {"String": "\u2581"}, "content": " "}, + {"type": "ByteFallback"}, + {"type": "Fuse"} + ] + }, + "model": { + "type": "BPE", + "dropout": None, + "unk_token": "", + "continuing_subword_prefix": None, + "end_of_word_suffix": None, + "fuse_unk": True, + "byte_fallback": True, + "ignore_merges": False, + "vocab": vocab, + "merges": merges + } + } + + json_string = json.dumps(tokenizer_dict, ensure_ascii=False) + + logging.info(f"Created tokenizer with vocab size of {len(vocab)}") + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="The given buffer is not writable") + return torch.frombuffer(json_string.encode('utf-8'), dtype=torch.uint8) + +def gguf_json_tokenizer_loader(path): + tenc_fname = os.path.basename(path) + tenc = os.path.splitext(tenc_fname)[0].lower() + tenc = strip_quant_suffix(tenc) + + target = [] + root = os.path.dirname(path) + for fname in os.listdir(root): + name, ext = os.path.splitext(fname) + if ext.lower() != ".json": + continue + if "tokenizer" not in name.lower(): + continue + if tenc in name.lower(): + target.append(fname) + + if len(target) == 0: + logging.info(f"Can't find tokenizer file for '{tenc_fname}' (matching:'{tenc}')!") + return None + if len(target) > 1: + logging.info(f"Ambiguous tokenizer for text encoder '{tenc_fname}', will use first match.") -def gguf_clip_loader(path): - sd, extra = gguf_sd_loader(path, is_text_model=True) + logging.info(f"Using tokenizer '{target[0]}' for text encoder '{tenc_fname}'.") + target = os.path.join(root, target[0]) + + with open(target, "rb") as f: + tokenizer_bytes = f.read() + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="The given buffer is not writable") + return torch.frombuffer(tokenizer_bytes, dtype=torch.uint8) + +def gguf_clip_loader(path, dynamic=False): + sd, extra = gguf_sd_loader(path, is_text_model=True, dynamic=dynamic) 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) + sd["spiece_model"] = gguf_tokenizer_loader(extra.pop("reader"), 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"}: + elif arch in {"llama", "qwen2vl", "qwen3", "qwen3vl", "gemma3", "gemma4"}: # 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) + sd["tekken_model"] = gguf_tekken_tokenizer_loader(extra.pop("reader"), 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) + sd["spiece_model"] = gguf_gemma3_tokenizer_loader(extra.pop("reader")) + if arch == "gemma4": + sd["tokenizer_json"] = gguf_gemma4_tokenizer_loader(extra.pop("reader")) + else: + # 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) + elif arch == "gemma4": + sd = sd_map_replace(sd, GEMMA4_SD_MAP) + + # temporary workaround + sd["model.embed_tokens.weight"] = dequantize_tensor(sd["model.embed_tokens.weight"], dtype=torch.bfloat16) + sd["model.embed_tokens_per_layer.weight"] = dequantize_tensor(sd["model.embed_tokens_per_layer.weight"], dtype=torch.bfloat16).as_subclass(torch.Tensor) + sd["model.norm.weight"] = dequantize_tensor(sd["model.norm.weight"], dtype=torch.bfloat16) 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) + if arch in {"qwen2vl", "qwen3vl", "gemma4"}: + vsd = gguf_mmproj_loader(path, dynamic=dynamic) + + if vsd: + # MiniMax-H3 uses the truncated Qwen3-VL-32B encoder. + # ComfyUI detects it by: + # visual.deepstack_merger_list... + # model.layers.49... + # + # The generic Qwen3-VL mmproj mapper produces model.visual.*, + # which makes ComfyUI incorrectly instantiate Qwen3-VL-8B. + is_minimax_h3 = ( + arch == "qwen3vl" + and "model.layers.49.self_attn.q_proj.weight" in sd + ) + + if is_minimax_h3: + vsd = { + ( + key.replace("model.visual.", "visual.", 1) + if key.startswith("model.visual.") + else key + ): value + for key, value in vsd.items() + } + sd.update(vsd) + + elif arch == "qwen3vl" and "model.norm.weight" in sd: + # Generic full-model fallback only. + weight = sd["model.norm.weight"].shape[0] + sd["model.visual.deepstack_merger_list.0.norm.weight"] = torch.zeros( + 4096 if weight < 4096 else 4608 + ) + sd["model.visual.merger.linear_fc2.weight"] = torch.zeros(weight) else: pass return sd + diff --git a/nodes.py b/nodes.py index 4683514..28bd645 100644 --- a/nodes.py +++ b/nodes.py @@ -11,6 +11,7 @@ import comfy.utils import comfy.model_patcher import comfy.model_management +import comfy.memory_management import folder_paths from .ops import GGMLOps, move_patch_to_device @@ -32,6 +33,19 @@ def update_folder_names_and_paths(key, targets=[]): update_folder_names_and_paths("unet_gguf", ["diffusion_models", "unet"]) update_folder_names_and_paths("clip_gguf", ["text_encoders", "clip"]) +def _clone_as_gguf_model_patcher(self, *args, model_override=None, **kwargs): + if model_override is None: + model_override = self.get_clone_model_override() + mmap_released = model_override[2] if len(model_override) > 2 else False + src_cls = self.__class__ + self.__class__ = GGUFModelPatcher + n = comfy.model_patcher.ModelPatcher.clone(self, *args, model_override=model_override, **kwargs) + n.__class__ = GGUFModelPatcher + self.__class__ = src_cls + n.patch_on_device = getattr(self, "patch_on_device", False) + n.mmap_released = mmap_released + return n + class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): patch_on_device = False @@ -89,6 +103,9 @@ def pin_weight_to_device(self, key): mmap_released = False named_modules_to_munmap = {} + def get_clone_model_override(self): + return (*super().get_clone_model_override(), self.mmap_released) + def load(self, *args, force_patch_weights=False, **kwargs): if not self.mmap_released: self.named_modules_to_munmap = dict(self.model.named_modules()) @@ -120,18 +137,104 @@ def load(self, *args, force_patch_weights=False, **kwargs): self.named_modules_to_munmap = {} def clone(self, *args, **kwargs): - src_cls = self.__class__ - self.__class__ = GGUFModelPatcher - n = super().clone(*args, **kwargs) - n.__class__ = GGUFModelPatcher - self.__class__ = src_cls - # GGUF specific clone values below - n.patch_on_device = getattr(self, "patch_on_device", False) - n.mmap_released = getattr(self, "mmap_released", False) - if src_cls != GGUFModelPatcher: + n = _clone_as_gguf_model_patcher(self, *args, **kwargs) + if self.__class__ != GGUFModelPatcher: n.size = 0 # force recalc return n +class GGUFModelPatcherDynamic(comfy.model_patcher.ModelPatcherDynamic): + patch_on_device = False + + def load(self, *args, **kwargs): + super().load(*args, **kwargs) + # GGML can't requantize after LoRA - demote lowvram_function to weight_function + for n, m in self.model.named_modules(): + for param_key in ("weight", "bias"): + attr = param_key + "_lowvram_function" + fn = getattr(m, attr, None) + if fn is not None: + setattr(m, attr, None) + fns = getattr(m, param_key + "_function", []) + fns.append(fn) + setattr(m, param_key + "_function", fns) + if self.patch_on_device: + for key in self.patches: + self.patches[key] = move_patch_to_device(self.patches[key], self.load_device) + + def clone(self, disable_dynamic=False, model_override=None): + if disable_dynamic: + if model_override is None: + temp = self.cached_patcher_init[0](*self.cached_patcher_init[1], disable_dynamic=True) + model_override = temp.get_clone_model_override() + n = _clone_as_gguf_model_patcher(self, model_override=model_override) + return n + n = super().clone(disable_dynamic=disable_dynamic, model_override=model_override) + n.patch_on_device = self.patch_on_device + return n + +def _clone_patcher_to_gguf(model_patcher): + if model_patcher.is_dynamic(): + src_cls = model_patcher.__class__ + model_patcher.__class__ = GGUFModelPatcherDynamic + n = model_patcher.clone() + model_patcher.__class__ = src_cls + return n + else: + return GGUFModelPatcher.clone(model_patcher) + +def _load_gguf_unet(unet_path, ops, disable_dynamic=False): + dynamic = not disable_dynamic and comfy.memory_management.aimdo_enabled + sd, extra = gguf_sd_loader(unet_path, dynamic=dynamic) + + 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={} if dynamic else { "custom_operations" : ops }, disable_dynamic=disable_dynamic, **kwargs, + ) + if model is None: + logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) + raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) + model = _clone_patcher_to_gguf(model) + + model.cached_patcher_init = (_load_gguf_unet, (unet_path, ops)) + + return model + +def _load_gguf_clip_patcher(clip_paths, clip_type, disable_dynamic=False): + return _load_gguf_clip(clip_paths, clip_type, disable_dynamic=disable_dynamic).patcher + +def _load_gguf_clip(clip_paths, clip_type, disable_dynamic=False): + dynamic = not disable_dynamic and comfy.memory_management.aimdo_enabled + + clip_data = [] + for p in clip_paths: + if p.endswith(".gguf"): + sd = gguf_clip_loader(p, dynamic=dynamic) + else: + sd = comfy.utils.load_torch_file(p, safe_load=True) + if not dynamic and "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) + + model_options = {"initial_device": comfy.model_management.text_encoder_offload_device()} + if not dynamic: + model_options["custom_operations"] = GGMLOps + + clip = comfy.sd.load_text_encoder_state_dicts( + clip_type = clip_type, + state_dicts = clip_data, + model_options = model_options, + embedding_directory = folder_paths.get_folder_paths("embeddings"), + disable_dynamic = disable_dynamic, + ) + clip.patcher = _clone_patcher_to_gguf(clip.patcher) + + clip.patcher.cached_patcher_init = (_load_gguf_clip_patcher, (clip_paths, clip_type)) + return clip + class UnetLoaderGGUF: @classmethod def INPUT_TYPES(s): @@ -164,22 +267,8 @@ 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", {}) - - model = comfy.sd.load_diffusion_model_state_dict( - sd, model_options={"custom_operations": ops}, **kwargs, - ) - if model is None: - logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) - raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) - model = GGUFModelPatcher.clone(model) + model = _load_gguf_unet(unet_path, ops) model.patch_on_device = patch_on_device return (model,) @@ -220,35 +309,10 @@ def get_filename_list(s): files += folder_paths.get_filename_list("clip_gguf") return sorted(files) - def load_data(self, ckpt_paths): - clip_data = [] - for p in ckpt_paths: - if p.endswith(".gguf"): - sd = gguf_clip_loader(p) - else: - sd = comfy.utils.load_torch_file(p, safe_load=True) - if "scaled_fp8" in sd: # NOTE: Scaled FP8 would require different custom ops, but only one can be active - raise NotImplementedError(f"Mixing scaled FP8 with GGUF is not supported! Use regular CLIP loader or switch model(s)\n({p})") - clip_data.append(sd) - return clip_data - - def load_patcher(self, clip_paths, clip_type, clip_data): - clip = comfy.sd.load_text_encoder_state_dicts( - clip_type = clip_type, - state_dicts = clip_data, - model_options = { - "custom_operations": GGMLOps, - "initial_device": comfy.model_management.text_encoder_offload_device() - }, - embedding_directory = folder_paths.get_folder_paths("embeddings"), - ) - clip.patcher = GGUFModelPatcher.clone(clip.patcher) - return clip - def load_clip(self, clip_name, type="stable_diffusion"): clip_path = folder_paths.get_full_path("clip", clip_name) clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),) + return (_load_gguf_clip([clip_path], clip_type),) class DualCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod @@ -270,7 +334,7 @@ def load_clip(self, clip_name1, clip_name2, type): clip_path2 = folder_paths.get_full_path("clip", clip_name2) clip_paths = (clip_path1, clip_path2) clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + return (_load_gguf_clip(clip_paths, clip_type),) class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod @@ -292,7 +356,7 @@ def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"): clip_path3 = folder_paths.get_full_path("clip", clip_name3) clip_paths = (clip_path1, clip_path2, clip_path3) clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + return (_load_gguf_clip(clip_paths, clip_type),) class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod @@ -316,7 +380,7 @@ def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable clip_path4 = folder_paths.get_full_path("clip", clip_name4) clip_paths = (clip_path1, clip_path2, clip_path3, clip_path4) clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + return (_load_gguf_clip(clip_paths, clip_type),) NODE_CLASS_MAPPINGS = { "UnetLoaderGGUF": UnetLoaderGGUF, @@ -326,4 +390,3 @@ def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable "QuadrupleCLIPLoaderGGUF": QuadrupleCLIPLoaderGGUF, "UnetLoaderGGUFAdvanced": UnetLoaderGGUFAdvanced, } - diff --git a/ops.py b/ops.py index 88a352e..6da399e 100644 --- a/ops.py +++ b/ops.py @@ -8,6 +8,11 @@ import comfy.model_management from .dequant import dequantize_tensor, is_quantized +try: + import comfy.weight_adapter as wadapter +except (ImportError, ModuleNotFoundError): + wadapter = None + def chained_hasattr(obj, chained_attr): probe = obj for attr in chained_attr.split('.'): @@ -253,7 +258,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) @@ -270,12 +275,28 @@ def forward_ggml_cast_weights(self, input): weight, bias = self.cast_bias_weight(input) return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps) -def move_patch_to_device(item, device): - if isinstance(item, torch.Tensor): - return item.to(device, non_blocking=True) - elif isinstance(item, tuple): - return tuple(move_patch_to_device(x, device) for x in item) - elif isinstance(item, list): - return [move_patch_to_device(x, device) for x in item] - else: +def move_patch_to_device(item, device, *, dtype=None): + if device is None: return item + if isinstance(item, torch.Tensor): + return comfy.model_management.cast_to_device(item, device, dtype, copy=True) + if isinstance(item, (tuple, list)): + return item.__class__( + move_patch_to_device(seqitem, device, dtype=dtype) for seqitem in item + ) + if ( + wadapter is not None + and isinstance(item, wadapter.WeightAdapterBase) + and hasattr(item, "loaded_keys") + and isinstance(getattr(item, "weights", None), (tuple, list)) + ): + return item.__class__( + item.loaded_keys, + item.weights.__class__( + wi + if not isinstance(wi, torch.Tensor) + else comfy.model_management.cast_to_device(wi, device, dtype, copy=True) + for wi in item.weights + ), + ) + return item diff --git a/quant_ops.py b/quant_ops.py new file mode 100644 index 0000000..a31c2ac --- /dev/null +++ b/quant_ops.py @@ -0,0 +1,95 @@ +# GGML QuantizedTensor support for dynamic VRAM loading +import gguf +import torch +from dataclasses import dataclass + +try: + from comfy_kitchen.tensor import ( + QuantizedTensor, + QuantizedLayout, + BaseLayoutParams, + register_layout_class, + ) + _CK_AVAILABLE = True +except ImportError: + _CK_AVAILABLE = False + + class QuantizedTensor: + pass + + class QuantizedLayout: + pass + + class BaseLayoutParams: + pass + + def register_layout_class(name, cls): + pass + +from .dequant import dequantize_functions, TORCH_COMPATIBLE_QTYPES, is_quantized + +COMPILED_DEQUANT_FUNCTIONS = {} +if hasattr(torch, 'compile'): + for k, func in dequantize_functions.items(): + COMPILED_DEQUANT_FUNCTIONS[k] = torch.compile(func) +else: + COMPILED_DEQUANT_FUNCTIONS = dequantize_functions.copy() + +if _CK_AVAILABLE: + @dataclass(frozen=True) + class GGMLLayoutParams(BaseLayoutParams): + tensor_type: int # gguf.GGMLQuantizationType stored as int + + class GGMLLayout(QuantizedLayout): + Params = GGMLLayoutParams + + @classmethod + def quantize(cls, tensor, **kwargs): + raise NotImplementedError("Quantization to GGML format is not supported") + + @classmethod + def dequantize(cls, qdata, params): + qtype = gguf.GGMLQuantizationType(params.tensor_type) + oshape = params.orig_shape + + if qtype in TORCH_COMPATIBLE_QTYPES: + return qdata.reshape(oshape).to(params.orig_dtype) + + if qtype not in dequantize_functions: + from tqdm import tqdm + tqdm.write(f"Falling back to numpy dequant for qtype: {qtype.name}") + new = gguf.quants.dequantize(qdata.cpu().numpy(), qtype) + return torch.from_numpy(new).reshape(oshape).to(device=qdata.device, dtype=params.orig_dtype) + + block_size, type_size = gguf.GGML_QUANT_SIZES[qtype] + raw = qdata.reshape(-1).view(torch.uint8) + n_blocks = raw.numel() // type_size + blocks = raw.reshape((n_blocks, type_size)) + + fn = COMPILED_DEQUANT_FUNCTIONS[qtype] + try: + blocks = fn(blocks, block_size, type_size, params.orig_dtype) + except Exception: + fn = COMPILED_DEQUANT_FUNCTIONS[qtype] = dequantize_functions[qtype] + blocks = fn(blocks, block_size, type_size, params.orig_dtype) + return blocks.reshape(oshape).to(params.orig_dtype) + + @classmethod + def get_plain_tensors(cls, qtensor): + return (qtensor._qdata,) + + @classmethod + def state_dict_tensors(cls, qdata, params): + return {"weight": qdata} + + register_layout_class("GGMLLayout", GGMLLayout) + +def make_quantized(qdata, tensor_type, tensor_shape, orig_dtype=torch.float16): + """Construct a GGML QuantizedTensor from raw packed data.""" + params = GGMLLayoutParams( + scale=torch.ones((), dtype=torch.float32), + orig_dtype=orig_dtype, + orig_shape=tuple(tensor_shape), + tensor_type=tensor_type.value if not isinstance(tensor_type, int) else tensor_type, + ) + return QuantizedTensor(qdata, "GGMLLayout", params) diff --git a/tools/convert.py b/tools/convert.py index 5029c87..4033377 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -145,8 +145,15 @@ class ModelLumina2(ModelTemplate): ("cap_embedder.1.weight", "context_refiner.0.attention.qkv.weight") ] +class ModelIdeogram4(ModelTemplate): + arch = "ideogram4" + keys_detect = [ + ("embed_image_indicator.weight", "llm_cond_proj.weight", "layers.0.adaln_modulation.weight"), + ("embed_image_indicator.weight", "input_proj.weight", "adaln_proj.weight"), + ] + arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, - ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2] + ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2, ModelIdeogram4] def is_model_arch(model, state_dict): # check if model is correct