From 599b71807de3ec68e359b4402bc020304d8f7dd4 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 8 Mar 2025 23:36:52 -0600 Subject: [PATCH] Added overlap-linear fuse method, marked some fuse methods as experimental to make the dropdown less overwhelming, made sigma-related fuse methods work with View Options --- animatediff/context.py | 30 ++++++++++++++++++++++++------ animatediff/motion_module_ad.py | 2 +- animatediff/sampling.py | 13 +++++-------- 3 files changed, 30 insertions(+), 15 deletions(-) diff --git a/animatediff/context.py b/animatediff/context.py index 61e20b5..9c6262f 100644 --- a/animatediff/context.py +++ b/animatediff/context.py @@ -20,6 +20,8 @@ class ContextFuseMethod: FLAT = "flat" PYRAMID = "pyramid" RELATIVE = "relative" + OVERLAP_LINEAR = "overlap-linear" + RANDOM = "🔬random" RANDOM_DEPR = "random" GAUSS_SIGMA = "🔬gauss-sigma" @@ -33,8 +35,8 @@ class ContextFuseMethod: PYRAMID_SIGMA_INV = "🔬pyramid-sigma inverse" PYRAMID_SIGMA_INV_DEPR = "pyramid-sigma inverse" - LIST = [PYRAMID, FLAT, DELAYED_REVERSE_SAWTOOTH, PYRAMID_SIGMA, PYRAMID_SIGMA_INV, GAUSS_SIGMA, GAUSS_SIGMA_INV, RANDOM] - LIST_STATIC = [PYRAMID, RELATIVE, FLAT, DELAYED_REVERSE_SAWTOOTH, PYRAMID_SIGMA, PYRAMID_SIGMA_INV, GAUSS_SIGMA, GAUSS_SIGMA_INV, RANDOM] + LIST = [PYRAMID, FLAT, OVERLAP_LINEAR, DELAYED_REVERSE_SAWTOOTH, PYRAMID_SIGMA, PYRAMID_SIGMA_INV, GAUSS_SIGMA, GAUSS_SIGMA_INV, RANDOM] + LIST_STATIC = [PYRAMID, RELATIVE, FLAT, OVERLAP_LINEAR, DELAYED_REVERSE_SAWTOOTH, PYRAMID_SIGMA, PYRAMID_SIGMA_INV, GAUSS_SIGMA, GAUSS_SIGMA_INV, RANDOM] class ContextType: @@ -360,11 +362,11 @@ CONTEXT_MAPPING = { } -def get_context_weights(num_frames: int, fuse_method: str, sigma: Tensor = None): - weights_func = FUSE_MAPPING.get(fuse_method, None) +def get_context_weights(length: int, full_length: int, idxs: list[int], ctx_opts: ContextOptions, sigma: Tensor=None): + weights_func = FUSE_MAPPING.get(ctx_opts.fuse_method, None) if not weights_func: - raise ValueError(f"Unknown fuse_method '{fuse_method}'.") - return weights_func(num_frames, sigma=sigma ) + raise ValueError(f"Unknown fuse_method '{ctx_opts.fuse_method}'.") + return weights_func(length, sigma=sigma, ctx_opts=ctx_opts, full_length=full_length, idxs=idxs) def create_weights_flat(length: int, **kwargs) -> list[float]: @@ -382,6 +384,20 @@ def create_weights_pyramid(length: int, **kwargs) -> list[float]: weight_sequence = list(range(1, max_weight, 1)) + [max_weight] + list(range(max_weight - 1, 0, -1)) return weight_sequence +def create_weights_overlap_linear(length: int, full_length: int, idxs: list[int], ctx_opts: ContextOptions, **kwargs): + # based on code in Kijai's WanVideoWrapper: https://github.com/kijai/ComfyUI-WanVideoWrapper/blob/dbb2523b37e4ccdf45127e5ae33e31362f755c8e/nodes.py#L1302 + # only expected overlap is given different weights + weights_torch = torch.ones((length)) + # blend left-side on all except first window + if min(idxs) > 0: + ramp_up = torch.linspace(1e-37, 1, ctx_opts.context_overlap) + weights_torch[:ctx_opts.context_overlap] = ramp_up + # blend right-side on all except last window + if max(idxs) < full_length-1: + ramp_down = torch.linspace(1, 1e-37, ctx_opts.context_overlap) + weights_torch[-ctx_opts.context_overlap:] = ramp_down + return weights_torch + def create_weights_random(length: int, **kwargs) -> list[float]: if length % 2 == 0: max_weight = length // 2 @@ -460,6 +476,8 @@ FUSE_MAPPING = { ContextFuseMethod.FLAT: create_weights_flat, ContextFuseMethod.PYRAMID: create_weights_pyramid, ContextFuseMethod.RELATIVE: create_weights_pyramid, + ContextFuseMethod.OVERLAP_LINEAR: create_weights_overlap_linear, + # experimental ContextFuseMethod.GAUSS_SIGMA: create_weights_gauss_sigma, ContextFuseMethod.GAUSS_SIGMA_DEPR: create_weights_gauss_sigma, ContextFuseMethod.GAUSS_SIGMA_INV: create_weights_gauss_sigma_inv, diff --git a/animatediff/motion_module_ad.py b/animatediff/motion_module_ad.py index 38c0323..c0ec964 100644 --- a/animatediff/motion_module_ad.py +++ b/animatediff/motion_module_ad.py @@ -1387,7 +1387,7 @@ class TemporalTransformerBlock(nn.Module): count += 1 sub_hidden_states = rearrange(sub_hidden_states, "(b f) d c -> b f d c", f=len(sub_idxs)) - weights = get_context_weights(len(sub_idxs), view_options.fuse_method) * batched_conds + weights = get_context_weights(len(sub_idxs), video_length, sub_idxs, view_options, sigma=transformer_options["sigmas"]) * batched_conds weights_tensor = torch.Tensor(weights).to(device=hidden_states.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) value_final[:, sub_idxs] += sub_hidden_states * weights_tensor count_final[:, sub_idxs] += weights_tensor diff --git a/animatediff/sampling.py b/animatediff/sampling.py index d3e6c98..ffabdf0 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -783,9 +783,6 @@ def sliding_calc_cond_batch(executor: Callable, model, conds: list[list[dict]], multigpu_windows = {} start_idx = 0 for device, work in ctxs_relative_work.items(): - # if device == x_in.device: - # continue - # multigpu_windows[device] = enumerated_context_windows if work == 0: continue end_idx = start_idx + work @@ -817,14 +814,14 @@ def sliding_calc_cond_batch(executor: Callable, model, conds: list[list[dict]], for results in combined_results: for result in results: - combine_context_window_results(x_in, result.sub_conds_out, result.sub_conds, result.ctx_idxs, result.window_idx, timestep, + combine_context_window_results(x_in, result.sub_conds_out, result.sub_conds, result.ctx_idxs, result.window_idx, len(enumerated_context_windows), timestep, ADGS, NAIVE, CREF, conds_final, counts_final, biases_final) else: for enum_window in enumerated_context_windows: results = evaluate_context_windows(executor, model, x_in, conds, timestep, [enum_window], model_options, CREF, ADGS) for result in results: - combine_context_window_results(x_in, result.sub_conds_out, result.sub_conds, result.ctx_idxs, result.window_idx, timestep, + combine_context_window_results(x_in, result.sub_conds_out, result.sub_conds, result.ctx_idxs, result.window_idx, len(enumerated_context_windows), timestep, ADGS, NAIVE, CREF, conds_final, counts_final, biases_final) finally: CREF.cleanup(model_options) @@ -834,7 +831,7 @@ def sliding_calc_cond_batch(executor: Callable, model, conds: list[list[dict]], # finalize conds if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE: - # already normalized, so return as is + # relative is already normalized, so return as is del counts_final return conds_final else: @@ -898,7 +895,7 @@ def evaluate_context_windows(executor, model: BaseModel, x_in: Tensor, conds, ti return results -def combine_context_window_results(x_in: Tensor, sub_conds_out, sub_conds, ctx_idxs: list[int], window_idx: int, timestep, +def combine_context_window_results(x_in: Tensor, sub_conds_out, sub_conds, ctx_idxs: list[int], window_idx: int, total_windows: int, timestep, ADGS: AnimateDiffGlobalState, NAIVE: NaiveReuseHandler, CREF: ContextRefHandler, conds_final: list[Tensor], counts_final: list[Tensor], biases_final: list[Tensor]): if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE: @@ -915,7 +912,7 @@ def combine_context_window_results(x_in: Tensor, sub_conds_out, sub_conds, ctx_i biases_final[i][idx] = bias_total + bias else: # add conds and counts based on weights of fuse method - weights = get_context_weights(len(ctx_idxs), ADGS.params.context_options.fuse_method, sigma=timestep) + weights = get_context_weights(len(ctx_idxs), x_in.shape[0], ctx_idxs, ADGS.params.context_options, sigma=timestep) weights_tensor = torch.Tensor(weights).to(device=x_in.device).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) for i in range(len(sub_conds_out)): conds_final[i][ctx_idxs] += sub_conds_out[i] * weights_tensor