Fix PAB cond/uncond cache cross-contamination

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
 
 

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."
+ )

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

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()