diff --git a/src/diffusers/hooks/pyramid_attention_broadcast.py b/src/diffusers/hooks/pyramid_attention_broadcast.py index ed5bd24..de77321 100644 --- a/src/diffusers/hooks/pyramid_attention_broadcast.py +++ b/src/diffusers/hooks/pyramid_attention_broadcast.py @@ -27,7 +27,7 @@ _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS, _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS, ) -from .hooks import HookRegistry, ModelHook +from .hooks import BaseState, HookRegistry, ModelHook, StateManager logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -105,7 +105,7 @@ def __repr__(self) -> str: ) -class PyramidAttentionBroadcastState: +class PyramidAttentionBroadcastState(BaseState): r""" State for Pyramid Attention Broadcast. @@ -148,33 +148,34 @@ def __init__( self.timestep_skip_range = timestep_skip_range self.block_skip_range = block_skip_range self.current_timestep_callback = current_timestep_callback - - def initialize_hook(self, module): - self.state = PyramidAttentionBroadcastState() - return module + self.state_manager = StateManager(PyramidAttentionBroadcastState, (), {}) def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any: + if self.state_manager._current_context is None: + self.state_manager.set_context("inference") + + state = self.state_manager.get_state() is_within_timestep_range = ( self.timestep_skip_range[0] < self.current_timestep_callback() < self.timestep_skip_range[1] ) should_compute_attention = ( - self.state.cache is None - or self.state.iteration == 0 + state.cache is None + or state.iteration == 0 or not is_within_timestep_range - or self.state.iteration % self.block_skip_range == 0 + or state.iteration % self.block_skip_range == 0 ) if should_compute_attention: output = self.fn_ref.original_forward(*args, **kwargs) else: - output = self.state.cache + output = state.cache - self.state.cache = output - self.state.iteration += 1 + state.cache = output + state.iteration += 1 return output def reset_state(self, module: torch.nn.Module) -> None: - self.state.reset() + self.state_manager.reset() return module diff --git a/tests/hooks/test_pyramid_attention_broadcast.py b/tests/hooks/test_pyramid_attention_broadcast.py new file mode 100644 index 0000000..938d12c --- /dev/null +++ b/tests/hooks/test_pyramid_attention_broadcast.py @@ -0,0 +1,71 @@ +# Copyright 2025 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 torch + +from diffusers.hooks.pyramid_attention_broadcast import _apply_pyramid_attention_broadcast_hook +from diffusers.models.cache_utils import CacheMixin +from diffusers.models.modeling_utils import ModelMixin + +from ..testing_utils import torch_device + + +class CountingBlock(torch.nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.proj = torch.nn.Linear(dim, dim) + self.call_count = 0 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self.call_count += 1 + return self.proj(hidden_states) + + +class DummyTransformer(ModelMixin, CacheMixin, torch.nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([CountingBlock(dim)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for block in self.transformer_blocks: + hidden_states = block(hidden_states) + return hidden_states + + +def test_pab_isolates_state_between_cond_and_uncond_contexts(): + """PAB must not reuse cond attention cache during the uncond forward pass.""" + dim = 4 + model = DummyTransformer(dim).to(torch_device) + block = model.transformer_blocks[0] + + _apply_pyramid_attention_broadcast_hook( + block, + timestep_skip_range=(100, 800), + block_skip_range=2, + current_timestep_callback=lambda: 500, + ) + + x = torch.randn(1, dim, device=torch_device) + + with model.cache_context("cond"): + model(x) + + assert block.call_count == 1 + + with model.cache_context("uncond"): + model(x) + + assert block.call_count == 2, ( + "Uncond pass must compute attention instead of reusing cond cache from a shared hook state." + ) diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index fcd8ab2..30aa825 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -2691,7 +2691,11 @@ def test_pyramid_attention_broadcast_layers(self): isinstance(hook, PyramidAttentionBroadcastHook), "Hook should be of type PyramidAttentionBroadcastHook.", ) - self.assertTrue(hook.state.cache is None, "Cache should be None at initialization.") + self.assertEqual( + len(hook.state_manager._state_cache), + 0, + "Cache should be None at initialization.", + ) self.assertEqual(count, expected_hooks, "Number of hooks should match the expected number.") # Perform dummy inference step to ensure state is updated @@ -2701,12 +2705,13 @@ def pab_state_check_callback(pipe, i, t, kwargs): hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast") if hook is None: continue + state = hook.state_manager.get_state() self.assertTrue( - hook.state.cache is not None, + state.cache is not None, "Cache should have updated during inference.", ) self.assertTrue( - hook.state.iteration == i + 1, + state.iteration == i + 1, "Hook iteration state should have updated during inference.", ) return {} @@ -2722,14 +2727,11 @@ def pab_state_check_callback(pipe, i, t, kwargs): hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast") if hook is None: continue - self.assertTrue( - hook.state.cache is None, + self.assertEqual( + len(hook.state_manager._state_cache), + 0, "Cache should be reset to None after inference.", ) - self.assertTrue( - hook.state.iteration == 0, - "Iteration should be reset to 0 after inference.", - ) def test_pyramid_attention_broadcast_inference(self, expected_atol: float = 0.2): # We need to use higher tolerance because we are using a random model. With a converged/trained diff --git a/utils/check_copies.py b/utils/check_copies.py index 001366c..338a25d 100644 --- a/utils/check_copies.py +++ b/utils/check_copies.py @@ -18,6 +18,7 @@ import os import re import subprocess +import sys # All paths are set with the intent you should run this script from the root of the repo with the command @@ -94,7 +95,7 @@ def get_indent(code): def run_ruff(code): - command = ["ruff", "format", "-", "--config", "pyproject.toml", "--silent"] + command = [sys.executable, "-m", "ruff", "format", "-", "--config", "pyproject.toml", "--silent"] process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) stdout, _ = process.communicate(input=code.encode()) return stdout.decode()