make linter happy

Author: lmxyy

nunchaku/caching/__init__.py

@@ -1 +0,0 @@
-

nunchaku/caching/diffusers_adapters/__init__.py

@@ -25,14 +25,14 @@
 
         # Load any supported pipeline
         pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")
-        
+
         # Apply caching automatically based on pipeline type
         cached_pipe = apply_cache_on_pipe(
-            pipe, 
+            pipe,
             residual_diff_threshold=0.1,
             use_double_fb_cache=True
         )
-        
+
         # Use the cached pipeline normally
         image = cached_pipe(prompt="A beautiful landscape")
 
@@ -48,12 +48,12 @@
 def apply_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
     """
     Apply caching to a diffusers pipeline with automatic type detection.
-    
+
     This function serves as a unified interface for applying Nunchaku caching
     to different types of diffusion pipelines. It automatically detects the
     pipeline type based on the class name and delegates to the appropriate
     caching implementation.
-    
+
     Args:
         pipe (DiffusionPipeline): The diffusers pipeline to apply caching to
         *args: Variable positional arguments passed to the specific caching function
@@ -63,34 +63,34 @@ def apply_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
             - use_double_fb_cache (bool): Whether to use double first-block caching
             - shallow_patch (bool): Whether to use shallow patching only
             - verbose (bool): Whether to enable verbose caching messages
-            
+
     Returns:
         DiffusionPipeline: The same pipeline instance with caching applied
-        
+
     Raises:
         ValueError: If the pipeline type is not supported (doesn't start with "Flux" or "Sana")
         AssertionError: If the input is not a DiffusionPipeline instance
-        
+
     Example:
         With a Flux pipeline::
-        
+
             from diffusers import FluxPipeline
             pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")
             cached_pipe = apply_cache_on_pipe(
-                pipe, 
+                pipe,
                 residual_diff_threshold=0.12,
                 use_double_fb_cache=True
             )
-            
+
         With a SANA pipeline::
-        
+
             from diffusers import SanaPipeline
             pipe = SanaPipeline.from_pretrained("Efficient-Large-Model/Sana_600M_512px")
             cached_pipe = apply_cache_on_pipe(
                 pipe,
                 residual_diff_threshold=0.1
             )
-            
+
     Note:
         The function modifies the pipeline in-place and returns the same instance.
         Currently supported pipeline types are those with class names starting

nunchaku/caching/diffusers_adapters/flux.py

@@ -71,11 +71,11 @@ def apply_cache_on_transformer(
 ):
     """
     Apply caching to a Flux transformer model.
-    
+
     This function modifies a FluxTransformer2DModel to use cached transformer blocks
     for improved inference performance. It supports both single and double first-block
     caching strategies with configurable similarity thresholds.
-    
+
     Args:
         transformer (FluxTransformer2DModel): The Flux transformer model to apply caching to
         use_double_fb_cache (bool, optional): Whether to use double first-block caching.
