From a5bd3c86c8ed6b83c55c2d0e7a59515b15a0137f Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Wed, 19 Mar 2025 10:13:08 +0200 Subject: [PATCH] Make TeaCache node error out instead of silently fail if ComfyUI isn't new enough to pass transformer_options --- nodes/model_optimization_nodes.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/nodes/model_optimization_nodes.py b/nodes/model_optimization_nodes.py index 5dfaf95..ac7c296 100644 --- a/nodes/model_optimization_nodes.py +++ b/nodes/model_optimization_nodes.py @@ -785,6 +785,9 @@ def teacache_wanvideo_forward_orig(self, x, t, context, clip_fea=None, freqs=Non cache['teacache_skipped_steps'] += 1 #print(f"TeaCache: Skipping {suffix} step") return should_calc, cache + + if not transformer_options: + raise RuntimeError("Can't access transformer_options, this requires ComfyUI nightly version from Mar 14, 2025 or later") teacache_enabled = transformer_options.get("teacache_enabled", False) if not teacache_enabled: @@ -938,16 +941,23 @@ Official recommended values https://github.com/ali-vilab/TeaCache/tree/main/TeaC with context: out = model_function(input, timestep, **c) if current_step_index+1 == last_step and hasattr(diffusion_model, "teacache_state"): - if cond_or_uncond[0] == 0: + if len(cond_or_uncond) == 1 and cond_or_uncond[0] == 0: skipped_steps_cond = diffusion_model.teacache_state["cond"]["teacache_skipped_steps"] skipped_steps_uncond = diffusion_model.teacache_state["uncond"]["teacache_skipped_steps"] - logging.info("-----------------------------------") logging.info(f"TeaCache skipped:") logging.info(f"{skipped_steps_cond} cond steps") logging.info(f"{skipped_steps_uncond} uncond step") logging.info(f"out of {last_step} steps") logging.info("-----------------------------------") + elif len(cond_or_uncond) == 2: + skipped_steps_cond = diffusion_model.teacache_state["uncond"]["teacache_skipped_steps"] + logging.info("-----------------------------------") + logging.info(f"TeaCache skipped:") + logging.info(f"{skipped_steps_cond} cond steps") + logging.info(f"out of {last_step} steps") + logging.info("-----------------------------------") + return out return unet_wrapper_function