add caching docs

Author: lmxyy

nunchaku/caching/__init__.py

@@ -0,0 +1 @@
+

nunchaku/caching/diffusers_adapters/__init__.py

@@ -1,7 +1,101 @@
+"""
+Diffusers Pipeline Adapters for Nunchaku Caching.
+
+This module provides adapter functions that integrate Nunchaku's caching capabilities
+with diffusers pipelines. The adapters automatically detect pipeline types and apply
+the appropriate caching strategy without requiring manual configuration.
+
+The module serves as a unified interface for applying caching to different types of
+diffusion pipelines, currently supporting:
+
+- Flux pipelines (FluxPipeline and related variants)
+- SANA pipelines (SanaPipeline and related variants)
+
+Key Features:
+    - Automatic pipeline type detection
+    - Unified interface for different caching strategies
+    - Seamless integration with existing diffusers workflows
+    - Support for both shallow and deep caching patches
+
+Example:
+    Basic usage with automatic pipeline detection::
+
+        from diffusers import FluxPipeline
+        from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
+
+        # 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, 
+            residual_diff_threshold=0.1,
+            use_double_fb_cache=True
+        )
+        
+        # Use the cached pipeline normally
+        image = cached_pipe(prompt="A beautiful landscape")
+
+Note:
+    The adapter functions modify the pipeline in-place, adding caching capabilities
+    while preserving the original API. The caching behavior is transparent to the
+    user and doesn't require changes to existing code.
+"""
+
 from diffusers import DiffusionPipeline
 
 
 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
+        **kwargs: Variable keyword arguments passed to the specific caching function.
+            Common arguments include:
+            - residual_diff_threshold (float): Similarity threshold for cache validity
+            - 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, 
+                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
+        with "Flux" or "Sana".
+    """
     assert isinstance(pipe, DiffusionPipeline)
 
     pipe_cls_name = pipe.__class__.__name__

nunchaku/caching/diffusers_adapters/flux.py

@@ -1,3 +1,57 @@
+"""
+Flux Pipeline Caching Adapter.
+
+This module provides caching adapters specifically for Flux diffusion pipelines.
+It implements both transformer-level and pipeline-level caching integration,
+enabling efficient inference through advanced first-block caching strategies.
+
+The module supports both single and double first-block caching for Flux models,
+with automatic context management to ensure proper cache lifecycle during
+inference.
+
+Key Functions:
+    apply_cache_on_transformer: Apply caching directly to a FluxTransformer2DModel
+    apply_cache_on_pipe: Apply caching to a complete Flux pipeline
+
+Caching Features:
+    - Single first-block caching: Caches the first transformer block only
+    - Double first-block caching: Caches both multi-head and single-head attention blocks
+    - Dynamic threshold adjustment: Automatically adjusts similarity thresholds
+    - Context management: Ensures proper cache setup and cleanup
+    - Shallow patching: Optional lightweight patching for testing
+
+Example:
+    Apply caching to a Flux transformer::
+
+        from diffusers import FluxTransformer2DModel
+        from nunchaku.caching.diffusers_adapters.flux import apply_cache_on_transformer
+
+        transformer = FluxTransformer2DModel.from_pretrained("model_name")
+        cached_transformer = apply_cache_on_transformer(
+            transformer,
+            use_double_fb_cache=True,
+            residual_diff_threshold_multi=0.12,
+            residual_diff_threshold_single=0.09
+        )
+
+    Apply caching to a complete pipeline::
+
+        from diffusers import FluxPipeline
+        from nunchaku.caching.diffusers_adapters.flux import apply_cache_on_pipe
+
+        pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")
+        cached_pipe = apply_cache_on_pipe(
+            pipe,
+            use_double_fb_cache=True,
+            residual_diff_threshold=0.12
+        )
+
+Note:
+    The caching is applied in-place and uses mock patching to temporarily replace
+    transformer components during inference. The original functionality is preserved
+    when not using caching context.
+"""
+
 import functools
 import unittest
 
@@ -15,6 +69,51 @@ def apply_cache_on_transformer(
     residual_diff_threshold_multi: float | None = None,
     residual_diff_threshold_single: float = 0.1,
 ):
+    """
+    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.
+            If True, caches both multi-head and single-head attention blocks. Defaults to False.
+        residual_diff_threshold (float, optional): Default similarity threshold for caching.
+            Used for residual_diff_threshold_multi if not explicitly provided. Defaults to 0.12.
+        residual_diff_threshold_multi (float, optional): Similarity threshold for multi-head
+            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
+        context is present.
+    """
     if residual_diff_threshold_multi is None:
         residual_diff_threshold_multi = residual_diff_threshold
 
@@ -60,6 +159,56 @@ 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
+            without modifying the transformer. Useful for testing. Defaults to False.
+        **kwargs: Additional keyword arguments passed to apply_cache_on_transformer,
+            including:
+            - use_double_fb_cache (bool): Whether to use double first-block caching
+            - 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,
+        it skips the pipeline-level patching but still applies transformer caching
+        unless shallow_patch is True.
+    """
     if not getattr(pipe, "_is_cached", False):
         original_call = pipe.__class__.__call__
 

nunchaku/caching/diffusers_adapters/sana.py

@@ -1,3 +1,55 @@
+"""
+SANA Pipeline Caching Adapter.
+
+This module provides caching adapters specifically for SANA diffusion pipelines.
+It implements single first-block caching for SANA models, enabling efficient
+inference through intelligent reuse of first transformer block computations.
+
+The SANA adapter uses a simpler caching strategy compared to Flux, focusing
+on single first-block caching with configurable similarity thresholds. This
+approach is optimized for the specific architecture and usage patterns of
+SANA models.
+
+Key Functions:
+    apply_cache_on_transformer: Apply caching directly to a SanaTransformer2DModel
+    apply_cache_on_pipe: Apply caching to a complete SANA pipeline
+
+Caching Features:
+    - Single first-block caching: Caches the first transformer block's output
+    - Configurable similarity thresholds: Adjust caching sensitivity
+    - Context management: Automatic cache setup and cleanup
+    - Batch size limitations: Optimized for batch sizes <= 2 (CFG support)
+
+Example:
+    Apply caching to a SANA transformer::
+
+        from diffusers import SanaTransformer2DModel
+        from nunchaku.caching.diffusers_adapters.sana import apply_cache_on_transformer
+
+        transformer = SanaTransformer2DModel.from_pretrained("model_name")
+        cached_transformer = apply_cache_on_transformer(
+            transformer,
+            residual_diff_threshold=0.12
+        )
+
+    Apply caching to a complete SANA pipeline::
+
+        from diffusers import SanaPipeline
+        from nunchaku.caching.diffusers_adapters.sana import apply_cache_on_pipe
+
+        pipe = SanaPipeline.from_pretrained("Efficient-Large-Model/Sana_600M_512px")
+        cached_pipe = apply_cache_on_pipe(
+            pipe,
+            residual_diff_threshold=0.1
+        )
+
+Note:
+    SANA caching is specifically designed for the SANA architecture and uses
+    mock patching to temporarily replace transformer blocks during inference.
+    The caching is automatically disabled for batch sizes > 2 to ensure
+    compatibility with classifier-free guidance.
+"""
+
 import functools
 import unittest