Apply style fixes

Author: github-actions[bot]

src/diffusers/__init__.py

@@ -163,12 +163,12 @@
     )
     _import_structure["hooks"].extend(
         [
+            "FLUX_MAG_RATIOS",
             "FasterCacheConfig",
             "FirstBlockCacheConfig",
             "HookRegistry",
             "LayerSkipConfig",
             "MagCacheConfig",
-            "FLUX_MAG_RATIOS",
             "PyramidAttentionBroadcastConfig",
             "SmoothedEnergyGuidanceConfig",
             "TaylorSeerCacheConfig",
@@ -898,9 +898,9 @@
             TangentialClassifierFreeGuidance,
         )
         from .hooks import (
+            FLUX_MAG_RATIOS,
             FasterCacheConfig,
             FirstBlockCacheConfig,
-            FLUX_MAG_RATIOS,
             HookRegistry,
             LayerSkipConfig,
             MagCacheConfig,

src/diffusers/hooks/_common.py

@@ -23,7 +23,13 @@
 _ATTENTION_CLASSES = (Attention, MochiAttention, AttentionModuleMixin)
 _FEEDFORWARD_CLASSES = (FeedForward, LuminaFeedForward)
 
-_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "single_transformer_blocks", "layers", "visual_transformer_blocks",)
+_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = (
+    "blocks",
+    "transformer_blocks",
+    "single_transformer_blocks",
+    "layers",
+    "visual_transformer_blocks",
+)
 _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("temporal_transformer_blocks",)
 _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "layers")
 

src/diffusers/hooks/_helpers.py

@@ -184,12 +184,12 @@ def _register_transformer_blocks_metadata():
         HunyuanImageSingleTransformerBlock,
         HunyuanImageTransformerBlock,
     )
+    from ..models.transformers.transformer_kandinsky import Kandinsky5TransformerDecoderBlock
     from ..models.transformers.transformer_ltx import LTXVideoTransformerBlock
     from ..models.transformers.transformer_mochi import MochiTransformerBlock
     from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock
     from ..models.transformers.transformer_wan import WanTransformerBlock
     from ..models.transformers.transformer_z_image import ZImageTransformerBlock
-    from ..models.transformers.transformer_kandinsky import Kandinsky5TransformerDecoderBlock
 
     # BasicTransformerBlock
     TransformerBlockRegistry.register(
@@ -332,7 +332,6 @@ def _register_transformer_blocks_metadata():
         ),
     )
 
-
     TransformerBlockRegistry.register(
         model_class=JointTransformerBlock,
         metadata=TransformerBlockMetadata(
@@ -341,7 +340,6 @@ def _register_transformer_blocks_metadata():
         ),
     )
 
