| from comfy import model_management as mm |
|
|
| class WanVideoTeaCache: |
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "rel_l1_thresh": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.001, |
| "tooltip": "Higher values will make TeaCache more aggressive, faster, but may cause artifacts. Good value range for 1.3B: 0.05 - 0.08, for other models 0.15-0.30"}), |
| "start_step": ("INT", {"default": 1, "min": 0, "max": 9999, "step": 1, "tooltip": "Start percentage of the steps to apply TeaCache"}), |
| "end_step": ("INT", {"default": -1, "min": -1, "max": 9999, "step": 1, "tooltip": "End steps to apply TeaCache"}), |
| "cache_device": (["main_device", "offload_device"], {"default": "offload_device", "tooltip": "Device to cache to"}), |
| "use_coefficients": ("BOOLEAN", {"default": True, "tooltip": "Use calculated coefficients for more accuracy. When enabled therel_l1_thresh should be about 10 times higher than without"}), |
| }, |
| "optional": { |
| "mode": (["e", "e0"], {"default": "e", "tooltip": "Choice between using e (time embeds, default) or e0 (modulated time embeds)"}), |
| }, |
| } |
| RETURN_TYPES = ("CACHEARGS",) |
| RETURN_NAMES = ("cache_args",) |
| FUNCTION = "process" |
| CATEGORY = "WanVideoWrapper" |
| DESCRIPTION = """ |
| Patch WanVideo model to use TeaCache. Speeds up inference by caching the output and |
| applying it instead of doing the step. Best results are achieved by choosing the |
| appropriate coefficients for the model. Early steps should never be skipped, with too |
| aggressive values this can happen and the motion suffers. Starting later can help with that too. |
| When NOT using coefficients, the threshold value should be |
| about 10 times smaller than the value used with coefficients. |
| |
| Official recommended values https://github.com/ali-vilab/TeaCache/tree/main/TeaCache4Wan2.1 |
| """ |
|
|
| def process(self, rel_l1_thresh, start_step, end_step, cache_device, use_coefficients, mode="e"): |
| if cache_device == "main_device": |
| cache_device = mm.get_torch_device() |
| else: |
| cache_device = mm.unet_offload_device() |
| cache_args = { |
| "cache_type": "TeaCache", |
| "rel_l1_thresh": rel_l1_thresh, |
| "start_step": start_step, |
| "end_step": end_step, |
| "cache_device": cache_device, |
| "use_coefficients": use_coefficients, |
| "mode": mode, |
| } |
| return (cache_args,) |
| |
| class WanVideoMagCache: |
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "magcache_thresh": ("FLOAT", {"default": 0.02, "min": 0.0, "max": 0.3, "step": 0.001, "tooltip": "How strongly to cache the output of diffusion model. This value must be non-negative."}), |
| "magcache_K": ("INT", {"default": 4, "min": 0, "max": 6, "step": 1, "tooltip": "The maxium skip steps of MagCache."}), |
| "start_step": ("INT", {"default": 1, "min": 0, "max": 9999, "step": 1, "tooltip": "Step to start applying MagCache"}), |
| "end_step": ("INT", {"default": -1, "min": -1, "max": 9999, "step": 1, "tooltip": "Step to end applying MagCache"}), |
| "cache_device": (["main_device", "offload_device"], {"default": "offload_device", "tooltip": "Device to cache to"}), |
| }, |
| } |
| RETURN_TYPES = ("CACHEARGS",) |
| RETURN_NAMES = ("cache_args",) |
| FUNCTION = "setargs" |
| CATEGORY = "WanVideoWrapper" |
| EXPERIMENTAL = True |
| DESCRIPTION = "MagCache for WanVideoWrapper, source https://github.com/Zehong-Ma/MagCache" |
|
|
| def setargs(self, magcache_thresh, magcache_K, start_step, end_step, cache_device): |
| if cache_device == "main_device": |
| cache_device = mm.get_torch_device() |
| else: |
| cache_device = mm.unet_offload_device() |
|
|
| cache_args = { |
| "cache_type": "MagCache", |
| "magcache_thresh": magcache_thresh, |
| "magcache_K": magcache_K, |
| "start_step": start_step, |
| "end_step": end_step, |
| "cache_device": cache_device, |
| } |
| return (cache_args,) |
| |
| class WanVideoEasyCache: |
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "easycache_thresh": ("FLOAT", {"default": 0.015, "min": 0.0, "max": 1.0, "step": 0.001, "tooltip": "How strongly to cache the output of diffusion model. This value must be non-negative."}), |
| "start_step": ("INT", {"default": 10, "min": 0, "max": 9999, "step": 1, "tooltip": "Step to start applying EasyCache"}), |
| "end_step": ("INT", {"default": -1, "min": -1, "max": 9999, "step": 1, "tooltip": "Step to end applying EasyCache"}), |
| "cache_device": (["main_device", "offload_device"], {"default": "offload_device", "tooltip": "Device to cache to"}), |
| }, |
| } |
| RETURN_TYPES = ("CACHEARGS",) |
| RETURN_NAMES = ("cache_args",) |
| FUNCTION = "setargs" |
| CATEGORY = "WanVideoWrapper" |
| EXPERIMENTAL = True |
| DESCRIPTION = "EasyCache for WanVideoWrapper, source https://github.com/H-EmbodVis/EasyCache" |
|
|
| def setargs(self, easycache_thresh, start_step, end_step, cache_device): |
| if cache_device == "main_device": |
| cache_device = mm.get_torch_device() |
| else: |
| cache_device = mm.unet_offload_device() |
|
|
| cache_args = { |
| "cache_type": "EasyCache", |
| "easycache_thresh": easycache_thresh, |
| "start_step": start_step, |
| "end_step": end_step, |
| "cache_device": cache_device, |
| } |
| return (cache_args,) |
|
|
| |
| NODE_CLASS_MAPPINGS = { |
| "WanVideoTeaCache": WanVideoTeaCache, |
| "WanVideoMagCache": WanVideoMagCache, |
| "WanVideoEasyCache": WanVideoEasyCache, |
| } |
| NODE_DISPLAY_NAME_MAPPINGS = { |
| "WanVideoTeaCache": "WanVideo TeaCache", |
| "WanVideoMagCache": "WanVideo MagCache", |
| "WanVideoEasyCache": "WanVideo EasyCache" |
| } |
