From e05ce931a25ebcf2a8d56f2a4cc26210e04b2048 Mon Sep 17 00:00:00 2001 From: wTechArtist Date: Wed, 21 Aug 2024 10:59:43 +0800 Subject: [PATCH] =?UTF-8?q?add=20node=EF=BC=9ALoad=20Lora=20With=20Shared?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- __init__.py | 3 +- libs/utils.py | 336 ++++++++++++++++++++++++++++++++++++++++++++++++++ node.py | 133 +++++++++++++++++++- 3 files changed, 470 insertions(+), 2 deletions(-) create mode 100644 libs/utils.py diff --git a/__init__.py b/__init__.py index 1ef33c2..22827df 100644 --- a/__init__.py +++ b/__init__.py @@ -6,7 +6,8 @@ NODE_CLASS_MAPPINGS = { "Load Image With Bool" : LoadImage_Bool, "IPAdapter Mad Scientist Weight_Type" : IPAdapter_Mad_Scientist_weight_type, "IPAdapter FaceID With Bool" : IPAdapter_FaceID_Bool, - "GPT4 WWL" : GPT4_WWL + "GPT4 WWL" : GPT4_WWL, + "Load Lora With Shared" : LoraLoaderShared } diff --git a/libs/utils.py b/libs/utils.py new file mode 100644 index 0000000..c69ecde --- /dev/null +++ b/libs/utils.py @@ -0,0 +1,336 @@ +import itertools +from typing import Optional +import numpy as np +import torch +from PIL import Image, ImageDraw +import math +import cv2 + + +def apply_variation_noise(latent_image, noise_device, variation_seed, variation_strength, mask=None, variation_method='linear'): + latent_size = latent_image.size() + latent_size_1batch = [1, latent_size[1], latent_size[2], latent_size[3]] + + if noise_device == "cpu": + variation_generator = torch.manual_seed(variation_seed) + else: + torch.cuda.manual_seed(variation_seed) + variation_generator = None + + variation_latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, + generator=variation_generator, device=noise_device) + + variation_noise = variation_latent.expand(latent_image.size()[0], -1, -1, -1) + + if variation_strength == 0: + return latent_image + elif mask is None: + result = (1 - variation_strength) * latent_image + variation_strength * variation_noise + else: + # this seems precision is not enough when variation_strength is 0.0 + mixed_noise = mix_noise(latent_image, variation_noise, variation_strength, variation_method=variation_method) + result = (mask == 1).float() * mixed_noise + (mask == 0).float() * latent_image + + return result + + +# CREDIT: https://github.com/BlenderNeko/ComfyUI_Noise/blob/afb14757216257b12268c91845eac248727a55e2/nodes.py#L68 +# https://discuss.pytorch.org/t/help-regarding-slerp-function-for-generative-model-sampling/32475/3 +def slerp(val, low, high): + dims = low.shape + + low = low.reshape(dims[0], -1) + high = high.reshape(dims[0], -1) + + low_norm = low/torch.norm(low, dim=1, keepdim=True) + high_norm = high/torch.norm(high, dim=1, keepdim=True) + + low_norm[low_norm != low_norm] = 0.0 + high_norm[high_norm != high_norm] = 0.0 + + omega = torch.acos((low_norm*high_norm).sum(1)) + so = torch.sin(omega) + res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high + + return res.reshape(dims) + + +def mix_noise(from_noise, to_noise, strength, variation_method): + if variation_method == 'slerp': + mixed_noise = slerp(strength, from_noise, to_noise) + else: + # linear + mixed_noise = (1 - strength) * from_noise + strength * to_noise + + # NOTE: Since the variance of the Gaussian noise in mixed_noise has changed, it must be corrected through scaling. + scale_factor = math.sqrt((1 - strength) ** 2 + strength ** 2) + mixed_noise /= scale_factor + + return mixed_noise + + +def prepare_noise(latent_image, seed, noise_inds=None, noise_device="cpu", incremental_seed_mode="comfy", variation_seed=None, variation_strength=None, variation_method="linear"): + """ + creates random noise given a latent image and a seed. + optional arg skip can be used to skip and discard x number of noise generations for a given seed + """ + + latent_size = latent_image.size() + latent_size_1batch = [1, latent_size[1], latent_size[2], latent_size[3]] + + if variation_strength is not None and variation_strength > 0 or incremental_seed_mode.startswith("variation str inc"): + if noise_device == "cpu": + variation_generator = torch.manual_seed(variation_seed) + else: + torch.cuda.manual_seed(variation_seed) + variation_generator = None + + variation_latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, + generator=variation_generator, device=noise_device) + else: + variation_latent = None + + def apply_variation(input_latent, strength_up=None): + if variation_latent is None: + return input_latent + else: + strength = variation_strength + + if strength_up is not None: + strength += strength_up + + variation_noise = variation_latent.expand(input_latent.size()[0], -1, -1, -1) + + mixed_noise = mix_noise(input_latent, variation_noise, strength, variation_method) + + return mixed_noise + + # method: incremental seed batch noise + if noise_inds is None and incremental_seed_mode == "incremental": + batch_cnt = latent_size[0] + + latents = None + for i in range(batch_cnt): + if noise_device == "cpu": + generator = torch.manual_seed(seed+i) + else: + torch.cuda.manual_seed(seed+i) + generator = None + + latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, + generator=generator, device=noise_device) + + latent = apply_variation(latent) + + if latents is None: + latents = latent + else: + latents = torch.cat((latents, latent), dim=0) + + return latents + + # method: incremental variation batch noise + elif noise_inds is None and incremental_seed_mode.startswith("variation str inc"): + batch_cnt = latent_size[0] + + latents = None + for i in range(batch_cnt): + if noise_device == "cpu": + generator = torch.manual_seed(seed) + else: + torch.cuda.manual_seed(seed) + generator = None + + latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, + generator=generator, device=noise_device) + + step = float(incremental_seed_mode[18:]) + latent = apply_variation(latent, step*i) + + if latents is None: + latents = latent + else: + latents = torch.cat((latents, latent), dim=0) + + return latents + + # method: comfy batch noise + if noise_device == "cpu": + generator = torch.manual_seed(seed) + else: + torch.cuda.manual_seed(seed) + generator = None + + if noise_inds is None: + latents = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, + generator=generator, device=noise_device) + latents = apply_variation(latents) + return latents + + unique_inds, inverse = np.unique(noise_inds, return_inverse=True) + noises = [] + for i in range(unique_inds[-1] + 1): + noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, + generator=generator, device=noise_device) + if i in unique_inds: + noises.append(noise) + noises = [noises[i] for i in inverse] + noises = torch.cat(noises, axis=0) + return noises + + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + + +def empty_pil_tensor(w=64, h=64): + image = Image.new("RGB", (w, h)) + draw = ImageDraw.Draw(image) + draw.rectangle((0, 0, w-1, h-1), fill=(0, 0, 0)) + return pil2tensor(image) + + +def try_install_custom_node(custom_node_url, msg): + try: + import cm_global + cm_global.try_call(api='cm.try-install-custom-node', + sender="Inspire Pack", custom_node_url=custom_node_url, msg=msg) + except Exception as e: + print(msg) + print(f"[Inspire Pack] ComfyUI-Manager is outdated. The custom node installation feature is not available.") + + +def empty_latent(): + return torch.zeros([1, 4, 8, 8]) + +# wildcard trick is taken from pythongossss's +class AnyType(str): + def __ne__(self, __value: object) -> bool: + return False + +any_typ = AnyType("*") + + +# author: Trung0246 ---> +class TautologyStr(str): + def __ne__(self, other): + return False + + +class ByPassTypeTuple(tuple): + def __getitem__(self, index): + if index > 0: + index = 0 + item = super().__getitem__(index) + if isinstance(item, str): + return TautologyStr(item) + return item + + +class TaggedCache: + def __init__(self, tag_settings: Optional[dict]=None): + self._tag_settings = tag_settings or {} # tag cache size + self._data = {} + + def __getitem__(self, key): + for tag_data in self._data.values(): + if key in tag_data: + return tag_data[key] + raise KeyError(f'Key `{key}` does not exist') + + def __setitem__(self, key, value: tuple): + # value: (tag: str, (islist: bool, data: *)) + + # if key already exists, pop old value + for tag_data in self._data.values(): + if key in tag_data: + tag_data.pop(key, None) + break + + tag = value[0] + if tag not in self._data: + + try: + from cachetools import LRUCache + + default_size = 20 + if 'ckpt' in tag: + default_size = 5 + elif tag in ['latent', 'image']: + default_size = 100 + + self._data[tag] = LRUCache(maxsize=self._tag_settings.get(tag, default_size)) + + except (ImportError, ModuleNotFoundError): + # TODO: implement a simple lru dict + self._data[tag] = {} + self._data[tag][key] = value + + def __delitem__(self, key): + for tag_data in self._data.values(): + if key in tag_data: + del tag_data[key] + return + raise KeyError(f'Key `{key}` does not exist') + + def __contains__(self, key): + return any(key in tag_data for tag_data in self._data.values()) + + def items(self): + yield from itertools.chain(*map(lambda x :x.items(), self._data.values())) + + def get(self, key, default=None): + """D.get(k[,d]) -> D[k] if k in D, else d. d defaults to None.""" + for tag_data in self._data.values(): + if key in tag_data: + return tag_data[key] + return default + + def clear(self): + # clear all cache + self._data = {} + + +def make_3d_mask(mask): + if len(mask.shape) == 4: + return mask.squeeze(0) + + elif len(mask.shape) == 2: + return mask.unsqueeze(0) + + return mask + + +def dilate_mask(mask: torch.Tensor, dilation_factor: float) -> torch.Tensor: + """Dilate a mask using a square kernel with a given dilation factor.""" + kernel_size = int(dilation_factor * 2) + 1 + kernel = np.ones((abs(kernel_size), abs(kernel_size)), np.uint8) + + masks = make_3d_mask(mask).numpy() + dilated_masks = [] + for m in masks: + if dilation_factor > 0: + m2 = cv2.dilate(m, kernel, iterations=1) + else: + m2 = cv2.erode(m, kernel, iterations=1) + + dilated_masks.append(torch.from_numpy(m2)) + + return torch.stack(dilated_masks) + + +def flatten_non_zero_override(masks: torch.Tensor): + """ + flatten multiple layer mask tensor to 1 layer mask tensor. + Override the lower layer with the tensor from the upper layer, but only override non-zero values. + + :param masks: 3d mask + :return: flatten mask + """ + final_mask = masks[0] + + for i in range(1, masks.size(0)): + non_zero_mask = masks[i] != 0 + final_mask[non_zero_mask] = masks[i][non_zero_mask] + + return final_mask diff --git a/node.py b/node.py index e8cf816..e3a2208 100644 --- a/node.py +++ b/node.py @@ -25,6 +25,8 @@ import time import random import logging import uuid +import folder_paths +import nodes # from ..ComfyUI_IPAdapter_plus import IPAdapterPlus from PIL import Image, ImageOps, ImageSequence, ImageFile,UnidentifiedImageError from PIL.PngImagePlugin import PngInfo @@ -51,6 +53,36 @@ import importlib import folder_paths import latent_preview import node_helpers +from .libs.utils import TaggedCache, any_typ + +root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +settings_file = os.path.join(root_dir, 'cache_settings.json') +try: + with open(settings_file) as f: + cache_settings = json.load(f) +except Exception as e: + print(e) + cache_settings = {} +cache = TaggedCache(cache_settings) +cache_count = {} + + +def update_cache(k, tag, v): + cache[k] = (tag, v) + cnt = cache_count.get(k) + if cnt is None: + cnt = 0 + cache_count[k] = cnt + else: + cache_count[k] += 1 + + +def cache_weak_hash(k): + cnt = cache_count.get(k) + if cnt is None: + cnt = 0 + + return k, cnt # 获取当前文件的目录 current_dir = os.path.dirname(os.path.abspath(__file__)) @@ -443,4 +475,103 @@ class IPAdapter_FaceID_Bool(): return (work_model, face_image, True) except Exception as e: print(f"Error: {e}") - return (model, image, False) \ No newline at end of file + return (model, image, False) + +class LoraLoaderShared(nodes.LoraLoader): + def __init__(self): + self.loaded_lora = None + + @classmethod + def INPUT_TYPES(s): + file_list = folder_paths.get_filename_list("loras") + file_list.insert(0, "None") + return { + "required": { + "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}), + "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}), + "lora_name": (file_list, {"tooltip": "The name of the LoRA."}), + "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}), + "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}), + "key_opt": ("STRING", {"multiline": False, "placeholder": "If empty, use 'lora_name' as the key."}), + }, + "optional": { + "mode": (['Auto', 'Override Cache', 'Read Only'],), + } + } + + RETURN_TYPES = ("MODEL", "CLIP", "STRING") + RETURN_NAMES = ("model", "clip", "cache key") + FUNCTION = "doit" + + CATEGORY = "InspirePack/Backend" + + def doit(self, model, clip, lora_name, strength_model, strength_clip, key_opt, mode='Auto'): + if mode == 'Read Only': + if key_opt.strip() == '': + raise Exception("[LoraLoaderShared] key_opt cannot be omit if mode is 'Read Only'") + key = key_opt.strip() + elif key_opt.strip() == '': + key = lora_name + else: + key = key_opt.strip() + + if strength_model == 0 and strength_clip == 0: + return (model, clip, key) + + if lora_name == "None": + return (model, clip, key) + + if key not in cache or mode == 'Override Cache': + res = self.load_lora(model, clip, lora_name, strength_model, strength_clip) + update_cache(key, "lora", (False, res)) + cache_kind = 'lora' + print(f"[Inspire Pack] LoraLoaderShared: Lora '{lora_name}' is cached to '{key}'.") + else: + cache_kind, (_, res) = cache[key] + print(f"[Inspire Pack] LoraLoaderShared: Cached lora '{key}' is loaded. (Loading skip)") + + if cache_kind == 'lora': + model, clip = res + else: + raise Exception(f"[LoraLoaderShared] Unexpected cache_kind '{cache_kind}'") + + return model, clip, key + + @staticmethod + def IS_CHANGED(model, clip, lora_name, strength_model, strength_clip, key_opt, mode='Auto'): + if mode == 'Read Only': + if key_opt.strip() == '': + raise Exception("[LoraLoaderShared] key_opt cannot be omit if mode is 'Read Only'") + key = key_opt.strip() + elif key_opt.strip() == '': + key = lora_name + else: + key = key_opt.strip() + + if mode == 'Read Only': + return (None, cache_weak_hash(key)) + elif mode == 'Override Cache': + return (lora_name, key) + + return (None, cache_weak_hash(key)) + + def load_lora(self, model, clip, lora_name, strength_model, strength_clip): + if strength_model == 0 and strength_clip == 0: + return (model, clip) + + 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) + + model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip) + return (model_lora, clip_lora) \ No newline at end of file