cache-dit 1.0.9__py3-none-any.whl → 1.0.10__py3-none-any.whl

cache_dit/cache_factory/cache_contexts/cache_config.py

@@ -56,6 +56,11 @@ class BasicCacheConfig:
     #     Compute separate diff values for CFG and non-CFG step, default True. If False, we will
     #     use the computed diff from current non-CFG transformer step for current CFG step.
     cfg_diff_compute_separate: bool = True
+    # num_inference_steps (`int`, *optional*, defaults to None):
+    #     num_inference_steps for DiffusionPipeline, used to adjust some internal settings
+    #     for better caching performance. For example, we will refresh the cache once the
+    #     executed steps exceed num_inference_steps if num_inference_steps is provided.
+    num_inference_steps: Optional[int] = None
 
     def update(self, **kwargs) -> "BasicCacheConfig":
         for key, value in kwargs.items():
@@ -108,9 +113,6 @@ class ExtraCacheConfig:
     # downsample_factor (`int`, *optional*, defaults to 1):
     #     Downsample factor for Fn buffer, in order the save GPU memory.
     downsample_factor: int = 1
-    # num_inference_steps (`int`, *optional*, defaults to -1):
-    #     num_inference_steps for DiffusionPipeline, for future use.
-    num_inference_steps: int = -1
 
 
 @dataclasses.dataclass

cache_dit/cache_factory/cache_contexts/cache_manager.py

@@ -7,6 +7,7 @@ import torch.distributed as dist
 
 from cache_dit.cache_factory.cache_contexts.calibrators import CalibratorBase
 from cache_dit.cache_factory.cache_contexts.cache_context import (
+    BasicCacheConfig,
     CachedContext,
 )
 from cache_dit.logger import init_logger
@@ -21,23 +22,143 @@ class ContextNotExistError(Exception):
 class CachedContextManager:
     # Each Pipeline should have it's own context manager instance.
 
-    def __init__(self, name: str = None):
+    def __init__(self, name: str = None, persistent_context: bool = False):
         self.name = name
         self._current_context: CachedContext = None
         self._cached_context_manager: Dict[str, CachedContext] = {}
+        # Whether to create new context automatically when setting
+        # a non-exist context name. Persistent context is useful when
+        # the pipeline class is not provided and users want to use
+        # cache-dit in a transformer-only way.
+        self._persistent_context = persistent_context
+        self._current_step_refreshed: bool = False
+
+    @property
+    def persistent_context(self) -> bool:
+        return self._persistent_context
+
+    @property
+    def current_context(self) -> CachedContext:
+        return self._current_context
+
+    @property
+    @torch.compiler.disable
+    def current_step_refreshed(self) -> bool:
+        return self._current_step_refreshed
+
+    @torch.compiler.disable
+    def is_pre_refreshed(self) -> bool:
+        _context = self._current_context
+        if _context is None:
+            return False
+
+        num_inference_steps = _context.cache_config.num_inference_steps
+        if num_inference_steps is not None:
+            current_step = _context.get_current_step()  # e.g, 0~49,50~99,...
+            return current_step == num_inference_steps - 1
+        return False
 
+    @torch.compiler.disable
     def new_context(self, *args, **kwargs) -> CachedContext:
+        if self._persistent_context:
+            cache_config: BasicCacheConfig = kwargs.get("cache_config", None)
+            assert (
+                cache_config is not None
+                and cache_config.num_inference_steps is not None
+            ), (
+                "When persistent_context is True, num_inference_steps "
+                "must be set in cache_config for proper cache refreshing."
+                f"\nkwargs: {kwargs}"
+            )
         _context = CachedContext(*args, **kwargs)
+        # NOTE: Patch args and kwargs for implicit refresh.
+        _context._init_args = args  # maybe empty tuple: ()
+        _context._init_kwargs = kwargs  # maybe empty dict: {}
         self._cached_context_manager[_context.name] = _context
         return _context
 
-    def set_context(self, cached_context: CachedContext | str) -> CachedContext:
+    @torch.compiler.disable
+    def maybe_refresh(
+        self,
+        cached_context: Optional[CachedContext | str] = None,
+    ) -> bool:
+        if cached_context is None:
+            _context = self._current_context
+            assert _context is not None, "Current context is not set!"
+
         if isinstance(cached_context, CachedContext):
