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

cache_dit/cache_factory/dynamic_block_prune/prune_context.py

@@ -1,14 +1,12 @@
-# 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, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, DefaultDict
 
 import torch
 
 import cache_dit.primitives as primitives
-from cache_dit.utils import is_diffusers_at_least_0_3_5
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -41,21 +39,56 @@ class DBPPruneContext:
     buffers: Dict[str, Any] = dataclasses.field(default_factory=dict)
 
     # Other settings
-    downsample_factor: int = 1
+    downsample_factor: int = 1  # un-used
     num_inference_steps: int = -1
     warmup_steps: int = 0  # DON'T pruned in warmup steps
     # DON'T prune if the number of pruned steps >= max_pruned_steps
     max_pruned_steps: int = -1
 
-    # Statistics
-    executed_steps: int = 0
+    # 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.
+    transformer_executed_steps: int = 0
+
+    # 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.
+    cfg_compute_first: bool = False
+    # Compute spearate 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
+
+    # CFG & non-CFG pruned steps
     pruned_blocks: List[int] = dataclasses.field(default_factory=list)
     actual_blocks: List[int] = dataclasses.field(default_factory=list)
-    # Residual diffs for each step, [step: list[float]]
-    residual_diffs: Dict[str, List[float]] = dataclasses.field(
+    residual_diffs: DefaultDict[str, list[float]] = dataclasses.field(
+        default_factory=lambda: defaultdict(list),
+    )
+    cfg_pruned_blocks: List[int] = dataclasses.field(default_factory=list)
+    cfg_actual_blocks: List[int] = dataclasses.field(default_factory=list)
+    cfg_residual_diffs: DefaultDict[str, list[float]] = dataclasses.field(
         default_factory=lambda: defaultdict(list),
     )
 
+    @torch.compiler.disable
+    def __post_init__(self):
+        # Some checks for settings
+        if self.do_separate_classifier_free_guidance:
+            assert (
+                self.cfg_diff_compute_separate
+            ), "cfg_diff_compute_separate must be True"
+            if self.cfg_diff_compute_separate:
+                assert self.cfg_compute_first is False, (
+                    "cfg_compute_first must set as False if "
+                    "cfg_diff_compute_separate is enabled."
+                )
+
     @torch.compiler.disable
     def get_residual_diff_threshold(self):
         residual_diff_threshold = self.residual_diff_threshold
@@ -119,42 +152,89 @@ class DBPPruneContext:
 
     @torch.compiler.disable
     def mark_step_begin(self):
-        self.executed_steps += 1
-        if self.get_current_step() == 0:
+        # Always increase transformer executed steps
+        # incr    step: prev 0 -> 1; prev 1 -> 2
+        # current step: incr step - 1
+        self.transformer_executed_steps += 1
+        if not self.do_separate_classifier_free_guidance:
+            self.executed_steps += 1
+        else:
+            # 0,1 -> 0 + 1, 2,3 -> 1 + 1, ...
+            if not self.cfg_compute_first:
+                if not self.is_separate_classifier_free_guidance_step():
+                    # transformer step: 0,2,4,...
+                    self.executed_steps += 1
+            else:
+                if self.is_separate_classifier_free_guidance_step():
+                    # transformer step: 0,2,4,...
+                    self.executed_steps += 1
+
+        # Reset the cached steps and residual diffs at the beginning
+        # of each inference.
+        if self.get_current_transformer_step() == 0:
             self.pruned_blocks.clear()
             self.actual_blocks.clear()
             self.residual_diffs.clear()
+            self.cfg_pruned_blocks.clear()
+            self.cfg_actual_blocks.clear()
+            self.cfg_residual_diffs.clear()
 
     @torch.compiler.disable
     def add_pruned_block(self, num_blocks):
-        self.pruned_blocks.append(num_blocks)
+        if not self.is_separate_classifier_free_guidance_step():
+            self.pruned_blocks.append(num_blocks)
+        else:
+            self.cfg_pruned_blocks.append(num_blocks)
 
     @torch.compiler.disable
     def add_actual_block(self, num_blocks):
-        self.actual_blocks.append(num_blocks)
+        if not self.is_separate_classifier_free_guidance_step():
+            self.actual_blocks.append(num_blocks)
+        else:
+            self.cfg_actual_blocks.append(num_blocks)
 
     @torch.compiler.disable
     def add_residual_diff(self, diff):
-        if isinstance(diff, torch.Tensor):
-            diff = diff.item()
-        step = self.get_current_step()
-        self.residual_diffs[step].append(diff)
-        max_num_block_diffs = 1000
-        # Avoid memory leak, keep only the last 1000 diffs
-        if len(self.residual_diffs[step]) > max_num_block_diffs:
-            self.residual_diffs[step] = self.residual_diffs[step][
-                -max_num_block_diffs:
-            ]
-        if logger.isEnabledFor(logging.DEBUG):
-            logger.debug(
-                f"Step {step}, block: {len(self.residual_diffs[step])}, "
-                f"residual diff: {diff:.6f}"
-            )
+        # step: executed_steps - 1, not transformer_steps - 1
+        step = str(self.get_current_step())
+        # Only add the diff if it is not already recorded for this step
+        if not self.is_separate_classifier_free_guidance_step():
+            if step not in self.residual_diffs:
+                self.residual_diffs[step] = [diff]
+            else:
+                self.residual_diffs[step].append(diff)
+        else:
+            if step not in self.cfg_residual_diffs:
+                self.cfg_residual_diffs[step] = [diff]
+            else:
+                self.cfg_residual_diffs[step].append(diff)
+
+    @torch.compiler.disable
+    def get_pruned_blocks(self):
+        return self.pruned_blocks.copy()
+
+    @torch.compiler.disable
+    def get_cfg_pruned_blocks(self):
+        return self.cfg_pruned_blocks.copy()
 
     @torch.compiler.disable
     def get_current_step(self):
         return self.executed_steps - 1
 
+    @torch.compiler.disable
+    def get_current_transformer_step(self):
+        return self.transformer_executed_steps - 1
+
+    @torch.compiler.disable
+    def is_separate_classifier_free_guidance_step(self):
+        if not self.do_separate_classifier_free_guidance:
+            return False
+        if self.cfg_compute_first:
+            # CFG steps: 0, 2, 4, 6, ...
+            return self.get_current_transformer_step() % 2 == 0
+        # CFG steps: 1, 3, 5, 7, ...
+        return self.get_current_transformer_step() % 2 != 0
+
     @torch.compiler.disable
     def is_in_warmup(self):
         return self.get_current_step() < self.warmup_steps
@@ -168,38 +248,35 @@ def get_residual_diff_threshold():
 
 
 @torch.compiler.disable
-def get_buffer(name):
-    prune_context = get_current_prune_context()
-    assert prune_context is not None, "prune_context must be set before"
-    return prune_context.get_buffer(name)
-
-
-@torch.compiler.disable
-def set_buffer(name, buffer):
+def mark_step_begin():
     prune_context = get_current_prune_context()
     assert prune_context is not None, "prune_context must be set before"
-    prune_context.set_buffer(name, buffer)
+    prune_context.mark_step_begin()
 
 
 @torch.compiler.disable
-def remove_buffer(name):
+def get_current_step():
     prune_context = get_current_prune_context()
     assert prune_context is not None, "prune_context must be set before"
-    prune_context.remove_buffer(name)
+    return prune_context.get_current_step()
 
 
 @torch.compiler.disable
-def mark_step_begin():
+def get_current_step_cfg_residual_diff():
     prune_context = get_current_prune_context()
     assert prune_context is not None, "prune_context must be set before"
-    prune_context.mark_step_begin()
+    step = str(get_current_step())
+    cfg_residual_diffs = get_cfg_residual_diffs()
+    if step in cfg_residual_diffs:
+        return cfg_residual_diffs[step]
+    return None
 
 
 @torch.compiler.disable
-def get_current_step():
+def get_current_transformer_step():
     prune_context = get_current_prune_context()
     assert prune_context is not None, "prune_context must be set before"
-    return prune_context.get_current_step()
+    return prune_context.get_current_transformer_step()
 
 
 @torch.compiler.disable
@@ -226,6 +303,13 @@ def get_pruned_blocks():
     return prune_context.pruned_blocks.copy()
 
 
+@torch.compiler.disable
+def get_cfg_pruned_blocks():
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    return prune_context.cfg_pruned_blocks.copy()
+
+
 @torch.compiler.disable
 def add_actual_block(num_blocks):
     assert (
@@ -243,6 +327,13 @@ def get_actual_blocks():
     return prune_context.actual_blocks.copy()
 
 
+@torch.compiler.disable
+def get_cfg_actual_blocks():
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    return prune_context.cfg_actual_blocks.copy()
+
+
 @torch.compiler.disable
 def get_pruned_steps():
     prune_context = get_current_prune_context()
@@ -252,6 +343,15 @@ def get_pruned_steps():
     return len(pruned_blocks)
 
 
+@torch.compiler.disable
+def get_cfg_pruned_steps():
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    cfg_pruned_blocks = get_cfg_pruned_blocks()
+    cfg_pruned_blocks = [x for x in cfg_pruned_blocks if x > 0]
+    return len(cfg_pruned_blocks)
+
+
 @torch.compiler.disable
 def is_in_warmup():
     prune_context = get_current_prune_context()
@@ -284,6 +384,14 @@ def get_residual_diffs():
     return prune_context.residual_diffs.copy()
 
 
+@torch.compiler.disable
+def get_cfg_residual_diffs():
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    # Return a copy of the residual diffs to avoid modification
+    return prune_context.cfg_residual_diffs.copy()
+
+
 @torch.compiler.disable
 def get_important_condition_threshold():
     prune_context = get_current_prune_context()
@@ -325,6 +433,27 @@ def get_non_prune_blocks_ids():
     return prune_context.non_prune_blocks_ids
 
 
+@torch.compiler.disable
+def do_separate_classifier_free_guidance():
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    return prune_context.do_separate_classifier_free_guidance
+
+
+@torch.compiler.disable
+def is_separate_classifier_free_guidance_step():
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    return prune_context.is_separate_classifier_free_guidance_step()
+
+
+@torch.compiler.disable
+def cfg_diff_compute_separate():
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    return prune_context.cfg_diff_compute_separate
+
+
 _current_prune_context: DBPPruneContext = None
 
 
@@ -463,6 +592,58 @@ def are_two_tensors_similar(
     return diff < threshold
 
 
+@torch.compiler.disable
+def _debugging_set_buffer(prefix):
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(
+            f"set {prefix}, "
+            f"transformer step: {get_current_transformer_step()}, "
+            f"executed step: {get_current_step()}"
+        )
+
+
+@torch.compiler.disable
+def _debugging_get_buffer(prefix):
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(
+            f"get {prefix}, "
+            f"transformer step: {get_current_transformer_step()}, "
+            f"executed step: {get_current_step()}"
+        )
+
+
+@torch.compiler.disable
+def set_buffer(name: str, buffer: torch.Tensor):
+    # Set hidden_states or residual for Fn blocks.
+    # This buffer is only use for L1 diff calculation.
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    if is_separate_classifier_free_guidance_step():
+        _debugging_set_buffer(f"{name}_buffer_cfg")
+        prune_context.set_buffer(f"{name}_buffer_cfg", buffer)
+    else:
+        _debugging_set_buffer(f"{name}_buffer")
+        prune_context.set_buffer(f"{name}_buffer", buffer)
+
+
+@torch.compiler.disable
+def get_buffer(name: str):
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    if is_separate_classifier_free_guidance_step():
+        _debugging_get_buffer(f"{name}_buffer_cfg")
+        return prune_context.get_buffer(f"{name}_buffer_cfg")
+    _debugging_get_buffer(f"{name}_buffer")
+    return prune_context.get_buffer(f"{name}_buffer")
+
+
+@torch.compiler.disable
+def remove_buffer(name: str):
+    prune_context = get_current_prune_context()
+    assert prune_context is not None, "prune_context must be set before"
+    prune_context.remove_buffer(name)
+
+
 @torch.compiler.disable
 def apply_hidden_states_residual(
     hidden_states: torch.Tensor,
@@ -506,7 +687,11 @@ def get_can_use_prune(
     if is_in_warmup():
         return False
 
-    pruned_steps = get_pruned_steps()
+    if not is_separate_classifier_free_guidance_step():
+        pruned_steps = get_pruned_steps()
+    else:
+        pruned_steps = get_cfg_pruned_steps()
+
     max_pruned_steps = get_max_pruned_steps()
     if max_pruned_steps >= 0 and (pruned_steps >= max_pruned_steps):
         if logger.isEnabledFor(logging.DEBUG):
@@ -521,16 +706,8 @@ def get_can_use_prune(
     if threshold <= 0.0:
         return False
 
-    downsample_factor = get_downsample_factor()
     prev_states_tensor = get_buffer(f"{name}")
 
-    if downsample_factor > 1:
-        states_tensor = states_tensor[..., ::downsample_factor]
-        states_tensor = states_tensor.contiguous()
-        if prev_states_tensor is not None:
-            prev_states_tensor = prev_states_tensor[..., ::downsample_factor]
-            prev_states_tensor = prev_states_tensor.contiguous()
-
     return prev_states_tensor is not None and are_two_tensors_similar(
         prev_states_tensor,
         states_tensor,
@@ -538,468 +715,3 @@ def get_can_use_prune(
         parallelized=parallelized,
         name=name,
     )
-
-
-class DBPrunedTransformerBlocks(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
-        self.pruned_blocks_step: int = 0
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        *args,
-        **kwargs,
-    ):
-        mark_step_begin()
-        self.pruned_blocks_step = 0
-        original_hidden_states = hidden_states
-
-        torch._dynamo.graph_break()
-        hidden_states, encoder_hidden_states = self.call_transformer_blocks(
-            hidden_states,
-            encoder_hidden_states,
-            *args,
-            **kwargs,
-        )
-
-        del original_hidden_states
-        torch._dynamo.graph_break()
-
-        add_pruned_block(self.pruned_blocks_step)
-        add_actual_block(self.num_transformer_blocks)
-        patch_pruned_stats(self.transformer)
-
-        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)
-            )
-        )
-
-    @property
-    @torch.compiler.disable
-    def num_transformer_blocks(self):
-        # Total number of transformer blocks, including single transformer blocks.
-        num_blocks = len(self.transformer_blocks)
-        if self.single_transformer_blocks is not None:
-            num_blocks += len(self.single_transformer_blocks)
-        return num_blocks
-
-    @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 _non_prune_blocks_ids(self):
-        # Never prune the first `Fn` and last `Bn` blocks.
-        num_blocks = self.num_transformer_blocks
-        Fn_compute_blocks_ = (
-            Fn_compute_blocks()
-            if Fn_compute_blocks() < num_blocks
-            else num_blocks
-        )
-        Fn_compute_blocks_ids = list(range(Fn_compute_blocks_))
-        Bn_compute_blocks_ = (
-            Bn_compute_blocks()
-            if Bn_compute_blocks() < num_blocks
-            else num_blocks
-        )
-        Bn_compute_blocks_ids = list(
-            range(
-                num_blocks - Bn_compute_blocks_,
-                num_blocks,
-            )
-        )
-        non_prune_blocks_ids = list(
-            set(
-                Fn_compute_blocks_ids
-                + Bn_compute_blocks_ids
-                + get_non_prune_blocks_ids()
-            )
-        )
-        non_prune_blocks_ids = [
-            d for d in non_prune_blocks_ids if d < num_blocks
-        ]
-        return sorted(non_prune_blocks_ids)
-
-    @torch.compiler.disable
-    def _compute_single_hidden_states_residual(
-        self,
-        single_hidden_states: torch.Tensor,
-        single_original_hidden_states: torch.Tensor,
-        # global original single hidden states
-        original_single_hidden_states: torch.Tensor,
-        original_single_encoder_hidden_states: torch.Tensor,
-    ):
-        single_hidden_states, single_encoder_hidden_states = (
-            self._split_single_hidden_states(
-                single_hidden_states,
-                original_single_hidden_states,
-                original_single_encoder_hidden_states,
-            )
-        )
-
-        single_original_hidden_states, single_original_encoder_hidden_states = (
-            self._split_single_hidden_states(
-                single_original_hidden_states,
-                original_single_hidden_states,
-                original_single_encoder_hidden_states,
-            )
-        )
-
-        single_hidden_states_residual = (
-            single_hidden_states - single_original_hidden_states
-        )
-        single_encoder_hidden_states_residual = (
-            single_encoder_hidden_states - single_original_encoder_hidden_states
-        )
-        return (
-            single_hidden_states_residual,
-            single_encoder_hidden_states_residual,
-        )
-
-    @torch.compiler.disable
-    def _split_single_hidden_states(
-        self,
-        single_hidden_states: torch.Tensor,
-        # global original single hidden states
-        original_single_hidden_states: torch.Tensor,
-        original_single_encoder_hidden_states: torch.Tensor,
-    ):
-        single_encoder_hidden_states, single_hidden_states = (
-            single_hidden_states.split(
-                [
-                    original_single_encoder_hidden_states.shape[1],
-                    single_hidden_states.shape[1]
-                    - original_single_encoder_hidden_states.shape[1],
-                ],
-                dim=1,
-            )
-        )
-        # Reshape the single_hidden_states and single_encoder_hidden_states
-        # to the original shape. This is necessary to ensure that the
-        # residuals are computed correctly.
-        single_hidden_states = (
-            single_hidden_states.reshape(-1)
-            .contiguous()
-            .reshape(original_single_hidden_states.shape)
-        )
-        single_encoder_hidden_states = (
-            single_encoder_hidden_states.reshape(-1)
-            .contiguous()
-            .reshape(original_single_encoder_hidden_states.shape)
-        )
-        return single_hidden_states, single_encoder_hidden_states
-
-    @torch.compiler.disable
-    def _should_update_residuals(self):
-        # Wrap for non compiled mode.
-        # Check if the current step is a multiple of
-        # the residual cache update interval.
-        return get_current_step() % residual_cache_update_interval() == 0
-
-    @torch.compiler.disable
-    def _get_can_use_prune(
-        self,
-        block_id: int,  # Block index in the transformer blocks
-        hidden_states: torch.Tensor,  # hidden_states or residual
-        name: str = "Bn_original",  # prev step name for single blocks
-    ):
-        # Wrap for non compiled mode.
-        can_use_prune = False
-        if block_id not in self._non_prune_blocks_ids():
-            can_use_prune = get_can_use_prune(
-                hidden_states,  # curr step
-                parallelized=self._is_parallelized(),
-                name=name,  # prev step
-            )
-        self.pruned_blocks_step += int(can_use_prune)
-        return can_use_prune
-
-    def _compute_or_prune_single_transformer_block(
-        self,
-        block_id: int,  # Block index in the transformer blocks
-        # Helper inputs for hidden states split and reshape
-        # Global original single hidden states
-        original_single_hidden_states: torch.Tensor,
-        original_single_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_transformer_blocks`
-        # block_id: global block index in the transformer blocks +
-        # single_transformer_blocks
-        can_use_prune = self._get_can_use_prune(
-            block_id,
-            hidden_states,  # hidden_states or residual
-            name=f"{block_id}_single_original",  # prev step
-        )
-
-        # Prune steps: Prune current block and reuse the cached
-        # residuals for hidden states approximate.
-        if can_use_prune:
-            single_original_hidden_states = hidden_states
-            (
-                single_original_hidden_states,
-                single_original_encoder_hidden_states,
-            ) = self._split_single_hidden_states(
-                single_original_hidden_states,
-                original_single_hidden_states,
-                original_single_encoder_hidden_states,
-            )
-            hidden_states, encoder_hidden_states = apply_hidden_states_residual(
-                single_original_hidden_states,
-                single_original_encoder_hidden_states,
-                name=f"{block_id}_single_residual",
-                encoder_name=f"{block_id}_single_encoder_residual",
-            )
-            hidden_states = torch.cat(
-                [encoder_hidden_states, hidden_states],
-                dim=1,
-            )
-            del single_original_hidden_states
-            del single_original_encoder_hidden_states
-
-        else:
-            # Normal steps: Compute the block and cache the residuals.
-            single_original_hidden_states = hidden_states
-            hidden_states = block(
-                hidden_states,
-                *args,
-                **kwargs,
-            )
-
-            # Save original_hidden_states for diff calculation.
-            # May not be necessary to update the hidden
-            # states and residuals each step?
-            if self._should_update_residuals():
-                # Cache residuals for the non-compute Bn blocks for
-                # subsequent prune steps.
-                single_hidden_states = hidden_states
-                (
-                    single_hidden_states_residual,
-                    single_encoder_hidden_states_residual,
-                ) = self._compute_single_hidden_states_residual(
-                    single_hidden_states,
-                    single_original_hidden_states,
-                    original_single_hidden_states,
-                    original_single_encoder_hidden_states,
-                )
-
-                set_buffer(
-                    f"{block_id}_single_original",
-                    single_original_hidden_states,
-                )
-
-                set_buffer(
-                    f"{block_id}_single_residual",
-                    single_hidden_states_residual,
-                )
-                set_buffer(
-                    f"{block_id}_single_encoder_residual",
-                    single_encoder_hidden_states_residual,
-                )
-
-                del single_hidden_states
-                del single_hidden_states_residual
-                del single_encoder_hidden_states_residual
-
-            del single_original_hidden_states
-
-        return hidden_states
-
-    def _compute_or_prune_transformer_block(
-        self,
-        block_id: int,  # Block index in the transformer blocks
-        # 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_transformer_blocks`
-        original_hidden_states = hidden_states
-        original_encoder_hidden_states = encoder_hidden_states
-
-        # block_id: global block index in the transformer blocks +
-        # single_transformer_blocks
-        can_use_prune = self._get_can_use_prune(
-            block_id,
-            hidden_states,  # hidden_states or residual
-            name=f"{block_id}_original",  # prev step
-        )
-
-        # Prune steps: Prune current block and reuse the cached
-        # residuals for hidden states approximate.
-        if can_use_prune:
-            hidden_states, encoder_hidden_states = apply_hidden_states_residual(
-                hidden_states,
-                encoder_hidden_states,
-                name=f"{block_id}_residual",
-                encoder_name=f"{block_id}_encoder_residual",
-            )
-        else:
-            # Normal steps: Compute the block and cache the residuals.
-            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,
-                    )
-
-            # Save original_hidden_states for diff calculation.
-            # May not be necessary to update the hidden
-            # states and residuals each step?
-            if self._should_update_residuals():
-                # Cache residuals for the non-compute Bn blocks for
-                # subsequent prune steps.
-                hidden_states_residual = hidden_states - original_hidden_states
-                encoder_hidden_states_residual = (
-                    encoder_hidden_states - original_encoder_hidden_states
-                )
-                set_buffer(
-                    f"{block_id}_original",
-                    original_hidden_states,
-                )
-
-                set_buffer(
-                    f"{block_id}_residual",
-                    hidden_states_residual,
-                )
-                set_buffer(
-                    f"{block_id}_encoder_residual",
-                    encoder_hidden_states_residual,
-                )
-                del hidden_states_residual
-                del encoder_hidden_states_residual
-
-        del original_hidden_states
-        del original_encoder_hidden_states
-
-        return hidden_states, encoder_hidden_states
-
-    def call_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
-
-        for i, block in enumerate(self.transformer_blocks):
-            hidden_states, encoder_hidden_states = (
-                self._compute_or_prune_transformer_block(
-                    i,
-                    block,
-                    hidden_states,
-                    encoder_hidden_states,
-                    *args,
-                    **kwargs,
-                )
-            )
-
-        if self.single_transformer_blocks is not None:
-            # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py#L380
-            if is_diffusers_at_least_0_3_5():
-                for j, block in enumerate(self.single_transformer_blocks):
-                    # NOTE: Reuse _compute_or_prune_transformer_block here.
-                    hidden_states, encoder_hidden_states = (
-                        self._compute_or_prune_transformer_block(
-                            j + len(self.transformer_blocks),
-                            block,
-                            hidden_states,
-                            encoder_hidden_states,
-                            *args,
-                            **kwargs,
-                        )
-                    )
-            else:
-                hidden_states = torch.cat(
-                    [encoder_hidden_states, hidden_states], dim=1
-                )
-                for j, block in enumerate(self.single_transformer_blocks):
-                    hidden_states = (
-                        self._compute_or_prune_single_transformer_block(
-                            j + len(self.transformer_blocks),
-                            original_hidden_states,
-                            original_encoder_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,
-                )
-
-        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_pruned_stats(
-    transformer,
-):
-    # Patch the pruned stats to the transformer, the pruned stats
-    # will be reset for each calling of pipe.__call__(**kwargs).
-    if transformer is None:
-        return
-
-    # TODO: Patch more pruned stats to the transformer
-    transformer._pruned_blocks = get_pruned_blocks()
-    transformer._pruned_steps = get_pruned_steps()
-    transformer._residual_diffs = get_residual_diffs()
-    transformer._actual_blocks = get_actual_blocks()