133 lines
4.4 KiB
Python
133 lines
4.4 KiB
Python
# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
|
import torch
|
|
import gguf
|
|
import logging
|
|
import numpy as np
|
|
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import comfy.model_management
|
|
import folder_paths
|
|
|
|
from .ops import GGMLTensor, GGMLOps
|
|
from .lora import load_lora_gguf
|
|
|
|
# TODO: This causes gguf files to show up in the main unet loader
|
|
folder_paths.folder_names_and_paths["unet"][1].add(".gguf")
|
|
|
|
def gguf_sd_loader(path):
|
|
"""
|
|
Read state dict as fake tensors
|
|
"""
|
|
reader = gguf.GGUFReader(path)
|
|
sd = {}
|
|
dt = {}
|
|
for tensor in reader.tensors:
|
|
sd[str(tensor.name)] = GGMLTensor(
|
|
torch.from_numpy(tensor.data), # mmap
|
|
tensor_type = tensor.tensor_type,
|
|
tensor_shape = torch.Size(
|
|
np.flip(list(tensor.shape))
|
|
)
|
|
)
|
|
dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1
|
|
|
|
# sanity check debug print
|
|
print("\nggml_sd_loader:")
|
|
for k,v in dt.items():
|
|
print(f" {k:30}{v:3}")
|
|
print("\n")
|
|
return sd
|
|
|
|
class UnetLoaderGGUF:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
unet_names = [x for x in folder_paths.get_filename_list("unet") if x.endswith(".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))
|
|
return (model,)
|
|
|
|
class LoraLoaderGGUFModelOnly:
|
|
def __init__(self):
|
|
self.loaded_lora = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
|
"lora_on_gpu": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load_lora"
|
|
CATEGORY = "bootleg"
|
|
TITLE = "LoRA Loader Model Only (GGUF)"
|
|
|
|
def load_lora(self, model, lora_name, strength_model, lora_on_gpu):
|
|
if strength_model == 0:
|
|
return (model,)
|
|
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
lora = None
|
|
if self.loaded_lora is not None:
|
|
if self.loaded_lora[0] == lora_path:
|
|
lora = self.loaded_lora[1]
|
|
else:
|
|
temp = self.loaded_lora
|
|
self.loaded_lora = None
|
|
del temp
|
|
|
|
if lora is None:
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
self.loaded_lora = (lora_path, lora)
|
|
|
|
# This is a super hacky method to get it to load alongside the model
|
|
param_dict_name = f"lora_state_dict"#_{lora_name}".replace('.', '_')
|
|
target_model = model.model.diffusion_model
|
|
if lora_on_gpu:
|
|
print(f"WARNING! Loading the LoRA to GPU memory may cause OOM issues or memory leaks!")
|
|
temp = {}
|
|
params = torch.nn.ParameterDict()
|
|
for k,v in lora.items():
|
|
param = torch.nn.Parameter(v, requires_grad=False)
|
|
temp[k] = param
|
|
params[k.replace('.', '_')] = param
|
|
params.to(model.load_device, dtype=torch.float16)
|
|
target_model.register_module(param_dict_name, params)
|
|
lora = temp
|
|
else:
|
|
params = getattr(target_model, param_dict_name, None)
|
|
if params:
|
|
params.to("cpu", dtype=torch.float16)
|
|
delattr(target_model, param_dict_name)
|
|
|
|
model_lora = load_lora_gguf(model, lora, strength_model)
|
|
return (model_lora,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UnetLoaderGGUF": UnetLoaderGGUF,
|
|
"LoraLoaderGGUFModelOnly": LoraLoaderGGUFModelOnly,
|
|
}
|