xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-29 10:47:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

Initial changes to calc_cond_batch to eventually support hook_patches

Browse Source
This commit is contained in:
kosinkadink1@gmail.com 2024-09-13 18:31:52 +09:00
parent 069ec7a64b
commit 3cbd40ada3
1 changed files with 88 additions and 80 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

168
comfy/samplers.py
View File

@ -1,11 +1,13 @@
from .k_diffusion import sampling as k_diffusion_sampling from .k_diffusion import sampling as k_diffusion_sampling
from .extra_samplers import uni_pc from .extra_samplers import uni_pc
from typing import Dict, List, Tuple
import torch import torch
import collections import collections
from comfy import model_management from comfy import model_management
import math import math
import logging import logging
import comfy.sampler_helpers import comfy.sampler_helpers
import comfy.hooks
import scipy.stats import scipy.stats
import numpy import numpy
@ -141,7 +143,9 @@ def cond_cat(c_list):
def calc_cond_batch(model, conds, x_in, timestep, model_options): def calc_cond_batch(model, conds, x_in, timestep, model_options):
out_conds = [] out_conds = []
out_counts = [] 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)): for i in range(len(conds)):
out_conds.append(torch.zeros_like(x_in)) 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] cond = conds[i]
if cond is not None: if cond is not None:
for x in cond: 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: if p is None:
continue 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: to_batch_temp.reverse()
first = to_run[0] to_batch = to_batch_temp[:1]
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() free_memory = model_management.get_free_memory(x_in.device)
to_batch = to_batch_temp[:1] 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) input_x = []
for i in range(1, len(to_batch_temp) + 1): mult = []
batch_amount = to_batch_temp[:len(to_batch_temp)//i] c = []
input_shape = [len(batch_amount) * first_shape[0]] + list(first_shape)[1:] cond_or_uncond = []
if model.memory_required(input_shape) * 1.5 < free_memory: area = []
to_batch = batch_amount control = None
break 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 = [] batch_chunks = len(cond_or_uncond)
mult = [] input_x = torch.cat(input_x)
c = [] c = cond_cat(c)
cond_or_uncond = [] timestep_ = torch.cat([timestep] * batch_chunks)
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) if control is not None:
input_x = torch.cat(input_x) c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond))
c = cond_cat(c)
timestep_ = torch.cat([timestep] * batch_chunks)
if control is not None: transformer_options = {}
c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond)) if 'transformer_options' in model_options:
transformer_options = model_options['transformer_options'].copy()
transformer_options = {} if patches is not None:
if 'transformer_options' in model_options: if "patches" in transformer_options:
transformer_options = model_options['transformer_options'].copy() 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: transformer_options["cond_or_uncond"] = cond_or_uncond[:]
if "patches" in transformer_options: transformer_options["sigmas"] = timestep
cur_patches = transformer_options["patches"].copy()
for p in patches: c['transformer_options'] = transformer_options
if p in cur_patches:
cur_patches[p] = cur_patches[p] + patches[p] if 'model_function_wrapper' in model_options:
else: 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)
cur_patches[p] = patches[p]
transformer_options["patches"] = cur_patches
else: else:
transformer_options["patches"] = patches output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
transformer_options["cond_or_uncond"] = cond_or_uncond[:] for o in range(batch_chunks):
transformer_options["sigmas"] = timestep cond_index = cond_or_uncond[o]
a = area[o]
c['transformer_options'] = transformer_options if a is None:
out_conds[cond_index] += output[o] * mult[o]
if 'model_function_wrapper' in model_options: out_counts[cond_index] += mult[o]
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:
else: out_c = out_conds[cond_index]
output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks) out_cts = out_counts[cond_index]
dims = len(a) // 2
for o in range(batch_chunks): for i in range(dims):
cond_index = cond_or_uncond[o] out_c = out_c.narrow(i + 2, a[i + dims], a[i])
a = area[o] out_cts = out_cts.narrow(i + 2, a[i + dims], a[i])
if a is None: out_c += output[o] * mult[o]
out_conds[cond_index] += output[o] * mult[o] out_cts += 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)): for i in range(len(out_conds)):
out_conds[i] /= out_counts[i] out_conds[i] /= out_counts[i]