mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-15 20:43:31 +08:00
Fixed math to make threshold work as expected, refactored code to use EasyCacheHolder instead of a dict wrapped by object
This commit is contained in:
Jedrzej Kosinski
parent
26d54b1de5
commit
6ac21a396c
@ -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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user