From 6ac21a396c2f1acc69905745907278c2d7e57812 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 19 Aug 2025 17:01:09 -0700 Subject: [PATCH] Fixed math to make threshold work as expected, refactored code to use EasyCacheHolder instead of a dict wrapped by object --- comfy_extras/nodes_easycache.py | 180 +++++++++++++++++++++----------- 1 file changed, 118 insertions(+), 62 deletions(-) diff --git a/comfy_extras/nodes_easycache.py b/comfy_extras/nodes_easycache.py index 002050221..a0ae3e7b2 100644 --- a/comfy_extras/nodes_easycache.py +++ b/comfy_extras/nodes_easycache.py @@ -4,73 +4,135 @@ import logging import torch import comfy.model_patcher +def easycache_forward_wrapper(executor, *args, **kwargs): + # get values from args + x: torch.Tensor = args[0] + timestep: torch.Tensor = args[1] + transformer_options: dict[str] = args[-1] + # x: torch.Tensor = args[0] + # timestep: torch.Tensor = args[4] + # transformer_options: dict[str] = args[-2] + easycache: EasyCacheHolder = transformer_options["easycache"] + if easycache.is_past_end_timestep(timestep): + return executor(*args, **kwargs) + # prepare next x_prev + next_x_prev = x.clone() + do_easycache = easycache.should_do_easycache(timestep) + logging.info(f"easycache_wrapper: do_easycache: {do_easycache}") + output_prev_norm = None + input_change = None + if do_easycache: + if easycache.has_x_prev(): + input_change = (x - easycache.x_prev).flatten().abs().mean() + if easycache.has_output_prev() and easycache.has_relative_transformation_rate(): + output_prev_norm = easycache.output_prev.flatten().abs().mean() + approx_output_change_rate = (easycache.relative_transformation_rate * input_change) / output_prev_norm + easycache.cumulative_change_rate += approx_output_change_rate + if easycache.cumulative_change_rate < easycache.reuse_threshold: + logging.info(f"easycache_wrapper: skipping step; cumulative_change_rate: {easycache.cumulative_change_rate}, reuse_threshold: {easycache.reuse_threshold}") + return x + easycache.cache_diff + else: + easycache.cumulative_change_rate = 0.0 + logging.info(f"easycache_wrapper: NOT skipping step; cumulative_change_rate: {easycache.cumulative_change_rate}, reuse_threshold: {easycache.reuse_threshold}") + logging.info(f"easycache_wrapper: approx_output_change_rate: {approx_output_change_rate}") + + output: torch.Tensor = executor(*args, **kwargs) + if easycache.has_output_prev(): + output_change = (output - easycache.output_prev).flatten().abs().mean() + if output_prev_norm is None: + output_prev_norm = easycache.output_prev.flatten().abs().mean() + output_change_rate = output_change / output_prev_norm + easycache.output_change_rates.append(output_change_rate.item()) + if easycache.has_relative_transformation_rate(): + approx_output_change_rate = (easycache.relative_transformation_rate * input_change) / output_prev_norm + easycache.approx_output_change_rates.append(approx_output_change_rate.item()) + logging.info(f"easycache_wrapper: approx_output_change_rate: {approx_output_change_rate}") + if input_change is not None: + easycache.relative_transformation_rate = output_change / input_change + logging.info(f"easycache_wrapper: output_change_rate: {output_change_rate}") + easycache.cache_diff = output - next_x_prev + easycache.x_prev = next_x_prev + easycache.output_prev = output.clone() + return output + def easycache_sample_wrapper(executor, *args, **kwargs): try: guider = executor.class_obj orig_model_options = guider.model_options guider.model_options = comfy.model_patcher.create_model_options_clone(orig_model_options) - if "easycache" in orig_model_options["transformer_options"]: - guider.model_options["transformer_options"]["easycache"] = guider.model_options["transformer_options"]["easycache"].clone() - guider.model_options["transformer_options"]["easycache"].dict["start_timestep"] = guider.model_patcher.model.model_sampling.percent_to_sigma(guider.model_options["transformer_options"]["easycache"].dict["start_percent"]) - guider.model_options["transformer_options"]["easycache"].dict["end_timestep"] = guider.model_patcher.model.model_sampling.percent_to_sigma(guider.model_options["transformer_options"]["easycache"].dict["end_percent"]) + # clone and prepare timesteps + guider.model_options["transformer_options"]["easycache"] = guider.model_options["transformer_options"]["easycache"].clone().prepare_timesteps(guider.model_patcher.model.model_sampling) return executor(*args, **kwargs) finally: + output_change_rates = guider.model_options['transformer_options']['easycache'].output_change_rates + approx_output_change_rates = guider.model_options['transformer_options']['easycache'].approx_output_change_rates + logging.info(f"easycache_sample_wrapper: output_change_rates {len(output_change_rates)}: {output_change_rates}") + logging.info(f"easycache_sample_wrapper: approx_output_change_rates {len(approx_output_change_rates)}: {approx_output_change_rates}") + guider.model_options["transformer_options"]["easycache"].reset() guider.model_options = orig_model_options -def easycache_forward_wrapper(executor, *args, **kwargs): - x: torch.Tensor = args[0] - timestep: torch.Tensor = args[1] - transformer_options = args[-1] - do_easycache = timestep < transformer_options["easycache"].dict["start_timestep"] and timestep > transformer_options["easycache"].dict["end_timestep"] - logging.info(f"easycache_wrapper: do_easycache: {do_easycache}") - x_prev = None - input_change = None - # input_data = x.flatten().abs().mean() - if do_easycache and "easycache" in transformer_options: - if "x_prev" in transformer_options["easycache"].dict: - x_prev = transformer_options["easycache"].dict["x_prev"] - else: - transformer_options["easycache"].dict["x_prev"] = x.clone() - if x_prev is not None: - input_change = (x_prev - x).flatten().abs().mean() - if do_easycache and transformer_options["easycache"].dict.get("change_rate", None) is not None: - change_rate = transformer_options["easycache"].dict["change_rate"] - output_prev = transformer_options["easycache"].dict["output_prev"] - pred_change = change_rate * (input_change / output_prev.flatten().abs().mean()) - accumulated_change = transformer_options["easycache"].dict["accumulated_change"] + pred_change - if transformer_options["easycache"].dict["reuse_threshold"] <= accumulated_change: - logging.info(f"easycache_wrapper: skipping step; accumulated_change: {accumulated_change}, reuse_threshold: {transformer_options['easycache'].dict['reuse_threshold']}") - transformer_options["easycache"].dict["accumulated_change"] = 0.0 - return x + transformer_options["easycache"].dict["cache_diff"] - else: - transformer_options["easycache"].dict["accumulated_change"] = accumulated_change - logging.info(f"easycache_wrapper: NOT skipping step; accumulated_change: {accumulated_change}, reuse_threshold: {transformer_options['easycache'].dict['reuse_threshold']}") - logging.info(f"easycache_wrapper pred_change: {pred_change}") - output: torch.Tensor = executor(*args, **kwargs) - if x_prev is not None: - # output_data = output.flatten().abs().mean() - output_prev = transformer_options["easycache"].dict["output_prev"] - output_change = (output_prev - output).flatten().abs().mean() - k = output_change / input_change - transformer_options["easycache"].dict["change_rate"] = k - logging.info(f"easycache_wrapper: {input_change} {output_change} {k}") - if do_easycache and "easycache" in transformer_options: - transformer_options["easycache"].dict["output_prev"] = output.clone() - transformer_options["easycache"].dict["cache_diff"] = output - x - if not do_easycache: - transformer_options["easycache"].dict["accumulated_change"] = 0.0 - transformer_options["easycache"].dict["change_rate"] = None - transformer_options["easycache"].dict["output_prev"] = None - transformer_options["easycache"].dict["cache_diff"] = None - return output +class EasyCacheHolder: + def __init__(self, reuse_threshold: float, start_percent: float, end_percent: float): + self.reuse_threshold = reuse_threshold + self.start_percent = start_percent + self.end_percent = end_percent + # timestep values + self.start_t = 0.0 + self.end_t = 0.0 + # control values + self.relative_transformation_rate: float = None + self.cumulative_change_rate = 0.0 + # cache values + self.x_prev = None + self.output_prev = None + self.cache_diff = None + self.output_change_rates = [] + self.approx_output_change_rates = [] -class EasyCacheStore: - def __init__(self, dict: dict): - self.dict = dict + def is_past_end_timestep(self, timestep: float) -> bool: + return not (timestep > self.end_t).item() + + def should_do_easycache(self, timestep: float) -> bool: + return (timestep <= self.start_t).item() + + def has_x_prev(self) -> bool: + return self.x_prev is not None + + def has_output_prev(self) -> bool: + return self.output_prev is not None + + def has_cache_diff(self) -> bool: + return self.cache_diff is not None + + def has_relative_transformation_rate(self) -> bool: + return self.relative_transformation_rate is not None + + def prepare_timesteps(self, model_sampling): + self.start_t = model_sampling.percent_to_sigma(self.start_percent) + self.end_t = model_sampling.percent_to_sigma(self.end_percent) + return self + + def apply_cache(self): + ... + + def accumulate_change(self): + ... + + def reset(self): + self.relative_transformation_rate = 0.0 + self.cumulative_change_rate = 0.0 + self.output_change_rates = [] + del self.x_prev + self.x_prev = None + del self.output_prev + self.output_prev = None + del self.cache_diff + self.cache_diff = None + return self def clone(self): - return EasyCacheStore(self.dict.copy()) + return EasyCacheHolder(self.reuse_threshold, self.start_percent, self.end_percent) class EasyCacheNode(io.ComfyNode): @@ -83,7 +145,7 @@ class EasyCacheNode(io.ComfyNode): category="advanced/debug/model", inputs=[ io.Model.Input("model", tooltip="The model to add EasyCache to."), - io.Float.Input("reuse_threshold", min=0.0, default=0.0, max=100.0, step=0.01, tooltip="The threshold for reusing cached steps."), + io.Float.Input("reuse_threshold", min=0.0, default=0.0, max=1.0, step=0.01, tooltip="The threshold for reusing cached steps."), io.Float.Input("start_percent", min=0.0, default=0.0, max=1.0, step=0.01, tooltip="The relative sampling step to begin use of EasyCache."), io.Float.Input("end_percent", min=0.0, default=1.0, max=1.0, step=0.01, tooltip="The relative sampling step to end use of EasyCache."), ], @@ -95,13 +157,7 @@ class EasyCacheNode(io.ComfyNode): @classmethod def execute(cls, model: io.Model.Type, reuse_threshold: float, start_percent: float, end_percent: float) -> io.NodeOutput: model = model.clone() - easycache_dict = { - "reuse_threshold": reuse_threshold, - "start_percent": start_percent, - "end_percent": end_percent, - "accumulated_change": 0.0, - } - model.model_options["transformer_options"]["easycache"] = EasyCacheStore(easycache_dict) + model.model_options["transformer_options"]["easycache"] = EasyCacheHolder(reuse_threshold, start_percent, end_percent) model.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, "easycache", easycache_forward_wrapper) model.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, "easycache", easycache_sample_wrapper) return io.NodeOutput(model)