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

cache_dit/cache_factory/cache_adapters/cache_adapter.py

@@ -5,10 +5,11 @@ import functools
 from contextlib import ExitStack
 from typing import Dict, List, Tuple, Any, Union, Callable, 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 FakeDiffusionPipeline
 from cache_dit.cache_factory.block_adapters import ParamsModifier
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 from cache_dit.cache_factory.cache_contexts import ContextManager
@@ -32,6 +33,9 @@ class CachedAdapter:
         pipe_or_adapter: Union[
             DiffusionPipeline,
             BlockAdapter,
+            # Transformer-only
+            torch.nn.Module,
+            ModelMixin,
         ],
         **context_kwargs,
     ) -> Union[
@@ -42,7 +46,9 @@ class CachedAdapter:
             pipe_or_adapter is not None
         ), "pipe or block_adapter can not both None!"
 
-        if isinstance(pipe_or_adapter, DiffusionPipeline):
+        if isinstance(
+            pipe_or_adapter, (DiffusionPipeline, torch.nn.Module, ModelMixin)
+        ):
             if BlockAdapterRegistry.is_supported(pipe_or_adapter):
                 logger.info(
                     f"{pipe_or_adapter.__class__.__name__} is officially "
@@ -52,16 +58,22 @@ class CachedAdapter:
                 block_adapter = BlockAdapterRegistry.get_adapter(
                     pipe_or_adapter
                 )
+                assert block_adapter is not None, (
+                    f"BlockAdapter for {pipe_or_adapter.__class__.__name__} "
+                    "should not be None!"
+                )
                 if params_modifiers := context_kwargs.pop(
                     "params_modifiers",
                     None,
                 ):
                     block_adapter.params_modifiers = params_modifiers
 
-                return cls.cachify(
-                    block_adapter,
-                    **context_kwargs,
-                ).pipe
+                block_adapter = cls.cachify(block_adapter, **context_kwargs)
+                if isinstance(pipe_or_adapter, DiffusionPipeline):
+                    return block_adapter.pipe
+
+                return block_adapter.transformer
+
             else:
                 raise ValueError(
                     f"{pipe_or_adapter.__class__.__name__} is not officially supported "
@@ -178,8 +190,6 @@ class CachedAdapter:
         context_kwargs = cls.check_context_kwargs(
             block_adapter, **context_kwargs
         )
-        # Apply cache on pipeline: wrap cache context
-        pipe_cls_name = block_adapter.pipe.__class__.__name__
 
         # Each Pipeline should have it's own context manager instance.
         # Different transformers (Wan2.2, etc) should shared the same
@@ -189,38 +199,58 @@ class CachedAdapter:
             "cache_config", None
         )
         assert cache_config is not None, "cache_config can not be None."
+        # Apply cache on pipeline: wrap cache context
+        pipe_cls_name = block_adapter.pipe.__class__.__name__
         context_manager = ContextManager(
             name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
             cache_type=cache_config.cache_type,
+            # Force use persistent_context for FakeDiffusionPipeline
+            persistent_context=isinstance(
+                block_adapter.pipe, FakeDiffusionPipeline
+            ),
         )
-        block_adapter.pipe._context_manager = context_manager  # instance level
-
         flatten_contexts, contexts_kwargs = cls.modify_context_params(
             block_adapter, **context_kwargs
         )
-        original_call = block_adapter.pipe.__class__.__call__
 
-        @functools.wraps(original_call)
-        def new_call(self, *args, **kwargs):
-            with ExitStack() as stack:
-                # cache context will be reset for each pipe inference
-                for context_name, context_kwargs in zip(
-                    flatten_contexts, contexts_kwargs
-                ):
-                    stack.enter_context(
-                        context_manager.enter_context(
-                            context_manager.reset_context(
-                                context_name,
-                                **context_kwargs,
-                            ),
+        block_adapter.pipe._context_manager = context_manager  # instance level
+
+        if not context_manager.persistent_context:
+
+            original_call = block_adapter.pipe.__class__.__call__
+
+            @functools.wraps(original_call)
+            def new_call(self, *args, **kwargs):
+                with ExitStack() as stack:
+                    # cache context will be reset for each pipe inference
+                    for context_name, context_kwargs in zip(
+                        flatten_contexts, contexts_kwargs
+                    ):
+                        stack.enter_context(
+                            context_manager.enter_context(
+                                context_manager.reset_context(
+                                    context_name,
+                                    **context_kwargs,
+                                ),
+                            )
                         )
-                    )
-                outputs = original_call(self, *args, **kwargs)
-                cls.apply_stats_hooks(block_adapter)
-                return outputs
+                    outputs = original_call(self, *args, **kwargs)
+                    cls.apply_stats_hooks(block_adapter)
+                    return outputs
+
+            block_adapter.pipe.__class__.__call__ = new_call
+            block_adapter.pipe.__class__._original_call = original_call
+
+        else:
+            # Init persistent cache context for transformer
+            for context_name, context_kwargs in zip(
+                flatten_contexts, contexts_kwargs
+            ):
+                context_manager.reset_context(
+                    context_name,
+                    **context_kwargs,
+                )
 
-        block_adapter.pipe.__class__.__call__ = new_call
-        block_adapter.pipe.__class__._original_call = original_call
         block_adapter.pipe.__class__._is_cached = True
 
         cls.apply_params_hooks(block_adapter, contexts_kwargs)
@@ -349,6 +379,7 @@ class CachedAdapter:
                 blocks_name,
                 unique_blocks_name,
                 dummy_blocks_names,
+                block_adapter,
             )
 
         return block_adapter.transformer
@@ -361,6 +392,7 @@ class CachedAdapter:
         blocks_name: List[str],
         unique_blocks_name: List[str],
         dummy_blocks_names: List[str],
+        block_adapter: BlockAdapter,
     ) -> torch.nn.Module:
         dummy_blocks = torch.nn.ModuleList()
 