-
     # Kandinsky 5.0 (Kandinsky5TransformerDecoderBlock)
     TransformerBlockRegistry.register(
         model_class=Kandinsky5TransformerDecoderBlock,

src/diffusers/hooks/mag_cache.py

@@ -93,18 +93,18 @@ class MagCacheConfig:
         max_skip_steps (`int`, defaults to `5`):
             The maximum number of consecutive steps that can be skipped (K in the paper).
         retention_ratio (`float`, defaults to `0.1`):
-            The fraction of initial steps during which skipping is disabled to ensure stability.
-            For example, if `num_inference_steps` is 28 and `retention_ratio` is 0.1, the first 3 steps will never be skipped.
+            The fraction of initial steps during which skipping is disabled to ensure stability. For example, if
+            `num_inference_steps` is 28 and `retention_ratio` is 0.1, the first 3 steps will never be skipped.
         num_inference_steps (`int`, defaults to `28`):
             The number of inference steps used in the pipeline. This is required to interpolate `mag_ratios` correctly.
         mag_ratios (`np.ndarray`, *optional*):
-            The pre-computed magnitude ratios for the model. These are checkpoint-dependent.
-            If not provided, you must set `calibrate=True` to calculate them for your specific model.
-            For Flux models, you can use `diffusers.hooks.mag_cache.FLUX_MAG_RATIOS`.
+            The pre-computed magnitude ratios for the model. These are checkpoint-dependent. If not provided, you must
+            set `calibrate=True` to calculate them for your specific model. For Flux models, you can use
+            `diffusers.hooks.mag_cache.FLUX_MAG_RATIOS`.
         calibrate (`bool`, defaults to `False`):
-            If True, enables calibration mode. In this mode, no blocks are skipped. Instead, the hook calculates
-            the magnitude ratios for the current run and logs them at the end. Use this to obtain `mag_ratios`
-            for new models or schedulers.
+            If True, enables calibration mode. In this mode, no blocks are skipped. Instead, the hook calculates the
+            magnitude ratios for the current run and logs them at the end. Use this to obtain `mag_ratios` for new
+            models or schedulers.
     """
 
     threshold: float = 0.24
@@ -335,10 +335,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 if diff == 0:
                     residual = out_hidden - in_hidden
                 else:
-                    residual = out_hidden - in_hidden # Fallback to matching tail
+                    residual = out_hidden - in_hidden  # Fallback to matching tail
             else:
-                 # Fallback for completely mismatched shapes
-                 residual = out_hidden # Invalid but prevents crash
+                # Fallback for completely mismatched shapes
+                residual = out_hidden  # Invalid but prevents crash
 
             if self.config.calibrate:
                 self._perform_calibration_step(state, residual)
@@ -429,9 +429,7 @@ def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None:
     _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True)
 
 
-def _apply_mag_cache_head_hook(
-    block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig
-) -> None:
+def _apply_mag_cache_head_hook(block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig) -> None:
     registry = HookRegistry.check_if_exists_or_initialize(block)
 
     # Automatically remove existing hook to allow re-application (e.g. switching modes)

tests/hooks/test_mag_cache.py

@@ -25,6 +25,7 @@
 
 logger = logging.get_logger(__name__)
 
+
 class DummyBlock(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -34,6 +35,7 @@ def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
         # This ensures Residual = 2*Input - Input = Input
         return hidden_states * 2.0
 
+
 class DummyTransformer(ModelMixin):
     def __init__(self):
         super().__init__()
@@ -44,6 +46,7 @@ def forward(self, hidden_states, encoder_hidden_states=None):
             hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states)
         return hidden_states
 
+
 class TupleOutputBlock(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -52,6 +55,7 @@ def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
         # Returns a tuple
         return hidden_states * 2.0, encoder_hidden_states
 
+
 class TupleTransformer(ModelMixin):
     def __init__(self):
         super().__init__()
@@ -65,23 +69,18 @@ def forward(self, hidden_states, encoder_hidden_states=None):
             encoder_hidden_states = output[1]
         return hidden_states, encoder_hidden_states
 
+
 class MagCacheTests(unittest.TestCase):
     def setUp(self):
         # Register standard dummy block
         TransformerBlockRegistry.register(
             DummyBlock,
-            TransformerBlockMetadata(
-                return_hidden_states_index=None,
-                return_encoder_hidden_states_index=None
-            )
+            TransformerBlockMetadata(return_hidden_states_index=None, return_encoder_hidden_states_index=None),
         )
         # Register tuple block (Flux style)
         TransformerBlockRegistry.register(
             TupleOutputBlock,
-            TransformerBlockMetadata(
-                return_hidden_states_index=0,
-                return_encoder_hidden_states_index=1
-            )
+            TransformerBlockMetadata(return_hidden_states_index=0, return_encoder_hidden_states_index=1),
         )
 
     def _set_context(self, model, context_name):
@@ -115,9 +114,9 @@ def test_mag_cache_skipping_logic(self):
         config = MagCacheConfig(
             threshold=100.0,
             num_inference_steps=2,
-            retention_ratio=0.0, # Enable immediate skipping
+            retention_ratio=0.0,  # Enable immediate skipping
             max_skip_steps=5,
-            mag_ratios=ratios
+            mag_ratios=ratios,
         )
 
         apply_mag_cache(model, config)
@@ -136,8 +135,7 @@ def test_mag_cache_skipping_logic(self):
         output_t1 = model(input_t1)
 
         self.assertTrue(
-            torch.allclose(output_t1, torch.tensor([[[41.0]]])),
-            f"Expected Skip (41.0), got {output_t1.item()}"
+            torch.allclose(output_t1, torch.tensor([[[41.0]]])), f"Expected Skip (41.0), got {output_t1.item()}"
         )
 
     def test_mag_cache_retention(self):
@@ -149,8 +147,8 @@ def test_mag_cache_retention(self):
         config = MagCacheConfig(
             threshold=100.0,
             num_inference_steps=2,
-            retention_ratio=1.0, # Force retention for ALL steps
-            mag_ratios=ratios
+            retention_ratio=1.0,  # Force retention for ALL steps
+            mag_ratios=ratios,
         )
 
         apply_mag_cache(model, config)
@@ -165,20 +163,15 @@ def test_mag_cache_retention(self):
 
         self.assertTrue(
             torch.allclose(output_t1, torch.tensor([[[44.0]]])),
-            f"Expected Compute (44.0) due to retention, got {output_t1.item()}"
+            f"Expected Compute (44.0) due to retention, got {output_t1.item()}",
         )
 
     def test_mag_cache_tuple_outputs(self):
         """Test compatibility with models returning (hidden, encoder_hidden) like Flux."""
         model = TupleTransformer()
         ratios = np.array([1.0, 1.0])
 
-        config = MagCacheConfig(
-            threshold=100.0,
-            num_inference_steps=2,
-            retention_ratio=0.0,
-            mag_ratios=ratios
-        )
+        config = MagCacheConfig(threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=ratios)
 
         apply_mag_cache(model, config)
         self._set_context(model, "test_context")
@@ -196,36 +189,29 @@ def test_mag_cache_tuple_outputs(self):
         out_1, _ = model(input_t1, encoder_hidden_states=enc_t0)
 
         self.assertTrue(
-            torch.allclose(out_1, torch.tensor([[[21.0]]])),
-            f"Tuple skip failed. Expected 21.0, got {out_1.item()}"
+            torch.allclose(out_1, torch.tensor([[[21.0]]])), f"Tuple skip failed. Expected 21.0, got {out_1.item()}"
         )
 
     def test_mag_cache_reset(self):
         """Test that state resets correctly after num_inference_steps."""
         model = DummyTransformer()
         config = MagCacheConfig(
-            threshold=100.0,
-            num_inference_steps=2,
-            retention_ratio=0.0,
-            mag_ratios=np.array([1.0, 1.0])
+            threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=np.array([1.0, 1.0])
         )
         apply_mag_cache(model, config)
         self._set_context(model, "test_context")
 
         input_t = torch.ones(1, 1, 1)
 
-        model(input_t) # Step 0
-        model(input_t) # Step 1 (Skipped)
+        model(input_t)  # Step 0
+        model(input_t)  # Step 1 (Skipped)
 
         # Step 2 (Reset -> Step 0) -> Should Compute
         # Input 2.0 -> Output 8.0
         input_t2 = torch.tensor([[[2.0]]])
         output_t2 = model(input_t2)
 
-        self.assertTrue(
-            torch.allclose(output_t2, torch.tensor([[[8.0]]])),
-            "State did not reset correctly"
-        )
+        self.assertTrue(torch.allclose(output_t2, torch.tensor([[[8.0]]])), "State did not reset correctly")
 
     def test_mag_cache_calibration(self):
         """Test that calibration mode records ratios."""