cache-dit 0.2.32__py3-none-any.whl → 0.2.34__py3-none-any.whl

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.32'
-__version_tuple__ = version_tuple = (0, 2, 32)
+__version__ = version = '0.2.34'
+__version_tuple__ = version_tuple = (0, 2, 34)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/block_adapters/__init__.py

@@ -153,7 +153,7 @@ def qwenimage_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("LTXVideo")
+@BlockAdapterRegistry.register("LTX")
 def ltxvideo_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import LTXVideoTransformer3DModel
 
@@ -248,7 +248,10 @@ def skyreelsv2_adapter(pipe, **kwargs) -> BlockAdapter:
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=pipe.transformer.blocks,
-        forward_pattern=ForwardPattern.Pattern_2,
+        # NOTE: Use Pattern_3 instead of Pattern_2 because the
+        # encoder_hidden_states will never change in the blocks
+        # forward loop.
+        forward_pattern=ForwardPattern.Pattern_3,
         has_separate_cfg=True,
         **kwargs,
     )
@@ -285,6 +288,7 @@ def consisid_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("DiT")
 def dit_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import DiTTransformer2DModel
+    from cache_dit.cache_factory.patch_functors import DiTPatchFunctor
 
     assert isinstance(pipe.transformer, DiTTransformer2DModel)
     return BlockAdapter(
@@ -292,6 +296,7 @@ def dit_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.transformer_blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        patch_functor=DiTPatchFunctor(),
         **kwargs,
     )
 
@@ -331,24 +336,13 @@ def bria_adapter(pipe, **kwargs) -> BlockAdapter:
 
 
 @BlockAdapterRegistry.register("Lumina")
-def lumina_adapter(pipe, **kwargs) -> BlockAdapter:
-    from diffusers import LuminaNextDiT2DModel
-
-    assert isinstance(pipe.transformer, LuminaNextDiT2DModel)
-    return BlockAdapter(
-        pipe=pipe,
-        transformer=pipe.transformer,
-        blocks=pipe.transformer.layers,
-        forward_pattern=ForwardPattern.Pattern_3,
-        **kwargs,
-    )
-
-
-@BlockAdapterRegistry.register("Lumina2")
 def lumina2_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import Lumina2Transformer2DModel
+    from diffusers import LuminaNextDiT2DModel
 
