| class Store:
|
| def __repr__(self):
|
| keys = sorted(self.__dict__)
|
| items = ("{}={!r}".format(k, self.__dict__[k]) for k in keys)
|
| return "{}({})".format(type(self).__name__, ", ".join(items))
|
|
|
| def __eq__(self, other):
|
| return self.__dict__ == other.__dict__
|
|
|
| store = Store()
|
|
|
|
|
|
|
| import comfy.samplers
|
|
|
| def KSAMPLER_sample(*args, **kwargs):
|
| orig_fn = store.KSAMPLER_sample
|
| extra_args = None
|
| model_options = None
|
| try:
|
| extra_args = kwargs['extra_args'] if 'extra_args' in kwargs else args[3]
|
| model_options = extra_args['model_options']
|
| except Exception: ...
|
| if model_options is not None and 'tiled_diffusion' in model_options and extra_args is not None:
|
| sigmas_ = kwargs['sigmas'] if 'sigmas' in kwargs else args[2]
|
| sigmas_all = model_options.pop('sigmas', None)
|
| sigmas = sigmas_all if sigmas_all is not None else sigmas_
|
| store.sigmas = sigmas
|
| store.model_options = model_options
|
| store.extra_args = extra_args
|
| else:
|
| for attr in ['sigmas', 'model_options', 'extra_args']:
|
| _delattr(store, attr)
|
| return orig_fn(*args, **kwargs)
|
|
|
| def KSampler_sample(*args, **kwargs):
|
| orig_fn = store.KSampler_sample
|
| self = args[0]
|
| model_patcher = getattr(self, 'model', None)
|
| model_options = getattr(model_patcher, 'model_options', None)
|
| if model_options is not None and 'tiled_diffusion' in model_options:
|
| sigmas = None
|
| try: sigmas = kwargs['sigmas'] if 'sigmas' in kwargs else args[10]
|
| except Exception: ...
|
| if sigmas is None:
|
| sigmas = getattr(self, 'sigmas', None)
|
| if sigmas is not None:
|
| model_options = model_options.copy()
|
| model_options['sigmas'] = sigmas
|
| self.model.model_options = model_options
|
| return orig_fn(*args, **kwargs)
|
|
|
| def get_area_and_mult(*args, **kwargs):
|
| conds = kwargs['conds'] if 'conds' in kwargs else args[0]
|
| if (model_options:=getattr(store, '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.get_area_and_mult(*args, **kwargs)
|
|
|
| def _delattr(obj, attr):
|
| try:
|
| if hasattr(obj, attr): delattr(obj, attr)
|
| except Exception: ...
|
|
|
| def register_hooks():
|
| patches = [
|
| (comfy.samplers.KSampler, 'sample', KSampler_sample),
|
| (comfy.samplers.KSAMPLER, 'sample', KSAMPLER_sample),
|
| (comfy.samplers, 'get_area_and_mult', get_area_and_mult),
|
| ]
|
| for parent, fn_name, fn_patch in patches:
|
| if not hasattr(parent, f"_{fn_name}"):
|
| setattr(store, f"_{fn_name}", getattr(parent, fn_name))
|
| setattr(store, fn_patch.__name__, getattr(parent, fn_name))
|
| setattr(parent, fn_name, fn_patch)
|
|
|
| register_hooks()
|
|
|
|
|
|
|
|
|
| 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
|
|
|
|
|
|
|
| from math import sqrt
|
| import torch.nn.functional as F
|
| import comfy_extras.nodes_sag
|
| from comfy_extras.nodes_sag import gaussian_blur_2d
|
|
|
| def calc_closest_factors(a):
|
| for b in range(int(sqrt(a)), 0, -1):
|
| if a % b == 0:
|
| c = a // b
|
| return (b, c)
|
|
|
| def create_blur_map(x0, attn, sigma=3.0, threshold=1.0):
|
|
|
| _, hw1, hw2 = attn.shape
|
| b, _, lh, lw = x0.shape
|
| attn = attn.reshape(b, -1, hw1, hw2)
|
|
|
| mask = attn.mean(1, keepdim=False).sum(1, keepdim=False) > threshold
|
| 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
|
|
|
|
|
| mask = (
|
| mask.reshape(b, *mid_shape)
|
| .unsqueeze(1)
|
| .type(attn.dtype)
|
| )
|
|
|
| mask = F.interpolate(mask, (lh, lw))
|
|
|
| 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 _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
|
|
|
