nodes/SipherAGI-comfyui-animatediff
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update sampling function to match new ComfyUI update (#58)

Browse Source
This commit is contained in:
Tung Nguyen (Blockchain)
2023-10-25 16:40:39 +07:00
committed by GitHub
parent dccfb2326e
commit 6416ffafa2
1 changed files with 47 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
animatediff/sliding_context_sampling.py
+47 -95
View File
@@ -74,47 +74,34 @@ def __sliding_sample_factory(ctx: SlidingContext):
ctx.current_step = start_step + step + 1
try:
return orig_comfy_sample(model, *args, **kwargs, callback=callback)
except RuntimeError as e:
if str(e).startswith("CUDA error: invalid configuration argument"):
raise RuntimeError(
f"An xformers bug was encountered in AnimateDiff - to run your workflow, \
disable xformers in ComfyUI using '--disable-xformers' startup argument."
)
raise
return orig_comfy_sample(model, *args, **kwargs, callback=callback)
def sampling_function(
model_function, x, timestep, uncond, cond, cond_scale, cond_concat=None, model_options={}, seed=None
):
def get_area_and_mult(cond, x_in, cond_concat_in, timestep_in):
def sampling_function(model_function, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
def get_area_and_mult(conds, x_in, timestep_in):
area = (x_in.shape[2], x_in.shape[3], 0, 0)
strength = 1.0
if "timestep_start" in cond[1]:
timestep_start = cond[1]["timestep_start"]
if "timestep_start" in conds:
timestep_start = conds["timestep_start"]
if timestep_in[0] > timestep_start:
return None
if "timestep_end" in cond[1]:
timestep_end = cond[1]["timestep_end"]
if "timestep_end" in conds:
timestep_end = conds["timestep_end"]
if timestep_in[0] < timestep_end:
return None
if "area" in cond[1]:
area = cond[1]["area"]
if "strength" in cond[1]:
strength = cond[1]["strength"]
adm_cond = None
if "adm_encoded" in cond[1]:
adm_cond = cond[1]["adm_encoded"]
if "area" in conds:
area = conds["area"]
if "strength" in conds:
strength = conds["strength"]
input_x = x_in[:, :, area[2] : area[0] + area[2], area[3] : area[1] + area[3]]
if "mask" in cond[1]:
if "mask" in conds:
# Scale the mask to the size of the input
# The mask should have been resized as we began the sampling process
mask_strength = 1.0
if "mask_strength" in cond[1]:
mask_strength = cond[1]["mask_strength"]
mask = cond[1]["mask"]
if "mask_strength" in conds:
mask_strength = conds["mask_strength"]
mask = conds["mask"]
assert mask.shape[1] == x_in.shape[2]
assert mask.shape[2] == x_in.shape[3]
mask = mask[:, area[2] : area[0] + area[2], area[3] : area[1] + area[3]] * mask_strength
@@ -123,7 +110,7 @@ def __sliding_sample_factory(ctx: SlidingContext):
mask = torch.ones_like(input_x)
mult = mask * strength
if "mask" not in cond[1]:
if "mask" not in conds:
rr = 8
if area[2] != 0:
for t in range(rr):
@@ -139,24 +126,17 @@ def __sliding_sample_factory(ctx: SlidingContext):
mult[:, :, :, area[1] - 1 - t : area[1] - t] *= (1.0 / rr) * (t + 1)
conditionning = {}
conditionning["c_crossattn"] = cond[0]
if cond_concat_in is not None and len(cond_concat_in) > 0:
cropped = []
for x in cond_concat_in:
cr = x[:, :, area[2] : area[0] + area[2], area[3] : area[1] + area[3]]
cropped.append(cr)
conditionning["c_concat"] = torch.cat(cropped, dim=1)
if adm_cond is not None:
conditionning["c_adm"] = adm_cond
model_conds = conds["model_conds"]
for c in model_conds:
conditionning[c] = model_conds[c].process_cond(batch_size=x_in.shape[0], device=x_in.device, area=area)
control = None
if "control" in cond[1]:
control = cond[1]["control"]
if "control" in conds:
control = conds["control"]
patches = None
if "gligen" in cond[1]:
gligen = cond[1]["gligen"]
if "gligen" in conds:
gligen = conds["gligen"]
patches = {}
gligen_type = gligen[0]
gligen_model = gligen[1]
@@ -174,24 +154,8 @@ def __sliding_sample_factory(ctx: SlidingContext):
return True
if c1.keys() != c2.keys():
return False
if "c_crossattn" in c1:
s1 = c1["c_crossattn"].shape
s2 = c2["c_crossattn"].shape
if s1 != s2:
if s1[0] != s2[0] or s1[2] != s2[2]: # these 2 cases should not happen
return False
mult_min = lcm(s1[1], s2[1])
diff = mult_min // min(s1[1], s2[1])
if (
diff > 4
): # arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
return False
if "c_concat" in c1:
if c1["c_concat"].shape != c2["c_concat"].shape:
return False
if "c_adm" in c1:
if c1["c_adm"].shape != c2["c_adm"].shape:
for k in c1:
if not c1[k].can_concat(c2[k]):
return False
return True
@@ -220,36 +184,22 @@ def __sliding_sample_factory(ctx: SlidingContext):
c_concat = []
c_adm = []
crossattn_max_len = 0
for x in c_list:
if "c_crossattn" in x:
c = x["c_crossattn"]
if crossattn_max_len == 0:
crossattn_max_len = c.shape[1]
else:
crossattn_max_len = lcm(crossattn_max_len, c.shape[1])
c_crossattn.append(c)
if "c_concat" in x:
c_concat.append(x["c_concat"])
if "c_adm" in x:
c_adm.append(x["c_adm"])
out = {}
c_crossattn_out = []
for c in c_crossattn:
if c.shape[1] < crossattn_max_len:
c = c.repeat(1, crossattn_max_len // c.shape[1], 1) # padding with repeat doesn't change result
c_crossattn_out.append(c)
if len(c_crossattn_out) > 0:
out["c_crossattn"] = torch.cat(c_crossattn_out)
if len(c_concat) > 0:
out["c_concat"] = torch.cat(c_concat)
if len(c_adm) > 0:
out["c_adm"] = torch.cat(c_adm)
temp = {}
for x in c_list:
for k in x:
cur = temp.get(k, [])
cur.append(x[k])
temp[k] = cur
out = {}
for k in temp:
conds = temp[k]
out[k] = conds[0].concat(conds[1:])
return out
def calc_cond_uncond_batch(
model_function, cond, uncond, x_in, timestep, max_total_area, cond_concat_in, model_options
):
def calc_cond_uncond_batch(model_function, cond, uncond, x_in, timestep, max_total_area, model_options):
out_cond = torch.zeros_like(x_in)
out_count = torch.ones_like(x_in) / 100000.0
@@ -261,14 +211,14 @@ def __sliding_sample_factory(ctx: SlidingContext):
to_run = []
for x in cond:
p = get_area_and_mult(x, x_in, cond_concat_in, timestep)
p = get_area_and_mult(x, x_in, timestep)
if p is None:
continue
to_run += [(p, COND)]
if uncond is not None:
for x in uncond:
p = get_area_and_mult(x, x_in, cond_concat_in, timestep)
p = get_area_and_mult(x, x_in, timestep)
if p is None:
continue
@@ -371,7 +321,7 @@ def __sliding_sample_factory(ctx: SlidingContext):
# sliding_calc_cond_uncond_batch inspired by ashen's initial hack for 16-frame sliding context:
# https://github.com/comfyanonymous/ComfyUI/compare/master...ashen-sensored:ComfyUI:master
def sliding_calc_cond_uncond_batch(
model_function, cond, uncond, x_in, timestep, max_total_area, cond_concat_in, model_options
model_function, cond, uncond, x_in, timestep, max_total_area, model_options
):
# figure out how input is split
axes_factor = x.size(0) // ctx.video_length
@@ -441,7 +391,6 @@ def __sliding_sample_factory(ctx: SlidingContext):
sub_timestep = timestep[full_idxs]
sub_cond = get_resized_cond(cond, full_idxs) if cond is not None else None
sub_uncond = get_resized_cond(uncond, full_idxs) if uncond is not None else None
sub_cond_concat = get_resized_cond(cond_concat, full_idxs) if cond_concat is not None else None
sub_cond_out, sub_uncond_out = calc_cond_uncond_batch(
model_function,
@@ -450,7 +399,6 @@ def __sliding_sample_factory(ctx: SlidingContext):
sub_x,
sub_timestep,
max_total_area,
sub_cond_concat,
model_options,
)
@@ -468,7 +416,7 @@ def __sliding_sample_factory(ctx: SlidingContext):
uncond = None
cond, uncond = sliding_calc_cond_uncond_batch(
model_function, cond, uncond, x, timestep, max_total_area, cond_concat, model_options
model_function, cond, uncond, x, timestep, max_total_area, model_options
)
if "sampler_cfg_function" in model_options:
@@ -481,6 +429,10 @@ def __sliding_sample_factory(ctx: SlidingContext):
def inject_sampling_function(ctx: SlidingContext):
global orig_comfy_sample, orig_sampling_function
orig_comfy_sample = comfy.sample.sample
orig_sampling_function = comfy_samplers.sampling_function
(sample, sampling_function) = __sliding_sample_factory(ctx)
comfy.sample.sample = sample
comfy_samplers.sampling_function = sampling_function