-    assert isinstance(pipe.transformer, Lumina2Transformer2DModel)
+    assert isinstance(
+        pipe.transformer, (Lumina2Transformer2DModel, LuminaNextDiT2DModel)
+    )
     return BlockAdapter(
         pipe=pipe,
         transformer=pipe.transformer,
@@ -386,12 +380,10 @@ def pixart_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("Sana", supported=False)
+@BlockAdapterRegistry.register("Sana")
 def sana_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import SanaTransformer2DModel
 
-    # TODO: fix -> got multiple values for argument 'encoder_hidden_states'
-
     assert isinstance(pipe.transformer, SanaTransformer2DModel)
     return BlockAdapter(
         pipe=pipe,
@@ -469,6 +461,7 @@ def auraflow_adapter(pipe, **kwargs) -> BlockAdapter:
 @BlockAdapterRegistry.register("Chroma")
 def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import ChromaTransformer2DModel
+    from cache_dit.cache_factory.patch_functors import ChromaPatchFunctor
 
     assert isinstance(pipe.transformer, ChromaTransformer2DModel)
     return BlockAdapter(
@@ -482,6 +475,7 @@ def chroma_adapter(pipe, **kwargs) -> BlockAdapter:
             ForwardPattern.Pattern_1,
             ForwardPattern.Pattern_3,
         ],
+        patch_functor=ChromaPatchFunctor(),
         has_separate_cfg=True,
         **kwargs,
     )

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -16,8 +16,52 @@ logger = init_logger(__name__)
 
 
 class ParamsModifier:
-    def __init__(self, **kwargs):
-        self._context_kwargs = kwargs.copy()
+    def __init__(
+        self,
+        # Cache context kwargs
+        Fn_compute_blocks: Optional[int] = None,
+        Bn_compute_blocks: Optional[int] = None,
+        max_warmup_steps: Optional[int] = None,
+        max_cached_steps: Optional[int] = None,
+        max_continuous_cached_steps: Optional[int] = None,
+        residual_diff_threshold: Optional[float] = None,
+        # Cache CFG or not
+        enable_separate_cfg: Optional[bool] = None,
+        cfg_compute_first: Optional[bool] = None,
+        cfg_diff_compute_separate: Optional[bool] = None,
+        # Hybird TaylorSeer
+        enable_taylorseer: Optional[bool] = None,
+        enable_encoder_taylorseer: Optional[bool] = None,
+        taylorseer_cache_type: Optional[str] = None,
+        taylorseer_order: Optional[int] = None,
+        **other_cache_context_kwargs,
+    ):
+        self._context_kwargs = other_cache_context_kwargs.copy()
+        self._maybe_update_param("Fn_compute_blocks", Fn_compute_blocks)
+        self._maybe_update_param("Bn_compute_blocks", Bn_compute_blocks)
+        self._maybe_update_param("max_warmup_steps", max_warmup_steps)
+        self._maybe_update_param("max_cached_steps", max_cached_steps)
+        self._maybe_update_param(
+            "max_continuous_cached_steps", max_continuous_cached_steps
+        )
+        self._maybe_update_param(
+            "residual_diff_threshold", residual_diff_threshold
+        )
+        self._maybe_update_param("enable_separate_cfg", enable_separate_cfg)
+        self._maybe_update_param("cfg_compute_first", cfg_compute_first)
+        self._maybe_update_param(
+            "cfg_diff_compute_separate", cfg_diff_compute_separate
+        )
+        self._maybe_update_param("enable_taylorseer", enable_taylorseer)
+        self._maybe_update_param(
+            "enable_encoder_taylorseer", enable_encoder_taylorseer
+        )
+        self._maybe_update_param("taylorseer_cache_type", taylorseer_cache_type)
+        self._maybe_update_param("taylorseer_order", taylorseer_order)
+
+    def _maybe_update_param(self, key: str, value: Any):
+        if value is not None:
+            self._context_kwargs[key] = value
 
 
 @dataclasses.dataclass
@@ -579,7 +623,7 @@ class BlockAdapter:
             assert isinstance(adapter[0], torch.nn.Module)
             return getattr(adapter[0], "_is_cached", False)
         else:
-            raise TypeError(f"Can't check this type: {type(adapter)}!")
+            return getattr(adapter, "_is_cached", False)
 
     @classmethod
     def nested_depth(cls, obj: Any):

cache_dit/cache_factory/block_adapters/block_registers.py

@@ -10,13 +10,14 @@ logger = init_logger(__name__)
 
 class BlockAdapterRegistry:
     _adapters: Dict[str, Callable[..., BlockAdapter]] = {}
-    _predefined_adapters_has_spearate_cfg: List[str] = [
+    _predefined_adapters_has_separate_cfg: List[str] = [
         "QwenImage",
         "Wan",
         "CogView4",
         "Cosmos",
         "SkyReelsV2",
         "Chroma",
+        "Lumina2",
     ]
 
     @classmethod
@@ -68,7 +69,7 @@ class BlockAdapterRegistry:
             return True
 
         pipe_cls_name = pipe_or_adapter.__class__.__name__
-        for name in cls._predefined_adapters_has_spearate_cfg:
+        for name in cls._predefined_adapters_has_separate_cfg:
             if pipe_cls_name.startswith(name):
                 return True
 

cache_dit/cache_factory/cache_adapters.py

@@ -114,27 +114,27 @@ class CachedAdapter:
         **cache_context_kwargs,
     ):
         # Check cache_context_kwargs
-        if cache_context_kwargs["enable_spearate_cfg"] is None:
+        if cache_context_kwargs["enable_separate_cfg"] is None:
             # Check cfg for some specific case if users don't set it as True
             if BlockAdapterRegistry.has_separate_cfg(block_adapter):
-                cache_context_kwargs["enable_spearate_cfg"] = True
+                cache_context_kwargs["enable_separate_cfg"] = True
                 logger.info(
-                    f"Use custom 'enable_spearate_cfg' from BlockAdapter: True. "
+                    f"Use custom 'enable_separate_cfg' from BlockAdapter: True. "
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
             else:
-                cache_context_kwargs["enable_spearate_cfg"] = (
+                cache_context_kwargs["enable_separate_cfg"] = (
                     BlockAdapterRegistry.has_separate_cfg(block_adapter.pipe)
                 )
                 logger.info(
-                    f"Use default 'enable_spearate_cfg' from block adapter "
-                    f"register: {cache_context_kwargs['enable_spearate_cfg']}, "
+                    f"Use default 'enable_separate_cfg' from block adapter "
+                    f"register: {cache_context_kwargs['enable_separate_cfg']}, "
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
         else:
             logger.info(
-                f"Use custom 'enable_spearate_cfg' from cache context "
-                f"kwargs: {cache_context_kwargs['enable_spearate_cfg']}. "
+                f"Use custom 'enable_separate_cfg' from cache context "
+                f"kwargs: {cache_context_kwargs['enable_separate_cfg']}. "
                 f"Pipeline: {block_adapter.pipe.__class__.__name__}."
             )
 

cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py

@@ -1,6 +1,5 @@
 import torch
 
-from typing import Dict, Any
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory.cache_blocks.pattern_base import (
     CachedBlocks_Pattern_Base,
@@ -24,14 +23,12 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         **kwargs,
     ):
         # Use it's own cache context.
-        self.cache_manager.set_context(
-            self.cache_context,
-        )
+        self.cache_manager.set_context(self.cache_context)
+        self._check_cache_params()
 
         original_hidden_states = hidden_states
         # Call first `n` blocks to process the hidden states for
         # more stable diff calculation.
-        # encoder_hidden_states: None Pattern 3, else 4, 5
         hidden_states, new_encoder_hidden_states = self.call_Fn_blocks(
             hidden_states,
             *args,
@@ -109,10 +106,7 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                 *args,
                 **kwargs,
             )
-            if new_encoder_hidden_states is not None:
-                new_encoder_hidden_states_residual = (
-                    new_encoder_hidden_states - old_encoder_hidden_states
-                )
+
             torch._dynamo.graph_break()
             if self.cache_manager.is_cache_residual():
                 self.cache_manager.set_Bn_buffer(
@@ -125,6 +119,10 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                     prefix=f"{self.cache_prefix}_Bn_hidden_states",
                 )
 
+            if new_encoder_hidden_states is not None:
+                new_encoder_hidden_states_residual = (
+                    new_encoder_hidden_states - old_encoder_hidden_states
+                )
             if self.cache_manager.is_encoder_cache_residual():
                 if new_encoder_hidden_states is not None:
                     self.cache_manager.set_Bn_encoder_buffer(
@@ -159,27 +157,12 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
             )
         )
 
-    @torch.compiler.disable
-    def maybe_update_kwargs(
-        self, encoder_hidden_states, kwargs: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        # if "encoder_hidden_states" in kwargs:
-        #     kwargs["encoder_hidden_states"] = encoder_hidden_states
-        # return kwargs
-        return kwargs
-
     def call_Fn_blocks(
         self,
         hidden_states: torch.Tensor,
         *args,
         **kwargs,
     ):
-        assert self.cache_manager.Fn_compute_blocks() <= len(
-            self.transformer_blocks
-        ), (
-            f"Fn_compute_blocks {self.cache_manager.Fn_compute_blocks()} must be less than "
-            f"the number of transformer blocks {len(self.transformer_blocks)}"
-        )
         new_encoder_hidden_states = None
         for block in self._Fn_blocks():
             hidden_states = block(
@@ -194,10 +177,6 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                         new_encoder_hidden_states,
                         hidden_states,
                     )
-            kwargs = self.maybe_update_kwargs(
-                new_encoder_hidden_states,
-                kwargs,
-            )
 
         return hidden_states, new_encoder_hidden_states
 
@@ -222,11 +201,6 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
                         new_encoder_hidden_states,
                         hidden_states,
                     )
-            kwargs = self.maybe_update_kwargs(
-                new_encoder_hidden_states,
-                kwargs,
-            )
-
         # compute hidden_states residual
         hidden_states = hidden_states.contiguous()
         hidden_states_residual = hidden_states - original_hidden_states
@@ -243,35 +217,22 @@ class CachedBlocks_Pattern_3_4_5(CachedBlocks_Pattern_Base):
         *args,
         **kwargs,
     ):
-        assert self.cache_manager.Bn_compute_blocks() <= len(
-            self.transformer_blocks
-        ), (
-            f"Bn_compute_blocks {self.cache_manager.Bn_compute_blocks()} must be less than "
-            f"the number of transformer blocks {len(self.transformer_blocks)}"
-        )
-        if len(self.cache_manager.Bn_compute_blocks_ids()) > 0:
-            raise ValueError(
-                f"Bn_compute_blocks_ids is not support for "
-                f"patterns: {self._supported_patterns}."
+        new_encoder_hidden_states = None
+        if self.cache_manager.Bn_compute_blocks() == 0:
+            return hidden_states, new_encoder_hidden_states
+
+        for block in self._Bn_blocks():
+            hidden_states = block(
+                hidden_states,
+                *args,
+                **kwargs,
             )
-        else:
-            # Compute all Bn blocks if no specific Bn compute blocks ids are set.
-            for block in self._Bn_blocks():
-                hidden_states = block(
-                    hidden_states,
-                    *args,
-                    **kwargs,
-                )
-                if not isinstance(hidden_states, torch.Tensor):  # Pattern 4,5
-                    hidden_states, new_encoder_hidden_states = hidden_states
-                    if not self.forward_pattern.Return_H_First:
-                        hidden_states, new_encoder_hidden_states = (
-                            new_encoder_hidden_states,
-                            hidden_states,
-                        )
-                kwargs = self.maybe_update_kwargs(
-                    new_encoder_hidden_states,
-                    kwargs,
-                )
+            if not isinstance(hidden_states, torch.Tensor):  # Pattern 4,5
+                hidden_states, new_encoder_hidden_states = hidden_states
+                if not self.forward_pattern.Return_H_First:
+                    hidden_states, new_encoder_hidden_states = (
+                        new_encoder_hidden_states,
+                        hidden_states,
+                    )
 
         return hidden_states, new_encoder_hidden_states

cache_dit/cache_factory/cache_blocks/pattern_base.py

@@ -93,6 +93,21 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                         required_param in forward_parameters
                     ), f"The input parameters must contains: {required_param}."
 
+    @torch.compiler.disable
+    def _check_cache_params(self):
+        assert self.cache_manager.Fn_compute_blocks() <= len(
+            self.transformer_blocks
+        ), (
+            f"Fn_compute_blocks {self.cache_manager.Fn_compute_blocks()} must be less than "
+            f"the number of transformer blocks {len(self.transformer_blocks)}"
+        )
+        assert self.cache_manager.Bn_compute_blocks() <= len(
+            self.transformer_blocks
+        ), (
+            f"Bn_compute_blocks {self.cache_manager.Bn_compute_blocks()} must be less than "
+            f"the number of transformer blocks {len(self.transformer_blocks)}"
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -100,7 +115,9 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
+        # Use it's own cache context.
         self.cache_manager.set_context(self.cache_context)
+        self._check_cache_params()
 
         original_hidden_states = hidden_states
         # Call first `n` blocks to process the hidden states for
@@ -191,18 +208,17 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                     prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
-                # TaylorSeer
                 self.cache_manager.set_Bn_buffer(
                     hidden_states,
                     prefix=f"{self.cache_prefix}_Bn_hidden_states",
                 )
+
             if self.cache_manager.is_encoder_cache_residual():
                 self.cache_manager.set_Bn_encoder_buffer(
                     encoder_hidden_states_residual,
                     prefix=f"{self.cache_prefix}_Bn_residual",
                 )
             else:
-                # TaylorSeer
                 self.cache_manager.set_Bn_encoder_buffer(
                     encoder_hidden_states,
                     prefix=f"{self.cache_prefix}_Bn_hidden_states",
@@ -296,12 +312,6 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         *args,
         **kwargs,
     ):
-        assert self.cache_manager.Fn_compute_blocks() <= len(
-            self.transformer_blocks
-        ), (
-            f"Fn_compute_blocks {self.cache_manager.Fn_compute_blocks()} must be less than "
-            f"the number of transformer blocks {len(self.transformer_blocks)}"
-        )
         for block in self._Fn_blocks():
             hidden_states = block(
                 hidden_states,
@@ -366,28 +376,17 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
             encoder_hidden_states_residual,
         )
 
-    def _compute_or_cache_block(
+    def call_Bn_blocks(
         self,
-        # Block index in the transformer blocks
-        # Bn: 8, block_id should be in [0, 8)
-        block_id: int,
-        # Below are the inputs to the block
-        block,  # The transformer block to be executed
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         *args,
         **kwargs,
     ):
-        # Helper function for `call_Bn_blocks`
-        # Skip the blocks by reuse residual cache if they are not
-        # in the Bn_compute_blocks_ids. NOTE: We should only skip
-        # the specific Bn blocks in cache steps. Compute the block
-        # and cache the residuals in non-cache steps.
-
-        # Normal steps: Compute the block and cache the residuals.
-        if not self._is_in_cache_step():
-            Bn_i_original_hidden_states = hidden_states
-            Bn_i_original_encoder_hidden_states = encoder_hidden_states
+        if self.cache_manager.Bn_compute_blocks() == 0:
+            return hidden_states, encoder_hidden_states
+
+        for block in self._Bn_blocks():
             hidden_states = block(
                 hidden_states,
                 encoder_hidden_states,
@@ -401,149 +400,5 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                         encoder_hidden_states,
                         hidden_states,
                     )
-            # Cache residuals for the non-compute Bn blocks for
-            # subsequent cache steps.
-            if block_id not in self.cache_manager.Bn_compute_blocks_ids():
-                Bn_i_hidden_states_residual = (
-                    hidden_states - Bn_i_original_hidden_states
-                )
-                if (
-                    encoder_hidden_states is not None
-                    and Bn_i_original_encoder_hidden_states is not None
-                ):
-                    Bn_i_encoder_hidden_states_residual = (
-                        encoder_hidden_states
-                        - Bn_i_original_encoder_hidden_states
-                    )
-                else:
-                    Bn_i_encoder_hidden_states_residual = None
-
-                # Save original_hidden_states for diff calculation.
-                self.cache_manager.set_Bn_buffer(
-                    Bn_i_original_hidden_states,
-                    prefix=f"{self.cache_prefix}_Bn_{block_id}_original",
-                )
-                self.cache_manager.set_Bn_encoder_buffer(
-                    Bn_i_original_encoder_hidden_states,
-                    prefix=f"{self.cache_prefix}_Bn_{block_id}_original",
-                )
-
-                self.cache_manager.set_Bn_buffer(
-                    Bn_i_hidden_states_residual,
-                    prefix=f"{self.cache_prefix}_Bn_{block_id}_residual",
-                )
-                self.cache_manager.set_Bn_encoder_buffer(
-                    Bn_i_encoder_hidden_states_residual,
-                    prefix=f"{self.cache_prefix}_Bn_{block_id}_residual",
-                )
-                del Bn_i_hidden_states_residual
-                del Bn_i_encoder_hidden_states_residual
-
-            del Bn_i_original_hidden_states
-            del Bn_i_original_encoder_hidden_states
-
-        else:
-            # Cache steps: Reuse the cached residuals.
-            # Check if the block is in the Bn_compute_blocks_ids.
-            if block_id in self.cache_manager.Bn_compute_blocks_ids():
-                hidden_states = block(
-                    hidden_states,
-                    encoder_hidden_states,
-                    *args,
-                    **kwargs,
-                )
-                if not isinstance(hidden_states, torch.Tensor):
-                    hidden_states, encoder_hidden_states = hidden_states
-                    if not self.forward_pattern.Return_H_First:
-                        hidden_states, encoder_hidden_states = (
-                            encoder_hidden_states,
-                            hidden_states,
-                        )
-            else:
-                # Skip the block if it is not in the Bn_compute_blocks_ids.
-                # Use the cached residuals instead.
-                # Check if can use the cached residuals.
-                if self.cache_manager.can_cache(
-                    hidden_states,  # curr step
-                    parallelized=self._is_parallelized(),
-                    threshold=self.cache_manager.non_compute_blocks_diff_threshold(),
-                    prefix=f"{self.cache_prefix}_Bn_{block_id}_original",  # prev step
-                ):
-                    hidden_states, encoder_hidden_states = (
-                        self.cache_manager.apply_cache(
-                            hidden_states,
-                            encoder_hidden_states,
-                            prefix=(
-                                f"{self.cache_prefix}_Bn_{block_id}_residual"
-                                if self.cache_manager.is_cache_residual()
-                                else f"{self.cache_prefix}_Bn_{block_id}_original"
-                            ),
-                            encoder_prefix=(
-                                f"{self.cache_prefix}_Bn_{block_id}_residual"
-                                if self.cache_manager.is_encoder_cache_residual()
-                                else f"{self.cache_prefix}_Bn_{block_id}_original"
-                            ),
-                        )
-                    )
-                else:
-                    hidden_states = block(
-                        hidden_states,
-                        encoder_hidden_states,
-                        *args,
-                        **kwargs,
-                    )
-                    if not isinstance(hidden_states, torch.Tensor):
-                        hidden_states, encoder_hidden_states = hidden_states
-                        if not self.forward_pattern.Return_H_First:
-                            hidden_states, encoder_hidden_states = (
-                                encoder_hidden_states,
-                                hidden_states,
-                            )
-        return hidden_states, encoder_hidden_states
-
-    def call_Bn_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        if self.cache_manager.Bn_compute_blocks() == 0:
-            return hidden_states, encoder_hidden_states
-
-        assert self.cache_manager.Bn_compute_blocks() <= len(
-            self.transformer_blocks
-        ), (
-            f"Bn_compute_blocks {self.cache_manager.Bn_compute_blocks()} must be less than "
-            f"the number of transformer blocks {len(self.transformer_blocks)}"
-        )
-        if len(self.cache_manager.Bn_compute_blocks_ids()) > 0:
-            for i, block in enumerate(self._Bn_blocks()):
-                hidden_states, encoder_hidden_states = (
-                    self._compute_or_cache_block(
-                        i,
-                        block,
-                        hidden_states,
-                        encoder_hidden_states,
-                        *args,
-                        **kwargs,
-                    )
-                )
-        else:
-            # Compute all Bn blocks if no specific Bn compute blocks ids are set.
-            for block in self._Bn_blocks():
-                hidden_states = block(
-                    hidden_states,
-                    encoder_hidden_states,
-                    *args,
-                    **kwargs,
-                )
-                if not isinstance(hidden_states, torch.Tensor):
-                    hidden_states, encoder_hidden_states = hidden_states
-                    if not self.forward_pattern.Return_H_First:
-                        hidden_states, encoder_hidden_states = (
-                            encoder_hidden_states,
-                            hidden_states,
-                        )
 
         return hidden_states, encoder_hidden_states