nodes/city96-ComfyUI-GGUF
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #49 from city96/gguf_tenc

Browse Source 
GGUF Text Encoder support
This commit is contained in:
City
2024-08-20 03:52:09 +02:00
committed by GitHub
parent 165212919e e7906b43eb
commit 5d7ffb582c
2 changed files with 187 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+177 -26
View File
@@ -12,12 +12,17 @@ import comfy.model_patcher
import folder_paths
from .ops import GGMLTensor, GGMLOps
from .dequant import dequantize_tensor
# Add a custom key for files ending in .gguf
# Add a custom keys for files ending in .gguf
if "unet_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["unet_gguf"] = (orig[0], {".gguf"})
if "clip_gguf" not in folder_paths.folder_names_and_paths:
orig = folder_paths.folder_names_and_paths.get("clip", [[], set()])
folder_paths.folder_names_and_paths["clip_gguf"] = (orig[0], {".gguf"})
def gguf_sd_loader(path):
"""
Read state dict as fake tensors
@@ -42,32 +47,40 @@ def gguf_sd_loader(path):
print("\n")
return sd
class UnetLoaderGGUF:
@classmethod
def INPUT_TYPES(s):
unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")]
return {
"required": {
"unet_name": (unet_names,),
}
}
# 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",
}
# weights that should be dequantized on load
clip_sd_dequant = {
"shared.weight",
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "load_unet"
CATEGORY = "bootleg"
TITLE = "Unet Loader (GGUF)"
def load_unet(self, unet_name):
unet_path = folder_paths.get_full_path("unet", unet_name)
sd = gguf_sd_loader(unet_path)
model = comfy.sd.load_diffusion_model_state_dict(
sd, model_options={"custom_operations": GGMLOps}
)
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)
return (model,)
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)
if k in clip_sd_dequant:
v = dequantize_tensor(v, torch.float32).to(torch.float16)
v = GGMLTensor(v, tensor_type=gguf.GGMLQuantizationType.F16, tensor_shape=v.shape)
sd[k] = v
return sd
# TODO: Temporary fix for now
class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
@@ -96,6 +109,144 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
n.object_patches_backup = self.object_patches_backup
return n
class UnetLoaderGGUF:
@classmethod
def INPUT_TYPES(s):
unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")]
return {
"required": {
"unet_name": (unet_names,),
}
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "load_unet"
CATEGORY = "bootleg"
TITLE = "Unet Loader (GGUF)"
def load_unet(self, unet_name):
unet_path = folder_paths.get_full_path("unet", unet_name)
sd = gguf_sd_loader(unet_path)
model = comfy.sd.load_diffusion_model_state_dict(
sd, model_options={"custom_operations": GGMLOps}
)
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)
return (model,)
clip_name_dict = {
"stable_diffusion": comfy.sd.CLIPType.STABLE_DIFFUSION,
"stable_cascade": comfy.sd.CLIPType.STABLE_CASCADE,
"stable_audio": comfy.sd.CLIPType.STABLE_AUDIO,
"sdxl": comfy.sd.CLIPType.STABLE_DIFFUSION,
"sd3": comfy.sd.CLIPType.SD3,
"flux": comfy.sd.CLIPType.FLUX,
}
class CLIPLoaderGGUF:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"clip_name": (s.get_filename_list(),),
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio"],),
}
}
RETURN_TYPES = ("CLIP",)
FUNCTION = "load_clip"
CATEGORY = "bootleg"
TITLE = "CLIPLoader (GGUF)"
@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"):
clip_data.append(gguf_clip_loader(p))
else:
sd = comfy.utils.load_torch_file(p, safe_load=True)
clip_data.append(
{k:GGMLTensor(v, tensor_type=gguf.GGMLQuantizationType.F16, tensor_shape=v.shape) for k,v in sd.items()}
)
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},
embedding_directory = folder_paths.get_folder_paths("embeddings"),
)
clip.patcher = GGUFModelPatcher.clone(clip.patcher)
# for some reason this is just missing in some SAI checkpoints
if hasattr(clip.cond_stage_model, "clip_l"):
if clip.cond_stage_model.clip_l.transformer.text_projection.weight.tensor_shape == None:
clip.cond_stage_model.clip_l.transformer.text_projection = comfy.ops.manual_cast.Linear(768, 768)
if hasattr(clip.cond_stage_model, "clip_g"):
if clip.cond_stage_model.clip_g.transformer.text_projection.weight.tensor_shape == None:
clip.cond_stage_model.clip_g.transformer.text_projection = comfy.ops.manual_cast.Linear(1280, 1280)
return clip
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])),)
class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"clip_name1": (s.get_filename_list(), ),
"clip_name2": (s.get_filename_list(), ),
"type": (["sdxl", "sd3", "flux"], ),
}
}
TITLE = "DualCLIPLoader (GGUF)"
def load_clip(self, clip_name1, clip_name2, type):
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)),)
class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"clip_name1": (s.get_filename_list(), ),
"clip_name2": (s.get_filename_list(), ),
"clip_name3": (s.get_filename_list(), ),
}
}
TITLE = "TripleCLIPLoader (GGUF)"
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]
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)),)
NODE_CLASS_MAPPINGS = {
"UnetLoaderGGUF": UnetLoaderGGUF,
"CLIPLoaderGGUF": CLIPLoaderGGUF,
"DualCLIPLoaderGGUF": DualCLIPLoaderGGUF,
"TripleCLIPLoaderGGUF": TripleCLIPLoaderGGUF,
}
ops.py
+10 -3
View File
@@ -32,6 +32,13 @@ class GGMLTensor(torch.Tensor):
def detach(self, *args, **kwargs):
return self
def copy_(self, *args, **kwargs):
# fixes .weight.copy_ in comfy/clip_model/CLIPTextModel
try:
return super().copy_(*args, **kwargs)
except Exception as e:
print(f"ignoring 'copy_' on tensor")
@property
def shape(self):
if not hasattr(self, "tensor_shape"):
@@ -44,7 +51,7 @@ class GGMLLayer(torch.nn.Module):
"""
def __init__(self, *args, **kwargs):
super().__init__()
self.weight = None
self.weight = GGMLTensor(1, tensor_type=None, tensor_shape=None)
self.bias = None
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
@@ -119,7 +126,7 @@ class GGMLLayer(torch.nn.Module):
bias = self.get_weight(self.bias, dtype)
return (weight, bias)
class GGMLOps(comfy.ops.disable_weight_init):
class GGMLOps(comfy.ops.manual_cast):
"""
Dequantize weights on the fly before doing the compute
"""
@@ -142,4 +149,4 @@ class GGMLOps(comfy.ops.disable_weight_init):
if device:
self.to(device)
return x
return x