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

Advanced loader with optimizations

Browse Source 
It may help with speed but it may also increase VRAM usage or degrade image quality.
Fastest would be target/target/true.
This commit is contained in:
City
2024-08-24 18:39:34 +02:00
parent 90a63c0c0e
commit 0ba6f64974
3 changed files with 72 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dequant.py
+4 -10
View File
@@ -1,24 +1,18 @@
# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
import gguf
import torch
import numpy as np
from tqdm import tqdm
def dequantize_tensor(tensor, dtype=torch.float16):
def dequantize_tensor(tensor, dtype=None, dequant_dtype=None):
data = torch.tensor(tensor.data)
qtype = getattr(tensor, "tensor_type", None)
oshape = getattr(tensor, "tensor_shape", tensor.data.shape)
if qtype == None:
return data.to(dtype)
elif qtype == gguf.GGMLQuantizationType.F32:
return data.to(dtype)
elif qtype == gguf.GGMLQuantizationType.F16:
if qtype in [None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16]:
return data.to(dtype)
elif qtype in dequantize_functions:
# dequantize in fp16 then convert instead of keeping FP32
out = dequantize(data, qtype, oshape, dtype=None)
return out.to(dtype) if out.dtype != dtype else out # why is .to() not a no-op?
dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype
return dequantize(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
+46 -7
View File
@@ -11,7 +11,7 @@ import comfy.model_management
import comfy.model_patcher
import folder_paths
from .ops import GGMLTensor, GGMLOps
from .ops import GGMLTensor, GGMLOps, move_patch_to_cuda
from .dequant import dequantize_tensor
# Add a custom keys for files ending in .gguf
@@ -85,6 +85,8 @@ def gguf_clip_loader(path):
# TODO: Temporary fix for now
import collections
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
@@ -92,20 +94,23 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
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)
qtype = getattr(weight, "tensor_type", None)
try:
from comfy.lora import calculate_weight
except:
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 device_to is not None:
out_weight = weight.to(device_to, copy=True)
else:
out_weight = weight.clone()
out_weight.patches.append((calculate_weight, self.patches[key], key))
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:
@@ -113,7 +118,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
else:
temp_weight = weight.to(torch.float32, copy=True)
out_weight = calculate_weight(self.patches[key], temp_weight, key)
out_weight = calculate_weight(patches, temp_weight, key)
out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype)
if inplace_update:
@@ -132,6 +137,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
n.model_options = copy.deepcopy(self.model_options)
n.backup = self.backup
n.object_patches_backup = self.object_patches_backup
n.patch_on_device = getattr(self, "patch_on_device", False)
return n
class UnetLoaderGGUF:
@@ -149,18 +155,50 @@ class UnetLoaderGGUF:
CATEGORY = "bootleg"
TITLE = "Unet Loader (GGUF)"
def load_unet(self, unet_name):
def load_unet(self, unet_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)
# init model
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}
sd, model_options={"custom_operations": ops}
)
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)
model.patch_on_device = patch_on_device
return (model,)
class UnetLoaderGGUFAdvanced(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,),
"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 = "Unet Loader (GGUF/Advanced)"
clip_name_dict = {
"stable_diffusion": comfy.sd.CLIPType.STABLE_DIFFUSION,
"stable_cascade": comfy.sd.CLIPType.STABLE_CASCADE,
@@ -274,4 +312,5 @@ NODE_CLASS_MAPPINGS = {
"CLIPLoaderGGUF": CLIPLoaderGGUF,
"DualCLIPLoaderGGUF": DualCLIPLoaderGGUF,
"TripleCLIPLoaderGGUF": TripleCLIPLoaderGGUF,
"UnetLoaderGGUFAdvanced": UnetLoaderGGUFAdvanced,
}
ops.py
+22 -16
View File
@@ -1,7 +1,5 @@
# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
import gguf
import torch
import numpy as np
import comfy.ops
from .dequant import dequantize_tensor
@@ -57,6 +55,9 @@ class GGMLLayer(torch.nn.Module):
"""
This (should) be responsible for de-quantizing on the fly
"""
dequant_dtype = None
patch_dtype = None
def __init__(self, *args, **kwargs):
super().__init__()
self.weight = GGMLTensor(1, tensor_type=None, tensor_shape=None)
@@ -100,16 +101,6 @@ class GGMLLayer(torch.nn.Module):
super()._apply(fn)
return self
def move_patch_to_cuda(self, item, device):
if isinstance(item, torch.Tensor):
return item.to(device, non_blocking=True)
elif isinstance(item, tuple):
return tuple(self.move_patch_to_cuda(x, device) for x in item)
elif isinstance(item, list):
return [self.move_patch_to_cuda(x, device) for x in item]
else:
return item
def get_weight(self, tensor, dtype):
if tensor is None:
return
@@ -118,15 +109,20 @@ class GGMLLayer(torch.nn.Module):
patch_list = []
device = tensor.device
t_move = lambda x: x.to(device) if torch.is_tensor(x) else x
for function, patches, _ in getattr(tensor, "patches", []):
patch_list += self.move_patch_to_cuda(patches, device)
for function, patches, key in getattr(tensor, "patches", []):
patch_list += move_patch_to_cuda(patches, device)
# dequantize tensor while patches load
weight = dequantize_tensor(tensor, dtype)
weight = dequantize_tensor(tensor, dtype, self.dequant_dtype)
# apply patches
if patch_list:
weight = function(patch_list, weight, "dequant.name.unknown")
if self.patch_dtype is None:
weight = function(patch_list, weight, key)
else:
# for testing, may degrade image quality
patch_dtype = dtype if self.patch_dtype == "target" else self.patch_dtype
weight = function(patch_list, weight, key, patch_dtype)
return weight
def get_weights(self, dtype=torch.float16):
@@ -158,3 +154,13 @@ class GGMLOps(comfy.ops.manual_cast):
if device:
self.to(device)
return x
def move_patch_to_cuda(item, device):
if isinstance(item, torch.Tensor):
return item.to(device, non_blocking=True)
elif isinstance(item, tuple):
return tuple(move_patch_to_cuda(x, device) for x in item)
elif isinstance(item, list):
return [move_patch_to_cuda(x, device) for x in item]
else:
return item