-            self._current_context = cached_context
+            _context = cached_context
         else:
             if cached_context not in self._cached_context_manager:
                 raise ContextNotExistError("Context not exist!")
-            self._current_context = self._cached_context_manager[cached_context]
+            _context = self._cached_context_manager[cached_context]
+
+        if self._persistent_context:
+            assert _context.cache_config.num_inference_steps is not None, (
+                "When persistent_context is True, num_inference_steps must be set "
+                "in cache_config for proper cache refreshing."
+            )
+
+        num_inference_steps = _context.cache_config.num_inference_steps
+        if num_inference_steps is not None:
+            current_step = _context.get_current_step()  # e.g, 0~49,50~99,...
+            # Another round of inference, need to refresh cache context.
+            if current_step >= num_inference_steps:
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(
+                        f"Refreshing cache context '{_context.name}' "
+                        f"as current step: {current_step} >= "
+                        f"num_inference_steps: {num_inference_steps}."
+                    )
+                return True
+        return False
+
+    @torch.compiler.disable
+    def set_context(
+        self,
+        cached_context: CachedContext | str,
+        *args,
+        **kwargs,
+    ) -> CachedContext:
+        if isinstance(cached_context, CachedContext):
+            self._current_context = cached_context
+        else:
+            if cached_context not in self._cached_context_manager:
+                if not self._persistent_context:
+                    raise ContextNotExistError(
+                        "Context not exist and persistent_context is False. Please "
+                        "create new context first or set persistent_context=True."
+                    )
+                else:
+                    # Create new context if not exist
+                    if any((bool(args), bool(kwargs))):
+                        kwargs["name"] = cached_context
+                        self._current_context = self.new_context(
+                            *args, **kwargs
+                        )
+                    else:
+                        raise ValueError(
+                            "To create new context, please provide args and kwargs."
+                        )
+            else:
+                self._current_context = self._cached_context_manager[
+                    cached_context
+                ]
+
+        if self.maybe_refresh(self._current_context):
+            if not any((bool(args), bool(kwargs))):
+                assert hasattr(self._current_context, "_init_args")
+                assert hasattr(self._current_context, "_init_kwargs")
+                args = self._current_context._init_args
+                kwargs = self._current_context._init_kwargs
+
+            self._current_context = self.reset_context(
+                self._current_context, *args, **kwargs
+            )
+            self._current_step_refreshed = True
+        else:
+            self._current_step_refreshed = False
+
         return self._current_context
 
     def get_context(self, name: str = None) -> CachedContext:

cache_dit/cache_factory/cache_contexts/context_manager.py

@@ -20,10 +20,17 @@ class ContextManager:
         cls,
         cache_type: CacheType,
         name: str = "default",
+        persistent_context: bool = False,
     ) -> CachedContextManager | PrunedContextManager:
         if cache_type == CacheType.DBCache:
-            return CachedContextManager(name)
+            return CachedContextManager(
+                name=name,
+                persistent_context=persistent_context,
+            )
         elif cache_type == CacheType.DBPrune:
-            return PrunedContextManager(name)
+            return PrunedContextManager(
+                name=name,
+                persistent_context=persistent_context,
+            )
         else:
             raise ValueError(f"Unsupported cache_type: {cache_type}.")

cache_dit/cache_factory/cache_contexts/prune_manager.py