@@ -387,6 +419,8 @@ class CachedAdapter:
         # re-apply hooks to transformer after cache applied.
         # from diffusers.hooks.hooks import HookFunctionReference, HookRegistry
         # from diffusers.hooks.group_offloading import apply_group_offloading
+        context_manager: ContextManager = block_adapter.pipe._context_manager
+        assert isinstance(context_manager, ContextManager._supported_managers)
 
         def new_forward(self, *args, **kwargs):
             with ExitStack() as stack:
@@ -406,6 +440,13 @@ class CachedAdapter:
                         )
                     )
                 outputs = original_forward(*args, **kwargs)
+
+                if (
+                    context_manager.persistent_context
+                    and context_manager.is_pre_refreshed()
+                ):
+                    cls.apply_stats_hooks(block_adapter)
+
             return outputs
 
         def new_forward_with_hf_hook(self, *args, **kwargs):
@@ -509,6 +550,7 @@ class CachedAdapter:
             params_shift += len(blocks)
 
     @classmethod
+    @torch.compiler.disable
     def apply_stats_hooks(
         cls,
         block_adapter: BlockAdapter,

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,13 +23,18 @@ def enable_cache(
     pipe_or_adapter: Union[
         DiffusionPipeline,
         BlockAdapter,
+        # Transformer-only
+        torch.nn.Module,
+        ModelMixin,
     ],
     # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
-    cache_config: Union[
-        BasicCacheConfig,
-        DBCacheConfig,
-        DBPruneConfig,
-    ] = DBCacheConfig(),
+    cache_config: Optional[
+        Union[
+            BasicCacheConfig,
+            DBCacheConfig,
+            DBPruneConfig,
+        ]
+    ] = None,
     # Calibrator config: TaylorSeerCalibratorConfig, etc.
     calibrator_config: Optional[CalibratorConfig] = None,
     # Modify cache context params for specific blocks.
@@ -45,6 +51,9 @@ def enable_cache(
     **kwargs,
 ) -> Union[
     DiffusionPipeline,
+    # Transformer-only
+    torch.nn.Module,
+    ModelMixin,
     BlockAdapter,
 ]:
     r"""
@@ -74,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.
@@ -117,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
@@ -135,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
@@ -154,13 +179,27 @@ def enable_cache(
     >>> stats = cache_dit.summary(pipe) # Then, get the summary of cache acceleration stats.
     >>> cache_dit.disable_cache(pipe) # Disable cache and run original pipe.
     """
+    # Precheck for compatibility of different configurations
+    if cache_config is None:
+        if parallelism_config is None:
+            # Set default cache config only when parallelism is not enabled
+            logger.info("cache_config is None, using default DBCacheConfig")
+            cache_config = DBCacheConfig()
+        else:
+            # Allow empty cache_config when parallelism is enabled
+            logger.warning(
+                "Parallelism is enabled and cache_config is None. Please manually "
+                "set cache_config to avoid potential compatibility issues. "
+                "Otherwise, cache will not be enabled."
+            )
+
     # Collect cache context kwargs
     context_kwargs = {}
     if (cache_type := context_kwargs.get("cache_type", None)) is not None:
         if cache_type == CacheType.NONE:
             return pipe_or_adapter
 
-    # WARNING: Deprecated cache config params. These parameters are now retained
+    # NOTE: Deprecated cache config params. These parameters are now retained
     # for backward compatibility but will be removed in the future.
     deprecated_kwargs = {
         "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
@@ -196,9 +235,9 @@ def enable_cache(
     if cache_config is not None:
         context_kwargs["cache_config"] = cache_config
 
-    # WARNING: Deprecated taylorseer params. These parameters are now retained
+    # NOTE: Deprecated taylorseer params. These parameters are now retained
     # for backward compatibility but will be removed in the future.
-    if (
+    if cache_config is not None and (
         kwargs.get("enable_taylorseer", None) is not None
         or kwargs.get("enable_encoder_taylorseer", None) is not None
     ):
@@ -226,16 +265,25 @@ def enable_cache(
     if params_modifiers is not None:
         context_kwargs["params_modifiers"] = params_modifiers
 
-    if isinstance(pipe_or_adapter, (DiffusionPipeline, BlockAdapter)):
-        pipe_or_adapter = CachedAdapter.apply(
+    if cache_config is not None:
+        if isinstance(
             pipe_or_adapter,
-            **context_kwargs,
-        )
+            (DiffusionPipeline, BlockAdapter, torch.nn.Module, ModelMixin),
+        ):
+            pipe_or_adapter = CachedAdapter.apply(
+                pipe_or_adapter,
+                **context_kwargs,
+            )
+        else:
+            raise ValueError(
+                f"type: {type(pipe_or_adapter)} is not valid, "
+                "Please pass DiffusionPipeline or BlockAdapter"
+                "for the 1's position param: pipe_or_adapter"
+            )
     else:
-        raise ValueError(
-            f"type: {type(pipe_or_adapter)} is not valid, "
-            "Please pass DiffusionPipeline or BlockAdapter"
-            "for the 1's position param: pipe_or_adapter"
+        logger.warning(
+            "cache_config is None, skip enabling cache for "
+            f"{pipe_or_adapter.__class__.__name__}."
         )
 
     # NOTE: Users should always enable parallelism after applying
@@ -244,19 +292,45 @@ def enable_cache(
         assert isinstance(
             parallelism_config, ParallelismConfig
         ), "parallelism_config should be of type ParallelismConfig."
+
+        transformers = []
         if isinstance(pipe_or_adapter, DiffusionPipeline):
-            transformer = pipe_or_adapter.transformer
+            adapter = BlockAdapterRegistry.get_adapter(pipe_or_adapter)
+            if adapter is None:
+                assert hasattr(pipe_or_adapter, "transformer"), (
+                    "The given DiffusionPipeline does not have "
+                    "a 'transformer' attribute, cannot enable "
+                    "parallelism."
+                )
+                transformers = [pipe_or_adapter.transformer]
+            else:
+                adapter = BlockAdapter.normalize(adapter, unique=False)
+                transformers = BlockAdapter.flatten(adapter.transformer)
         else:
-            assert BlockAdapter.assert_normalized(pipe_or_adapter)
-            assert (
-                len(BlockAdapter.flatten(pipe_or_adapter.transformer)) == 1
-            ), (
-                "Only single transformer is supported to enable parallelism "
-                "currently for BlockAdapter."
+            if not BlockAdapter.is_normalized(pipe_or_adapter):
+                pipe_or_adapter = BlockAdapter.normalize(
+                    pipe_or_adapter, unique=False
+                )
+            transformers = BlockAdapter.flatten(pipe_or_adapter.transformer)
+
+        if len(transformers) == 0:
+            logger.warning(
+                "No transformer is detected in the "
+                "BlockAdapter, skip enabling parallelism."
+            )
+            return pipe_or_adapter
+
+        if len(transformers) > 1:
+            logger.warning(
+                "Multiple transformers are detected in the "
+                "BlockAdapter, all transfomers will be "
+                "enabled for parallelism."
+            )
+        for i, transformer in enumerate(transformers):
+            # Enable parallelism for the transformer inplace
+            transformers[i] = enable_parallelism(
+                transformer, parallelism_config
             )
-            transformer = BlockAdapter.flatten(pipe_or_adapter.transformer)[0]
-        # Enable parallelism for the transformer inplace
-        transformer = enable_parallelism(transformer, parallelism_config)
     return pipe_or_adapter