refactor

Author: a-r-r-o-w

src/diffusers/__init__.py

@@ -78,8 +78,10 @@
 else:
     _import_structure["hooks"].extend(
         [
+            "FasterCacheConfig",
             "HookRegistry",
             "PyramidAttentionBroadcastConfig",
+            "apply_faster_cache",
             "apply_pyramid_attention_broadcast",
         ]
     )
@@ -287,7 +289,6 @@
             "CogView3PlusPipeline",
             "ConsisIDPipeline",
             "CycleDiffusionPipeline",
-            "FasterCacheConfig",
             "FluxControlImg2ImgPipeline",
             "FluxControlInpaintPipeline",
             "FluxControlNetImg2ImgPipeline",
@@ -434,7 +435,6 @@
             "WuerstchenCombinedPipeline",
             "WuerstchenDecoderPipeline",
             "WuerstchenPriorPipeline",
-            "apply_fastercache",
         ]
     )
 
@@ -599,7 +599,13 @@
     except OptionalDependencyNotAvailable:
         from .utils.dummy_pt_objects import *  # noqa F403
     else:
-        from .hooks import HookRegistry, PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
+        from .hooks import (
+            FasterCacheConfig,
+            HookRegistry,
+            PyramidAttentionBroadcastConfig,
+            apply_faster_cache,
+            apply_pyramid_attention_broadcast,
+        )
         from .models import (
             AllegroTransformer3DModel,
             AsymmetricAutoencoderKL,
@@ -782,7 +788,6 @@
             CogView3PlusPipeline,
             ConsisIDPipeline,
             CycleDiffusionPipeline,
-            FasterCacheConfig,
             FluxControlImg2ImgPipeline,
             FluxControlInpaintPipeline,
             FluxControlNetImg2ImgPipeline,
@@ -927,7 +932,6 @@
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,
             WuerstchenPriorPipeline,
-            apply_fastercache,
         )
 
     try:

src/diffusers/hooks/__init__.py

@@ -2,6 +2,7 @@
 
 
 if is_torch_available():
+    from .faster_cache import FasterCacheConfig, apply_faster_cache
     from .hooks import HookRegistry, ModelHook
     from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast

src/diffusers/pipelines/fastercache_utils.py renamed to src/diffusers/hooks/faster_cache.py

@@ -87,7 +87,7 @@ class PyramidAttentionBroadcastConfig:
 
     def __repr__(self) -> str:
         return (
-            f"PyramidAttentionBroadcastConfig("
+            f"PyramidAttentionBroadcastConfig(\n"
             f"  spatial_attention_block_skip_range={self.spatial_attention_block_skip_range},\n"
             f"  temporal_attention_block_skip_range={self.temporal_attention_block_skip_range},\n"
             f"  cross_attention_block_skip_range={self.cross_attention_block_skip_range},\n"

src/diffusers/hooks/pyramid_attention_broadcast.py

@@ -58,7 +58,6 @@
         "StableDiffusionMixin",
         "ImagePipelineOutput",
     ]
-    _import_structure["faster_cache_utils"] = ["FasterCacheConfig", "apply_fastercache"]
     _import_structure["deprecated"].extend(
         [
             "PNDMPipeline",
@@ -451,7 +450,6 @@
         from .ddpm import DDPMPipeline
         from .deprecated import KarrasVePipeline, LDMPipeline, PNDMPipeline, RePaintPipeline, ScoreSdeVePipeline
         from .dit import DiTPipeline
-        from .fastercache_utils import FasterCacheConfig, apply_fastercache
         from .latent_diffusion import LDMSuperResolutionPipeline
         from .pipeline_utils import (
             AudioPipelineOutput,

src/diffusers/pipelines/__init__.py

@@ -2,6 +2,21 @@
 from ..utils import DummyObject, requires_backends
 
 
+class FasterCacheConfig(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 HookRegistry(metaclass=DummyObject):
     _backends = ["torch"]
 
@@ -32,6 +47,10 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["torch"])
 
 
+def apply_faster_cache(*args, **kwargs):
+    requires_backends(apply_faster_cache, ["torch"])
+
+
 def apply_pyramid_attention_broadcast(*args, **kwargs):
     requires_backends(apply_pyramid_attention_broadcast, ["torch"])
 

src/diffusers/utils/dummy_pt_objects.py

@@ -392,21 +392,6 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["torch", "transformers"])
 
 
-class FasterCacheConfig(metaclass=DummyObject):
-    _backends = ["torch", "transformers"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["torch", "transformers"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-
 class FluxControlImg2ImgPipeline(metaclass=DummyObject):
     _backends = ["torch", "transformers"]
 
@@ -2565,7 +2550,3 @@ def from_config(cls, *args, **kwargs):
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["torch", "transformers"])
-
-
-def apply_fastercache(*args, **kwargs):
-    requires_backends(apply_fastercache, ["torch", "transformers"])

src/diffusers/utils/dummy_torch_and_transformers_objects.py

@@ -29,7 +29,7 @@
     StableDiffusionPipeline,
     StableDiffusionXLPipeline,
     UNet2DConditionModel,
-    apply_fastercache,
+    apply_faster_cache,
 )
 from diffusers.hooks.pyramid_attention_broadcast import PyramidAttentionBroadcastHook
 from diffusers.image_processor import VaeImageProcessor
@@ -2479,21 +2479,21 @@ def test_fastercache_basic_warning_or_errors_raised(self):
         # Check if warning is raised when no FasterCacheConfig is provided
         pipe = self.pipeline_class(**components)
         with CaptureLogger(logger) as cap_logger:
-            apply_fastercache(pipe)
+            apply_faster_cache(pipe)
         self.assertTrue("No FasterCacheConfig provided" in cap_logger.out)
 
         # Check if warning is raise when no attention_weight_callback is provided
         pipe = self.pipeline_class(**components)
         with CaptureLogger(logger) as cap_logger:
             config = FasterCacheConfig(spatial_attention_block_skip_range=2, attention_weight_callback=None)
-            apply_fastercache(pipe, config)
+            apply_faster_cache(pipe, config)
         self.assertTrue("No `attention_weight_callback` provided when enabling FasterCache" in cap_logger.out)
 
         # Check if error raised when unsupported tensor format used
         pipe = self.pipeline_class(**components)
         with self.assertRaises(ValueError):
             config = FasterCacheConfig(spatial_attention_block_skip_range=2, tensor_format="BFHWC")
-            apply_fastercache(pipe, config)
+            apply_faster_cache(pipe, config)
 
     def test_fastercache_inference(self, expected_atol: float = 0.1):
         device = "cpu"  # ensure determinism for the device-dependent torch.Generator
@@ -2509,7 +2509,7 @@ def test_fastercache_inference(self, expected_atol: float = 0.1):
         original_image_slice = output.flatten()
         original_image_slice = np.concatenate((original_image_slice[:8], original_image_slice[-8:]))
 
-        apply_fastercache(pipe, self.fastercache_config)
+        apply_faster_cache(pipe, self.fastercache_config)
 
         inputs = self.get_dummy_inputs(device)
         inputs["num_inference_steps"] = 4
@@ -2541,7 +2541,7 @@ def test_fastercache_state(self):
         pipe = self.pipeline_class(**components)
         pipe.set_progress_bar_config(disable=None)
 
-        apply_fastercache(pipe, self.fastercache_config)
+        apply_faster_cache(pipe, self.fastercache_config)
 
         expected_hooks = 0
         if self.fastercache_config.spatial_attention_block_skip_range is not None: