From d86b16cabb3ad19996825a47f3d6cbb9d72fc4d3 Mon Sep 17 00:00:00 2001 From: shiimizu Date: Sun, 25 Aug 2024 02:39:09 -0700 Subject: [PATCH] Simpler hooks. --- __init__.py | 9 +- utils.py | 326 +++++++++++++++------------------------------------- 2 files changed, 96 insertions(+), 239 deletions(-) diff --git a/__init__.py b/__init__.py index 6affb6c..433b6f8 100644 --- a/__init__.py +++ b/__init__.py @@ -1,9 +1,6 @@ from .tiled_diffusion import NODE_CLASS_MAPPINGS as TD_NCM, NODE_DISPLAY_NAME_MAPPINGS as TD_NDCM from .tiled_vae import NODE_CLASS_MAPPINGS as TV_NCM, NODE_DISPLAY_NAME_MAPPINGS as TV_NDCM -NODE_CLASS_MAPPINGS = {} -NODE_DISPLAY_NAME_MAPPINGS = {} -NODE_CLASS_MAPPINGS.update(TD_NCM) -NODE_DISPLAY_NAME_MAPPINGS.update(TD_NDCM) -NODE_CLASS_MAPPINGS.update(TV_NCM) -NODE_DISPLAY_NAME_MAPPINGS.update(TV_NDCM) +from .utils import store as _ +NODE_CLASS_MAPPINGS = {**TD_NCM, **TV_NCM} +NODE_DISPLAY_NAME_MAPPINGS = {**TD_NDCM, **TV_NDCM} __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/utils.py b/utils.py index b9309fd..c01deb7 100644 --- a/utils.py +++ b/utils.py @@ -1,249 +1,109 @@ -import inspect -import importlib -from textwrap import dedent, indent -from copy import copy -import types -import functools -import os -import sys -import binascii -from typing import List, NamedTuple +store = {} -class Hook(NamedTuple): - fn: object - module_name: str - target: str - orig_key: str - module_name_path: str +# ==================== Hook into sampling functions for ControlNet ==================== -def gen_id(): - return binascii.hexlify(os.urandom(1024))[64:72].decode("utf-8") +import comfy.samplers -def hook_calc_cond_uncond_batch(): - try: - from comfy.samplers import calc_cond_batch - calc_cond_batch_ = calc_cond_batch - except Exception: - from comfy.samplers import calc_cond_uncond_batch - calc_cond_batch_ = calc_cond_uncond_batch - # this function should only be run by us - orig_key = f"{calc_cond_batch_.__name__}_original_tiled_diffusion_{gen_id()}" - payload = [{ - "mode": "replace", - "target_line": 'control.get_control', - "code_to_insert": """control if 'tiled_diffusion' in model_options else control.get_control""" - }, - { - "dedent": False, - "target_line": calc_cond_batch_.__name__, - "code_to_insert": f""" - if 'tiled_diffusion' not in model_options: - return {orig_key}{inspect.signature(calc_cond_batch_)}""" - }] - fn = inject_code(calc_cond_batch_, payload, 'w') - return create_hook(fn, 'comfy.samplers', orig_key=orig_key) +def patch1(fn_name): + def calc_cond_batch(*args, **kwargs): + x_in = kwargs['x_in'] if 'x_in' in kwargs else args[2] + model_options = kwargs['model_options'] if 'model_options' in kwargs else args[4] + if not hasattr(x_in, 'model_options'): + x_in.model_options = model_options + return store[fn_name](*args, **kwargs) + return calc_cond_batch -def hook_sag_create_blur_map(): - imported = False - try: - import comfy_extras - from comfy_extras import nodes_sag - imported = True - except Exception: ... - if not imported: return - import comfy_extras - from comfy_extras import nodes_sag - import re - source=inspect.getsource(nodes_sag.create_blur_map) - replace_str=""" +def patch2(fn_name): + def get_area_and_mult(*args, **kwargs): + x_in = kwargs['x_in'] if 'x_in' in kwargs else args[1] + conds = kwargs['conds'] if 'conds' in kwargs else args[0] + if (model_options:=getattr(x_in, 'model_options', None)) is not None and 'tiled_diffusion' in model_options: + if 'control' in conds: + control = conds['control'] + if not hasattr(control, 'get_control_orig'): + control.get_control_orig = control.get_control + control.get_control = lambda *a, **kw: control + else: + if 'control' in conds: + control = conds['control'] + if hasattr(control, 'get_control_orig') and control.get_control != control.get_control_orig: + control.get_control = control.get_control_orig + return store[fn_name](*args, **kwargs) + return get_area_and_mult + +patches = [ + (comfy.samplers, 'calc_cond_batch', patch1), + (comfy.samplers, 'get_area_and_mult', patch2), +] +for parent, fn_name, create_patch in patches: + store[fn_name] = getattr(parent, fn_name) + setattr(parent, fn_name, create_patch(fn_name)) + +# ==================== Patch pre_run_control ==================== + +# Is this necessary anymore? +def pre_run_control(model, conds): + s = model.model_sampling + for t in range(len(conds)): + x = conds[t] + + timestep_start = None + timestep_end = None + percent_to_timestep_function = lambda a: s.percent_to_sigma(a) + if 'control' in x: + try: x['control'].cleanup() + except Exception: ... + x['control'].pre_run(model, percent_to_timestep_function) +comfy.samplers.pre_run_control = pre_run_control + +# ==================== Patch SAG ==================== + +import math +import torch.nn.functional as F +import comfy_extras.nodes_sag +from comfy_extras.nodes_sag import gaussian_blur_2d +def create_blur_map(x0, attn, sigma=3.0, threshold=1.0): + # reshape and GAP the attention map + _, hw1, hw2 = attn.shape + b, _, lh, lw = x0.shape + attn = attn.reshape(b, -1, hw1, hw2) + # Global Average Pool + mask = attn.mean(1, keepdim=False).sum(1, keepdim=False) > threshold def calc_closest_factors(a): for b in range(int(math.sqrt(a)), 0, -1): - if a%b == 0: + if a % b == 0: c = a // b return (b,c) m = calc_closest_factors(hw1) mh = max(m) if lh > lw else min(m) mw = m[1] if mh == m[0] else m[0] - mid_shape = mh, mw""" - modified_source = re.sub(r"ratio =.*\s+mid_shape =.*", replace_str, source, flags=re.MULTILINE) - fn = write_to_file_and_return_fn(nodes_sag.create_blur_map, modified_source) - return create_hook(fn, 'comfy_extras.nodes_sag') + mid_shape = mh, mw -def hook_samplers_pre_run_control(): - from comfy.samplers import pre_run_control - payload = [{ - "dedent": False, - "target_line": "if 'control' in x:", - "code_to_insert": """ try: x['control'].cleanup()\n except Exception: ...""" - }, - { - "target_line": "s = model.model_sampling", - "code_to_insert": """ - def find_outer_instance(target:str, target_type): - import inspect - frame = inspect.currentframe() - i = 0 - while frame and i < 7: - if (found:=frame.f_locals.get(target, None)) is not None: - if isinstance(found, target_type): - return found - frame = frame.f_back - i += 1 - return None - from comfy.model_patcher import ModelPatcher - if (_model:=find_outer_instance('model', ModelPatcher)) is not None: - if (model_function_wrapper:=_model.model_options.get('model_function_wrapper', None)) is not None: - import sys - tiled_diffusion = sys.modules.get('ComfyUI-TiledDiffusion.tiled_diffusion', None) - if tiled_diffusion is None: - for key in sys.modules: - if 'tiled_diffusion' in key: - tiled_diffusion = sys.modules[key] - break - if (AbstractDiffusion:=getattr(tiled_diffusion, 'AbstractDiffusion', None)) is not None: - if isinstance(model_function_wrapper, AbstractDiffusion): - model_function_wrapper.reset() - """}] - fn = inject_code(pre_run_control, payload) - return create_hook(fn, 'comfy.samplers') + # Reshape + mask = ( + mask.reshape(b, *mid_shape) + .unsqueeze(1) + .type(attn.dtype) + ) + # Upsample + mask = F.interpolate(mask, (lh, lw)) -def hook_gligen__set_position(): - from comfy.gligen import Gligen - import re - source=inspect.getsource(Gligen._set_position) - replace_str=""" - nonlocal objs - if x.shape[0] > objs.shape[0]: - _objs = objs.repeat(-(x.shape[0] // -objs.shape[0]),1,1) - else: - _objs = objs - return module(x, _objs.to(device=x.device, dtype=x.dtype))""" - modified_source = dedent(re.sub(r" return module\(.*", replace_str, source, 1, re.MULTILINE) or source) - fn = write_to_file_and_return_fn(Gligen._set_position, modified_source) - return create_hook(fn, 'comfy.gligen', 'Gligen._set_position') + blurred = gaussian_blur_2d(x0, kernel_size=9, sigma=sigma) + blurred = blurred * mask + x0 * (1 - mask) + return blurred +comfy_extras.nodes_sag.create_blur_map = create_blur_map -def create_hook(fn, module_name:str, target = None, orig_key = None): - if target is None: target = fn.__name__ - if orig_key is None: orig_key = f'{target}_original' - module_name_path = os.path.normpath(module_name.replace('.', '/')) - return Hook(fn, module_name, target, orig_key, module_name_path) +# ==================== Patch Gligen ==================== -def _getattr(obj, name:str, default=None): - """multi-level getattr""" - for attr in name.split('.'): - obj = getattr(obj, attr, default) - return obj - -def _hasattr(obj, name:str): - """multi-level hasattr""" - return _getattr(obj, name) is not None - -def _setattr(obj, name:str, value=None): - """multi-level setattr""" - split = name.split('.') - if not split[:-1]: - return setattr(obj, name, value) - else: - name = split[-1] - for attr in split[:-1]: - obj = getattr(obj, attr, None) - return setattr(obj, name, value) - -def hook_all(restore=False, hooks=None): - if hooks is None: - hooks: List[Hook] = [ - hook_calc_cond_uncond_batch(), - hook_sag_create_blur_map(), - hook_samplers_pre_run_control(), - hook_gligen__set_position(), - ] - for key, module in sys.modules.items(): - for hook in hooks: - if hook is None: - continue - if key == hook.module_name or key.endswith(hook.module_name_path): - if _hasattr(module, hook.target): - if not _hasattr(module, hook.orig_key): - if (orig_fn:=_getattr(module, hook.target, None)) is not None: - _setattr(module, hook.orig_key, orig_fn) - if restore: - _setattr(module, hook.target, _getattr(module, hook.orig_key, None)) - else: - _setattr(module, hook.target, hook.fn) - -def inject_code(original_func, data, mode='a'): - # Get the source code of the original function - original_source = inspect.getsource(original_func) - - # Split the source code into lines - lines = original_source.split("\n") - - for item in data: - # Find the line number of the target line - target_line_number = None - for i, line in enumerate(lines): - if item['target_line'] not in line: continue - target_line_number = i + 1 - if item.get("mode","insert") == "replace": - lines[i] = lines[i].replace(item['target_line'], item['code_to_insert']) - break - - # Find the indentation of the line where the new code will be inserted - indentation = '' - for char in line: - if char == ' ': - indentation += char - else: - break - - # Indent the new code to match the original - code_to_insert = item['code_to_insert'] - if item.get("dedent",True): - code_to_insert = dedent(item['code_to_insert']) - code_to_insert = indent(code_to_insert, indentation) - - break - - # Insert the code to be injected after the target line - if item.get("mode","insert") == "insert" and target_line_number is not None: - lines.insert(target_line_number, code_to_insert) - - # Recreate the modified source code - modified_source = "\n".join(lines) - modified_source = dedent(modified_source.strip("\n")) - return write_to_file_and_return_fn(original_func, modified_source, mode) - -def write_to_file_and_return_fn(original_func, source:str, mode='a'): - # Write the modified source code to a temporary file so the - # source code and stack traces can still be viewed when debugging. - custom_name = ".patches.py" - current_dir = os.path.dirname(os.path.abspath(__file__)) - temp_file_path = os.path.join(current_dir, custom_name) - with open(temp_file_path, mode) as temp_file: - temp_file.write(source) - temp_file.write("\n") - temp_file.flush() - - MODULE_PATH = temp_file.name - MODULE_NAME = __name__.split('.')[0].replace('-','_') + "_patch_modules" - spec = importlib.util.spec_from_file_location(MODULE_NAME, MODULE_PATH) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - - # Retrieve the modified function from the module - modified_function = getattr(module, original_func.__name__) - - # Adapted from https://stackoverflow.com/a/49077211 - def copy_func(f, globals=None, module=None, code=None, update_wrapper=True): - if globals is None: globals = f.__globals__ - if code is None: code = f.__code__ - g = types.FunctionType(code, globals, name=f.__name__, - argdefs=f.__defaults__, closure=f.__closure__) - if update_wrapper: g = functools.update_wrapper(g, f) - if module is not None: g.__module__ = module - g.__kwdefaults__ = copy(f.__kwdefaults__) - return g - - return copy_func(original_func, code=modified_function.__code__, update_wrapper=False) +def _set_position(self, boxes, masks, positive_embeddings): + objs = self.position_net(boxes, masks, positive_embeddings) + def func(x, extra_options): + key = extra_options["transformer_index"] + module = self.module_list[key] + nonlocal objs + _objs = objs.repeat(-(x.shape[0] // -objs.shape[0]),1,1) if x.shape[0] > objs.shape[0] else objs + return module(x, _objs.to(device=x.device, dtype=x.dtype)) + return func +import comfy.gligen +comfy.gligen.Gligen._set_position = _set_position