import comfy.sd import comfy.ops import comfy.utils import comfy.model_patcher import comfy.model_management import torch, numpy, os, json, logging, collections, nodes, folder_paths from safetensors.torch import load_file, save_file from typing import Dict, Tuple from tqdm import tqdm as loading from .gguf_connector import reader as gr from .gguf_connector.writer import GGUFWriter, GGMLQuantizationType from .gguf_connector.const import GGML_QUANT_VERSION, LlamaFileType from .gguf_connector.quant import quantize, dequantize, QuantError from .gguf_connector.quant2c import dequantize_tensor, is_quantized, is_torch_compatible from .gguf_connector.tkn import get_field, tokenizer_builder pig = os.path.join(os.path.dirname(__file__), 'version.json') with open(pig, 'r') as file: data = json.load(file) arrays = {} for key, value in data[0].items(): arrays[key] = value class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): patch_on_device = False def patch_weight_to_device(self, key, device_to=None, inplace_update=False): if key not in self.patches: return weight = comfy.utils.get_attr(self.model, key) try: from comfy.lora import calculate_weight except Exception: calculate_weight = self.calculate_weight patches = self.patches[key] if is_quantized(weight): out_weight = weight.to(device_to) patches = load_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device) out_weight.patches = [(calculate_weight, patches, key)] else: inplace_update = self.weight_inplace_update or inplace_update if key not in self.backup: self.backup[key] = collections.namedtuple('Dimension', [ 'weight', 'inplace_update'])(weight.to(device=self. offload_device, copy=inplace_update), inplace_update) if device_to is not None: temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True) else: temp_weight = weight.to(torch.float32, copy=True) out_weight = calculate_weight(patches, temp_weight, key) out_weight = comfy.float.stochastic_rounding(out_weight, weight .dtype) if inplace_update: comfy.utils.copy_to_param(self.model, key, out_weight) else: comfy.utils.set_attr_param(self.model, key, out_weight) def unpatch_model(self, device_to=None, unpatch_weights=True): if unpatch_weights: for p in self.model.parameters(): if is_torch_compatible(p): continue patches = getattr(p, 'patches', []) if len(patches) > 0: p.patches = [] return super().unpatch_model(device_to=device_to, unpatch_weights= unpatch_weights) mmap_released = False def load(self, *args, force_patch_weights=False, **kwargs): super().load(*args, force_patch_weights=True, **kwargs) if not self.mmap_released: linked = [] if kwargs.get('lowvram_model_memory', 0) > 0: for n, m in self.model.named_modules(): if hasattr(m, 'weight'): device = getattr(m.weight, 'device', None) if device == self.offload_device: linked.append((n, m)) continue if hasattr(m, 'bias'): device = getattr(m.bias, 'device', None) if device == self.offload_device: linked.append((n, m)) continue if linked: print(f'Attempting to release mmap ({len(linked)})') for n, m in linked: m.to(self.load_device).to(self.offload_device) self.mmap_released = True def clone(self, *args, **kwargs): src_cls = self.__class__ self.__class__ = GGUFModelPatcher n = super().clone(*args, **kwargs) n.__class__ = GGUFModelPatcher self.__class__ = src_cls n.patch_on_device = getattr(self, 'patch_on_device', False) return n class GGMLTensor(torch.Tensor): def __init__(self, *args, tensor_type, tensor_shape, patches=[], **kwargs): super().__init__() self.tensor_type = tensor_type self.tensor_shape = tensor_shape self.patches = patches def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): return super().__new__(cls, *args, **kwargs) def to(self, *args, **kwargs): new = super().to(*args, **kwargs) new.tensor_type = getattr(self, 'tensor_type', None) new.tensor_shape = getattr(self, 'tensor_shape', new.data.shape) new.patches = getattr(self, 'patches', []).copy() return new def clone(self, *args, **kwargs): return self def detach(self, *args, **kwargs): return self def copy_(self, *args, **kwargs): try: return super().copy_(*args, **kwargs) except Exception as e: print(f"ignoring 'copy_' on tensor: {e}") def empty_(self, size, *args, **kwargs): new_tensor = super().empty_(size, *args, **kwargs) return GGMLTensor(new_tensor, tensor_type=getattr(self, 'tensor_type', None), tensor_shape=size, patches=getattr(self, 'patches', []).copy()) @property def shape(self): if not hasattr(self, 'tensor_shape'): self.tensor_shape = self.size() return self.tensor_shape if hasattr(torch, 'compiler') and hasattr(torch.compiler, 'disable'): torch_compiler_disable = torch.compiler.disable else: def torch_compiler_disable(*args, **kwargs): def noop(x): return x return noop class GGMLLayer(torch.nn.Module): comfy_cast_weights = True dequant_dtype = None patch_dtype = None largest_layer = False torch_compatible_tensor_types = {None, gr.GGMLQuantizationType.F32, gr. GGMLQuantizationType.F16} def is_ggml_quantized(self, *, weight=None, bias=None): if weight is None: weight = self.weight if bias is None: bias = self.bias return is_quantized(weight) or is_quantized(bias) def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): weight, bias = state_dict.get(f'{prefix}weight'), state_dict.get( f'{prefix}bias') if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(self, torch.nn.Linear): return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) return super()._load_from_state_dict(state_dict, prefix, *args, ** kwargs) def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): prefix_len = len(prefix) for k, v in state_dict.items(): if k[prefix_len:] == 'weight': self.weight = torch.nn.Parameter(v, requires_grad=False) elif k[prefix_len:] == 'bias' and v is not None: self.bias = torch.nn.Parameter(v, requires_grad=False) else: missing_keys.append(k) if self.weight is None and isinstance(self, torch.nn.Linear): v = torch.zeros(self.in_features, self.out_features) self.weight = torch.nn.Parameter(v, requires_grad=False) missing_keys.append(prefix + 'weight') if getattr(self.weight, 'is_largest_weight', False): self.largest_layer = True def _save_to_state_dict(self, *args, **kwargs): if self.is_ggml_quantized(): return self.ggml_save_to_state_dict(*args, **kwargs) return super()._save_to_state_dict(*args, **kwargs) def ggml_save_to_state_dict(self, destination, prefix, keep_vars): weight = torch.zeros_like(self.weight, device=torch.device('meta')) destination[prefix + 'weight'] = weight if self.bias is not None: bias = torch.zeros_like(self.bias, device=torch.device('meta')) destination[prefix + 'bias'] = bias if self.largest_layer: shape = getattr(self.weight, 'tensor_shape', self.weight.shape) dtype = self.dequant_dtype or torch.float16 temp = torch.empty(*shape, device=torch.device('meta'), dtype=dtype ) destination[prefix + 'temp.weight'] = temp return destination[prefix + 'weight'] = self.get_weight(self.weight) if bias is not None: destination[prefix + 'bias'] = self.get_weight(self.bias) def get_weight(self, tensor, dtype): if tensor is None: return patch_list = [] device = tensor.device for function, patches, key in getattr(tensor, 'patches', []): patch_list += load_patch_to_device(patches, device) weight = dequantize_tensor(tensor, dtype, self.dequant_dtype) if isinstance(weight, GGMLTensor): weight = torch.Tensor(weight) if patch_list: if self.patch_dtype is None: weight = function(patch_list, weight, key) else: patch_dtype = (dtype if self.patch_dtype == 'target' else self.patch_dtype) weight = function(patch_list, weight, key, patch_dtype) return weight @torch_compiler_disable() def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype =None): if input is not None: if dtype is None: dtype = getattr(input, 'dtype', torch.float32) if bias_dtype is None: bias_dtype = dtype if device is None: device = input.device bias = None non_blocking = comfy.model_management.device_supports_non_blocking( device) if s.bias is not None: bias = s.get_weight(s.bias.to(device), dtype) bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking =non_blocking, copy=False) weight = s.get_weight(s.weight.to(device), dtype) weight = comfy.ops.cast_to(weight, dtype, device, non_blocking= non_blocking, copy=False) return weight, bias def forward_comfy_cast_weights(self, input, *args, **kwargs): if self.is_ggml_quantized(): out = self.forward_ggml_cast_weights(input, *args, **kwargs) else: out = super().forward_comfy_cast_weights(input, *args, **kwargs) if isinstance(out, GGMLTensor): out = torch.Tensor(out) return out def forward_ggml_cast_weights(self, input): raise NotImplementedError class GGMLOps(comfy.ops.manual_cast): class Linear(GGMLLayer, comfy.ops.manual_cast.Linear): def __init__(self, in_features, out_features, bias=True, device= None, dtype=None): torch.nn.Module.__init__(self) self.in_features = in_features self.out_features = out_features self.weight = None self.bias = None def forward_ggml_cast_weights(self, input): weight, bias = self.cast_bias_weight(input) return torch.nn.functional.linear(input, weight, bias) class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d): def forward_ggml_cast_weights(self, input): weight, bias = self.cast_bias_weight(input) return self._conv_forward(input, weight, bias) class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding): 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) 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) class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm): def forward_ggml_cast_weights(self, input): if self.weight is None: return super().forward_comfy_cast_weights(input) weight, bias = self.cast_bias_weight(input) return torch.nn.functional.layer_norm(input, self. normalized_shape, weight, bias, self.eps) class GroupNorm(GGMLLayer, comfy.ops.manual_cast.GroupNorm): 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 load_patch_to_device(item, device): if isinstance(item, torch.Tensor): return item.to(device, non_blocking=True) elif isinstance(item, tuple): return tuple(load_patch_to_device(x, device) for x in item) elif isinstance(item, list): return [load_patch_to_device(x, device) for x in item] else: return item def get_folder_names_and_paths(key, targets=[]): base = folder_paths.folder_names_and_paths.get(key, ([], {})) base = base[0] if isinstance(base[0], (list, set, tuple)) else [] target = next((x for x in targets if x in folder_paths. folder_names_and_paths), targets[0]) orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {})) folder_paths.folder_names_and_paths[key] = orig or base, {'.gguf'} if base and base != orig: logging.warning( f'Unknown file list already present on key {key}: {base}') get_folder_names_and_paths('model_gguf', ['diffusion_models', 'unet']) get_folder_names_and_paths('clip_gguf', ['text_encoders', 'clip']) get_folder_names_and_paths('vae_gguf', ['vae']) 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 if len(field.types) != 2 or field.types[0 ] != gr.GGUFValueType.ARRAY or field.types[1 ] != gr.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 load_gguf_sd(path, handle_prefix='model.diffusion_model.', return_arch= False): reader = gr.GGUFReader(path) 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)) compat = None arch_str = get_field(reader, 'general.architecture', str) if arch_str is None: compat = 'sd.cpp' elif arch_str not in arrays['PIG_ARCH_LIST'] and arch_str not in arrays['TXT_ARCH_LIST']: raise ValueError(f"Unknown architecture: {arch_str!r}") state_dict, qtype_dict = {}, {} for sd_key, tensor in tensors: tensor_name = tensor.name torch_tensor = torch.from_numpy(tensor.data) shape = get_orig_shape(reader, tensor_name) if shape is None: shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) 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] if tensor.tensor_type in {gr.GGMLQuantizationType.F32, gr. GGMLQuantizationType.F16}: torch_tensor = torch_tensor.view(*shape) state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor. tensor_type, tensor_shape=shape) 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('gguf qtypes: ' + ', '.join(f'{k} ({v})' for k, v in qtype_dict.items())) 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 if return_arch: return state_dict, arch_str return state_dict def tensor_swap(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 pig_work(raw_sd): sd = {} for k, v in raw_sd.items(): sd[k] = v return sd def llama_permute(raw_sd, n_head, n_head_kv): 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 load_gguf_clip(path): sd, arch = load_gguf_sd(path, return_arch=True) if arch in {'t5', 't5encoder'}: temb_key = 'token_embd.weight' if temb_key in sd and (sd[temb_key].shape == (256384, 4096) or sd[temb_key].shape == (256384, 768)): sd['spiece_model'] = tokenizer_builder(path) sd = tensor_swap(sd, arrays['T5']) elif temb_key in sd and sd[temb_key].shape == (32128, 768): sd = tensor_swap(sd, arrays['B5']) else: sd = tensor_swap(sd, arrays['T5']) elif arch in {'llama'}: sd = tensor_swap(sd, arrays['L3']) sd = llama_permute(sd, 32, 8) elif arch in {'gemma2'}: sd["spiece_model"] = tokenizer_builder(path) sd = tensor_swap(sd, arrays['G2']) elif arch in {'pig'}: sd = pig_work(sd) else: pass return sd class LoaderGGUF: @classmethod def INPUT_TYPES(s): gguf_names = [x for x in folder_paths.get_filename_list('model_gguf')] return {'required': {'gguf_name': (gguf_names,)}} RETURN_TYPES = 'MODEL', FUNCTION = 'load_model' CATEGORY = 'gguf' TITLE = 'GGUF Loader' def load_model(self, gguf_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): ops = GGMLOps() if dequant_dtype in ('default', None): ops.Linear.dequant_dtype = None elif dequant_dtype in ['target']: ops.Linear.dequant_dtype = dequant_dtype else: ops.Linear.dequant_dtype = getattr(torch, dequant_dtype) if patch_dtype in ('default', None): ops.Linear.patch_dtype = None elif patch_dtype in ['target']: ops.Linear.patch_dtype = patch_dtype else: ops.Linear.patch_dtype = getattr(torch, patch_dtype) model_path = folder_paths.get_full_path('unet', gguf_name) sd = load_gguf_sd(model_path) model = comfy.sd.load_diffusion_model_state_dict(sd, model_options= {'custom_operations': ops}) if model is None: logging.error('ERROR UNSUPPORTED MODEL {}'.format(model_path)) raise RuntimeError('ERROR: Could not detect model type of: {}'. format(model_path)) model = GGUFModelPatcher.clone(model) model.patch_on_device = patch_on_device return model, class LoaderGGUFAdvanced(LoaderGGUF): @classmethod def INPUT_TYPES(s): model_names = [x for x in folder_paths.get_filename_list('model_gguf')] return {'required': {'gguf_name': (model_names,), 'dequant_dtype': (['default', 'target', 'float32', 'float16', 'bfloat16'], { 'default': 'default'}), 'patch_dtype': (['default', 'target', 'float32', 'float16', 'bfloat16'], {'default': 'default'}), 'patch_on_device': ('BOOLEAN', {'default': False})}} TITLE = 'GGUF Loader (Advanced)' def get_clip_type(type): clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) return clip_type def get_device(device): model_options = {} if device == "cpu": model_options["load_device"] = model_options["offload_device"] = torch.device("cpu") return model_options class ClipLoaderGGUF: @classmethod def INPUT_TYPES(s): base = nodes.CLIPLoader.INPUT_TYPES() return {'required': {'clip_name': (s.get_filename_list(),), 'type': base['required']['type']}, 'optional':{'device':(['default','cpu'],{'advanced':True}),}} RETURN_TYPES = 'CLIP', FUNCTION = 'load_clip' CATEGORY = 'gguf' TITLE = 'GGUF CLIP Loader' @classmethod def get_filename_list(s): files = [] files += folder_paths.get_filename_list('clip') files += folder_paths.get_filename_list('clip_gguf') return sorted(files) def load_data(self, ckpt_paths): clip_data = [] for p in ckpt_paths: if p.endswith('.gguf'): sd = load_gguf_clip(p) else: sd = comfy.utils.load_torch_file(p, safe_load=True) 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', device='default'): clip_path = folder_paths.get_full_path('clip', clip_name) if clip_name.endswith('.safetensors') and device != 'default': clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=get_clip_type(type), model_options=get_device(device)) return (clip,) else: return (self.load_patcher([clip_path], get_clip_type(type), self.load_data([clip_path])), get_device(device)) class DualClipLoaderGGUF(ClipLoaderGGUF): @classmethod def INPUT_TYPES(s): base = nodes.DualCLIPLoader.INPUT_TYPES() file_options = s.get_filename_list(), return {'required': {'clip_name1':file_options, 'clip_name2':file_options, 'type': base['required']['type']}, 'optional':{'device':(['default','cpu'],{'advanced':True}),}} TITLE = 'GGUF DualCLIP Loader' def load_clip(self, clip_name1, clip_name2, type, device='default'): clip_path1 = folder_paths.get_full_path('clip', clip_name1) clip_path2 = folder_paths.get_full_path('clip', clip_name2) clip_paths = clip_path1, clip_path2 if device != 'default' and (clip_name1.endswith('.safetensors') and clip_name2.endswith('.safetensors')): clip = comfy.sd.load_clip(ckpt_paths=clip_paths, embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=get_clip_type(type), model_options=get_device(device)) return (clip,) else: return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)), get_device(device)) class TripleClipLoaderGGUF(ClipLoaderGGUF): @classmethod def INPUT_TYPES(s): file_options = s.get_filename_list(), return {'required': {'clip_name1': file_options, 'clip_name2': file_options, 'clip_name3': file_options}} TITLE = 'GGUF TripleCLIP Loader' def load_clip(self, clip_name1, clip_name2, clip_name3, type='sd3'): clip_path1 = folder_paths.get_full_path('clip', clip_name1) clip_path2 = folder_paths.get_full_path('clip', clip_name2) clip_path3 = folder_paths.get_full_path('clip', clip_name3) clip_paths = clip_path1, clip_path2, clip_path3 return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),) class QuadrupleClipLoaderGGUF(ClipLoaderGGUF): @classmethod def INPUT_TYPES(s): file_options = s.get_filename_list(), return {'required': {'clip_name1': file_options, 'clip_name2': file_options, 'clip_name3': file_options, 'clip_name4': file_options}} TITLE = 'GGUF QuadrupleCLIP Loader' def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type='hunyuan_video'): clip_path1 = folder_paths.get_full_path('clip', clip_name1) clip_path2 = folder_paths.get_full_path('clip', clip_name2) clip_path3 = folder_paths.get_full_path('clip', clip_name3) clip_path4 = folder_paths.get_full_path('clip', clip_name4) clip_paths = clip_path1, clip_path2, clip_path3, clip_path4 return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),) QUANTIZATION_THRESHOLD = 1024 REARRANGE_THRESHOLD = 512 MAX_TENSOR_NAME_LENGTH = 127 class ModelTemplate: arch = 'invalid' keys_detect = [] keys_banned = [] 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 ModelSDXL(ModelTemplate): arch = 'sdxl' 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'), ('label_emb.0.0.weight',)] class ModelSD1(ModelTemplate): arch = 'sd1' 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')] arch_list = [ModelFlux, ModelSD3, ModelSDXL, ModelSD1] def is_model_arch(model, state_dict): 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 supported for alpha (please opt to use zero)' 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 detected!' return model_arch def load_state_dict(path): state_dict = load_file(path) prefix = None for pfx in ['model.diffusion_model.', 'model.', 'net.']: if any([x.startswith(pfx) for x in state_dict.keys()]): prefix = pfx break sd = {} for k, v in state_dict.items(): if prefix and prefix not in k: continue if prefix: k = k.replace(prefix, '') sd[k] = v return sd def load_model(path): state_dict = load_state_dict(path) model_arch = detect_arch(state_dict) print(f'* Architecture detected from input: {model_arch.arch}') writer = GGUFWriter(path=None, arch=model_arch.arch) return writer, state_dict, model_arch def load_pig_state(path): pig_state = load_file(path) sd = {} for k, v in pig_state.items(): sd[k] = v return sd def load_pig(path): state_dict = load_pig_state(path) model_arch = arrays['TXT_ARCH_LIST'][0] writer = GGUFWriter(path=None, arch=model_arch) return writer, state_dict, model_arch def handle_tensors(args, 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 loading(state_dict.items()): old_dtype = data.dtype if data.dtype == torch.bfloat16: data = data.to(torch.float32).numpy() 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 data_qtype = getattr(GGMLQuantizationType, 'BF16' if old_dtype == torch.bfloat16 else 'F16') n_params = 1 for dim_size in data_shape: n_params *= dim_size blacklist = {'time_embedding.', 'add_embedding.', 'time_in.', 'txt_in.', 'vector_in.', 'img_in.', 'guidance_in.', 'final_layer.'} if old_dtype in (torch.float32, torch.bfloat16): if n_dims == 1: data_qtype = GGMLQuantizationType.F32 elif n_params <= QUANTIZATION_THRESHOLD: data_qtype = GGMLQuantizationType.F32 elif '.weight' in key and any(x in key for x in blacklist): data_qtype = GGMLQuantizationType.F32 try: data = quantize(data, data_qtype) except (AttributeError, QuantError) as e: loading.write(f'falling back to F16: {e}') data_qtype = GGMLQuantizationType.F16 data = quantize(data, data_qtype) new_name = key shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" loading.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) if 'select_safetensors' not in folder_paths.folder_names_and_paths: orig = folder_paths.folder_names_and_paths.get('diffusion_models', folder_paths.folder_names_and_paths.get('checkpoints', [[], set()])) folder_paths.folder_names_and_paths['select_safetensors'] = orig[0], { '.safetensors'} class GGUFSave: def __init__(self): self.output_dir = folder_paths.get_output_directory() @classmethod def INPUT_TYPES(s): return {'required': {'select_safetensors': (s.get_filename_list(),)}} RETURN_TYPES = () FUNCTION = 'save' OUTPUT_NODE = True CATEGORY = 'gguf' TITLE = 'GGUF Convertor (Alpha)' @classmethod def get_filename_list(s): files = [] files += folder_paths.get_filename_list('select_safetensors') return sorted(files) def save(self, select_safetensors): path = folder_paths.get_full_path('select_safetensors', select_safetensors) writer, state_dict, model_arch = load_model(path) writer.add_quantization_version(GGML_QUANT_VERSION) if next(iter(state_dict.values())).dtype == torch.bfloat16: output_path = ( f'{self.output_dir}/{os.path.splitext(select_safetensors)[0]}-bf16.gguf' ) writer.add_file_type(LlamaFileType.MOSTLY_BF16) else: output_path = ( f'{self.output_dir}/{os.path.splitext(select_safetensors)[0]}-f16.gguf' ) writer.add_file_type(LlamaFileType.MOSTLY_F16) if os.path.isfile(output_path): input('Output exists enter to continue or ctrl+c to abort!') handle_tensors(output_path, writer, state_dict, model_arch) writer.write_header_to_file(path=output_path) writer.write_kv_data_to_file() writer.write_tensors_to_file(progress=True) writer.close() return {} class GGUFRun: def __init__(self): self.output_dir = folder_paths.get_output_directory() @classmethod def INPUT_TYPES(s): return {'required': {'select_safetensors': (s.get_filename_list(),)}} RETURN_TYPES = () FUNCTION = 'run' OUTPUT_NODE = True CATEGORY = 'gguf' TITLE = 'GGUF Convertor (Zero)' @classmethod def get_filename_list(s): files = [] files += folder_paths.get_filename_list('select_safetensors') return sorted(files) def run(self, select_safetensors): path = folder_paths.get_full_path('select_safetensors', select_safetensors) writer, state_dict, model_arch = load_pig(path) writer.add_quantization_version(GGML_QUANT_VERSION) if next(iter(state_dict.values())).dtype == torch.bfloat16: output_path = ( f'{self.output_dir}/{os.path.splitext(select_safetensors)[0]}-bf16.gguf' ) writer.add_file_type(LlamaFileType.MOSTLY_BF16) else: output_path = ( f'{self.output_dir}/{os.path.splitext(select_safetensors)[0]}-f16.gguf' ) writer.add_file_type(LlamaFileType.MOSTLY_F16) if os.path.isfile(output_path): input('Output exists enter to continue or ctrl+c to abort!') handle_tensors(output_path, writer, state_dict, model_arch) writer.write_header_to_file(path=output_path) writer.write_kv_data_to_file() writer.write_tensors_to_file(progress=True) writer.close() return {} def quantize_to_fp8(tensor): if tensor.dtype != torch.bfloat16: raise ValueError('Input tensor must be in BF16 format.') tensor = tensor.to(torch.float16) fp8_max = 240.0 fp8_min = -fp8_max clamped_tensor = tensor.clamp(min=fp8_min, max=fp8_max) scale = fp8_max / torch.max(torch.abs(clamped_tensor)) quantized_tensor = torch.round(clamped_tensor * scale) / scale return quantized_tensor.to(torch.float8_e4m3fn) class TENSORCut: def __init__(self): self.output_dir = folder_paths.get_output_directory() @classmethod def INPUT_TYPES(s): return {'required': {'select_safetensors': (s.get_filename_list(),)}} RETURN_TYPES = () FUNCTION = 'cut' OUTPUT_NODE = True CATEGORY = 'gguf' TITLE = 'TENSOR Cutter (Beta)' @classmethod def get_filename_list(s): files = [] files += folder_paths.get_filename_list('select_safetensors') return sorted(files) def cut(self, select_safetensors): input_file = folder_paths.get_full_path('select_safetensors', select_safetensors) output_file = ( f'{self.output_dir}/{os.path.splitext(select_safetensors)[0]}_fp8_e4m3fn.safetensors' ) data = load_file(input_file) quantized_data = {} print('Starting quantization process...') for key, tensor in loading(data.items(), desc='Quantizing tensors', unit='tensor'): tensor = tensor.to(dtype=torch.bfloat16, device='cuda') quantized_tensor = quantize_to_fp8(tensor) quantized_data[key] = quantized_tensor.cpu() save_file(quantized_data, output_file) print(f'Quantized safetensors saved to {output_file}.') return {} class TENSORBoost: def __init__(self): self.output_dir = folder_paths.get_output_directory() @classmethod def INPUT_TYPES(s): return {'required': {'select_safetensors': (s.get_filename_list(),)}} RETURN_TYPES = () FUNCTION = 'boost' OUTPUT_NODE = True CATEGORY = 'gguf' TITLE = 'TENSOR Booster' @classmethod def get_filename_list(s): files = [] files += folder_paths.get_filename_list('select_safetensors') return sorted(files) def boost(self, select_safetensors): input_file = folder_paths.get_full_path('select_safetensors', select_safetensors) output_file = ( f'{self.output_dir}/{os.path.splitext(select_safetensors)[0]}_fp32.safetensors' ) data = load_file(input_file) quantized_data = {} print('Starting quantization process...') for key, tensor in loading(data.items(), desc="Converting tensors", unit="tensor"): quantized_data[key] = tensor.to(torch.float32) save_file(quantized_data, output_file) print(f'Quantized safetensors saved to {output_file}.') return {} def load_gguf_and_extract_metadata(gguf_path): reader = gr.GGUFReader(gguf_path) tensors_metadata = [] for tensor in reader.tensors: tensor_metadata = {'name': tensor.name, 'shape': tuple(tensor.shape .tolist()), 'n_elements': tensor.n_elements, 'n_bytes': tensor. n_bytes, 'data_offset': tensor.data_offset, 'type': tensor. tensor_type} tensors_metadata.append(tensor_metadata) return reader, tensors_metadata def convert_gguf_to_safetensors(gguf_path, output_path, use_bf16): reader, tensors_metadata = load_gguf_and_extract_metadata(gguf_path) print(f'Extracted {len(tensors_metadata)} tensors from GGUF file') tensors_dict: dict[str, torch.Tensor] = {} for i, tensor_info in enumerate(loading(tensors_metadata, desc= 'Converting tensors', unit='tensor')): tensor_name = tensor_info['name'] tensor_data = reader.get_tensor(i) weights = dequantize(tensor_data.data, tensor_data.tensor_type).copy() try: if use_bf16: weights_tensor = torch.from_numpy(weights).to(dtype=torch. float32) weights_tensor = weights_tensor.to(torch.bfloat16) else: weights_tensor = torch.from_numpy(weights).to(dtype=torch. float16) weights_hf = weights_tensor except Exception as e: print( f"Error during BF16 conversion for tensor '{tensor_name}': {e}" ) weights_tensor = torch.from_numpy(weights.astype(numpy.float32) ).to(torch.float16) weights_hf = weights_tensor tensors_dict[tensor_name] = weights_hf metadata = {key: str(reader.get_field(key)) for key in reader.fields} save_file(tensors_dict, output_path, metadata=metadata) print('Conversion complete!') if 'select_gguf' not in folder_paths.folder_names_and_paths: orig = folder_paths.folder_names_and_paths.get('diffusion_models', folder_paths.folder_names_and_paths.get('unet', [[], set()])) folder_paths.folder_names_and_paths['select_gguf'] = orig[0], {'.gguf'} class GGUFUndo: def __init__(self): self.output_dir = folder_paths.get_output_directory() @classmethod def INPUT_TYPES(s): return {'required': {'select_gguf': (s.get_filename_list(),)}} RETURN_TYPES = () FUNCTION = 'undo' OUTPUT_NODE = True CATEGORY = 'gguf' TITLE = 'GGUF Convertor (Reverse)' @classmethod def get_filename_list(s): files = [] files += folder_paths.get_filename_list('select_gguf') return sorted(files) def undo(self, select_gguf): in_file = folder_paths.get_full_path('select_gguf', select_gguf) out_file = (f'{self.output_dir}/{os.path.splitext(select_gguf)[0]}_fp16.safetensors') use_bf16 = False convert_gguf_to_safetensors(in_file, out_file, use_bf16) return {} class VaeGGUF: @staticmethod def vae_list(): vaes = [] vaes += folder_paths.get_filename_list('vae') vaes += folder_paths.get_filename_list('vae_gguf') approx_vaes = folder_paths.get_filename_list('vae_approx') sdxl_taesd_enc = False sdxl_taesd_dec = False sd1_taesd_enc = False sd1_taesd_dec = False sd3_taesd_enc = False sd3_taesd_dec = False f1_taesd_enc = False f1_taesd_dec = False for v in approx_vaes: if v.startswith('taesd_decoder.'): sd1_taesd_dec = True elif v.startswith('taesd_encoder.'): sd1_taesd_enc = True elif v.startswith('taesdxl_decoder.'): sdxl_taesd_dec = True elif v.startswith('taesdxl_encoder.'): sdxl_taesd_enc = True elif v.startswith('taesd3_decoder.'): sd3_taesd_dec = True elif v.startswith('taesd3_encoder.'): sd3_taesd_enc = True elif v.startswith('taef1_encoder.'): f1_taesd_dec = True elif v.startswith('taef1_decoder.'): f1_taesd_enc = True if sd1_taesd_dec and sd1_taesd_enc: vaes.append('taesd') if sdxl_taesd_dec and sdxl_taesd_enc: vaes.append('taesdxl') if sd3_taesd_dec and sd3_taesd_enc: vaes.append('taesd3') if f1_taesd_dec and f1_taesd_enc: vaes.append('taef1') return vaes @staticmethod def load_taesd(name): sd = {} approx_vaes = folder_paths.get_filename_list('vae_approx') encoder = next(filter(lambda a: a.startswith('{}_encoder.'.format(name)), approx_vaes)) decoder = next(filter(lambda a: a.startswith('{}_decoder.'.format(name)), approx_vaes)) enc = comfy.utils.load_torch_file(folder_paths. get_full_path_or_raise('vae_approx', encoder)) for k in enc: sd['taesd_encoder.{}'.format(k)] = enc[k] dec = comfy.utils.load_torch_file(folder_paths. get_full_path_or_raise('vae_approx', decoder)) for k in dec: sd['taesd_decoder.{}'.format(k)] = dec[k] if name == 'taesd': sd['vae_scale'] = torch.tensor(0.18215) sd['vae_shift'] = torch.tensor(0.0) elif name == 'taesdxl': sd['vae_scale'] = torch.tensor(0.13025) sd['vae_shift'] = torch.tensor(0.0) elif name == 'taesd3': sd['vae_scale'] = torch.tensor(1.5305) sd['vae_shift'] = torch.tensor(0.0609) elif name == 'taef1': sd['vae_scale'] = torch.tensor(0.3611) sd['vae_shift'] = torch.tensor(0.1159) return sd @classmethod def INPUT_TYPES(s): return {'required': {'vae_name': (s.vae_list(),)}} RETURN_TYPES = 'VAE', FUNCTION = 'load_vae' CATEGORY = 'gguf' TITLE = 'GGUF VAE Loader' def load_vae(self, vae_name): if vae_name.endswith('.gguf'): vae_path = folder_paths.get_full_path_or_raise('vae_gguf', vae_name) sd = load_gguf_clip(vae_path) elif vae_name in ['taesd', 'taesdxl', 'taesd3', 'taef1']: sd = self.load_taesd(vae_name) else: vae_path = folder_paths.get_full_path_or_raise('vae', vae_name) sd = comfy.utils.load_torch_file(vae_path) vae = comfy.sd.VAE(sd=sd) return vae, NODE_CLASS_MAPPINGS = { "LoaderGGUF": LoaderGGUF, "ClipLoaderGGUF": ClipLoaderGGUF, "DualClipLoaderGGUF": DualClipLoaderGGUF, "TripleClipLoaderGGUF": TripleClipLoaderGGUF, "QuadrupleClipLoaderGGUF": QuadrupleClipLoaderGGUF, "LoaderGGUFAdvanced": LoaderGGUFAdvanced, "VaeGGUF": VaeGGUF, "GGUFUndo": GGUFUndo, "GGUFSave": GGUFSave, "GGUFRun": GGUFRun, "TENSORCut": TENSORCut, "TENSORBoost": TENSORBoost, }