diff --git a/comfy/samplers.py b/comfy/samplers.py index a20f65d4a..578e29dc8 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -1,11 +1,13 @@ from .k_diffusion import sampling as k_diffusion_sampling from .extra_samplers import uni_pc +from typing import Dict, List, Tuple import torch import collections from comfy import model_management import math import logging import comfy.sampler_helpers +import comfy.hooks import scipy.stats import numpy @@ -141,7 +143,9 @@ def cond_cat(c_list): def calc_cond_batch(model, conds, x_in, timestep, model_options): out_conds = [] out_counts = [] - to_run = [] + # separate conds by matching hooks + # TODO: implement default_conds support + hooked_to_run: Dict[comfy.hooks.HookWeightGroup,List[Tuple[Tuple,int]]] = {} for i in range(len(conds)): out_conds.append(torch.zeros_like(x_in)) @@ -150,98 +154,102 @@ def calc_cond_batch(model, conds, x_in, timestep, model_options): cond = conds[i] if cond is not None: for x in cond: - p = get_area_and_mult(x, x_in, timestep) + p = comfy.samplers.get_area_and_mult(x, x_in, timestep) if p is None: continue + hook: comfy.hooks.HookWeightGroup = x.get('hooks', None) + hooked_to_run.setdefault(hook, list()) + hooked_to_run[hook] += [(p, i)] - to_run += [(p, i)] + # run every hooked_to_run separately + for hooks, to_run in hooked_to_run.items(): + while len(to_run) > 0: + first = to_run[0] + first_shape = first[0][0].shape + to_batch_temp = [] + for x in range(len(to_run)): + if can_concat_cond(to_run[x][0], first[0]): + to_batch_temp += [x] - while len(to_run) > 0: - first = to_run[0] - first_shape = first[0][0].shape - to_batch_temp = [] - for x in range(len(to_run)): - if can_concat_cond(to_run[x][0], first[0]): - to_batch_temp += [x] + to_batch_temp.reverse() + to_batch = to_batch_temp[:1] - to_batch_temp.reverse() - to_batch = to_batch_temp[:1] + free_memory = model_management.get_free_memory(x_in.device) + for i in range(1, len(to_batch_temp) + 1): + batch_amount = to_batch_temp[:len(to_batch_temp)//i] + input_shape = [len(batch_amount) * first_shape[0]] + list(first_shape)[1:] + if model.memory_required(input_shape) * 1.5 < free_memory: + to_batch = batch_amount + break + # TODO: add apply_hooks call here, once a ModelPatcher ref is added to BaseModel - free_memory = model_management.get_free_memory(x_in.device) - for i in range(1, len(to_batch_temp) + 1): - batch_amount = to_batch_temp[:len(to_batch_temp)//i] - input_shape = [len(batch_amount) * first_shape[0]] + list(first_shape)[1:] - if model.memory_required(input_shape) * 1.5 < free_memory: - to_batch = batch_amount - break + input_x = [] + mult = [] + c = [] + cond_or_uncond = [] + area = [] + control = None + patches = None + for x in to_batch: + o = to_run.pop(x) + p = o[0] + input_x.append(p.input_x) + mult.append(p.mult) + c.append(p.conditioning) + area.append(p.area) + cond_or_uncond.append(o[1]) + control = p.control + patches = p.patches - input_x = [] - mult = [] - c = [] - cond_or_uncond = [] - area = [] - control = None - patches = None - for x in to_batch: - o = to_run.pop(x) - p = o[0] - input_x.append(p.input_x) - mult.append(p.mult) - c.append(p.conditioning) - area.append(p.area) - cond_or_uncond.append(o[1]) - control = p.control - patches = p.patches + batch_chunks = len(cond_or_uncond) + input_x = torch.cat(input_x) + c = cond_cat(c) + timestep_ = torch.cat([timestep] * batch_chunks) - batch_chunks = len(cond_or_uncond) - input_x = torch.cat(input_x) - c = cond_cat(c) - timestep_ = torch.cat([timestep] * batch_chunks) + if control is not None: + c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond)) - if control is not None: - c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond)) + transformer_options = {} + if 'transformer_options' in model_options: + transformer_options = model_options['transformer_options'].copy() - transformer_options = {} - if 'transformer_options' in model_options: - transformer_options = model_options['transformer_options'].copy() + if patches is not None: + if "patches" in transformer_options: + cur_patches = transformer_options["patches"].copy() + for p in patches: + if p in cur_patches: + cur_patches[p] = cur_patches[p] + patches[p] + else: + cur_patches[p] = patches[p] + transformer_options["patches"] = cur_patches + else: + transformer_options["patches"] = patches - if patches is not None: - if "patches" in transformer_options: - cur_patches = transformer_options["patches"].copy() - for p in patches: - if p in cur_patches: - cur_patches[p] = cur_patches[p] + patches[p] - else: - cur_patches[p] = patches[p] - transformer_options["patches"] = cur_patches + transformer_options["cond_or_uncond"] = cond_or_uncond[:] + transformer_options["sigmas"] = timestep + + c['transformer_options'] = transformer_options + + if 'model_function_wrapper' in model_options: + output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks) else: - transformer_options["patches"] = patches + output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks) - transformer_options["cond_or_uncond"] = cond_or_uncond[:] - transformer_options["sigmas"] = timestep - - c['transformer_options'] = transformer_options - - if 'model_function_wrapper' in model_options: - output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks) - else: - output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks) - - for o in range(batch_chunks): - cond_index = cond_or_uncond[o] - a = area[o] - if a is None: - out_conds[cond_index] += output[o] * mult[o] - out_counts[cond_index] += mult[o] - else: - out_c = out_conds[cond_index] - out_cts = out_counts[cond_index] - dims = len(a) // 2 - for i in range(dims): - out_c = out_c.narrow(i + 2, a[i + dims], a[i]) - out_cts = out_cts.narrow(i + 2, a[i + dims], a[i]) - out_c += output[o] * mult[o] - out_cts += mult[o] + for o in range(batch_chunks): + cond_index = cond_or_uncond[o] + a = area[o] + if a is None: + out_conds[cond_index] += output[o] * mult[o] + out_counts[cond_index] += mult[o] + else: + out_c = out_conds[cond_index] + out_cts = out_counts[cond_index] + dims = len(a) // 2 + for i in range(dims): + out_c = out_c.narrow(i + 2, a[i + dims], a[i]) + out_cts = out_cts.narrow(i + 2, a[i + dims], a[i]) + out_c += output[o] * mult[o] + out_cts += mult[o] for i in range(len(out_conds)): out_conds[i] /= out_counts[i]