cache-dit 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

cache_dit/cache_factory/dual_block_cache/cache_context.py

@@ -1,16 +1,13 @@
-# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/context.py
-
 import logging
 import contextlib
 import dataclasses
 from collections import defaultdict
-from typing import Any, DefaultDict, Dict, List, Optional, Union
+from typing import Any, DefaultDict, Dict, List, Optional, Union, Tuple
 
 import torch
 
 import cache_dit.primitives as primitives
 from cache_dit.cache_factory.taylorseer import TaylorSeer
-from cache_dit.utils import is_diffusers_at_least_0_3_5
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -55,8 +52,7 @@ class DBCacheContext:
     # DON'T Cache if the number of cached steps >= max_cached_steps
     max_cached_steps: int = -1  # for both CFG and non-CFG
 
-    # Statistics for botch alter cache and non-alter cache
-    # Record the steps that have been cached, both alter cache and non-alter cache
+    # Record the steps that have been cached, both cached and non-cache
     executed_steps: int = 0  # cache + non-cache steps pippeline
     # steps for transformer, for CFG, transformer_executed_steps will
     # be double of executed_steps.
@@ -73,10 +69,10 @@ class DBCacheContext:
     taylorseer: Optional[TaylorSeer] = None
     encoder_tarlorseer: Optional[TaylorSeer] = None
 
-    # Support do_separate_classifier_free_guidance, such as Wan 2.1
-    # For model that fused CFG and non-CFG into single forward step,
-    # should set do_separate_classifier_free_guidance as False. For
-    # example: CogVideoX, HunyuanVideo, Mochi.
+    # Support do_separate_classifier_free_guidance, such as Wan 2.1,
+    # Qwen-Image. For model that fused CFG and non-CFG into single
+    # forward step, should set do_separate_classifier_free_guidance
+    # as False. For example: CogVideoX, HunyuanVideo, Mochi.
     do_separate_classifier_free_guidance: bool = False
     # Compute cfg forward first or not, default False, namely,
     # 0, 2, 4, ..., -> non-CFG step; 1, 3, 5, ... -> CFG step.
@@ -98,13 +94,6 @@ class DBCacheContext:
         default_factory=lambda: defaultdict(float),
     )
 
-    # TODO: Support SLG in Dual Block Cache
-    # Skip Layer Guidance, SLG
-    # https://github.com/huggingface/candle/issues/2588
-    slg_layers: Optional[List[int]] = None
-    slg_start: float = 0.0
-    slg_end: float = 0.1
-
     @torch.compiler.disable
     def __post_init__(self):
         # Some checks for settings
@@ -144,18 +133,6 @@ class DBCacheContext:
                     **self.taylorseer_kwargs
                 )
 
-    @torch.compiler.disable
-    def get_incremental_name(self, name=None):
-        if name is None:
-            name = "default"
-        idx = self.incremental_name_counters[name]
-        self.incremental_name_counters[name] += 1
-        return f"{name}_{idx}"
-
-    @torch.compiler.disable
-    def reset_incremental_names(self):
-        self.incremental_name_counters.clear()
-
     @torch.compiler.disable
     def get_residual_diff_threshold(self):
         if self.enable_alter_cache:
@@ -222,7 +199,6 @@ class DBCacheContext:
             self.residual_diffs.clear()
             self.cfg_cached_steps.clear()
             self.cfg_residual_diffs.clear()
-            self.reset_incremental_names()
             # Reset the TaylorSeers cache at the beginning of each inference.
             # reset_cache will set the current step to -1 for TaylorSeer,
             if self.enable_taylorseer or self.enable_encoder_taylorseer:
@@ -264,12 +240,10 @@ class DBCacheContext:
                 if encoder_taylorseer is not None:
                     encoder_taylorseer.mark_step_begin()
 
-    @torch.compiler.disable
-    def get_taylorseers(self):
+    def get_taylorseers(self) -> Tuple[TaylorSeer, TaylorSeer]:
         return self.taylorseer, self.encoder_tarlorseer
 
