Spaces:
Running
on
Zero
Running
on
Zero
| store = {} | |
| # ==================== Hook into sampling functions for ControlNet ==================== | |
| import comfy.samplers | |
| 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 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: | |
| 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 | |
| # Reshape | |
| mask = ( | |
| mask.reshape(b, *mid_shape) | |
| .unsqueeze(1) | |
| .type(attn.dtype) | |
| ) | |
| # Upsample | |
| 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 | |
| # ==================== Patch Gligen ==================== | |
| 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 | |