diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py index c4de82795..0a77fa097 100644 --- a/comfy/ldm/flux/model.py +++ b/comfy/ldm/flux/model.py @@ -6,6 +6,7 @@ import torch from torch import Tensor, nn from einops import rearrange, repeat import comfy.ldm.common_dit +import comfy.patcher_extension from .layers import ( DoubleStreamBlock, @@ -214,6 +215,13 @@ class Flux(nn.Module): return img, repeat(img_ids, "h w c -> b (h w) c", b=bs) def forward(self, x, timestep, context, y=None, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs): + return comfy.patcher_extension.WrapperExecutor.new_class_executor( + self._forward, + self, + comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options) + ).execute(x, timestep, context, y, guidance, ref_latents, control, transformer_options, **kwargs) + + def _forward(self, x, timestep, context, y=None, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs): bs, c, h_orig, w_orig = x.shape patch_size = self.patch_size diff --git a/comfy_extras/nodes_easycache.py b/comfy_extras/nodes_easycache.py new file mode 100644 index 000000000..002050221 --- /dev/null +++ b/comfy_extras/nodes_easycache.py @@ -0,0 +1,116 @@ +from comfy_api.latest import io, ComfyExtension +import comfy.patcher_extension +import logging +import torch +import comfy.model_patcher + +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"]) + return executor(*args, **kwargs) + finally: + 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 EasyCacheStore: + def __init__(self, dict: dict): + self.dict = dict + + def clone(self): + return EasyCacheStore(self.dict.copy()) + + +class EasyCacheNode(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="EasyCache", + display_name="Easy Cache", + description="Easy Cache", + 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("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."), + ], + outputs=[ + io.Model.Output(tooltip="The model with EasyCache."), + ], + ) + + @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.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) + +class EasyCacheExtension(ComfyExtension): + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + EasyCacheNode, + ] + +def comfy_entrypoint(): + return EasyCacheExtension() diff --git a/nodes.py b/nodes.py index 860a236aa..bcabd2d4d 100644 --- a/nodes.py +++ b/nodes.py @@ -2321,6 +2321,7 @@ async def init_builtin_extra_nodes(): "nodes_edit_model.py", "nodes_tcfg.py", "nodes_context_windows.py", + "nodes_easycache.py", ] import_failed = []