-    @torch.compiler.disable
-    def get_cfg_taylorseers(self):
+    def get_cfg_taylorseers(self) -> Tuple[TaylorSeer, TaylorSeer]:
         return self.cfg_taylorseer, self.cfg_encoder_taylorseer
 
     @torch.compiler.disable
@@ -464,15 +438,13 @@ def is_encoder_taylorseer_enabled():
     return cache_context.enable_encoder_taylorseer
 
 
-@torch.compiler.disable
-def get_taylorseers():
+def get_taylorseers() -> Tuple[TaylorSeer, TaylorSeer]:
     cache_context = get_current_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_taylorseers()
 
 
-@torch.compiler.disable
-def get_cfg_taylorseers():
+def get_cfg_taylorseers() -> Tuple[TaylorSeer, TaylorSeer]:
     cache_context = get_current_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_cfg_taylorseers()
@@ -1105,825 +1077,3 @@ def get_can_use_cache(
             and is_alter_cache()
         )
     return can_use_cache
-
-
-class DBCachedTransformerBlocks(torch.nn.Module):
-    def __init__(
-        self,
-        transformer_blocks,
-        single_transformer_blocks=None,
-        *,
-        transformer=None,
-        return_hidden_states_first=True,
-        return_hidden_states_only=False,
-    ):
-        super().__init__()
-
-        self.transformer = transformer
-        self.transformer_blocks = transformer_blocks
-        self.single_transformer_blocks = single_transformer_blocks
-        self.return_hidden_states_first = return_hidden_states_first
-        self.return_hidden_states_only = return_hidden_states_only
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        original_hidden_states = hidden_states
-        # Call first `n` blocks to process the hidden states for
-        # more stable diff calculation.
-        hidden_states, encoder_hidden_states = self.call_Fn_transformer_blocks(
-            hidden_states,
-            encoder_hidden_states,
-            *args,
-            **kwargs,
-        )
-
-        Fn_hidden_states_residual = hidden_states - original_hidden_states
-        del original_hidden_states
-
-        mark_step_begin()
-        # Residual L1 diff or Hidden States L1 diff
-        can_use_cache = get_can_use_cache(
-            (
-                Fn_hidden_states_residual
-                if not is_l1_diff_enabled()
-                else hidden_states
-            ),
-            parallelized=self._is_parallelized(),
-            prefix=(
-                "Fn_residual"
-                if not is_l1_diff_enabled()
-                else "Fn_hidden_states"
-            ),
-        )
-
-        torch._dynamo.graph_break()
-        if can_use_cache:
-            add_cached_step()
-            del Fn_hidden_states_residual
-            hidden_states, encoder_hidden_states = apply_hidden_states_residual(
-                hidden_states,
-                encoder_hidden_states,
-                prefix=(
-                    "Bn_residual" if is_cache_residual() else "Bn_hidden_states"
-                ),
-                encoder_prefix=(
-                    "Bn_residual"
-                    if is_encoder_cache_residual()
-                    else "Bn_hidden_states"
-                ),
-            )
-            torch._dynamo.graph_break()
-            # Call last `n` blocks to further process the hidden states
-            # for higher precision.
-            hidden_states, encoder_hidden_states = (
-                self.call_Bn_transformer_blocks(
-                    hidden_states,
-                    encoder_hidden_states,
-                    *args,
-                    **kwargs,
-                )
-            )
-        else:
-            set_Fn_buffer(Fn_hidden_states_residual, prefix="Fn_residual")
-            if is_l1_diff_enabled():
-                # for hidden states L1 diff
-                set_Fn_buffer(hidden_states, "Fn_hidden_states")
-            del Fn_hidden_states_residual
-            torch._dynamo.graph_break()
-            (
-                hidden_states,
-                encoder_hidden_states,
-                hidden_states_residual,
-                encoder_hidden_states_residual,
-            ) = self.call_Mn_transformer_blocks(  # middle
-                hidden_states,
-                encoder_hidden_states,
-                *args,
-                **kwargs,
-            )
-            torch._dynamo.graph_break()
-            if is_cache_residual():
-                set_Bn_buffer(
-                    hidden_states_residual,
-                    prefix="Bn_residual",
-                )
-            else:
-                # TaylorSeer
-                set_Bn_buffer(
-                    hidden_states,
-                    prefix="Bn_hidden_states",
-                )
-            if is_encoder_cache_residual():
-                set_Bn_encoder_buffer(
-                    encoder_hidden_states_residual,
-                    prefix="Bn_residual",
-                )
-            else:
-                # TaylorSeer
-                set_Bn_encoder_buffer(
-                    encoder_hidden_states,
-                    prefix="Bn_hidden_states",
-                )
-            torch._dynamo.graph_break()
-            # Call last `n` blocks to further process the hidden states
-            # for higher precision.
-            hidden_states, encoder_hidden_states = (
-                self.call_Bn_transformer_blocks(
-                    hidden_states,
-                    encoder_hidden_states,
-                    *args,
-                    **kwargs,
-                )
-            )
-
-        patch_cached_stats(self.transformer)
-        torch._dynamo.graph_break()
-
-        return (
-            hidden_states
-            if self.return_hidden_states_only
-            else (
-                (hidden_states, encoder_hidden_states)
-                if self.return_hidden_states_first
-                else (encoder_hidden_states, hidden_states)
-            )
-        )
-
-    @torch.compiler.disable
-    def _is_parallelized(self):
-        # Compatible with distributed inference.
-        return all(
-            (
-                self.transformer is not None,
-                getattr(self.transformer, "_is_parallelized", False),
-            )
-        )
-
-    @torch.compiler.disable
-    def _is_in_cache_step(self):
-        # Check if the current step is in cache steps.
-        # If so, we can skip some Bn blocks and directly
-        # use the cached values.
-        return (get_current_step() in get_cached_steps()) or (
-            get_current_step() in get_cfg_cached_steps()
-        )
-
-    @torch.compiler.disable
-    def _Fn_transformer_blocks(self):
-        # Select first `n` blocks to process the hidden states for
-        # more stable diff calculation.
-        # Fn: [0,...,n-1]
-        selected_Fn_transformer_blocks = self.transformer_blocks[
-            : Fn_compute_blocks()
-        ]
-        return selected_Fn_transformer_blocks
-
-    @torch.compiler.disable
-    def _Mn_single_transformer_blocks(self):  # middle blocks
-        # M(N-2n): transformer_blocks [n,...] + single_transformer_blocks [0,...,N-n]
-        selected_Mn_single_transformer_blocks = []
-        if self.single_transformer_blocks is not None:
-            if Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
-                selected_Mn_single_transformer_blocks = (
-                    self.single_transformer_blocks
-                )
-            else:
-                selected_Mn_single_transformer_blocks = (
-                    self.single_transformer_blocks[: -Bn_compute_blocks()]
-                )
-        return selected_Mn_single_transformer_blocks
-
-    @torch.compiler.disable
-    def _Mn_transformer_blocks(self):  # middle blocks
-        # M(N-2n): only transformer_blocks [n,...,N-n], middle
-        if Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
-            selected_Mn_transformer_blocks = self.transformer_blocks[
-                Fn_compute_blocks() :
-            ]
-        else:
-            selected_Mn_transformer_blocks = self.transformer_blocks[
-                Fn_compute_blocks() : -Bn_compute_blocks()
-            ]
-        return selected_Mn_transformer_blocks
-
-    @torch.compiler.disable
-    def _Bn_single_transformer_blocks(self):
-        # Bn: single_transformer_blocks [N-n+1,...,N-1]
-        selected_Bn_single_transformer_blocks = []
-        if self.single_transformer_blocks is not None:
-            selected_Bn_single_transformer_blocks = (
-                self.single_transformer_blocks[-Bn_compute_blocks() :]
-            )
-        return selected_Bn_single_transformer_blocks
-
-    @torch.compiler.disable
-    def _Bn_transformer_blocks(self):
-        # Bn: transformer_blocks [N-n+1,...,N-1]
-        selected_Bn_transformer_blocks = self.transformer_blocks[
-            -Bn_compute_blocks() :
-        ]
-        return selected_Bn_transformer_blocks
-
-    def call_Fn_transformer_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        assert Fn_compute_blocks() <= len(self.transformer_blocks), (
-            f"Fn_compute_blocks {Fn_compute_blocks()} must be less than "
-            f"the number of transformer blocks {len(self.transformer_blocks)}"
-        )
-        for block in self._Fn_transformer_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.return_hidden_states_first:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
-                        hidden_states,
-                    )
-
-        return hidden_states, encoder_hidden_states
-
-    def call_Mn_transformer_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
-        # This condition branch is mainly for FLUX series.
-        if self.single_transformer_blocks is not None:
-            for block in self.transformer_blocks[Fn_compute_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.return_hidden_states_first:
-                        hidden_states, encoder_hidden_states = (
-                            encoder_hidden_states,
-                            hidden_states,
-                        )
-
-            # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py#L380
-            if is_diffusers_at_least_0_3_5():
-                for block in self._Mn_single_transformer_blocks():
-                    encoder_hidden_states, hidden_states = block(
-                        hidden_states,
-                        encoder_hidden_states,
-                        *args,
-                        **kwargs,
-                    )
-            else:
-                hidden_states = torch.cat(
-                    [encoder_hidden_states, hidden_states], dim=1
-                )
-                for block in self._Mn_single_transformer_blocks():
-                    hidden_states = block(
-                        hidden_states,
-                        *args,
-                        **kwargs,
-                    )
-                encoder_hidden_states, hidden_states = hidden_states.split(
-                    [
-                        encoder_hidden_states.shape[1],
-                        hidden_states.shape[1] - encoder_hidden_states.shape[1],
-                    ],
-                    dim=1,
-                )
-        else:
-            for block in self._Mn_transformer_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.return_hidden_states_first:
-                        hidden_states, encoder_hidden_states = (
-                            encoder_hidden_states,
-                            hidden_states,
-                        )
-
-        # hidden_states_shape = hidden_states.shape
-        # encoder_hidden_states_shape = encoder_hidden_states.shape
-        hidden_states = (
-            hidden_states.reshape(-1)
-            .contiguous()
-            .reshape(original_hidden_states.shape)
-        )
-        encoder_hidden_states = (
-            encoder_hidden_states.reshape(-1)
-            .contiguous()
-            .reshape(original_encoder_hidden_states.shape)
-        )
-
-        # hidden_states = hidden_states.contiguous()
-        # encoder_hidden_states = encoder_hidden_states.contiguous()
-
-        hidden_states_residual = hidden_states - original_hidden_states
-        encoder_hidden_states_residual = (
-            encoder_hidden_states - original_encoder_hidden_states
-        )
-
-        hidden_states_residual = (
-            hidden_states_residual.reshape(-1)
-            .contiguous()
-            .reshape(original_hidden_states.shape)
-        )
-        encoder_hidden_states_residual = (
-            encoder_hidden_states_residual.reshape(-1)
-            .contiguous()
-            .reshape(original_encoder_hidden_states.shape)
-        )
-
-        return (
-            hidden_states,
-            encoder_hidden_states,
-            hidden_states_residual,
-            encoder_hidden_states_residual,
-        )
-
-    @torch.compiler.disable
-    def _Bn_i_single_hidden_states_residual(
-        self,
-        Bn_i_hidden_states: torch.Tensor,
-        Bn_i_original_hidden_states: torch.Tensor,
-        original_hidden_states: torch.Tensor,
-        original_encoder_hidden_states: torch.Tensor,
-    ):
-        # Split the Bn_i_hidden_states and Bn_i_original_hidden_states
-        # into encoder_hidden_states and hidden_states.
-        Bn_i_hidden_states, Bn_i_encoder_hidden_states = (
-            self._split_Bn_i_single_hidden_states(
-                Bn_i_hidden_states,
-                original_hidden_states,
-                original_encoder_hidden_states,
-            )
-        )
-        # Split the Bn_i_original_hidden_states into encoder_hidden_states
-        # and hidden_states.
-        Bn_i_original_hidden_states, Bn_i_original_encoder_hidden_states = (
-            self._split_Bn_i_single_hidden_states(
-                Bn_i_original_hidden_states,
-                original_hidden_states,
-                original_encoder_hidden_states,
-            )
-        )
-
-        # Compute the residuals for the Bn_i_hidden_states and
-        # Bn_i_encoder_hidden_states.
-        Bn_i_hidden_states_residual = (
-            Bn_i_hidden_states - Bn_i_original_hidden_states
-        )
-        Bn_i_encoder_hidden_states_residual = (
-            Bn_i_encoder_hidden_states - Bn_i_original_encoder_hidden_states
-        )
-        return (
-            Bn_i_hidden_states_residual,
-            Bn_i_encoder_hidden_states_residual,
-        )
-
-    @torch.compiler.disable
-    def _split_Bn_i_single_hidden_states(
-        self,
-        Bn_i_hidden_states: torch.Tensor,
-        original_hidden_states: torch.Tensor,
-        original_encoder_hidden_states: torch.Tensor,
-    ):
-        # Split the Bn_i_hidden_states into encoder_hidden_states and hidden_states.
-        Bn_i_encoder_hidden_states, Bn_i_hidden_states = (
-            Bn_i_hidden_states.split(
-                [
-                    original_encoder_hidden_states.shape[1],
-                    Bn_i_hidden_states.shape[1]
-                    - original_encoder_hidden_states.shape[1],
-                ],
-                dim=1,
-            )
-        )
-        # Reshape the Bn_i_hidden_states and Bn_i_encoder_hidden_states
-        # to the original shape. This is necessary to ensure that the
-        # residuals are computed correctly.
-        Bn_i_hidden_states = (
-            Bn_i_hidden_states.reshape(-1)
-            .contiguous()
-            .reshape(original_hidden_states.shape)
-        )
-        Bn_i_encoder_hidden_states = (
-            Bn_i_encoder_hidden_states.reshape(-1)
-            .contiguous()
-            .reshape(original_encoder_hidden_states.shape)
-        )
-        return Bn_i_hidden_states, Bn_i_encoder_hidden_states
-
-    def _compute_and_cache_single_transformer_block(
-        self,
-        # Block index in the transformer blocks
-        # Bn: 8, block_id should be in [0, 8)
-        block_id: int,
-        # Helper inputs for hidden states split and reshape
-        original_hidden_states: torch.Tensor,
-        original_encoder_hidden_states: torch.Tensor,
-        # Below are the inputs to the block
-        block,  # The transformer block to be executed
-        hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        # Helper function for `call_Bn_transformer_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
-            hidden_states = block(
-                hidden_states,
-                *args,
-                **kwargs,
-            )
-            # Cache residuals for the non-compute Bn blocks for
-            # subsequent cache steps.
-            if block_id not in Bn_compute_blocks_ids():
-                Bn_i_hidden_states = hidden_states
-                (
-                    Bn_i_hidden_states_residual,
-                    Bn_i_encoder_hidden_states_residual,
-                ) = self._Bn_i_single_hidden_states_residual(
-                    Bn_i_hidden_states,
-                    Bn_i_original_hidden_states,
-                    original_hidden_states,
-                    original_encoder_hidden_states,
-                )
-
-                # Save original_hidden_states for diff calculation.
-                set_Bn_buffer(
-                    Bn_i_original_hidden_states,
-                    prefix=f"Bn_{block_id}_single_original",
-                )
-                set_Bn_encoder_buffer(
-                    Bn_i_original_hidden_states,
-                    prefix=f"Bn_{block_id}_single_original",
-                )
-
-                set_Bn_buffer(
-                    Bn_i_hidden_states_residual,
-                    prefix=f"Bn_{block_id}_single_residual",
-                )
-                set_Bn_encoder_buffer(
-                    Bn_i_encoder_hidden_states_residual,
-                    prefix=f"Bn_{block_id}_single_residual",
-                )
-                del Bn_i_hidden_states
-                del Bn_i_hidden_states_residual
-                del Bn_i_encoder_hidden_states_residual
-
-            del Bn_i_original_hidden_states
-
-        else:
-            # Cache steps: Reuse the cached residuals.
-            # Check if the block is in the Bn_compute_blocks_ids.
-            if block_id in Bn_compute_blocks_ids():
-                hidden_states = block(
-                    hidden_states,
-                    *args,
-                    **kwargs,
-                )
-            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 get_can_use_cache(
-                    hidden_states,  # curr step
-                    parallelized=self._is_parallelized(),
-                    threshold=non_compute_blocks_diff_threshold(),
-                    prefix=f"Bn_{block_id}_single_original",  # prev step
-                ):
-                    Bn_i_original_hidden_states = hidden_states
-                    (
-                        Bn_i_original_hidden_states,
-                        Bn_i_original_encoder_hidden_states,
-                    ) = self._split_Bn_i_single_hidden_states(
-                        Bn_i_original_hidden_states,
-                        original_hidden_states,
-                        original_encoder_hidden_states,
-                    )
-                    hidden_states, encoder_hidden_states = (
-                        apply_hidden_states_residual(
-                            Bn_i_original_hidden_states,
-                            Bn_i_original_encoder_hidden_states,
-                            prefix=(
-                                f"Bn_{block_id}_single_residual"
-                                if is_cache_residual()
-                                else f"Bn_{block_id}_single_original"
-                            ),
-                            encoder_prefix=(
-                                f"Bn_{block_id}_single_residual"
-                                if is_encoder_cache_residual()
-                                else f"Bn_{block_id}_single_original"
-                            ),
-                        )
-                    )
-                    hidden_states = torch.cat(
-                        [encoder_hidden_states, hidden_states],
-                        dim=1,
-                    )
-                    del Bn_i_original_hidden_states
-                    del Bn_i_original_encoder_hidden_states
-                else:
-                    hidden_states = block(
-                        hidden_states,
-                        *args,
-                        **kwargs,
-                    )
-        return hidden_states
-
-    def _compute_and_cache_transformer_block(
-        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_transformer_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
-            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.return_hidden_states_first:
-                    hidden_states, encoder_hidden_states = (
-                        encoder_hidden_states,
-                        hidden_states,
-                    )
-            # Cache residuals for the non-compute Bn blocks for
-            # subsequent cache steps.
-            if block_id not in Bn_compute_blocks_ids():
-                Bn_i_hidden_states_residual = (
-                    hidden_states - Bn_i_original_hidden_states
-                )
-                Bn_i_encoder_hidden_states_residual = (
-                    encoder_hidden_states - Bn_i_original_encoder_hidden_states
-                )
-
-                # Save original_hidden_states for diff calculation.
-                set_Bn_buffer(
-                    Bn_i_original_hidden_states,
-                    prefix=f"Bn_{block_id}_original",
-                )
-                set_Bn_encoder_buffer(
-                    Bn_i_original_encoder_hidden_states,
-                    prefix=f"Bn_{block_id}_original",
-                )
-
-                set_Bn_buffer(
-                    Bn_i_hidden_states_residual,
-                    prefix=f"Bn_{block_id}_residual",
-                )
-                set_Bn_encoder_buffer(
-                    Bn_i_encoder_hidden_states_residual,
-                    prefix=f"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 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.return_hidden_states_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 get_can_use_cache(
-                    hidden_states,  # curr step
-                    parallelized=self._is_parallelized(),
-                    threshold=non_compute_blocks_diff_threshold(),
-                    prefix=f"Bn_{block_id}_original",  # prev step
-                ):
-                    hidden_states, encoder_hidden_states = (
-                        apply_hidden_states_residual(
-                            hidden_states,
-                            encoder_hidden_states,
-                            prefix=(
-                                f"Bn_{block_id}_residual"
-                                if is_cache_residual()
-                                else f"Bn_{block_id}_original"
-                            ),
-                            encoder_prefix=(
-                                f"Bn_{block_id}_residual"
-                                if is_encoder_cache_residual()
-                                else f"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.return_hidden_states_first:
-                            hidden_states, encoder_hidden_states = (
-                                encoder_hidden_states,
-                                hidden_states,
-                            )
-        return hidden_states, encoder_hidden_states
-
-    def call_Bn_transformer_blocks(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        if Bn_compute_blocks() == 0:
-            return hidden_states, encoder_hidden_states
-
-        original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
-        # This condition branch is mainly for FLUX series.
-        if self.single_transformer_blocks is not None:
-            assert Bn_compute_blocks() <= len(self.single_transformer_blocks), (
-                f"Bn_compute_blocks {Bn_compute_blocks()} must be less than "
-                f"the number of single transformer blocks {len(self.single_transformer_blocks)}"
-            )
-            if is_diffusers_at_least_0_3_5():
-                if len(Bn_compute_blocks_ids()) > 0:
-                    # NOTE: Reuse _compute_and_cache_transformer_block here.
-                    for i, block in enumerate(
-                        self._Bn_single_transformer_blocks()
-                    ):
-                        hidden_states, encoder_hidden_states = (
-                            self._compute_and_cache_transformer_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_single_transformer_blocks():
-                        encoder_hidden_states, hidden_states = block(
-                            hidden_states,
-                            encoder_hidden_states,
-                            *args,
-                            **kwargs,
-                        )
-            else:
-                hidden_states = torch.cat(
-                    [encoder_hidden_states, hidden_states], dim=1
-                )
-                if len(Bn_compute_blocks_ids()) > 0:
-                    for i, block in enumerate(
-                        self._Bn_single_transformer_blocks()
-                    ):
-                        hidden_states = (
-                            self._compute_and_cache_single_transformer_block(
-                                i,
-                                original_hidden_states,
-                                original_encoder_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_single_transformer_blocks():
-                        hidden_states = block(
-                            hidden_states,
-                            *args,
-                            **kwargs,
-                        )
-                encoder_hidden_states, hidden_states = hidden_states.split(
-                    [
-                        encoder_hidden_states.shape[1],
-                        hidden_states.shape[1] - encoder_hidden_states.shape[1],
-                    ],
-                    dim=1,
-                )
-        else:
-            assert Bn_compute_blocks() <= len(self.transformer_blocks), (
-                f"Bn_compute_blocks {Bn_compute_blocks()} must be less than "
-                f"the number of transformer blocks {len(self.transformer_blocks)}"
-            )
-            if len(Bn_compute_blocks_ids()) > 0:
-                for i, block in enumerate(self._Bn_transformer_blocks()):
-                    hidden_states, encoder_hidden_states = (
-                        self._compute_and_cache_transformer_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_transformer_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.return_hidden_states_first:
-                            hidden_states, encoder_hidden_states = (
-                                encoder_hidden_states,
-                                hidden_states,
-                            )
-
-        hidden_states = (
-            hidden_states.reshape(-1)
-            .contiguous()
-            .reshape(original_hidden_states.shape)
-        )
-        encoder_hidden_states = (
-            encoder_hidden_states.reshape(-1)
-            .contiguous()
-            .reshape(original_encoder_hidden_states.shape)
-        )
-        return hidden_states, encoder_hidden_states
-
-
-@torch.compiler.disable
-def patch_cached_stats(
-    transformer,
-):
-    # Patch the cached stats to the transformer, the cached stats
-    # will be reset for each calling of pipe.__call__(**kwargs).
-    if transformer is None:
-        return
-
-    # TODO: Patch more cached stats to the transformer
-    transformer._cached_steps = get_cached_steps()
-    transformer._residual_diffs = get_residual_diffs()
-    transformer._cfg_cached_steps = get_cfg_cached_steps()
-    transformer._cfg_residual_diffs = get_cfg_residual_diffs()