@@ -86,29 +86,29 @@ def apply_cache_on_transformer(
             attention blocks. If None, uses residual_diff_threshold. Defaults to None.
         residual_diff_threshold_single (float, optional): Similarity threshold for single-head
             attention blocks. Defaults to 0.1.
-            
+
     Returns:
         FluxTransformer2DModel: The same transformer instance with caching applied
-        
+
     Example:
         Basic caching setup::
-        
+
             transformer = FluxTransformer2DModel.from_pretrained("model_name")
             cached_transformer = apply_cache_on_transformer(
                 transformer,
                 use_double_fb_cache=True,
                 residual_diff_threshold=0.12
             )
-            
+
         Advanced configuration::
-        
+
             cached_transformer = apply_cache_on_transformer(
                 transformer,
                 use_double_fb_cache=True,
                 residual_diff_threshold_multi=0.15,
                 residual_diff_threshold_single=0.08
             )
-            
+
     Note:
         If the transformer is already cached, the function updates the thresholds
         instead of reapplying caching. The caching only activates when a cache
@@ -161,11 +161,11 @@ def new_forward(self, *args, **kwargs):
 def apply_cache_on_pipe(pipe: DiffusionPipeline, *, shallow_patch: bool = False, **kwargs):
     """
     Apply caching to a complete Flux diffusion pipeline.
-    
+
     This function modifies a Flux diffusion pipeline to use caching during inference.
     It wraps the pipeline's __call__ method to automatically create and manage cache
     contexts, and optionally applies transformer-level caching.
-    
+
     Args:
         pipe (DiffusionPipeline): The Flux diffusion pipeline to apply caching to
         shallow_patch (bool, optional): If True, only applies pipeline-level caching
@@ -176,33 +176,33 @@ def apply_cache_on_pipe(pipe: DiffusionPipeline, *, shallow_patch: bool = False,
             - residual_diff_threshold (float): Similarity threshold for caching
             - residual_diff_threshold_multi (float): Multi-head attention threshold
             - residual_diff_threshold_single (float): Single-head attention threshold
-            
+
     Returns:
         DiffusionPipeline: The same pipeline instance with caching applied
-        
+
     Example:
         Basic usage::
-        
+
             from diffusers import FluxPipeline
             pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")
             cached_pipe = apply_cache_on_pipe(pipe)
-            
+
             # Use normally - caching is transparent
             image = cached_pipe(prompt="A beautiful landscape")
-            
+
         Advanced configuration::
-        
+
             cached_pipe = apply_cache_on_pipe(
                 pipe,
                 use_double_fb_cache=True,
                 residual_diff_threshold=0.1,
                 residual_diff_threshold_single=0.05
             )
-            
+
         Shallow patching for testing::
-        
+
             cached_pipe = apply_cache_on_pipe(pipe, shallow_patch=True)
-            
+
     Note:
         The function modifies the pipeline class's __call__ method, affecting all
         instances of the same pipeline class. If the pipeline is already cached,

nunchaku/caching/teacache.py

@@ -28,7 +28,7 @@
         from diffusers import FluxTransformer2DModel
 
         model = FluxTransformer2DModel.from_pretrained("black-forest-labs/FLUX.1-dev")
-        
+
         with TeaCache(model, num_steps=50, rel_l1_thresh=0.6, skip_steps=10):
             # Model forward passes will use temporal caching
             for step in range(50):
@@ -59,33 +59,34 @@
 def make_teacache_forward(num_steps: int = 50, rel_l1_thresh: float = 0.6, skip_steps: int = 0) -> Callable:
     """
     Create a cached forward method for Flux transformers using TeaCache.
-    
+
     This factory function creates a modified forward method that implements temporal
     caching based on the relative L1 distance of modulated inputs. The caching
     decision is made by analyzing the first transformer block's modulated input
     and comparing it to the previous timestep.
-    
+
     Args:
         num_steps (int, optional): Total number of inference steps. Used to determine
             when to reset the counter. Defaults to 50.
         rel_l1_thresh (float, optional): Relative L1 distance threshold for caching.
             Lower values mean more aggressive caching. Defaults to 0.6.
         skip_steps (int, optional): Number of initial steps to skip caching.
             Useful for allowing the model to stabilize. Defaults to 0.
-            
+
     Returns:
         Callable: A cached forward method that can be bound to a transformer model
-        
+
     Example:
         >>> model = FluxTransformer2DModel.from_pretrained("model_name")
         >>> cached_forward = make_teacache_forward(num_steps=50, rel_l1_thresh=0.6)
         >>> model.forward = cached_forward.__get__(model, type(model))
-        
+
     Note:
         The rescaling function uses polynomial coefficients optimized for Flux models.
         The accumulated distance is reset when it exceeds the threshold or at the
         beginning/end of the inference sequence.
     """
+
     def teacache_forward(
         self: Union[FluxTransformer2DModel, NunchakuFluxTransformer2dModel],
         hidden_states: torch.Tensor,
@@ -405,14 +406,14 @@ def custom_forward(*inputs):  # type: ignore
 class TeaCache:
     """
     Context manager for applying TeaCache temporal caching to transformer models.
-    
+
     This class provides a context manager that temporarily modifies a Flux transformer
     model to use TeaCache temporal caching. When entering the context, the model's
     forward method is replaced with a cached version that tracks temporal changes
     and skips computation when appropriate.
-    
+
     Args:
-        model (Union[FluxTransformer2DModel, NunchakuFluxTransformer2dModel]): 
+        model (Union[FluxTransformer2DModel, NunchakuFluxTransformer2dModel]):
             The transformer model to apply caching to
         num_steps (int, optional): Total number of inference steps. Defaults to 50.
         rel_l1_thresh (float, optional): Relative L1 distance threshold for caching.
@@ -421,32 +422,33 @@ class TeaCache:
             Useful for model stabilization. Defaults to 0.
         enabled (bool, optional): Whether caching is enabled. If False, the model
             behaves normally. Defaults to True.
-            
+
     Attributes:
         model: Reference to the transformer model
         num_steps (int): Total number of inference steps
         rel_l1_thresh (float): Caching threshold
         skip_steps (int): Number of steps to skip caching
         enabled (bool): Caching enabled flag
         previous_model_forward: Original forward method (for restoration)
-        
+
     Example:
         Basic usage::
-        
+
             with TeaCache(model, num_steps=50, rel_l1_thresh=0.6):
                 for step in range(50):
                     output = model(inputs[step])
-                    
+
         Disabling caching conditionally::
-        
+
             with TeaCache(model, enabled=use_caching):
                 # Model will use caching only if use_caching is True
                 output = model(inputs)
-                
+
     Note:
         The context manager automatically restores the original forward method
         when exiting, ensuring the model can be used normally afterward.
     """
+
     def __init__(
         self,
         model: Union[FluxTransformer2DModel, NunchakuFluxTransformer2dModel],
@@ -465,14 +467,14 @@ def __init__(
     def __enter__(self) -> "TeaCache":
         """
         Enter the TeaCache context and apply caching to the model.
-        
+
         This method is called when entering the 'with' block. It replaces the
         model's forward method with a cached version and initializes the
         necessary state variables for tracking temporal changes.
-        
+
         Returns:
             TeaCache: Self reference for context manager protocol
-            
+
         Note:
             If caching is disabled (enabled=False), the model is left unchanged.
         """
@@ -490,16 +492,16 @@ def __enter__(self) -> "TeaCache":
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         """
         Exit the TeaCache context and restore the original model.
-        
+
         This method is called when exiting the 'with' block. It restores the
         model's original forward method and cleans up the state variables
         that were added for caching.
-        
+
         Args:
             exc_type: Exception type (if any occurred)
-            exc_value: Exception value (if any occurred)  
+            exc_value: Exception value (if any occurred)
             traceback: Exception traceback (if any occurred)
-            
+
         Note:
             If caching was disabled (enabled=False), no cleanup is performed.
         """