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

Merge pull request #80 from blepping/feat_load_stable_diffusion.cpp

Browse Source 
Support loading u-net from stable-diffusion.cpp GGUF files
This commit is contained in:
City
2024-09-02 02:51:40 +02:00
committed by GitHub
parent ba8ce44b4c fa1671d63e
commit cd343d1426
4 changed files with 113 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dequant.py
+11 -2
View File
@@ -3,15 +3,24 @@ import gguf
import torch
from tqdm import tqdm
TORCH_COMPATIBLE_QTYPES = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
def is_torch_compatible(tensor):
return tensor is None or getattr(tensor, "tensor_type", None) in TORCH_COMPATIBLE_QTYPES
def is_quantized(tensor):
return not is_torch_compatible(tensor)
def dequantize_tensor(tensor, dtype=None, dequant_dtype=None):
qtype = getattr(tensor, "tensor_type", None)
oshape = getattr(tensor, "tensor_shape", tensor.shape)
if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
if qtype in TORCH_COMPATIBLE_QTYPES:
return tensor.to(dtype)
elif qtype in dequantize_functions:
dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype
return dequantize(tensor, qtype, oshape, dtype=dequant_dtype).to(dtype)
return dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype).to(dtype)
else:
# this is incredibly slow
tqdm.write(f"Falling back to numpy dequant for qtype: {qtype}")
nodes.py
+50 -17
View File
@@ -11,7 +11,7 @@ import comfy.model_patcher
import folder_paths
from .ops import GGMLTensor, GGMLOps, move_patch_to_device
from .dequant import dequantize_tensor
from .tools.convert import detect_arch
# Add a custom keys for files ending in .gguf
if "unet_gguf" not in folder_paths.folder_names_and_paths:
@@ -32,27 +32,58 @@ def gguf_sd_loader_get_orig_shape(reader, tensor_name):
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):
def gguf_sd_loader(path, handle_prefix="model.diffusion_model."):
"""
Read state dict as fake tensors
"""
reader = gguf.GGUFReader(path)
sd = {}
dt = {}
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", "t5encoder"}:
raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}")
else:
arch_str = None
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)
else:
has_prefix = False
tensors = []
for tensor in reader.tensors:
tensor_name = str(tensor.name)
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))
detected_arch = detect_arch(set(val[0] for val in tensors)) if arch_str is None else None
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)))
elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
if arch_str is None and detected_arch == "sdxl" and \
(tensor_name.endswith(".proj_in.weight") or tensor_name.endswith(".proj_out.weight")):
# Workaround for stable-diffusion.cpp SDXL detection.
# Impossible to land here for our own GGUF files as we set architecture.
while len(shape) > 2 and shape[-1] == 1:
shape = shape[:-1]
if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
torch_tensor = torch_tensor.view(*shape)
sd[tensor_name] = GGMLTensor(
sd[sd_key] = GGMLTensor(
torch_tensor,
tensor_type = tensor.tensor_type,
tensor_shape = shape
)
dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1
dt[tensor_type_str] = dt.get(tensor_type_str, 0) + 1
# sanity check debug print
print("\nggml_sd_loader:")
@@ -133,7 +164,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
def unpatch_model(self, device_to=None, unpatch_weights=True):
if unpatch_weights:
for p in self.model.parameters():
qtype = getattr(p, "tensor_type", None)
qtype = getattr(p, "tensor_type", None)
patches = getattr(p, "patches", [])
if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
continue
@@ -177,14 +208,14 @@ class UnetLoaderGGUF:
def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None):
ops = GGMLOps()
if dequant_dtype in ["default", None]:
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]:
if patch_dtype in ("default", None):
ops.Linear.patch_dtype = None
elif patch_dtype in ["target"]:
ops.Linear.patch_dtype = patch_dtype
@@ -288,11 +319,12 @@ class CLIPLoaderGGUF:
class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
@classmethod
def INPUT_TYPES(s):
file_options = (s.get_filename_list(), )
return {
"required": {
"clip_name1": (s.get_filename_list(), ),
"clip_name2": (s.get_filename_list(), ),
"type": (["sdxl", "sd3", "flux"], ),
"clip_name1": file_options,
"clip_name2": file_options,
"type": (("sdxl", "sd3", "flux"), ),
}
}
@@ -301,18 +333,19 @@ class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
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_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):
file_options = (s.get_filename_list(), )
return {
"required": {
"clip_name1": (s.get_filename_list(), ),
"clip_name2": (s.get_filename_list(), ),
"clip_name3": (s.get_filename_list(), ),
"clip_name1": file_options,
"clip_name2": file_options,
"clip_name3": file_options,
}
}
@@ -322,7 +355,7 @@ class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
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_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)),)
ops.py
+49 -9
View File
@@ -3,7 +3,7 @@ import gguf
import torch
import comfy.ops
from .dequant import dequantize_tensor
from .dequant import dequantize_tensor, is_quantized
class GGMLTensor(torch.Tensor):
"""
@@ -59,17 +59,37 @@ class GGMLLayer(torch.nn.Module):
comfy_cast_weights = True
dequant_dtype = None
patch_dtype = None
torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
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):
return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
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_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[len(prefix):] == "weight":
if k[prefix_len:] == "weight":
self.weight = torch.nn.Parameter(v, requires_grad=False)
elif k[len(prefix):] == "bias" and v is not None:
elif k[prefix_len:] == "bias" and v is not None:
self.bias = torch.nn.Parameter(v, requires_grad=False)
else:
missing_keys.append(k)
def _save_to_state_dict(self, destination, prefix, keep_vars):
def ggml_save_to_state_dict(self, destination, prefix, keep_vars):
# This is a fake state dict for vram estimation
weight = torch.zeros_like(self.weight, device=torch.device("meta"))
destination[prefix + "weight"] = weight
@@ -125,30 +145,50 @@ class GGMLLayer(torch.nn.Module):
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():
return self.ggml_forward(input, *args, **kwargs)
return super().forward_comfy_cast_weights(input, *args, **kwargs)
def ggml_forward(self, input):
raise NotImplementedError
class GGMLOps(comfy.ops.manual_cast):
"""
Dequantize weights on the fly before doing the compute
"""
class Linear(GGMLLayer, comfy.ops.manual_cast.Linear):
def forward_comfy_cast_weights(self, input):
def ggml_forward(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_comfy_cast_weights(self, input):
def ggml_forward(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_comfy_cast_weights(self, input, out_dtype=None):
def ggml_forward(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(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 ggml_forward(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 ggml_forward(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)
tools/convert.py
+3
View File
@@ -109,6 +109,9 @@ def handle_tensors(args, writer, state_dict):
if data.dtype == torch.bfloat16:
data = data.to(torch.float32).numpy()
# this is so we don't break torch 2.0.X
elif data.dtype in [getattr(torch, "float8_e4m3fn", "_invalid"), getattr(torch, "float8_e5m2", "_invalid")]:
data = data.to(torch.float16).numpy()
else:
data = data.numpy()