diff --git a/docs/source/en/api/cache.md b/docs/source/en/api/cache.md index a5ed8751118d..bf6acb801662 100644 --- a/docs/source/en/api/cache.md +++ b/docs/source/en/api/cache.md @@ -41,6 +41,12 @@ Cache methods speedup diffusion transformers by storing and reusing intermediate [[autodoc]] apply_taylorseer_cache +## ResilPhaseCacheConfig + +[[autodoc]] ResilPhaseCacheConfig + +[[autodoc]] apply_resilphase_cache + ## MagCacheConfig [[autodoc]] MagCacheConfig diff --git a/docs/source/en/optimization/cache.md b/docs/source/en/optimization/cache.md index 04f23758b88b..e9df2800c7e5 100644 --- a/docs/source/en/optimization/cache.md +++ b/docs/source/en/optimization/cache.md @@ -112,6 +112,43 @@ config = TaylorSeerCacheConfig( pipe.transformer.enable_cache(config) ``` +## ResilPhase Cache + +[ResilPhase](https://github.com/zqc214/ResilPhase) accelerates diffusion inference by mapping denoising steps to a +normalized phase axis and interpolating the residual produced by the transformer's block stack. On intermediate +steps, the predicted residual is added to the stack input while the transformer blocks are skipped. The method is +supported for FLUX and HunyuanVideo transformers. + +FLUX ControlNet inference automatically falls back to full transformer block computation because ControlNet residuals +are injected between individual blocks. + +Create a [`ResilPhaseCacheConfig`] and pass it to the pipeline transformer. `cache_interval` controls the interval +between full block-stack computations, `warmup_steps` initializes the interpolation history, and `max_order` controls +the number of historical residuals used by the interpolation. + +```python +import torch + +from diffusers import FluxPipeline, ResilPhaseCacheConfig + + +pipe = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + dtype=torch.bfloat16, +).to("cuda") + +config = ResilPhaseCacheConfig( + cache_interval=6, + warmup_steps=3, + max_order=1, + mapping_method="balanced", + balance_alpha=0.55, +) +pipe.transformer.enable_cache(config) + +image = pipe("A cat playing chess", num_inference_steps=50).images[0] +``` + ## MagCache [MagCache](https://github.com/Zehong-Ma/MagCache) accelerates inference by skipping transformer blocks based on the magnitude of the residual update. It observes that the magnitude of updates (Output - Input) decays predictably over the diffusion process. By accumulating an "error budget" based on pre-computed magnitude ratios, it dynamically decides when to skip computation and reuse the previous residual. diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2deff4b71cb2..7e019f615bd9 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -203,6 +203,7 @@ "LayerSkipConfig", "MagCacheConfig", "PyramidAttentionBroadcastConfig", + "ResilPhaseCacheConfig", "SmoothedEnergyGuidanceConfig", "TaylorSeerCacheConfig", "TextKVCacheConfig", @@ -211,6 +212,7 @@ "apply_layer_skip", "apply_mag_cache", "apply_pyramid_attention_broadcast", + "apply_resilphase_cache", "apply_taylorseer_cache", "apply_text_kv_cache", ] @@ -1052,6 +1054,7 @@ LayerSkipConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, + ResilPhaseCacheConfig, SmoothedEnergyGuidanceConfig, TaylorSeerCacheConfig, TextKVCacheConfig, @@ -1060,6 +1063,7 @@ apply_layer_skip, apply_mag_cache, apply_pyramid_attention_broadcast, + apply_resilphase_cache, apply_taylorseer_cache, apply_text_kv_cache, ) diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 2a9aa81608e7..dc894e6ed4c3 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -25,6 +25,7 @@ from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook from .mag_cache import MagCacheConfig, apply_mag_cache from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast + from .resilphase_cache import ResilPhaseCacheConfig, apply_resilphase_cache from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache from .text_kv_cache import TextKVCacheConfig, apply_text_kv_cache diff --git a/src/diffusers/hooks/resilphase_cache.py b/src/diffusers/hooks/resilphase_cache.py new file mode 100644 index 000000000000..2ccb4e0bf41f --- /dev/null +++ b/src/diffusers/hooks/resilphase_cache.py @@ -0,0 +1,375 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import math +from dataclasses import dataclass +from typing import Literal + +import torch + +from ..utils import get_logger +from ..utils.torch_utils import unwrap_module +from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS +from ._helpers import TransformerBlockMetadata, TransformerBlockRegistry +from .hooks import BaseState, HookRegistry, ModelHook, StateManager + + +logger = get_logger(__name__) # pylint: disable=invalid-name + +_RESILPHASE_LEADER_BLOCK_HOOK = "resilphase_leader_block_hook" +_RESILPHASE_BLOCK_HOOK = "resilphase_block_hook" +_RESILPHASE_DENOISER_HOOK = "resilphase_denoiser_hook" +_CONTROL_RESIDUAL_ARGUMENTS = ("controlnet_block_samples", "controlnet_single_block_samples") + + +@dataclass +class ResilPhaseCacheConfig: + r""" + Configuration for ResilPhase cache. + + ResilPhase approximates the residual produced by a denoiser's transformer block stack with barycentric Lagrange + interpolation on a normalized phase axis. On prediction steps, the expensive transformer blocks are skipped and the + predicted residual is added to the block stack input. + + Args: + cache_interval (`int`, defaults to `6`): + Number of denoising steps between full transformer block computations. The intermediate `cache_interval - + 1` steps use ResilPhase predictions. + warmup_steps (`int`, defaults to `3`): + Number of initial denoising steps that always perform full computations to initialize the interpolation + history. + max_order (`int`, defaults to `1`): + Maximum interpolation order. ResilPhase retains at most `max_order + 1` fully computed residuals. + mapping_method (`str`, defaults to `"balanced"`): + Phase-axis mapping. Must be `"balanced"` for a hyperbolic-tangent mapping or `"chebyshev"` for Chebyshev + nodes. + balance_alpha (`float`, defaults to `0.55`): + Scale of the hyperbolic-tangent phase mapping. Only used when `mapping_method="balanced"`. + """ + + cache_interval: int = 6 + warmup_steps: int = 3 + max_order: int = 1 + mapping_method: Literal["balanced", "chebyshev"] = "balanced" + balance_alpha: float = 0.55 + + def __post_init__(self) -> None: + if self.cache_interval < 1: + raise ValueError("`cache_interval` must be greater than zero.") + if self.warmup_steps < 0: + raise ValueError("`warmup_steps` must be greater than or equal to zero.") + if self.max_order < 0: + raise ValueError("`max_order` must be greater than or equal to zero.") + if self.mapping_method not in {"balanced", "chebyshev"}: + raise ValueError('`mapping_method` must be either "balanced" or "chebyshev".') + if self.balance_alpha <= 0: + raise ValueError("`balance_alpha` must be greater than zero.") + + +class ResilPhaseState(BaseState): + def __init__(self, config: ResilPhaseCacheConfig) -> None: + self.config = config + self.step_index = -1 + self.skipped_steps = 0 + self.should_compute = True + self.bypass = False + + self.stack_input: tuple[torch.Tensor, torch.Tensor | None] | None = None + self.history_steps: list[int] = [] + self.history_residuals: list[tuple[torch.Tensor, torch.Tensor | None]] = [] + + def reset(self) -> None: + self.step_index = -1 + self.skipped_steps = 0 + self.should_compute = True + self.bypass = False + self.stack_input = None + self.history_steps = [] + self.history_residuals = [] + + def start_step(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None) -> None: + self.step_index += 1 + has_history = len(self.history_residuals) > 0 + is_warmup = self.step_index < self.config.warmup_steps + reached_refresh = self.skipped_steps >= self.config.cache_interval - 1 + self.should_compute = is_warmup or not has_history or reached_refresh + + if self.should_compute: + self.skipped_steps = 0 + self.stack_input = (hidden_states, encoder_hidden_states) + else: + self.skipped_steps += 1 + self.stack_input = None + + def update(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None) -> None: + if self.stack_input is None: + raise ValueError("Cannot update ResilPhase state without a transformer block stack input.") + + input_hidden_states, input_encoder_hidden_states = self.stack_input + hidden_states_residual = (hidden_states - input_hidden_states).detach().clone() + encoder_hidden_states_residual = None + if encoder_hidden_states is not None: + encoder_hidden_states_residual = (encoder_hidden_states - input_encoder_hidden_states).detach().clone() + + self.history_steps.append(self.step_index) + self.history_residuals.append((hidden_states_residual, encoder_hidden_states_residual)) + history_size = self.config.max_order + 1 + self.history_steps = self.history_steps[-history_size:] + self.history_residuals = self.history_residuals[-history_size:] + self.stack_input = None + + @torch.compiler.disable + def predict(self) -> tuple[torch.Tensor, torch.Tensor | None]: + if len(self.history_residuals) == 1: + return self.history_residuals[0] + + if self.config.mapping_method == "chebyshev": + num_nodes = len(self.history_steps) + phase_nodes = [math.cos((2 * index + 1) * math.pi / (2 * num_nodes)) for index in range(num_nodes)] + phase_nodes.sort(reverse=True) + + left_step, right_step = self.history_steps[-2:] + left_node, right_node = phase_nodes[-2:] + slope = (right_node - left_node) / (right_step - left_step) + target_phase = right_node + slope * (self.step_index - right_step) + target_phase = min(1.0, target_phase) + else: + mean_step = sum(self.history_steps) / len(self.history_steps) + max_distance = max(abs(step - mean_step) for step in self.history_steps) + phase_nodes = [ + math.tanh(self.config.balance_alpha * (step - mean_step) / max_distance) for step in self.history_steps + ] + target_phase = math.tanh(self.config.balance_alpha * (self.step_index - mean_step) / max_distance) + + barycentric_weights = [] + for index, node in enumerate(phase_nodes): + log_weight = 0.0 + sign = 1 + for other_index, other_node in enumerate(phase_nodes): + if index != other_index: + difference = node - other_node + log_weight -= math.log(abs(difference)) + if difference < 0: + sign *= -1 + log_weight = max(-700, min(700, log_weight)) + barycentric_weights.append(sign * math.exp(log_weight)) + + total_absolute_weight = sum(abs(weight) for weight in barycentric_weights) + barycentric_weights = [ + weight / total_absolute_weight * len(barycentric_weights) for weight in barycentric_weights + ] + + hidden_states_numerator = None + encoder_hidden_states_numerator = None + denominator = 0.0 + for index, (hidden_states_residual, encoder_hidden_states_residual) in enumerate(self.history_residuals): + distance = target_phase - phase_nodes[index] + if abs(distance) < 1e-12: + return hidden_states_residual, encoder_hidden_states_residual + + coefficient = barycentric_weights[index] / distance + hidden_states_term = coefficient * hidden_states_residual + hidden_states_numerator = ( + hidden_states_term if hidden_states_numerator is None else hidden_states_numerator + hidden_states_term + ) + if encoder_hidden_states_residual is not None: + encoder_hidden_states_term = coefficient * encoder_hidden_states_residual + encoder_hidden_states_numerator = ( + encoder_hidden_states_term + if encoder_hidden_states_numerator is None + else encoder_hidden_states_numerator + encoder_hidden_states_term + ) + denominator += coefficient + + hidden_states_prediction = hidden_states_numerator / denominator + encoder_hidden_states_prediction = ( + None if encoder_hidden_states_numerator is None else encoder_hidden_states_numerator / denominator + ) + return hidden_states_prediction, encoder_hidden_states_prediction + + +def _get_block_inputs( + metadata: TransformerBlockMetadata, args: tuple, kwargs: dict +) -> tuple[torch.Tensor, torch.Tensor | None]: + hidden_states = metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs) + encoder_hidden_states = None + if metadata.return_encoder_hidden_states_index is not None: + encoder_hidden_states = metadata._get_parameter_from_args_kwargs("encoder_hidden_states", args, kwargs) + return hidden_states, encoder_hidden_states + + +def _get_block_outputs( + metadata: TransformerBlockMetadata, output: torch.Tensor | tuple[torch.Tensor, ...] +) -> tuple[torch.Tensor, torch.Tensor | None]: + if isinstance(output, tuple): + hidden_states = output[metadata.return_hidden_states_index] + encoder_hidden_states = ( + None + if metadata.return_encoder_hidden_states_index is None + else output[metadata.return_encoder_hidden_states_index] + ) + return hidden_states, encoder_hidden_states + return output, None + + +def _pack_block_outputs( + metadata: TransformerBlockMetadata, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, +) -> torch.Tensor | tuple[torch.Tensor, ...]: + if metadata.return_encoder_hidden_states_index is None: + return hidden_states + + output = [None, None] + output[metadata.return_hidden_states_index] = hidden_states + output[metadata.return_encoder_hidden_states_index] = encoder_hidden_states + return tuple(output) + + +class ResilPhaseHeadBlockHook(ModelHook): + _is_stateful = True + + def __init__(self, state_manager: StateManager) -> None: + super().__init__() + self.state_manager = state_manager + self._metadata = None + + def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: + self._metadata = TransformerBlockRegistry.get(unwrap_module(module).__class__) + return module + + def reset_state(self, module: torch.nn.Module) -> torch.nn.Module: + self.state_manager.reset() + return module + + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + state: ResilPhaseState = self.state_manager.get_state() + if state.bypass: + return self.fn_ref.original_forward(*args, **kwargs) + + hidden_states, encoder_hidden_states = _get_block_inputs(self._metadata, args, kwargs) + state.start_step(hidden_states, encoder_hidden_states) + + if state.should_compute: + return self.fn_ref.original_forward(*args, **kwargs) + + hidden_states_residual, encoder_hidden_states_residual = state.predict() + hidden_states = hidden_states + hidden_states_residual + if encoder_hidden_states is not None: + encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_residual + return _pack_block_outputs(self._metadata, hidden_states, encoder_hidden_states) + + +class ResilPhaseBlockHook(ModelHook): + def __init__(self, state_manager: StateManager, is_tail: bool = False) -> None: + super().__init__() + self.state_manager = state_manager + self.is_tail = is_tail + self._metadata = None + + def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: + self._metadata = TransformerBlockRegistry.get(unwrap_module(module).__class__) + return module + + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + state: ResilPhaseState = self.state_manager.get_state() + if state.bypass: + return self.fn_ref.original_forward(*args, **kwargs) + + if state.should_compute: + output = self.fn_ref.original_forward(*args, **kwargs) + if self.is_tail: + hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output) + state.update(hidden_states, encoder_hidden_states) + return output + + hidden_states, encoder_hidden_states = _get_block_inputs(self._metadata, args, kwargs) + return _pack_block_outputs(self._metadata, hidden_states, encoder_hidden_states) + + +class ResilPhaseDenoiserHook(ModelHook): + def __init__(self, state_manager: StateManager) -> None: + super().__init__() + self.state_manager = state_manager + self._control_argument_indices = {} + + def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: + parameters = list(inspect.signature(unwrap_module(module).__class__.forward).parameters)[1:] + self._control_argument_indices = { + name: parameters.index(name) for name in _CONTROL_RESIDUAL_ARGUMENTS if name in parameters + } + return module + + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + state: ResilPhaseState = self.state_manager.get_state() + state.bypass = any(kwargs.get(name) is not None for name in _CONTROL_RESIDUAL_ARGUMENTS) + if not state.bypass: + state.bypass = any( + index < len(args) and args[index] is not None for index in self._control_argument_indices.values() + ) + + try: + return self.fn_ref.original_forward(*args, **kwargs) + finally: + state.bypass = False + + +def apply_resilphase_cache(module: torch.nn.Module, config: ResilPhaseCacheConfig) -> None: + r"""Apply ResilPhase cache to the transformer blocks of a denoiser. + + Args: + module (`torch.nn.Module`): + The denoiser module whose transformer block stack should be cached. + config (`ResilPhaseCacheConfig`): + Configuration for ResilPhase cache. + + Example: + ```python + >>> from diffusers import FluxPipeline, ResilPhaseCacheConfig + + >>> pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev") + >>> pipe.transformer.enable_cache(ResilPhaseCacheConfig()) + ``` + """ + + transformer_blocks = [] + for name, submodule in module.named_children(): + if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList): + continue + for index, block in enumerate(submodule): + transformer_blocks.append((f"{name}.{index}", block)) + + if len(transformer_blocks) < 2: + raise ValueError("ResilPhase cache requires a denoiser with at least two transformer blocks.") + + state_manager = StateManager(ResilPhaseState, init_args=(config,)) + head_block_name, head_block = transformer_blocks.pop(0) + tail_block_name, tail_block = transformer_blocks.pop(-1) + + registry = HookRegistry.check_if_exists_or_initialize(module) + registry.register_hook(ResilPhaseDenoiserHook(state_manager), _RESILPHASE_DENOISER_HOOK) + + logger.debug(f"Applying ResilPhaseHeadBlockHook to '{head_block_name}'") + registry = HookRegistry.check_if_exists_or_initialize(head_block) + registry.register_hook(ResilPhaseHeadBlockHook(state_manager), _RESILPHASE_LEADER_BLOCK_HOOK) + + for name, block in transformer_blocks: + logger.debug(f"Applying ResilPhaseBlockHook to '{name}'") + registry = HookRegistry.check_if_exists_or_initialize(block) + registry.register_hook(ResilPhaseBlockHook(state_manager), _RESILPHASE_BLOCK_HOOK) + + logger.debug(f"Applying ResilPhaseBlockHook to tail block '{tail_block_name}'") + registry = HookRegistry.check_if_exists_or_initialize(tail_block) + registry.register_hook(ResilPhaseBlockHook(state_manager, is_tail=True), _RESILPHASE_BLOCK_HOOK) diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index 5aa189987ba2..eea3a9fe4e95 100644 --- a/src/diffusers/models/cache_utils.py +++ b/src/diffusers/models/cache_utils.py @@ -28,6 +28,7 @@ class CacheMixin: - [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) - [FasterCache](https://huggingface.co/papers/2410.19355) - [FirstBlockCache](https://github.com/chengzeyi/ParaAttention/blob/7a266123671b55e7e5a2fe9af3121f07a36afc78/README.md#first-block-cache-our-dynamic-caching) + - [ResilPhase](https://github.com/zqc214/ResilPhase) """ _cache_config = None @@ -41,11 +42,12 @@ def enable_cache(self, config) -> None: Enable caching techniques on the model. Args: - config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | TextKVCacheConfig`): + config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | ResilPhaseCacheConfig | TextKVCacheConfig`): The configuration for applying the caching technique. Currently supported caching techniques are: - [`~hooks.PyramidAttentionBroadcastConfig`] - [`~hooks.FasterCacheConfig`] - [`~hooks.FirstBlockCacheConfig`] + - [`~hooks.ResilPhaseCacheConfig`] - [`~hooks.TextKVCacheConfig`] Example: @@ -71,12 +73,14 @@ def enable_cache(self, config) -> None: FirstBlockCacheConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, + ResilPhaseCacheConfig, TaylorSeerCacheConfig, TextKVCacheConfig, apply_faster_cache, apply_first_block_cache, apply_mag_cache, apply_pyramid_attention_broadcast, + apply_resilphase_cache, apply_taylorseer_cache, apply_text_kv_cache, ) @@ -96,6 +100,8 @@ def enable_cache(self, config) -> None: apply_text_kv_cache(self, config) elif isinstance(config, PyramidAttentionBroadcastConfig): apply_pyramid_attention_broadcast(self, config) + elif isinstance(config, ResilPhaseCacheConfig): + apply_resilphase_cache(self, config) elif isinstance(config, TaylorSeerCacheConfig): apply_taylorseer_cache(self, config) else: @@ -110,6 +116,7 @@ def disable_cache(self) -> None: HookRegistry, MagCacheConfig, PyramidAttentionBroadcastConfig, + ResilPhaseCacheConfig, TaylorSeerCacheConfig, TextKVCacheConfig, ) @@ -117,6 +124,11 @@ def disable_cache(self) -> None: from ..hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK from ..hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK from ..hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK + from ..hooks.resilphase_cache import ( + _RESILPHASE_BLOCK_HOOK, + _RESILPHASE_DENOISER_HOOK, + _RESILPHASE_LEADER_BLOCK_HOOK, + ) from ..hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK from ..hooks.text_kv_cache import _TEXT_KV_CACHE_BLOCK_HOOK, _TEXT_KV_CACHE_TRANSFORMER_HOOK @@ -136,6 +148,10 @@ def disable_cache(self) -> None: registry.remove_hook(_MAG_CACHE_BLOCK_HOOK, recurse=True) elif isinstance(self._cache_config, PyramidAttentionBroadcastConfig): registry.remove_hook(_PYRAMID_ATTENTION_BROADCAST_HOOK, recurse=True) + elif isinstance(self._cache_config, ResilPhaseCacheConfig): + registry.remove_hook(_RESILPHASE_DENOISER_HOOK, recurse=True) + registry.remove_hook(_RESILPHASE_LEADER_BLOCK_HOOK, recurse=True) + registry.remove_hook(_RESILPHASE_BLOCK_HOOK, recurse=True) elif isinstance(self._cache_config, TextKVCacheConfig): registry.remove_hook(_TEXT_KV_CACHE_TRANSFORMER_HOOK, recurse=True) registry.remove_hook(_TEXT_KV_CACHE_BLOCK_HOOK, recurse=True) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index f7a8b6ec9de7..eee719bc6120 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -257,6 +257,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ResilPhaseCacheConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class SmoothedEnergyGuidanceConfig(metaclass=DummyObject): _backends = ["torch"] @@ -322,6 +337,10 @@ def apply_pyramid_attention_broadcast(*args, **kwargs): requires_backends(apply_pyramid_attention_broadcast, ["torch"]) +def apply_resilphase_cache(*args, **kwargs): + requires_backends(apply_resilphase_cache, ["torch"]) + + def apply_taylorseer_cache(*args, **kwargs): requires_backends(apply_taylorseer_cache, ["torch"]) diff --git a/tests/hooks/test_resilphase_cache.py b/tests/hooks/test_resilphase_cache.py new file mode 100644 index 000000000000..ffbbfa059fe4 --- /dev/null +++ b/tests/hooks/test_resilphase_cache.py @@ -0,0 +1,174 @@ +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import ResilPhaseCacheConfig, apply_resilphase_cache +from diffusers.hooks import HookRegistry +from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry +from diffusers.hooks.resilphase_cache import ResilPhaseState + + +class DummyBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor): + self.calls += 1 + return hidden_states + 1, encoder_hidden_states + 2 + + +class DummyTransformer(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([DummyBlock(), DummyBlock(), DummyBlock()]) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + controlnet_block_samples: torch.Tensor | None = None, + controlnet_single_block_samples: torch.Tensor | None = None, + ): + for block in self.transformer_blocks: + hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states) + if controlnet_block_samples is not None: + hidden_states = hidden_states + controlnet_block_samples + return hidden_states, encoder_hidden_states + + +@pytest.fixture(autouse=True) +def register_dummy_block(): + TransformerBlockRegistry.register( + DummyBlock, + TransformerBlockMetadata(return_hidden_states_index=0, return_encoder_hidden_states_index=1), + ) + + +@pytest.mark.parametrize( + ("kwargs", "error"), + [ + ({"cache_interval": 0}, "cache_interval"), + ({"warmup_steps": -1}, "warmup_steps"), + ({"max_order": -1}, "max_order"), + ({"mapping_method": "linear"}, "mapping_method"), + ({"balance_alpha": 0}, "balance_alpha"), + ], +) +def test_resilphase_config_validation(kwargs, error): + with pytest.raises(ValueError, match=error): + ResilPhaseCacheConfig(**kwargs) + + +@pytest.mark.parametrize( + ("mapping_method", "expected_hidden_states", "expected_encoder_hidden_states"), + [ + ("balanced", 4.670691622065688, 9.341383244131376), + ("chebyshev", 5.0, 10.0), + ], +) +def test_resilphase_barycentric_prediction(mapping_method, expected_hidden_states, expected_encoder_hidden_states): + state = ResilPhaseState(ResilPhaseCacheConfig(mapping_method=mapping_method)) + state.step_index = 4 + state.history_steps = [0, 3] + state.history_residuals = [ + (torch.tensor([1.0]), torch.tensor([2.0])), + (torch.tensor([4.0]), torch.tensor([8.0])), + ] + + hidden_states, encoder_hidden_states = state.predict() + + assert torch.allclose(hidden_states, torch.tensor([expected_hidden_states])) + assert torch.allclose(encoder_hidden_states, torch.tensor([expected_encoder_hidden_states])) + + +@pytest.mark.parametrize("mapping_method", ["balanced", "chebyshev"]) +def test_resilphase_skips_blocks_and_predicts_both_streams(mapping_method): + model = DummyTransformer() + config = ResilPhaseCacheConfig( + cache_interval=3, + warmup_steps=2, + max_order=1, + mapping_method=mapping_method, + ) + apply_resilphase_cache(model, config) + HookRegistry.check_if_exists_or_initialize(model)._set_context("cond") + + hidden_states = torch.tensor([[[0.0]]]) + encoder_hidden_states = torch.tensor([[[0.0]]]) + + output_0 = model(hidden_states, encoder_hidden_states) + output_1 = model(hidden_states + 1, encoder_hidden_states + 1) + output_2 = model(hidden_states + 2, encoder_hidden_states + 2) + + assert torch.allclose(output_0[0], torch.tensor([[[3.0]]])) + assert torch.allclose(output_0[1], torch.tensor([[[6.0]]])) + assert torch.allclose(output_1[0], torch.tensor([[[4.0]]])) + assert torch.allclose(output_1[1], torch.tensor([[[7.0]]])) + assert torch.allclose(output_2[0], torch.tensor([[[5.0]]])) + assert torch.allclose(output_2[1], torch.tensor([[[8.0]]])) + assert [block.calls for block in model.transformer_blocks] == [2, 2, 2] + + +def test_resilphase_contexts_keep_independent_state(): + model = DummyTransformer() + apply_resilphase_cache(model, ResilPhaseCacheConfig(cache_interval=3, warmup_steps=1, max_order=0)) + registry = HookRegistry.check_if_exists_or_initialize(model) + + hidden_states = torch.tensor([[[1.0]]]) + encoder_hidden_states = torch.tensor([[[1.0]]]) + + registry._set_context("cond") + cond_output = model(hidden_states, encoder_hidden_states) + registry._set_context("uncond") + uncond_output = model(hidden_states, encoder_hidden_states) + + assert torch.equal(cond_output[0], uncond_output[0]) + assert torch.equal(cond_output[1], uncond_output[1]) + assert [block.calls for block in model.transformer_blocks] == [2, 2, 2] + + +def test_resilphase_refreshes_after_cache_interval(): + model = DummyTransformer() + apply_resilphase_cache(model, ResilPhaseCacheConfig(cache_interval=3, warmup_steps=2, max_order=1)) + HookRegistry.check_if_exists_or_initialize(model)._set_context("cond") + + hidden_states = torch.tensor([[[0.0]]]) + encoder_hidden_states = torch.tensor([[[0.0]]]) + for _ in range(6): + model(hidden_states, encoder_hidden_states) + + assert [block.calls for block in model.transformer_blocks] == [3, 3, 3] + + +def test_resilphase_bypasses_cache_for_controlnet_residuals(): + model = DummyTransformer() + apply_resilphase_cache(model, ResilPhaseCacheConfig(cache_interval=3, warmup_steps=1, max_order=0)) + HookRegistry.check_if_exists_or_initialize(model)._set_context("cond") + + hidden_states = torch.tensor([[[0.0]]]) + encoder_hidden_states = torch.tensor([[[0.0]]]) + model(hidden_states, encoder_hidden_states) + model(hidden_states, encoder_hidden_states) + output = model( + hidden_states, + encoder_hidden_states, + controlnet_block_samples=torch.tensor([[[10.0]]]), + ) + + assert torch.equal(output[0], torch.tensor([[[33.0]]])) + assert torch.equal(output[1], torch.tensor([[[6.0]]])) + assert [block.calls for block in model.transformer_blocks] == [2, 2, 2] diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 23aa871b80a2..01dacbec3061 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -9,6 +9,8 @@ MagCacheTesterMixin, PyramidAttentionBroadcastConfigMixin, PyramidAttentionBroadcastTesterMixin, + ResilPhaseCacheConfigMixin, + ResilPhaseCacheTesterMixin, TaylorSeerCacheConfigMixin, TaylorSeerCacheTesterMixin, ) @@ -87,6 +89,8 @@ "NunchakuLiteTesterMixin", "PyramidAttentionBroadcastConfigMixin", "PyramidAttentionBroadcastTesterMixin", + "ResilPhaseCacheConfigMixin", + "ResilPhaseCacheTesterMixin", "TaylorSeerCacheConfigMixin", "TaylorSeerCacheTesterMixin", "QuantizationCompileTesterMixin", diff --git a/tests/models/testing_utils/cache.py b/tests/models/testing_utils/cache.py index 8357d34f3077..99c4349b3bd3 100644 --- a/tests/models/testing_utils/cache.py +++ b/tests/models/testing_utils/cache.py @@ -23,12 +23,18 @@ FirstBlockCacheConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, + ResilPhaseCacheConfig, TaylorSeerCacheConfig, ) from diffusers.hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK from diffusers.hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK from diffusers.hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK from diffusers.hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK +from diffusers.hooks.resilphase_cache import ( + _RESILPHASE_BLOCK_HOOK, + _RESILPHASE_DENOISER_HOOK, + _RESILPHASE_LEADER_BLOCK_HOOK, +) from diffusers.hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK from diffusers.models.cache_utils import CacheMixin @@ -631,6 +637,89 @@ def test_mag_cache_reset_stateful_cache(self): self._test_reset_stateful_cache() +@is_cache +class ResilPhaseCacheConfigMixin: + """Base mixin providing ResilPhase cache config.""" + + RESILPHASE_CACHE_CONFIG = { + "cache_interval": 3, + "warmup_steps": 1, + "max_order": 1, + } + + def _get_cache_config(self): + return ResilPhaseCacheConfig(**self.RESILPHASE_CACHE_CONFIG) + + def _get_hook_names(self): + return [_RESILPHASE_DENOISER_HOOK, _RESILPHASE_LEADER_BLOCK_HOOK, _RESILPHASE_BLOCK_HOOK] + + +@is_cache +class ResilPhaseCacheTesterMixin(ResilPhaseCacheConfigMixin, CacheTesterMixin): + """Mixin class for testing ResilPhase cache on models.""" + + @torch.no_grad() + def _test_cache_inference(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + + model.enable_cache(self._get_cache_config()) + with model.cache_context("resilphase_test"): + _ = model(**inputs_dict, return_dict=False)[0] + + inputs_dict_step2 = inputs_dict.copy() + if self.cache_input_key in inputs_dict_step2: + inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like( + inputs_dict_step2[self.cache_input_key] + ) + output_with_cache = model(**inputs_dict_step2, return_dict=False)[0] + + assert output_with_cache is not None + assert not torch.isnan(output_with_cache).any() + + model.disable_cache() + output_without_cache = model(**inputs_dict_step2, return_dict=False)[0] + assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5) + + @torch.no_grad() + def _test_reset_stateful_cache(self): + model = self.model_class(**self.get_init_dict()).to(torch_device) + model.eval() + model.enable_cache(self._get_cache_config()) + + with model.cache_context("resilphase_test"): + _ = model(**self.get_dummy_inputs(), return_dict=False)[0] + + model._reset_stateful_cache() + model.disable_cache() + + @require_cache_mixin + def test_resilphase_cache_enable_disable_state(self): + self._test_cache_enable_disable_state() + + @require_cache_mixin + def test_resilphase_cache_double_enable_raises_error(self): + self._test_cache_double_enable_raises_error() + + @require_cache_mixin + def test_resilphase_cache_hooks_registered(self): + self._test_cache_hooks_registered() + + @require_cache_mixin + def test_resilphase_cache_inference(self): + self._test_cache_inference() + + @require_cache_mixin + def test_resilphase_cache_context_manager(self): + self._test_cache_context_manager() + + @require_cache_mixin + def test_resilphase_cache_reset_stateful_cache(self): + self._test_reset_stateful_cache() + + @is_cache class TaylorSeerCacheConfigMixin: """ diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 719429526945..7b5f46d6c4d6 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -47,6 +47,7 @@ ModelTesterMixin, QuantoCompileTesterMixin, QuantoTesterMixin, + ResilPhaseCacheTesterMixin, SDNQCompileTesterMixin, SDNQTesterMixin, SingleFileTesterMixin, @@ -571,5 +572,9 @@ class TestFluxTransformerMagCache(FluxTransformerTesterConfig, MagCacheTesterMix """MagCache tests for Flux Transformer.""" +class TestFluxTransformerResilPhaseCache(FluxTransformerTesterConfig, ResilPhaseCacheTesterMixin): + """ResilPhase cache tests for Flux Transformer.""" + + class TestFluxTransformerTaylorSeerCache(FluxTransformerTesterConfig, TaylorSeerCacheTesterMixin): """TaylorSeerCache tests for Flux Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_hunyuan_video.py b/tests/models/transformers/test_models_transformer_hunyuan_video.py index 145ad4b42481..a35c823eaaee 100644 --- a/tests/models/transformers/test_models_transformer_hunyuan_video.py +++ b/tests/models/transformers/test_models_transformer_hunyuan_video.py @@ -22,6 +22,7 @@ BaseModelTesterConfig, BitsAndBytesTesterMixin, ModelTesterMixin, + ResilPhaseCacheTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, @@ -138,6 +139,10 @@ class TestHunyuanVideoTransformerCompile(HunyuanVideoTransformerTesterConfig, To pass +class TestHunyuanVideoTransformerResilPhaseCache(HunyuanVideoTransformerTesterConfig, ResilPhaseCacheTesterMixin): + """ResilPhase cache tests for HunyuanVideo Transformer.""" + + class TestHunyuanVideoTransformerBitsAndBytes(HunyuanVideoTransformerTesterConfig, BitsAndBytesTesterMixin): """BitsAndBytes quantization tests for HunyuanVideo Transformer.""" diff --git a/tests/pipelines/flux/test_pipeline_flux.py b/tests/pipelines/flux/test_pipeline_flux.py index 132a0f636181..4e5188595597 100644 --- a/tests/pipelines/flux/test_pipeline_flux.py +++ b/tests/pipelines/flux/test_pipeline_flux.py @@ -36,6 +36,7 @@ MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, + ResilPhaseCacheTesterMixin, TaylorSeerCacheTesterMixin, ) @@ -353,6 +354,10 @@ class TestFluxPipelineTaylorSeerCache(FluxPipelineTesterConfig, TaylorSeerCacheT """TaylorSeer cache tests for the Flux pipeline.""" +class TestFluxPipelineResilPhaseCache(FluxPipelineTesterConfig, ResilPhaseCacheTesterMixin): + """ResilPhase cache tests for the Flux pipeline.""" + + class TestFluxPipelineMagCache(FluxPipelineTesterConfig, MagCacheTesterMixin): """MagCache tests for the Flux pipeline.""" diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video.py b/tests/pipelines/hunyuan_video/test_hunyuan_video.py index b16c8d7158a6..ccc148a2b39a 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_video.py @@ -33,6 +33,7 @@ FirstBlockCacheTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, + ResilPhaseCacheTesterMixin, TaylorSeerCacheTesterMixin, to_np, ) @@ -46,6 +47,7 @@ class HunyuanVideoPipelineFastTests( PyramidAttentionBroadcastTesterMixin, FasterCacheTesterMixin, FirstBlockCacheTesterMixin, + ResilPhaseCacheTesterMixin, TaylorSeerCacheTesterMixin, unittest.TestCase, ): diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index cb35fc46ee7d..7fd860dd7e3f 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -32,6 +32,7 @@ from diffusers.hooks.first_block_cache import FirstBlockCacheConfig from diffusers.hooks.mag_cache import MagCacheConfig from diffusers.hooks.pyramid_attention_broadcast import PyramidAttentionBroadcastHook +from diffusers.hooks.resilphase_cache import ResilPhaseCacheConfig from diffusers.hooks.taylorseer_cache import TaylorSeerCacheConfig from diffusers.image_processor import VaeImageProcessor from diffusers.loaders import FluxIPAdapterMixin, IPAdapterMixin @@ -2827,6 +2828,46 @@ def run_forward(pipe): ) +class ResilPhaseCacheTesterMixin: + resilphase_cache_config = ResilPhaseCacheConfig(cache_interval=2, warmup_steps=3, max_order=1) + + def test_resilphase_cache_inference(self, expected_atol: float = 0.1): + device = "cpu" + + def create_pipe(): + torch.manual_seed(0) + components = self.get_dummy_components(num_layers=2) + pipe = self.pipeline_class(**components).to(device) + pipe.set_progress_bar_config(disable=None) + return pipe + + def run_forward(pipe): + torch.manual_seed(0) + inputs = self.get_dummy_inputs(device) + inputs["num_inference_steps"] = 4 + return pipe(**inputs)[0] + + pipe = create_pipe() + output = run_forward(pipe).flatten() + original_output_slice = np.concatenate((output[:8], output[-8:])) + + pipe = create_pipe() + pipe.transformer.enable_cache(self.resilphase_cache_config) + output = run_forward(pipe).flatten() + cached_output_slice = np.concatenate((output[:8], output[-8:])) + + pipe.transformer.disable_cache() + output = run_forward(pipe).flatten() + disabled_output_slice = np.concatenate((output[:8], output[-8:])) + + assert np.allclose(original_output_slice, cached_output_slice, atol=expected_atol), ( + "ResilPhase cache outputs should not differ much." + ) + assert np.allclose(original_output_slice, disabled_output_slice, atol=1e-4), ( + "Outputs from normal inference and after disabling cache should not differ." + ) + + class TaylorSeerCacheTesterMixin: taylorseer_cache_config = TaylorSeerCacheConfig( cache_interval=5, diff --git a/tests/pipelines/testing_utils/__init__.py b/tests/pipelines/testing_utils/__init__.py index 029e7c7e4f4d..e15ab1080852 100644 --- a/tests/pipelines/testing_utils/__init__.py +++ b/tests/pipelines/testing_utils/__init__.py @@ -4,6 +4,7 @@ FirstBlockCacheTesterMixin, MagCacheTesterMixin, PyramidAttentionBroadcastTesterMixin, + ResilPhaseCacheTesterMixin, TaylorSeerCacheTesterMixin, ) from .common import BasePipelineTesterConfig, PipelineTesterMixin @@ -32,6 +33,7 @@ "PyramidAttentionBroadcastTesterMixin", "FasterCacheTesterMixin", "FirstBlockCacheTesterMixin", + "ResilPhaseCacheTesterMixin", "TaylorSeerCacheTesterMixin", "MagCacheTesterMixin", "check_qkv_fused_layers_exist", diff --git a/tests/pipelines/testing_utils/cache.py b/tests/pipelines/testing_utils/cache.py index 82980d57fb3c..60e091ed588d 100644 --- a/tests/pipelines/testing_utils/cache.py +++ b/tests/pipelines/testing_utils/cache.py @@ -21,6 +21,7 @@ from diffusers.hooks.first_block_cache import FirstBlockCacheConfig from diffusers.hooks.mag_cache import MagCacheConfig from diffusers.hooks.pyramid_attention_broadcast import PyramidAttentionBroadcastHook +from diffusers.hooks.resilphase_cache import ResilPhaseCacheConfig from diffusers.hooks.taylorseer_cache import TaylorSeerCacheConfig from ...testing_utils import assert_tensors_close, is_cache, torch_device @@ -309,6 +310,21 @@ def test_first_block_cache_inference(self, expected_atol: float = 0.1): self._test_cache_inference(self._get_cache_config(), num_inference_steps=4, expected_atol=expected_atol) +@is_cache +class ResilPhaseCacheTesterMixin(CacheTesterMixin): + RESILPHASE_CACHE_CONFIG = { + "cache_interval": 2, + "warmup_steps": 3, + "max_order": 1, + } + + def _get_cache_config(self): + return ResilPhaseCacheConfig(**self.RESILPHASE_CACHE_CONFIG) + + def test_resilphase_cache_inference(self, expected_atol: float = 0.1): + self._test_cache_inference(self._get_cache_config(), num_inference_steps=4, expected_atol=expected_atol) + + @is_cache class TaylorSeerCacheTesterMixin(CacheTesterMixin): TAYLORSEER_CACHE_CONFIG = {