add imports

Author: AlanPonnachan

src/diffusers/__init__.py

@@ -168,6 +168,7 @@
             "HookRegistry",
             "LayerSkipConfig",
             "MagCacheConfig",
+            "FLUX_MAG_RATIOS",
             "PyramidAttentionBroadcastConfig",
             "SmoothedEnergyGuidanceConfig",
             "apply_faster_cache",
@@ -897,6 +898,7 @@
         from .hooks import (
             FasterCacheConfig,
             FirstBlockCacheConfig,
+            FLUX_MAG_RATIOS,
             HookRegistry,
             LayerSkipConfig,
             MagCacheConfig,

src/diffusers/hooks/__init__.py

@@ -23,6 +23,6 @@
     from .hooks import HookRegistry, ModelHook
     from .layer_skip import LayerSkipConfig, apply_layer_skip
     from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
-    from .mag_cache import MagCacheConfig, apply_mag_cache
+    from .mag_cache import FLUX_MAG_RATIOS, MagCacheConfig, apply_mag_cache
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
     from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig

src/diffusers/hooks/_common.py

@@ -23,7 +23,7 @@
 _ATTENTION_CLASSES = (Attention, MochiAttention, AttentionModuleMixin)
 _FEEDFORWARD_CLASSES = (FeedForward, LuminaFeedForward)
 
-_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "single_transformer_blocks", "layers")
+_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")