Split loader into separate file
This commit is contained in:
City
@@ -0,0 +1,122 @@
|
||||
# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
||||
import torch
|
||||
import gguf
|
||||
|
||||
from .ops import GGMLTensor
|
||||
|
||||
IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "ltxv", "hyvid"}
|
||||
TXT_ARCH_LIST = {"t5", "t5encoder"}
|
||||
|
||||
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 gguf_sd_loader(path, handle_prefix="model.diffusion_model.", return_arch=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 = None
|
||||
arch_field = reader.get_field("general.architecture")
|
||||
if arch_field is not None:
|
||||
if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING:
|
||||
raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}")
|
||||
arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8")
|
||||
if arch_str not in IMG_ARCH_LIST and arch_str not in TXT_ARCH_LIST:
|
||||
raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}")
|
||||
else: # stable-diffusion.cpp
|
||||
# import here to avoid changes to convert.py breaking regular models
|
||||
from .tools.convert import detect_arch
|
||||
arch_str = detect_arch(set(val[0] for val in tensors)).arch
|
||||
compat = "sd.cpp"
|
||||
|
||||
# main loading loop
|
||||
state_dict = {}
|
||||
qtype_dict = {}
|
||||
for sd_key, tensor in tensors:
|
||||
tensor_name = tensor.name
|
||||
tensor_type_str = str(tensor.tensor_type)
|
||||
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)
|
||||
qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1
|
||||
|
||||
# sanity check debug print
|
||||
print("\nggml_sd_loader:")
|
||||
for k,v in qtype_dict.items():
|
||||
print(f" {k:30}{v:3}")
|
||||
|
||||
if return_arch:
|
||||
return (state_dict, arch_str)
|
||||
return state_dict
|
||||
|
||||
# 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",
|
||||
}
|
||||
|
||||
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 gguf_clip_loader(path):
|
||||
sd, arch = gguf_sd_loader(path, return_arch=True)
|
||||
if "enc.blk.23.ffn_up.weight" in sd or arch in {"t5", "t5encoder"}:
|
||||
sd = sd_map_replace(sd, T5_SD_MAP)
|
||||
else:
|
||||
pass
|
||||
return sd
|
||||
@@ -1,16 +1,16 @@
|
||||
# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
||||
import torch
|
||||
import gguf
|
||||
import copy
|
||||
import logging
|
||||
import collections
|
||||
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import comfy.model_management
|
||||
import comfy.model_patcher
|
||||
import comfy.model_management
|
||||
import folder_paths
|
||||
|
||||
from .ops import GGMLTensor, GGMLOps, move_patch_to_device
|
||||
from .ops import GGMLOps, move_patch_to_device
|
||||
from .loader import gguf_sd_loader, gguf_clip_loader
|
||||
from .dequant import is_quantized, is_torch_compatible
|
||||
|
||||
# Add a custom keys for files ending in .gguf
|
||||
@@ -22,114 +22,6 @@ if "clip_gguf" not in folder_paths.folder_names_and_paths:
|
||||
orig = folder_paths.folder_names_and_paths.get("text_encoders", folder_paths.folder_names_and_paths.get("clip", [[], set()]))
|
||||
folder_paths.folder_names_and_paths["clip_gguf"] = (orig[0], {".gguf"})
|
||||
|
||||
def gguf_sd_loader_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 gguf_sd_loader(path, handle_prefix="model.diffusion_model."):
|
||||
"""
|
||||
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 = None
|
||||
arch_field = reader.get_field("general.architecture")
|
||||
if arch_field is not None:
|
||||
if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING:
|
||||
raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}")
|
||||
arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8")
|
||||
if arch_str not in {"flux", "sd1", "sdxl", "sd3", "aura", "ltxv", "hyvid", "t5", "t5encoder"}:
|
||||
raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}")
|
||||
else: # stable-diffusion.cpp
|
||||
# import here to avoid changes to convert.py breaking regular models
|
||||
from .tools.convert import detect_arch
|
||||
arch_str = detect_arch(set(val[0] for val in tensors)).arch
|
||||
compat = "sd.cpp"
|
||||
|
||||
# main loading loop
|
||||
state_dict = {}
|
||||
qtype_dict = {}
|
||||
for sd_key, tensor in tensors:
|
||||
tensor_name = tensor.name
|
||||
tensor_type_str = str(tensor.tensor_type)
|
||||
torch_tensor = torch.from_numpy(tensor.data) # mmap
|
||||
|
||||
shape = gguf_sd_loader_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)
|
||||
qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1
|
||||
|
||||
# sanity check debug print
|
||||
print("\nggml_sd_loader:")
|
||||
for k,v in qtype_dict.items():
|
||||
print(f" {k:30}{v:3}")
|
||||
|
||||
return state_dict
|
||||
|
||||
# for remapping llama.cpp -> original key names
|
||||
clip_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",
|
||||
}
|
||||
|
||||
def gguf_clip_loader(path):
|
||||
raw_sd = gguf_sd_loader(path)
|
||||
assert "enc.blk.23.ffn_up.weight" in raw_sd, "Invalid Text Encoder!"
|
||||
sd = {}
|
||||
for k,v in raw_sd.items():
|
||||
for s,d in clip_sd_map.items():
|
||||
k = k.replace(s,d)
|
||||
sd[k] = v
|
||||
return sd
|
||||
|
||||
# TODO: Temporary fix for now
|
||||
import collections
|
||||
class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
|
||||
patch_on_device = False
|
||||
|
||||
@@ -276,7 +168,8 @@ class UnetLoaderGGUFAdvanced(UnetLoaderGGUF):
|
||||
}
|
||||
TITLE = "Unet Loader (GGUF/Advanced)"
|
||||
|
||||
clip_name_dict = {
|
||||
# Mapping from name to CLIPType enum with fallback
|
||||
CLIP_TYPE_MAP = {
|
||||
"stable_diffusion": comfy.sd.CLIPType.STABLE_DIFFUSION,
|
||||
"stable_cascade": comfy.sd.CLIPType.STABLE_CASCADE,
|
||||
"stable_audio": comfy.sd.CLIPType.STABLE_AUDIO,
|
||||
@@ -287,6 +180,14 @@ clip_name_dict = {
|
||||
"ltxv": getattr(comfy.sd.CLIPType, "LTXV", None),
|
||||
}
|
||||
|
||||
def get_clip_type(type):
|
||||
if type not in CLIP_TYPE_MAP:
|
||||
raise ValueError(f"Unknown CLIP model type {type}")
|
||||
clip_type = CLIP_TYPE_MAP[type]
|
||||
if clip_type is None:
|
||||
raise ValueError(f"Unsupported CLIP model type {type} (Update ComfyUI)")
|
||||
return clip_type
|
||||
|
||||
class CLIPLoaderGGUF:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -343,8 +244,7 @@ class CLIPLoaderGGUF:
|
||||
|
||||
def load_clip(self, clip_name, type="stable_diffusion"):
|
||||
clip_path = folder_paths.get_full_path("clip", clip_name)
|
||||
clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||
return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),)
|
||||
return (self.load_patcher([clip_path], get_clip_type(type), self.load_data([clip_path])),)
|
||||
|
||||
class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
|
||||
@classmethod
|
||||
@@ -364,8 +264,7 @@ class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
|
||||
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)
|
||||
clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||
return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
|
||||
return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),)
|
||||
|
||||
class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
|
||||
@classmethod
|
||||
@@ -386,8 +285,7 @@ class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
|
||||
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)
|
||||
clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||
return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
|
||||
return (self.load_patcher(clip_paths, get_clip_type(type), self.load_data(clip_paths)),)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"UnetLoaderGGUF": UnetLoaderGGUF,
|
||||
|
||||
Reference in New Issue
Block a user