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

Cleanups, refactoring

Browse Source
This commit is contained in:
blepping
2024-08-25 07:53:19 -06:00
parent 2832969aea
commit f3cf802dd6
3 changed files with 132 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+10 -8
View File
@@ -3,7 +3,6 @@ import torch
import gguf
import copy
import logging
import numpy as np
import comfy.sd
import comfy.utils
@@ -40,10 +39,13 @@ def gguf_sd_loader(path):
for tensor in reader.tensors:
tensor_name = str(tensor.name)
shape = tensor_fixup.get(tensor_name)
torch_tensor = torch.from_numpy(tensor.data) # mmap
if shape is None:
shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
torch_tensor = torch_tensor.view(*shape)
sd[tensor_name] = GGMLTensor(
torch.from_numpy(tensor.data), # mmap
torch_tensor,
tensor_type = tensor.tensor_type,
tensor_shape = shape
)
@@ -106,21 +108,21 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
try:
from comfy.lora import calculate_weight
except:
except Exception:
calculate_weight = self.calculate_weight
patches = self.patches[key]
qtype = getattr(weight, "tensor_type", None)
if qtype not in [None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16]:
if qtype not in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
if device_to is not None:
out_weight = weight.to(device_to, copy=True)
else:
out_weight = weight.clone()
if self.patch_on_device:
patches = move_patch_to_cuda(patches, self.load_device)
out_weight.patches.append((calculate_weight, patches, key))
else:
if device_to is not None:
temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
@@ -262,10 +264,10 @@ class CLIPLoaderGGUF:
# for some reason this is just missing in some SAI checkpoints
if getattr(clip.cond_stage_model, "clip_l", None) is not None:
if getattr(clip.cond_stage_model.clip_l.transformer.text_projection.weight, "tensor_shape", None) == None:
if getattr(clip.cond_stage_model.clip_l.transformer.text_projection.weight, "tensor_shape", None) is None:
clip.cond_stage_model.clip_l.transformer.text_projection = comfy.ops.manual_cast.Linear(768, 768)
if getattr(clip.cond_stage_model, "clip_g", None) is not None:
if getattr(clip.cond_stage_model.clip_g.transformer.text_projection.weight, "tensor_shape", None) == None:
if getattr(clip.cond_stage_model.clip_g.transformer.text_projection.weight, "tensor_shape", None) is None:
clip.cond_stage_model.clip_g.transformer.text_projection = comfy.ops.manual_cast.Linear(1280, 1280)
return clip
ops.py
+31 -40
View File
@@ -1,5 +1,6 @@
# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
import torch
from functools import partial
import comfy.ops
from .dequant import dequantize_tensor
@@ -14,7 +15,7 @@ class GGMLTensor(torch.Tensor):
self.tensor_shape = tensor_shape
self.patches = patches
def __new__(cls, *args, tensor_type, tensor_shape, tensor_data=None, patches=[], **kwargs):
def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs):
return super().__new__(cls, *args, **kwargs)
def to(self, *args, **kwargs):
@@ -35,7 +36,7 @@ class GGMLTensor(torch.Tensor):
try:
return super().copy_(*args, **kwargs)
except Exception as e:
print(f"ignoring 'copy_' on tensor")
print(f"ignoring 'copy_' on tensor: {e}")
def __deepcopy__(self, *args, **kwargs):
# Intel Arc fix, ref#50
@@ -130,53 +131,43 @@ class GGMLLayer(torch.nn.Module):
bias = self.get_weight(self.bias, dtype)
return (weight, bias)
class GGMLOpsLayer(GGMLLayer):
comfy_cast_weights = True
def __init__(self, *args, device=None, dtype=None, **kwargs):
super().__init__(device=device, dtype=dtype)
def _forward_operation(self, x, weight, bias):
raise NotImplementedError
def forward(self, x):
# lowvram hack
device = None
if self.weight.device != x.device:
device = self.weight.device
self.to(x.device)
weight, bias = self.get_weights(x.dtype)
x = self._forward_operation(x, weight, bias)
del weight, bias
if device:
self.to(device)
return x
class GGMLOps(comfy.ops.manual_cast):
"""
Dequantize weights on the fly before doing the compute
"""
class Linear(GGMLLayer):
comfy_cast_weights = True
def __init__(self, *args, device=None, dtype=None, **kwargs):
super().__init__(device=device, dtype=dtype)
def forward(self, x):
# lowvram hack
device = None
if self.weight.device != x.device:
device = self.weight.device
self.to(x.device)
weight, bias = self.get_weights(x.dtype)
x = torch.nn.functional.linear(x, weight, bias)
del weight, bias
if device:
self.to(device)
return x
class Conv2d(GGMLLayer):
comfy_cast_weights = True
class Linear(GGMLOpsLayer):
_forward_operation = staticmethod(torch.nn.functional.linear)
class Conv2d(GGMLOpsLayer):
def __init__(self, *args, device=None, dtype=None, **kwargs):
super().__init__(device=device, dtype=dtype)
_ = kwargs.pop("kernel_size", None)
self.kwargs=kwargs
self._forward_operation = partial(staticmethod(torch.nn.functional.conv2d), **kwargs)
def forward(self, x):
# lowvram hack
device = None
if self.weight.device != x.device:
device = self.weight.device
self.to(x.device)
weight, bias = self.get_weights(x.dtype)
x = torch.nn.functional.conv2d(x, weight, bias, **self.kwargs)
del weight, bias
if device:
self.to(device)
return x
def move_patch_to_cuda(item, device):
if isinstance(item, torch.Tensor):
tools/convert.py
+91 -62
View File
@@ -1,24 +1,55 @@
# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
import os
import sys
import torch
import numpy as np
import gguf # This needs to be the llama.cpp one specifically!
import argparse
from tqdm import tqdm
from safetensors.torch import load_file
QUANTIZATION_THRESHOLD = 1024
REARRANGE_THRESHOLD = 512
MAX_TENSOR_NAME_LENGTH = 63
# Tuple of arch_name, match_lists.
# Each item in match_lists is a tuple of keys that must match.
# All keys in a match_lists item must exist for the architecture to match.
# The architectures are checked in order and the first successful match terminates the search.
MODEL_DETECTION = (
("flux", (
("transformer_blocks.0.attn.norm_added_k.weight",),
("double_blocks.0.img_attn.proj.weight",),
)),
("sd3", (
("transformer_blocks.0.attn.add_q_proj.weight",),
)),
("sdxl", (
("down_blocks.0.downsamplers.0.conv.weight", "add_embedding.linear_1.weight",),
(
"input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight",
"output_blocks.2.2.conv.weight", "output_blocks.5.2.conv.weight",
), # Non-diffusers
("label_emb.0.0.weight",),
)),
("sd1", (
("down_blocks.0.downsamplers.0.conv.weight",),
(
"input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", "input_blocks.9.0.op.weight",
"output_blocks.2.1.conv.weight", "output_blocks.5.2.conv.weight", "output_blocks.8.2.conv.weight"
), # Non-diffusers
)),
)
def parse_args():
parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET")
parser.add_argument("--src", required=True, help="Source model ckpt file.")
parser.add_argument("--dst", help="Output unet gguf file.")
parser.add_argument("--force-arch", type=str, choices=("flux", "sd1", "sdxl"), default=None, help="Advanced option, allows forcing the model architecture. May be necessary for non-Flux models.")
parser.add_argument("--dst", help="Output unet gguf file.")
args = parser.parse_args()
if not os.path.isfile(args.src):
parser.error("No input provided!")
return args
def load_state_dict(path):
@@ -27,7 +58,7 @@ def load_state_dict(path):
state_dict = state_dict.get("model", state_dict)
else:
state_dict = load_file(path)
# only keep unet with no prefix!
sd = {}
has_prefix = any(["model.diffusion_model." in x for x in state_dict.keys()])
@@ -40,27 +71,21 @@ def load_state_dict(path):
return sd
def load_model(path, force_arch=None):
state_dict = load_state_dict(path)
# from ComfyUI model detection
if force_arch is not None:
arch = force_arch
elif "transformer_blocks.0.attn.norm_added_k.weight" in state_dict:
arch = "flux"
raise ValueError(f"The Diffusers UNET can not be used for this!")
elif "double_blocks.0.img_attn.proj.weight" in state_dict:
arch = "flux" # mmdit ...?
elif "transformer_blocks.0.attn.add_q_proj.weight" in state_dict:
arch = "sd3"
elif "down_blocks.0.downsamplers.0.conv.weight" in state_dict:
if "add_embedding.linear_1.weight" in state_dict:
arch = "sdxl"
else:
arch = "sd1"
else:
breakpoint()
raise ValueError(f"Unknown model architecture!")
def detect_arch(state_dict):
for arch, match_lists in MODEL_DETECTION:
for match_list in match_lists:
if all(key in state_dict for key in match_list):
return arch
breakpoint()
raise ValueError("Unknown model architecture!")
def load_model(path):
state_dict = load_state_dict(path)
arch = detect_arch(state_dict)
print(f"* Architecture detected from input: {arch}")
if arch == "flux" and "transformer_blocks.0.attn.norm_added_k.weight" in state_dict:
raise ValueError("The Diffusers UNET can not be used for this!")
writer = gguf.GGUFWriter(path=None, arch=arch)
return (writer, state_dict)
@@ -68,9 +93,17 @@ def handle_tensors(args, writer, state_dict):
# TODO list:
# - do something about this being awful and hacky
max_name_len = max([len(s) for s in state_dict.keys()]) + 4
if (max_name_len - 4) >= 64:
raise ValueError("Can only handle tensor names up to 63 characters!")
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}")
tensor_fixup = {}
for key, data in tqdm(state_dict.items()):
old_dtype = data.dtype
@@ -93,7 +126,7 @@ def handle_tensors(args, writer, state_dict):
n_params *= dim_size
# keys to keep as max precision
blacklist = [
blacklist = {
"time_embedding.",
"add_embedding.",
"time_in.",
@@ -102,41 +135,33 @@ def handle_tensors(args, writer, state_dict):
"img_in.",
"guidance_in.",
"final_layer.",
]
}
if any([x in key for x in blacklist]) and ".weight" in key:
data_qtype = gguf.GGMLQuantizationType.F32
if old_dtype in (torch.float32, torch.bfloat16):
if n_dims == 1:
# one-dimensional tensors should be kept in F32
# also speeds up inference due to not dequantizing
data_qtype = gguf.GGMLQuantizationType.F32
if n_dims == 1:
# one-dimensional tensors should be kept in F32
# also speeds up inference due to not dequantizing
data_qtype = gguf.GGMLQuantizationType.F32
elif n_params <= 1024:
# very small tensors
data_qtype = gguf.GGMLQuantizationType.F32
elif n_dims == 4:
if min(data.shape[:2]) == 4: # output tensor
data_qtype = gguf.GGMLQuantizationType.F16
elif data_shape[-1] == 3: # 3x3 kernel
data_qtype = gguf.GGMLQuantizationType.F16
elif data_shape[-1] == 1: # 1x1 kernel
#data = np.squeeze(data) # don't do this
data_qtype = gguf.GGMLQuantizationType.F16
elif n_params <= QUANTIZATION_THRESHOLD:
# very small tensors
data_qtype = gguf.GGMLQuantizationType.F32
if data_qtype != gguf.GGMLQuantizationType.F32 and n_dims >= 2 and n_params >= 256*256 and (n_params / 256).is_integer() and data.shape[0] >= 256 and data.shape[1] >= 256:
orig_shape=data.shape
orig_n_params=n_params
data = data.reshape(256+((n_params - 256*256) // 256),256)
elif ".weight" in key and any(x in key for x in blacklist):
data_qtype = gguf.GGMLQuantizationType.F32
if ( n_dims > 1 # Skip one-dimensional tensors
and n_params >= REARRANGE_THRESHOLD # Only rearrange tensors meeting the size requirement
and (n_params / 256).is_integer() # Rearranging only makes sense if total elements is divisible by 256
and not (data.shape[-1] / 256).is_integer() # Only need to rearrange if the last dimension is not divisible by 256
):
orig_shape = data.shape
data = data.reshape(n_params // 256, 256)
tensor_fixup[key] = orig_shape
try:
data = gguf.quants.quantize(data, data_qtype)
except gguf.QuantError as e:
tqdm.write(f"falling back to F16: {e}")
data_qtype = gguf.GGMLQuantizationType.F16
data = gguf.quants.quantize(data, data_qtype)
except AttributeError as e:
except (AttributeError, gguf.QuantError) as e:
tqdm.write(f"falling back to F16: {e}")
data_qtype = gguf.GGMLQuantizationType.F16
data = gguf.quants.quantize(data, data_qtype)
@@ -144,16 +169,20 @@ def handle_tensors(args, writer, state_dict):
new_name = key # do we need to rename?
shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}"
tqdm.write(f"{f'%-{max_name_len}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
writer.add_tensor(new_name, data, raw_dtype=data_qtype)
if not tensor_fixup:
return
fixup_str = "\n".join(" ".join((k,)+tuple(str(dimval) for dimval in v)) for k,v in tensor_fixup.items())
writer.add_string("comfy.gguf.tensor_fixup", fixup_str)
if __name__ == "__main__":
args = parse_args()
path = args.src
writer, state_dict = load_model(path, force_arch=args.force_arch)
writer, state_dict = load_model(path)
writer.add_quantization_version(gguf.GGML_QUANT_VERSION)
if next(iter(state_dict.values())).dtype == torch.bfloat16:
@@ -162,7 +191,7 @@ if __name__ == "__main__":
else:
out_path = f"{os.path.splitext(path)[0]}-F16.gguf"
writer.add_file_type(gguf.LlamaFileType.MOSTLY_F16)
out_path = args.dst or out_path
if os.path.isfile(out_path):
input("Output exists enter to continue or ctrl+c to abort!")