@@ -3,6 +3,7 @@ import functools
 from typing import Dict, List, Tuple, Union
 
 from cache_dit.cache_factory.cache_contexts.cache_manager import (
+    BasicCacheConfig,
     CachedContextManager,
 )
 from cache_dit.cache_factory.cache_contexts.prune_context import (
@@ -16,15 +17,27 @@ logger = init_logger(__name__)
 class PrunedContextManager(CachedContextManager):
     # Reuse CachedContextManager for Dynamic Block Prune
 
-    def __init__(self, name: str = None):
-        super().__init__(name)
+    def __init__(self, name: str = None, **kwargs):
+        super().__init__(name, **kwargs)
         # Overwrite for Dynamic Block Prune
         self._current_context: PrunedContext = None
         self._cached_context_manager: Dict[str, PrunedContext] = {}
 
     # Overwrite for Dynamic Block Prune
     def new_context(self, *args, **kwargs) -> PrunedContext:
+        if self._persistent_context:
+            cache_config: BasicCacheConfig = kwargs.get("cache_config", None)
+            assert (
+                cache_config is not None
+                and cache_config.num_inference_steps is not None
+            ), (
+                "When persistent_context is True, num_inference_steps "
+                "must be set in cache_config for proper cache refreshing."
+            )
         _context = PrunedContext(*args, **kwargs)
+        # NOTE: Patch args and kwargs for implicit refresh.
+        _context._init_args = args  # maybe empty tuple: ()
+        _context._init_kwargs = kwargs  # maybe empty dict: {}
         self._cached_context_manager[_context.name] = _context
         return _context
 

cache_dit/cache_factory/cache_interface.py

@@ -1,5 +1,6 @@
+import torch
 from typing import Any, Tuple, List, Union, Optional
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, ModelMixin
 from cache_dit.cache_factory.cache_types import CacheType
 from cache_dit.cache_factory.block_adapters import BlockAdapter
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
@@ -22,6 +23,9 @@ def enable_cache(
     pipe_or_adapter: Union[
         DiffusionPipeline,
         BlockAdapter,
+        # Transformer-only
+        torch.nn.Module,
+        ModelMixin,
     ],
     # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
     cache_config: Optional[
@@ -47,6 +51,9 @@ def enable_cache(
     **kwargs,
 ) -> Union[
     DiffusionPipeline,
+    # Transformer-only
+    torch.nn.Module,
+    ModelMixin,
     BlockAdapter,
 ]:
     r"""
@@ -76,7 +83,7 @@ def enable_cache(
     with minimal code changes.
 
     Args:
-        pipe_or_adapter (`DiffusionPipeline` or `BlockAdapter`, *required*):
+        pipe_or_adapter (`DiffusionPipeline`, `BlockAdapter` or `Transformer`, *required*):
             The standard Diffusion Pipeline or custom BlockAdapter (from cache-dit or user-defined).
             For example: cache_dit.enable_cache(FluxPipeline(...)). Please check https://github.com/vipshop/cache-dit/blob/main/docs/BlockAdapter.md
             for the usgae of BlockAdapter.
@@ -119,6 +126,10 @@ def enable_cache(
                     Whether to compute separate difference values for CFG and non-CFG steps, default is True.
                     If False, we will use the computed difference from the current non-CFG transformer step
                     for the current CFG step.
+                num_inference_steps (`int`, *optional*, defaults to None):
+                    num_inference_steps for DiffusionPipeline, used to adjust some internal settings
+                    for better caching performance. For example, we will refresh the cache once the
+                    executed steps exceed num_inference_steps if num_inference_steps is provided.
 
         calibrator_config (`CalibratorConfig`, *optional*, defaults to None):
             Config for calibrator. If calibrator_config is not None, it means the user wants to use DBCache
@@ -137,10 +148,22 @@ def enable_cache(
             Config for Parallelism. If parallelism_config is not None, it means the user wants to enable
             parallelism for cache-dit. Please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/parallelism/parallel_config.py
             for more details of ParallelismConfig.
+                backend: (`ParallelismBackend`, *required*, defaults to "ParallelismBackend.NATIVE_DIFFUSER"):
+                    Parallelism backend, currently only NATIVE_DIFFUSER and NVTIVE_PYTORCH are supported.
+                    For context parallelism, only NATIVE_DIFFUSER backend is supported, for tensor parallelism,
+                    only NATIVE_PYTORCH backend is supported.
                 ulysses_size: (`int`, *optional*, defaults to None):
                     The size of Ulysses cluster. If ulysses_size is not None, enable Ulysses style parallelism.
+                    This setting is only valid when backend is NATIVE_DIFFUSER.
                 ring_size: (`int`, *optional*, defaults to None):
                     The size of ring for ring parallelism. If ring_size is not None, enable ring attention.
+                    This setting is only valid when backend is NATIVE_DIFFUSER.
+                tp_size: (`int`, *optional*, defaults to None):
+                    The size of tensor parallelism. If tp_size is not None, enable tensor parallelism.
+                    This setting is only valid when backend is NATIVE_PYTORCH.
+                parallel_kwargs: (`dict`, *optional*, defaults to {}):
+                    Additional kwargs for parallelism backends. For example, for NATIVE_DIFFUSER backend,
+                    it can include `cp_plan` and `attention_backend` arguments for `Context Parallelism`.
 
         kwargs (`dict`, *optional*, defaults to {})
             Other cache context kwargs, please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/cache_factory/cache_contexts/cache_context.py
@@ -243,7 +266,10 @@ def enable_cache(
         context_kwargs["params_modifiers"] = params_modifiers
 
     if cache_config is not None:
-        if isinstance(pipe_or_adapter, (DiffusionPipeline, BlockAdapter)):
+        if isinstance(
+            pipe_or_adapter,
+            (DiffusionPipeline, BlockAdapter, torch.nn.Module, ModelMixin),
+        ):
             pipe_or_adapter = CachedAdapter.apply(
                 pipe_or_adapter,
                 **context_kwargs,

cache_dit/cache_factory/forward_pattern.py

@@ -20,33 +20,33 @@ class ForwardPattern(Enum):
 
     Pattern_0 = (
         True,  # Return_H_First
-        False, # Return_H_Only
-        False, # Forward_H_only
+        False,  # Return_H_Only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
         ("hidden_states", "encoder_hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_1 = (
-        False, # Return_H_First
-        False, # Return_H_Only
-        False, # Forward_H_only
+        False,  # Return_H_First
+        False,  # Return_H_Only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
         ("encoder_hidden_states", "hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_2 = (
-        False, # Return_H_First
+        False,  # Return_H_First
         True,  # Return_H_Only
-        False, # Forward_H_only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
-        ("hidden_states",),                          # Out
+        ("hidden_states",),  # Out
         True,  # Supported
     )
 
     Pattern_3 = (
-        False, # Return_H_First
+        False,  # Return_H_First
         True,  # Return_H_Only
         True,  # Forward_H_only
         ("hidden_states",),  # In
@@ -56,18 +56,18 @@ class ForwardPattern(Enum):
 
     Pattern_4 = (
         True,  # Return_H_First
-        False, # Return_H_Only
+        False,  # Return_H_Only
         True,  # Forward_H_only
-        ("hidden_states",),                          # In
+        ("hidden_states",),  # In
         ("hidden_states", "encoder_hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_5 = (
-        False, # Return_H_First
-        False, # Return_H_Only
+        False,  # Return_H_First
+        False,  # Return_H_Only
         True,  # Forward_H_only
-        ("hidden_states",),                          # In
+        ("hidden_states",),  # In
         ("encoder_hidden_states", "hidden_states"),  # Out
         True,  # Supported
     )

cache_dit/parallelism/backends/native_diffusers/__init__.py

@@ -1,6 +1,3 @@
 from cache_dit.parallelism.backends.native_diffusers.parallel_difffusers import (
     maybe_enable_parallelism,
 )
-from cache_dit.parallelism.backends.native_diffusers.parallel_difffusers import (
-    native_diffusers_parallelism_available,
-)

cache_dit/parallelism/backends/native_diffusers/context_parallelism/__init__.py

@@ -0,0 +1,95 @@
+import torch
+from typing import Optional
+
+from diffusers.models.modeling_utils import ModelMixin
+from cache_dit.parallelism.parallel_backend import ParallelismBackend
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+from cache_dit.logger import init_logger
+from ..utils import (
+    native_diffusers_parallelism_available,
+    ContextParallelConfig,
+)
+from .cp_planners import *
+
+logger = init_logger(__name__)
+
+
+def maybe_enable_context_parallelism(
+    transformer: torch.nn.Module,
+    parallelism_config: Optional[ParallelismConfig],
+) -> torch.nn.Module:
+    assert isinstance(transformer, ModelMixin), (
+        "transformer must be an instance of diffusers' ModelMixin, "
+        f"but got {type(transformer)}"
+    )
+    if parallelism_config is None:
+        return transformer
+
+    assert isinstance(parallelism_config, ParallelismConfig), (
+        "parallelism_config must be an instance of ParallelismConfig"
+        f" but got {type(parallelism_config)}"
+    )
+
+    if (
+        parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER
+        and native_diffusers_parallelism_available()
+    ):
+        cp_config = None
+        if (
+            parallelism_config.ulysses_size is not None
+            or parallelism_config.ring_size is not None
+        ):
+            cp_config = ContextParallelConfig(
+                ulysses_degree=parallelism_config.ulysses_size,
+                ring_degree=parallelism_config.ring_size,
+            )
+        if cp_config is not None:
+            attention_backend = parallelism_config.parallel_kwargs.get(
+                "attention_backend", None
+            )
+            if hasattr(transformer, "enable_parallelism"):
+                if hasattr(transformer, "set_attention_backend"):
+                    # _native_cudnn, flash, etc.
+                    if attention_backend is None:
+                        # Now only _native_cudnn is supported for parallelism
+                        # issue: https://github.com/huggingface/diffusers/pull/12443
+                        transformer.set_attention_backend("_native_cudnn")
+                        logger.warning(
+                            "attention_backend is None, set default attention backend "
+                            "to _native_cudnn for parallelism because of the issue: "
+                            "https://github.com/huggingface/diffusers/pull/12443"
+                        )
+                    else:
+                        transformer.set_attention_backend(attention_backend)
+                        logger.info(
+                            "Found attention_backend from config, set attention "
+                            f"backend to: {attention_backend}"
+                        )
+                # Prefer custom cp_plan if provided
+                cp_plan = parallelism_config.parallel_kwargs.get(
+                    "cp_plan", None
+                )
+                if cp_plan is not None:
+                    logger.info(
+                        f"Using custom context parallelism plan: {cp_plan}"
+                    )
+                else:
+                    # Try get context parallelism plan from register if not provided
+                    extra_parallel_kwargs = {}
+                    if parallelism_config.parallel_kwargs is not None:
+                        extra_parallel_kwargs = (
+                            parallelism_config.parallel_kwargs
+                        )
+                    cp_plan = ContextParallelismPlannerRegister.get_planner(
+                        transformer
+                    )().apply(transformer=transformer, **extra_parallel_kwargs)
+
+                transformer.enable_parallelism(
+                    config=cp_config, cp_plan=cp_plan
+                )
+            else:
+                raise ValueError(
+                    f"{transformer.__class__.__name__} does not support context parallelism."
+                )
+
+    return transformer

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_plan_registers.py

@@ -0,0 +1,74 @@
+import torch
+import logging
+from abc import abstractmethod
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+
+try:
+    from diffusers.models._modeling_parallel import (
+        ContextParallelModelPlan,
+    )
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+__all__ = [
+    "ContextParallelismPlanner",
+    "ContextParallelismPlannerRegister",
+]
+
+
+class ContextParallelismPlanner:
+    @abstractmethod
+    def apply(
+        self,
+        # NOTE: Keep this kwarg for future extensions
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        # NOTE: This method should only return the CP plan dictionary.
+        raise NotImplementedError(
+            "apply method must be implemented by subclasses"
+        )
+
+
+class ContextParallelismPlannerRegister:
+    _cp_planner_registry: dict[str, ContextParallelismPlanner] = {}
+
+    @classmethod
+    def register(cls, name: str):
+        def decorator(planner_cls: type[ContextParallelismPlanner]):
+            assert (
+                name not in cls._cp_planner_registry
+            ), f"ContextParallelismPlanner with name {name} is already registered."
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(f"Registering ContextParallelismPlanner: {name}")
+            cls._cp_planner_registry[name] = planner_cls
+            return planner_cls
+
+        return decorator
+
+    @classmethod
+    def get_planner(
+        cls, transformer: str | torch.nn.Module | ModelMixin
+    ) -> type[ContextParallelismPlanner]:
+        if isinstance(transformer, (torch.nn.Module, ModelMixin)):
+            name = transformer.__class__.__name__
+        else:
+            name = transformer
+        planner_cls = None
+        for planner_name in cls._cp_planner_registry:
+            if name.startswith(planner_name):
+                planner_cls = cls._cp_planner_registry.get(planner_name)
+                break
+        if planner_cls is None:
+            raise ValueError(f"No planner registered under name: {name}")
+        return planner_cls