cache-dit 0.2.30__py3-none-any.whl → 0.2.32__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.30'
-__version_tuple__ = version_tuple = (0, 2, 30)
+__version__ = version = '0.2.32'
+__version_tuple__ = version_tuple = (0, 2, 32)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/block_adapters/__init__.py

@@ -254,7 +254,7 @@ def skyreelsv2_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("SD3")
+@BlockAdapterRegistry.register("StableDiffusion3")
 def sd3_adapter(pipe, **kwargs) -> BlockAdapter:
     from diffusers import SD3Transformer2DModel
 
@@ -501,7 +501,7 @@ def shape_adapter(pipe, **kwargs) -> BlockAdapter:
     )
 
 
-@BlockAdapterRegistry.register("HiDream", supported=True)
+@BlockAdapterRegistry.register("HiDream")
 def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
     # NOTE: Need to patch Transformer forward to fully support
     # double_stream_blocks and single_stream_blocks, namely, need
@@ -509,29 +509,32 @@ def hidream_adapter(pipe, **kwargs) -> BlockAdapter:
     # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py#L893
     # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py#L927
     from diffusers import HiDreamImageTransformer2DModel
+    from cache_dit.cache_factory.patch_functors import HiDreamPatchFunctor
 
     assert isinstance(pipe.transformer, HiDreamImageTransformer2DModel)
     return BlockAdapter(
         pipe=pipe,
         transformer=pipe.transformer,
         blocks=[
-            # pipe.transformer.double_stream_blocks,
+            pipe.transformer.double_stream_blocks,
             pipe.transformer.single_stream_blocks,
         ],
         forward_pattern=[
-            # ForwardPattern.Pattern_4,
+            ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_3,
         ],
-        # The type hint in diffusers is wrong
-        check_num_outputs=False,
+        patch_functor=HiDreamPatchFunctor(),
+        # NOTE: The type hint in diffusers is wrong
+        check_forward_pattern=True,
+        check_num_outputs=True,
         **kwargs,
     )
 
 
-@BlockAdapterRegistry.register("HunyuanDiT", supported=False)
+@BlockAdapterRegistry.register("HunyuanDiT")
 def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
-    # TODO: Patch Transformer forward
     from diffusers import HunyuanDiT2DModel, HunyuanDiT2DControlNetModel
+    from cache_dit.cache_factory.patch_functors import HunyuanDiTPatchFunctor
 
     assert isinstance(
         pipe.transformer,
@@ -542,14 +545,15 @@ def hunyuandit_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        patch_functor=HunyuanDiTPatchFunctor(),
         **kwargs,
     )
 
 
-@BlockAdapterRegistry.register("HunyuanDiTPAG", supported=False)
+@BlockAdapterRegistry.register("HunyuanDiTPAG")
 def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
-    # TODO: Patch Transformer forward
     from diffusers import HunyuanDiT2DModel
+    from cache_dit.cache_factory.patch_functors import HunyuanDiTPatchFunctor
 
     assert isinstance(pipe.transformer, HunyuanDiT2DModel)
     return BlockAdapter(
@@ -557,5 +561,6 @@ def hunyuanditpag_adapter(pipe, **kwargs) -> BlockAdapter:
         transformer=pipe.transformer,
         blocks=pipe.transformer.blocks,
         forward_pattern=ForwardPattern.Pattern_3,
+        patch_functor=HunyuanDiTPatchFunctor(),
         **kwargs,
     )

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -75,6 +75,7 @@ class BlockAdapter:
         List[List[ParamsModifier]],
     ] = None
 
+    check_forward_pattern: bool = True
     check_num_outputs: bool = False
 
     # Pipeline Level Flags
@@ -391,11 +392,20 @@ class BlockAdapter:
         forward_pattern: ForwardPattern,
         **kwargs,
     ) -> bool:
+
+        if not kwargs.get("check_forward_pattern", True):
+            return True
+
         assert (
             forward_pattern.Supported
             and forward_pattern in ForwardPattern.supported_patterns()
         ), f"Pattern {forward_pattern} is not support now!"
 
+        # NOTE: Special case for HiDreamBlock
+        if hasattr(block, "block"):
+            if isinstance(block.block, torch.nn.Module):
+                block = block.block
+
         forward_parameters = set(
             inspect.signature(block.forward).parameters.keys()
         )
@@ -425,6 +435,14 @@ class BlockAdapter:
         logging: bool = True,
         **kwargs,
     ) -> bool:
+
+        if not kwargs.get("check_forward_pattern", True):
+            if logging:
+                logger.warning(
+                    f"Skipped Forward Pattern Check: {forward_pattern}"
+                )
+            return True
+
         assert (
             forward_pattern.Supported
             and forward_pattern in ForwardPattern.supported_patterns()
@@ -531,6 +549,7 @@ class BlockAdapter:
                     blocks,
                     forward_pattern=forward_pattern,
                     check_num_outputs=adapter.check_num_outputs,
+                    check_forward_pattern=adapter.check_forward_pattern,
                 ), (
                     "No block forward pattern matched, "
                     f"supported lists: {ForwardPattern.supported_patterns()}"

cache_dit/cache_factory/cache_adapters.py

@@ -114,7 +114,7 @@ class CachedAdapter:
         **cache_context_kwargs,
     ):
         # Check cache_context_kwargs
-        if not cache_context_kwargs["enable_spearate_cfg"]:
+        if cache_context_kwargs["enable_spearate_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
@@ -131,6 +131,12 @@ class CachedAdapter:
                     f"register: {cache_context_kwargs['enable_spearate_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"Pipeline: {block_adapter.pipe.__class__.__name__}."
+            )
 
         if (
             cache_type := cache_context_kwargs.pop("cache_type", None)
@@ -345,6 +351,7 @@ class CachedAdapter:
                             block_adapter.blocks[i][j],
                             transformer=block_adapter.transformer[i],
                             forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
                             check_num_outputs=block_adapter.check_num_outputs,
                             # 1. Cache context configuration
                             cache_prefix=block_adapter.blocks_name[i][j],

cache_dit/cache_factory/cache_blocks/__init__.py

@@ -25,6 +25,7 @@ class CachedBlocks:
         transformer_blocks: torch.nn.ModuleList,
         transformer: torch.nn.Module = None,
         forward_pattern: ForwardPattern = None,
+        check_forward_pattern: bool = True,
         check_num_outputs: bool = True,
         # 1. Cache context configuration
         # 'transformer_blocks', 'blocks', 'single_transformer_blocks',
@@ -45,6 +46,7 @@ class CachedBlocks:
                 transformer_blocks,
                 transformer=transformer,
                 forward_pattern=forward_pattern,
+                check_forward_pattern=check_forward_pattern,
                 check_num_outputs=check_num_outputs,
                 # 1. Cache context configuration
                 cache_prefix=cache_prefix,
@@ -58,6 +60,7 @@ class CachedBlocks:
                 transformer_blocks,
                 transformer=transformer,
                 forward_pattern=forward_pattern,
+                check_forward_pattern=check_forward_pattern,
                 check_num_outputs=check_num_outputs,
                 # 1. Cache context configuration
                 cache_prefix=cache_prefix,

cache_dit/cache_factory/cache_blocks/pattern_base.py

@@ -25,6 +25,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         transformer_blocks: torch.nn.ModuleList,
         transformer: torch.nn.Module = None,
         forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
+        check_forward_pattern: bool = True,
         check_num_outputs: bool = True,
         # 1. Cache context configuration
         cache_prefix: str = None,  # maybe un-need.
@@ -38,6 +39,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         self.transformer = transformer
         self.transformer_blocks = transformer_blocks
         self.forward_pattern = forward_pattern
+        self.check_forward_pattern = check_forward_pattern
         self.check_num_outputs = check_num_outputs
         # 1. Cache context configuration
         self.cache_prefix = cache_prefix
@@ -52,6 +54,12 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         )
 
     def _check_forward_pattern(self):
+        if not self.check_forward_pattern:
+            logger.warning(
+                f"Skipped Forward Pattern Check: {self.forward_pattern}"
+            )
+            return
+
         assert (
             self.forward_pattern.Supported
             and self.forward_pattern in self._supported_patterns
@@ -59,6 +67,11 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
 
         if self.transformer_blocks is not None:
             for block in self.transformer_blocks:
+                # Special case for HiDreamBlock
+                if hasattr(block, "block"):
+                    if isinstance(block.block, torch.nn.Module):
+                        block = block.block
+
                 forward_parameters = set(
                     inspect.signature(block.forward).parameters.keys()
                 )
@@ -332,12 +345,19 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
 
         # compute hidden_states residual
         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
-        )
+
+        if (
+            encoder_hidden_states is not None
+            and original_encoder_hidden_states is not None
+        ):
+            encoder_hidden_states = encoder_hidden_states.contiguous()
+            encoder_hidden_states_residual = (
+                encoder_hidden_states - original_encoder_hidden_states
+            )
+        else:
+            encoder_hidden_states_residual = None
 
         return (
             hidden_states,
@@ -387,9 +407,16 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
                 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
-                )
+                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(

cache_dit/cache_factory/cache_interface.py

@@ -24,7 +24,7 @@ def enable_cache(
     max_continuous_cached_steps: int = -1,
     residual_diff_threshold: float = 0.08,
     # Cache CFG or not
-    enable_spearate_cfg: bool = False,
+    enable_spearate_cfg: bool | None = None,
     cfg_compute_first: bool = False,
     cfg_diff_compute_separate: bool = True,
     # Hybird TaylorSeer
@@ -70,7 +70,7 @@ def enable_cache(
         residual_diff_threshold (`float`, *required*, defaults to 0.08):
             he value of residual diff threshold, a higher value leads to faster performance at the
             cost of lower precision.
-        enable_spearate_cfg (`bool`, *required*,  defaults to False):
+        enable_spearate_cfg (`bool`, *required*,  defaults to None):
             Whether to do separate cfg or not, such as Wan 2.1, Qwen-Image. For model that fused CFG
             and non-CFG into single forward step, should set enable_spearate_cfg as False, for example:
             CogVideoX, HunyuanVideo, Mochi, etc.

cache_dit/cache_factory/patch_functors/__init__.py

@@ -3,3 +3,9 @@ from cache_dit.cache_factory.patch_functors.functor_flux import FluxPatchFunctor
 from cache_dit.cache_factory.patch_functors.functor_chroma import (
     ChromaPatchFunctor,
 )
+from cache_dit.cache_factory.patch_functors.functor_hidream import (
+    HiDreamPatchFunctor,
+)
+from cache_dit.cache_factory.patch_functors.functor_hunyuan_dit import (
+    HunyuanDiTPatchFunctor,
+)

cache_dit/cache_factory/patch_functors/functor_chroma.py

@@ -46,8 +46,10 @@ class ChromaPatchFunctor(PatchFunctor):
                     block.forward = __patch_single_forward__.__get__(block)
                     is_patched = True
 
+        cls_name = transformer.__class__.__name__
+
         if is_patched:
-            logger.warning("Patched Chroma for cache-dit.")
+            logger.warning(f"Patched {cls_name} for cache-dit.")
             assert not getattr(transformer, "_is_parallelized", False), (
                 "Please call `cache_dit.enable_cache` before Parallelize, "
                 "the __patch_transformer_forward__ will overwrite the "
@@ -59,7 +61,6 @@ class ChromaPatchFunctor(PatchFunctor):
 
         transformer._is_patched = is_patched  # True or False
 
-        cls_name = transformer.__class__.__name__
         logger.info(
             f"Applied {self.__class__.__name__} for {cls_name}, "
             f"Patch: {is_patched}."

cache_dit/cache_factory/patch_functors/functor_flux.py

@@ -47,8 +47,10 @@ class FluxPatchFunctor(PatchFunctor):
                     block.forward = __patch_single_forward__.__get__(block)
                     is_patched = True
 
+        cls_name = transformer.__class__.__name__
+
         if is_patched:
-            logger.warning("Patched Flux for cache-dit.")
+            logger.warning(f"Patched {cls_name} for cache-dit.")
             assert not getattr(transformer, "_is_parallelized", False), (
                 "Please call `cache_dit.enable_cache` before Parallelize, "
                 "the __patch_transformer_forward__ will overwrite the "
@@ -60,7 +62,6 @@ class FluxPatchFunctor(PatchFunctor):
 
         transformer._is_patched = is_patched  # True or False
 
-        cls_name = transformer.__class__.__name__
         logger.info(
             f"Applied {self.__class__.__name__} for {cls_name}, "
             f"Patch: {is_patched}."

cache_dit/cache_factory/patch_functors/functor_hidream.py

@@ -0,0 +1,412 @@
+import torch
+from typing import Tuple, Optional, Dict, Any, Union, List
+from diffusers import HiDreamImageTransformer2DModel
+from diffusers.models.transformers.transformer_hidream_image import (
+    HiDreamBlock,
+    HiDreamImageTransformerBlock,
+    HiDreamImageSingleTransformerBlock,
+    Transformer2DModelOutput,
+)
+from diffusers.utils import (
+    deprecate,
+    USE_PEFT_BACKEND,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from cache_dit.cache_factory.patch_functors.functor_base import (
+    PatchFunctor,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class HiDreamPatchFunctor(PatchFunctor):
+
+    def apply(
+        self,
+        transformer: HiDreamImageTransformer2DModel,
+        **kwargs,
+    ) -> HiDreamImageTransformer2DModel:
+        if hasattr(transformer, "_is_patched"):
+            return transformer
+
+        is_patched = False
+
+        _block_id = 0
+        for block in transformer.double_stream_blocks:
+            assert isinstance(block, HiDreamBlock)
+            block.forward = __patch_block_forward__.__get__(block)
+            # NOTE: Patch Inner block and block_id
+            _block = block.block
+            assert isinstance(_block, HiDreamImageTransformerBlock)
+            _block._block_id = _block_id
+            _block.forward = __patch_double_forward__.__get__(_block)
+            _block_id += 1
+
+        for block in transformer.single_stream_blocks:
+            assert isinstance(block, HiDreamBlock)
+            block.forward = __patch_block_forward__.__get__(block)
+            # NOTE: Patch Inner block and block_id
+            _block = block.block
+            assert isinstance(_block, HiDreamImageSingleTransformerBlock)
+            _block._block_id = _block_id
+            _block.forward = __patch_single_forward__.__get__(_block)
+            _block_id += 1
+
+        is_patched = True
+        cls_name = transformer.__class__.__name__
+
+        if is_patched:
+            logger.warning(f"Patched {cls_name} for cache-dit.")
+            assert not getattr(transformer, "_is_parallelized", False), (
+                "Please call `cache_dit.enable_cache` before Parallelize, "
+                "the __patch_transformer_forward__ will overwrite the "
+                "parallized forward and cause a downgrade of performance."
+            )
+            transformer.forward = __patch_transformer_forward__.__get__(
+                transformer
+            )
+
+        transformer._is_patched = is_patched  # True or False
+
+        logger.info(
+            f"Applied {self.__class__.__name__} for {cls_name}, "
+            f"Patch: {is_patched}."
+        )
+
+        return transformer
+
+
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py
+def __patch_double_forward__(
+    self: HiDreamImageTransformerBlock,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,  # initial_encoder_hidden_states
+    hidden_states_masks: Optional[torch.Tensor] = None,
+    temb: Optional[torch.Tensor] = None,
+    image_rotary_emb: torch.Tensor = None,
+    llama31_encoder_hidden_states: List[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # Assume block_id was patched in transformer forward:
+    # for i, block in enumerate(blocks): block._block_id = i;
+    block_id = self._block_id
+    initial_encoder_hidden_states_seq_len = encoder_hidden_states.shape[1]
+    cur_llama31_encoder_hidden_states = llama31_encoder_hidden_states[block_id]
+    cur_encoder_hidden_states = torch.cat(
+        [encoder_hidden_states, cur_llama31_encoder_hidden_states],
+        dim=1,
+    )
+    encoder_hidden_states = cur_encoder_hidden_states
+
+    wtype = hidden_states.dtype
+    (
+        shift_msa_i,
+        scale_msa_i,
+        gate_msa_i,
+        shift_mlp_i,
+        scale_mlp_i,
+        gate_mlp_i,
+        shift_msa_t,
+        scale_msa_t,
+        gate_msa_t,
+        shift_mlp_t,
+        scale_mlp_t,
+        gate_mlp_t,
+    ) = self.adaLN_modulation(temb)[:, None].chunk(12, dim=-1)
+
+    # 1. MM-Attention
+    norm_hidden_states = self.norm1_i(hidden_states).to(dtype=wtype)
+    norm_hidden_states = norm_hidden_states * (1 + scale_msa_i) + shift_msa_i
+    norm_encoder_hidden_states = self.norm1_t(encoder_hidden_states).to(
+        dtype=wtype
+    )
+    norm_encoder_hidden_states = (
+        norm_encoder_hidden_states * (1 + scale_msa_t) + shift_msa_t
+    )
+
+    attn_output_i, attn_output_t = self.attn1(
+        norm_hidden_states,
+        hidden_states_masks,
+        norm_encoder_hidden_states,
+        image_rotary_emb=image_rotary_emb,
+    )
+
+    hidden_states = gate_msa_i * attn_output_i + hidden_states
+    encoder_hidden_states = gate_msa_t * attn_output_t + encoder_hidden_states
+
+    # 2. Feed-forward
+    norm_hidden_states = self.norm3_i(hidden_states).to(dtype=wtype)
+    norm_hidden_states = norm_hidden_states * (1 + scale_mlp_i) + shift_mlp_i
+    norm_encoder_hidden_states = self.norm3_t(encoder_hidden_states).to(
+        dtype=wtype
+    )
+    norm_encoder_hidden_states = (
+        norm_encoder_hidden_states * (1 + scale_mlp_t) + shift_mlp_t
+    )
+
+    ff_output_i = gate_mlp_i * self.ff_i(norm_hidden_states)
+    ff_output_t = gate_mlp_t * self.ff_t(norm_encoder_hidden_states)
+    hidden_states = ff_output_i + hidden_states
+    encoder_hidden_states = ff_output_t + encoder_hidden_states
+
+    initial_encoder_hidden_states = encoder_hidden_states[
+        :, :initial_encoder_hidden_states_seq_len
+    ]
+
+    return hidden_states, initial_encoder_hidden_states
+
+
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py
+def __patch_single_forward__(
+    self: HiDreamImageSingleTransformerBlock,
+    hidden_states: torch.Tensor,
+    hidden_states_masks: Optional[torch.Tensor] = None,
+    temb: Optional[torch.Tensor] = None,
+    image_rotary_emb: torch.Tensor = None,
+    llama31_encoder_hidden_states: List[torch.Tensor] = None,
+) -> torch.Tensor:
+    # Assume block_id was patched in transformer forward:
+    # for i, block in enumerate(blocks): block._block_id = i;
+    block_id = self._block_id
+    hidden_states_seq_len = hidden_states.shape[1]
+    cur_llama31_encoder_hidden_states = llama31_encoder_hidden_states[block_id]
+    hidden_states = torch.cat(
+        [hidden_states, cur_llama31_encoder_hidden_states], dim=1
+    )
+
+    wtype = hidden_states.dtype
+    (
+        shift_msa_i,
+        scale_msa_i,
+        gate_msa_i,
+        shift_mlp_i,
+        scale_mlp_i,
+        gate_mlp_i,
+    ) = self.adaLN_modulation(temb)[:, None].chunk(6, dim=-1)
+
+    # 1. MM-Attention
+    norm_hidden_states = self.norm1_i(hidden_states).to(dtype=wtype)
+    norm_hidden_states = norm_hidden_states * (1 + scale_msa_i) + shift_msa_i
+    attn_output_i = self.attn1(
+        norm_hidden_states,
+        hidden_states_masks,
+        image_rotary_emb=image_rotary_emb,
+    )
+    hidden_states = gate_msa_i * attn_output_i + hidden_states
+
+    # 2. Feed-forward
+    norm_hidden_states = self.norm3_i(hidden_states).to(dtype=wtype)
+    norm_hidden_states = norm_hidden_states * (1 + scale_mlp_i) + shift_mlp_i
+    ff_output_i = gate_mlp_i * self.ff_i(norm_hidden_states.to(dtype=wtype))
+    hidden_states = ff_output_i + hidden_states
+
+    hidden_states = hidden_states[:, :hidden_states_seq_len]
+
+    return hidden_states
+
+
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py
+def __patch_block_forward__(
+    self: HiDreamBlock,
+    hidden_states: torch.Tensor,
+    *args,
+    **kwargs,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    return self.block(hidden_states, *args, **kwargs)
+
+
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_hidream_image.py
+def __patch_transformer_forward__(
+    self: HiDreamImageTransformer2DModel,
+    hidden_states: torch.Tensor,
+    timesteps: torch.LongTensor = None,
+    encoder_hidden_states_t5: torch.Tensor = None,
+    encoder_hidden_states_llama3: torch.Tensor = None,
+    pooled_embeds: torch.Tensor = None,
+    img_ids: Optional[torch.Tensor] = None,
+    img_sizes: Optional[List[Tuple[int, int]]] = None,
+    hidden_states_masks: Optional[torch.Tensor] = None,
+    attention_kwargs: Optional[Dict[str, Any]] = None,
+    return_dict: bool = True,
+    **kwargs,
+) -> Union[torch.Tensor, Transformer2DModelOutput]:
+    encoder_hidden_states = kwargs.get("encoder_hidden_states", None)
+
+    if encoder_hidden_states is not None:
+        deprecation_message = "The `encoder_hidden_states` argument is deprecated. Please use `encoder_hidden_states_t5` and `encoder_hidden_states_llama3` instead."
+        deprecate("encoder_hidden_states", "0.35.0", deprecation_message)
+        encoder_hidden_states_t5 = encoder_hidden_states[0]
+        encoder_hidden_states_llama3 = encoder_hidden_states[1]
+
+    if (
+        img_ids is not None
+        and img_sizes is not None
+        and hidden_states_masks is None
+    ):
+        deprecation_message = "Passing `img_ids` and `img_sizes` with unpachified `hidden_states` is deprecated and will be ignored."
+        deprecate("img_ids", "0.35.0", deprecation_message)
+
+    if hidden_states_masks is not None and (
+        img_ids is None or img_sizes is None
+    ):
+        raise ValueError(
+            "if `hidden_states_masks` is passed, `img_ids` and `img_sizes` must also be passed."
+        )
+    elif hidden_states_masks is not None and hidden_states.ndim != 3:
+        raise ValueError(
+            "if `hidden_states_masks` is passed, `hidden_states` must be a 3D tensors with shape (batch_size, patch_height * patch_width, patch_size * patch_size * channels)"
+        )
+
+    if attention_kwargs is not None:
+        attention_kwargs = attention_kwargs.copy()
+        lora_scale = attention_kwargs.pop("scale", 1.0)
+    else:
+        lora_scale = 1.0
+
+    if USE_PEFT_BACKEND:
+        # weight the lora layers by setting `lora_scale` for each PEFT layer
+        scale_lora_layers(self, lora_scale)
+    else:
+        if (
+            attention_kwargs is not None
+            and attention_kwargs.get("scale", None) is not None
+        ):
+            logger.warning(
+                "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+            )
+
+    # spatial forward
+    batch_size = hidden_states.shape[0]
+    hidden_states_type = hidden_states.dtype
+
+    # Patchify the input
+    if hidden_states_masks is None:
+        hidden_states, hidden_states_masks, img_sizes, img_ids = self.patchify(
+            hidden_states
+        )
+
+    # Embed the hidden states
+    hidden_states = self.x_embedder(hidden_states)
+
+    # 0. time
+    timesteps = self.t_embedder(timesteps, hidden_states_type)
+    p_embedder = self.p_embedder(pooled_embeds)
+    temb = timesteps + p_embedder
+
+    encoder_hidden_states = [
+        encoder_hidden_states_llama3[k] for k in self.config.llama_layers
+    ]
+
+    if self.caption_projection is not None:
+        new_encoder_hidden_states = []
+        for i, enc_hidden_state in enumerate(encoder_hidden_states):
+            enc_hidden_state = self.caption_projection[i](enc_hidden_state)
+            enc_hidden_state = enc_hidden_state.view(
+                batch_size, -1, hidden_states.shape[-1]
+            )
+            new_encoder_hidden_states.append(enc_hidden_state)
+        encoder_hidden_states = new_encoder_hidden_states
+        encoder_hidden_states_t5 = self.caption_projection[-1](
+            encoder_hidden_states_t5
+        )
+        encoder_hidden_states_t5 = encoder_hidden_states_t5.view(
+            batch_size, -1, hidden_states.shape[-1]
+        )
+        encoder_hidden_states.append(encoder_hidden_states_t5)
+
+    txt_ids = torch.zeros(
+        batch_size,
+        encoder_hidden_states[-1].shape[1]
+        + encoder_hidden_states[-2].shape[1]
+        + encoder_hidden_states[0].shape[1],
+        3,
+        device=img_ids.device,
+        dtype=img_ids.dtype,
+    )
+    ids = torch.cat((img_ids, txt_ids), dim=1)
+    image_rotary_emb = self.pe_embedder(ids)
+
+    # 2. Blocks
+    # NOTE: block_id is no-need anymore.
+    initial_encoder_hidden_states = torch.cat(
+        [encoder_hidden_states[-1], encoder_hidden_states[-2]], dim=1
+    )
+    llama31_encoder_hidden_states = encoder_hidden_states
+    for bid, block in enumerate(self.double_stream_blocks):
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states, initial_encoder_hidden_states = (
+                self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    initial_encoder_hidden_states,
+                    hidden_states_masks,
+                    temb,
+                    image_rotary_emb,
+                    llama31_encoder_hidden_states,
+                )
+            )
+        else:
+            hidden_states, initial_encoder_hidden_states = block(
+                hidden_states,
+                initial_encoder_hidden_states,  # encoder_hidden_states
+                hidden_states_masks=hidden_states_masks,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                llama31_encoder_hidden_states=llama31_encoder_hidden_states,
+            )
+
+    image_tokens_seq_len = hidden_states.shape[1]
+    hidden_states = torch.cat(
+        [hidden_states, initial_encoder_hidden_states], dim=1
+    )
+    if hidden_states_masks is not None:
+        # NOTE: Patched
+        cur_llama31_encoder_hidden_states = llama31_encoder_hidden_states[
+            self.double_stream_blocks[-1].block._block_id
+        ]
+        encoder_attention_mask_ones = torch.ones(
+            (
+                batch_size,
+                initial_encoder_hidden_states.shape[1]
+                + cur_llama31_encoder_hidden_states.shape[1],
+            ),
+            device=hidden_states_masks.device,
+            dtype=hidden_states_masks.dtype,
+        )
+        hidden_states_masks = torch.cat(
+            [hidden_states_masks, encoder_attention_mask_ones], dim=1
+        )
+
+    for bid, block in enumerate(self.single_stream_blocks):
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(
+                block,
+                hidden_states,
+                hidden_states_masks,
+                temb,
+                image_rotary_emb,
+                llama31_encoder_hidden_states,
+            )
+        else:
+            hidden_states = block(
+                hidden_states,
+                hidden_states_masks=hidden_states_masks,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                llama31_encoder_hidden_states=llama31_encoder_hidden_states,
+            )
+
+    hidden_states = hidden_states[:, :image_tokens_seq_len, ...]
+    output = self.final_layer(hidden_states, temb)
+    output = self.unpatchify(output, img_sizes, self.training)
+    if hidden_states_masks is not None:
+        hidden_states_masks = hidden_states_masks[:, :image_tokens_seq_len]
+
+    if USE_PEFT_BACKEND:
+        # remove `lora_scale` from each PEFT layer
+        unscale_lora_layers(self, lora_scale)
+
+    if not return_dict:
+        return (output,)
+    return Transformer2DModelOutput(sample=output)

cache_dit/cache_factory/patch_functors/functor_hunyuan_dit.py

@@ -0,0 +1,213 @@
+import torch
+from typing import Optional, Union, List
+from diffusers import HunyuanDiT2DModel
+from diffusers.models.transformers.hunyuan_transformer_2d import (
+    HunyuanDiTBlock,
+    Transformer2DModelOutput,
+)
+from cache_dit.cache_factory.patch_functors.functor_base import (
+    PatchFunctor,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class HunyuanDiTPatchFunctor(PatchFunctor):
+
+    def apply(
+        self,
+        transformer: HunyuanDiT2DModel,
+        **kwargs,
+    ) -> HunyuanDiT2DModel:
+        if hasattr(transformer, "_is_patched"):
+            return transformer
+
+        is_patched = False
+
+        num_layers = transformer.config.num_layers
+        layer_id = 0
+        for block in transformer.blocks:
+            assert isinstance(block, HunyuanDiTBlock)
+            block._num_layers = num_layers
+            block._layer_id = layer_id
+            block.forward = __patch_block_forward__.__get__(block)
+            layer_id += 1
+
+        is_patched = True
+
+        cls_name = transformer.__class__.__name__
+
+        if is_patched:
+            logger.warning(f"Patched {cls_name} for cache-dit.")
+            assert not getattr(transformer, "_is_parallelized", False), (
+                "Please call `cache_dit.enable_cache` before Parallelize, "
+                "the __patch_transformer_forward__ will overwrite the "
+                "parallized forward and cause a downgrade of performance."
+            )
+            transformer.forward = __patch_transformer_forward__.__get__(
+                transformer
+            )
+
+        transformer._is_patched = is_patched  # True or False
+
+        logger.info(
+            f"Applied {self.__class__.__name__} for {cls_name}, "
+            f"Patch: {is_patched}."
+        )
+
+        return transformer
+
+
+def __patch_block_forward__(
+    self: HunyuanDiTBlock,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: Optional[torch.Tensor] = None,
+    temb: Optional[torch.Tensor] = None,
+    image_rotary_emb: torch.Tensor = None,
+    controlnet_block_samples: torch.Tensor = None,
+    skips: List[torch.Tensor] = [],
+) -> torch.Tensor:
+    # Notice that normalization is always applied before the real computation in the following blocks.
+    # 0. Long Skip Connection
+    num_layers = self._num_layers
+    layer_id = self._layer_id
+
+    if layer_id > num_layers // 2:
+        if controlnet_block_samples is not None:
+            skip = skips.pop() + controlnet_block_samples.pop()
+        else:
+            skip = skips.pop()
+    else:
+        skip = None
+
+    if self.skip_linear is not None:
+        cat = torch.cat([hidden_states, skip], dim=-1)
+        cat = self.skip_norm(cat)
+        hidden_states = self.skip_linear(cat)
+
+    # 1. Self-Attention
+    norm_hidden_states = self.norm1(
+        hidden_states, temb
+    )  # checked: self.norm1 is correct
+    attn_output = self.attn1(
+        norm_hidden_states,
+        image_rotary_emb=image_rotary_emb,
+    )
+    hidden_states = hidden_states + attn_output
+
+    # 2. Cross-Attention
+    hidden_states = hidden_states + self.attn2(
+        self.norm2(hidden_states),
+        encoder_hidden_states=encoder_hidden_states,
+        image_rotary_emb=image_rotary_emb,
+    )
+
+    # FFN Layer
+    mlp_inputs = self.norm3(hidden_states)
+    hidden_states = hidden_states + self.ff(mlp_inputs)
+
+    if layer_id < (num_layers // 2 - 1):
+        skips.append(hidden_states)
+
+    return hidden_states
+
+
+def __patch_transformer_forward__(
+    self: HunyuanDiT2DModel,
+    hidden_states,
+    timestep,
+    encoder_hidden_states=None,
+    text_embedding_mask=None,
+    encoder_hidden_states_t5=None,
+    text_embedding_mask_t5=None,
+    image_meta_size=None,
+    style=None,
+    image_rotary_emb=None,
+    controlnet_block_samples=None,
+    return_dict=True,
+) -> Union[torch.Tensor, Transformer2DModelOutput]:
+    height, width = hidden_states.shape[-2:]
+
+    hidden_states = self.pos_embed(hidden_states)
+
+    temb = self.time_extra_emb(
+        timestep,
+        encoder_hidden_states_t5,
+        image_meta_size,
+        style,
+        hidden_dtype=timestep.dtype,
+    )  # [B, D]
+
+    # text projection
+    batch_size, sequence_length, _ = encoder_hidden_states_t5.shape
+    encoder_hidden_states_t5 = self.text_embedder(
+        encoder_hidden_states_t5.view(-1, encoder_hidden_states_t5.shape[-1])
+    )
+    encoder_hidden_states_t5 = encoder_hidden_states_t5.view(
+        batch_size, sequence_length, -1
+    )
+
+    encoder_hidden_states = torch.cat(
+        [encoder_hidden_states, encoder_hidden_states_t5], dim=1
+    )
+    text_embedding_mask = torch.cat(
+        [text_embedding_mask, text_embedding_mask_t5], dim=-1
+    )
+    text_embedding_mask = text_embedding_mask.unsqueeze(2).bool()
+
+    encoder_hidden_states = torch.where(
+        text_embedding_mask, encoder_hidden_states, self.text_embedding_padding
+    )
+
+    skips = []
+    for layer, block in enumerate(self.blocks):
+        hidden_states = block(
+            hidden_states,
+            temb=temb,
+            encoder_hidden_states=encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            controlnet_block_samples=controlnet_block_samples,
+            skips=skips,
+        )  # (N, L, D)
+
+    if (
+        controlnet_block_samples is not None
+        and len(controlnet_block_samples) != 0
+    ):
+        raise ValueError(
+            "The number of controls is not equal to the number of skip connections."
+        )
+
+    # final layer
+    hidden_states = self.norm_out(hidden_states, temb.to(torch.float32))
+    hidden_states = self.proj_out(hidden_states)
+    # (N, L, patch_size ** 2 * out_channels)
+
+    # unpatchify: (N, out_channels, H, W)
+    patch_size = self.pos_embed.patch_size
+    height = height // patch_size
+    width = width // patch_size
+
+    hidden_states = hidden_states.reshape(
+        shape=(
+            hidden_states.shape[0],
+            height,
+            width,
+            patch_size,
+            patch_size,
+            self.out_channels,
+        )
+    )
+    hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+    output = hidden_states.reshape(
+        shape=(
+            hidden_states.shape[0],
+            self.out_channels,
+            height * patch_size,
+            width * patch_size,
+        )
+    )
+    if not return_dict:
+        return (output,)
+    return Transformer2DModelOutput(sample=output)

{cache_dit-0.2.30.dist-info → cache_dit-0.2.32.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.30
+Version: 0.2.32
 Summary: 🤗 A Unified and Training-free Cache Acceleration Toolbox for Diffusion Transformers
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
@@ -59,29 +59,37 @@ Dynamic: requires-python
   🔥<b><a href="#unified">Unified Cache APIs</a> | <a href="#dbcache">DBCache</a> | <a href="#taylorseer">Hybrid TaylorSeer</a> | <a href="#cfg">Hybrid Cache CFG</a></b>🔥
   </p>
   <p align="center">
-  🎉Now, <b>cache-dit</b> covers <b>mainstream</b> Diffusers' <b>DiT-based</b> Pipelines🎉<br>
-  🔥<b><a href="#supported">Qwen-Image</a> | <a href="#supported">FLUX.1</a> | <a href="#supported">Wan 2.1/2.2</a> | <a href="#supported"> ... </a> | <a href="#supported">CogVideoX</a></b>🔥
+    🎉Now, <b>cache-dit</b> covers <b>most</b> mainstream Diffusers' <b>DiT</b> Pipelines🎉<br>
+    🔥<a href="#supported">Qwen-Image</a> | <a href="#supported">FLUX.1</a> | <a href="#supported">Qwen-Image-Lightning</a> | <a href="#supported"> Wan 2.1/2.2 </a>🔥<br>
+    🔥<a href="#supported">HunyuanVideo</a> | <a href="#supported">HunyuanDiT</a> | <a href="#supported">HiDream</a> | <a href="#supported">Mochi</a> | <a href="#supported">CogVideoX 1/1.5</a>🔥<br>
+    🔥<a href="#supported">CogView3Plus</a> | <a href="#supported">CogView4</a> | <a href="#supported">Chroma</a> | <a href="#supported"> LTXVideo </a> | <a href="#supported">PixArt</a>🔥<br>
+    🔥<a href="#supported">Cosmos</a> | <a href="#supported">SkyReelsV2</a> | <a href="#supported">VisualCloze</a> | <a href="#supported"> OmniGen </a> | <a href="#supported">Lumina 1/2</a>🔥<br>
+    🔥<a href="#supported">Allegro</a> | <a href="#supported">EasyAnimate</a> | <a href="#supported">SD 3/3.5</a> | <a href="#supported"> ... </a> | <a href="#supported">DiT-XL</a>🔥
   </p>
 </div>
 <div align='center'>
-  <img src=./assets/gifs/wan2.2.C0_Q0_NONE.gif width=160px>
-  <img src=./assets/gifs/wan2.2.C1_Q0_DBCACHE_F1B0_W2M8MC2_T1O2_R0.08.gif width=160px>
-  <img src=./assets/gifs/wan2.2.C1_Q1_fp8_w8a8_dq_DBCACHE_F1B0_W2M8MC2_T1O2_R0.08.gif width=160px>
-  <p><b>🔥Wan2.2 MoE</b> Baseline | <b><a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:~2.0x↑🎉</b> | +FP8 DQ:<b>~2.4x↑🎉</b></p>
-  <img src=./assets/qwen-image.C0_Q0_NONE.png width=160px>
-  <img src=./assets/qwen-image.C1_Q0_DBCACHE_F8B0_W8M0MC0_T1O4_R0.12_S23.png width=160px>
-  <img src=./assets/qwen-image.C1_Q1_fp8_w8a8_dq_DBCACHE_F8B0_W8M0MC0_T1O4_R0.12_S18.png width=160px>
-  <p><b>🔥Qwen-Image</b> Baseline | <b><a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:~1.8x↑🎉</b> | +FP8 DQ:<b>~2.2x↑🎉</b><br>♥️ Please consider to leave a <b>⭐️ Star</b> to support us ~ ♥️</p>
-  </p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/wan2.2.C0_Q0_NONE.gif width=160px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/wan2.2.C1_Q0_DBCACHE_F1B0_W2M8MC2_T1O2_R0.08.gif width=160px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/gifs/wan2.2.C1_Q1_fp8_w8a8_dq_DBCACHE_F1B0_W2M8MC2_T1O2_R0.08.gif width=160px>
+  <p><b>🔥Wan2.2 MoE</b> | <b><a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:~2.0x↑🎉</b> | +FP8 DQ:<b>~2.4x↑🎉</b></p>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image.C0_Q0_NONE.png width=160px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image.C1_Q0_DBCACHE_F8B0_W8M0MC0_T1O4_R0.12_S23.png width=160px>
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/qwen-image.C1_Q1_fp8_w8a8_dq_DBCACHE_F8B0_W8M0MC0_T1O4_R0.12_S18.png width=160px>
+  <p><b>🔥Qwen-Image</b> | <b><a href="https://github.com/vipshop/cache-dit">+cache-dit</a>:~1.8x↑🎉</b> | +FP8 DQ:<b>~2.2x↑🎉</b></p>
+  <img src=./assets/qwen-image-lightning.4steps.C0_L1_Q0_NONE.png width=200px>
+  <img src=./assets/qwen-image-lightning.4steps.C0_L1_Q0_DBCACHE_F16B16_W2M1MC1_T0O2_R0.9_S1.png width=200px>
+  <p><b>🔥Qwen-Image-Lightning</b> 4 steps | <b><a href="https://github.com/vipshop/cache-dit">+cache-dit</a></b> 3.5 steps:<b>~1.14x↑🎉</b>
+  <br>♥️ Please consider to leave a <b>⭐️ Star</b> to support us ~ ♥️</p>
 </div>
 
 ## 🔥News  
 
+- [2025-09-08] 🔥[**Qwen-Image-Lightning**](./examples/pipeline/run_qwen_image_lightning.py) **7.1/3.5 steps🎉** inference with **[DBCache: F16B16](https://github.com/vipshop/cache-dit)**.
 - [2025-09-03] 🎉[**Wan2.2-MoE**](https://github.com/Wan-Video) **2.4x↑🎉** speedup! Please refer to [run_wan_2.2.py](./examples/pipeline/run_wan_2.2.py) as an example.
 - [2025-08-19] 🔥[**Qwen-Image-Edit**](https://github.com/QwenLM/Qwen-Image) **2x↑🎉** speedup! Check the example: [run_qwen_image_edit.py](./examples/pipeline/run_qwen_image_edit.py).
 - [2025-08-12] 🎉First caching mechanism in [QwenLM/Qwen-Image](https://github.com/QwenLM/Qwen-Image) with **[cache-dit](https://github.com/vipshop/cache-dit)**, check this [PR](https://github.com/QwenLM/Qwen-Image/pull/61). 
 - [2025-08-11] 🔥[**Qwen-Image**](https://github.com/QwenLM/Qwen-Image) **1.8x↑🎉** speedup! Please refer to [run_qwen_image.py](./examples/pipeline/run_qwen_image.py) as an example.
-- [2025-07-13] 🎉[**FLUX.1-Dev**](https://github.com/xlite-dev/flux-faster) **3.3x↑🎉** speedup! NVIDIA L20 with **[cache-dit](https://github.com/vipshop/cache-dit)** + `compile + FP8 DQ`.
+- [2025-07-13] 🎉[**FLUX.1-dev**](https://github.com/xlite-dev/flux-faster) **3.3x↑🎉** speedup! NVIDIA L20 with **[cache-dit](https://github.com/vipshop/cache-dit)** + **compile + FP8 DQ**.
 
 <details>
 <summary> Previous News </summary>  
@@ -131,6 +139,7 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 
 Currently, **cache-dit** library supports almost **Any** Diffusion Transformers (with **Transformer Blocks** that match the specific Input and Output **patterns**). Please check [🎉Unified Cache APIs](#unified) for more details. Here are just some of the tested models listed:
 
+- [🚀Qwen-Image-Lightning](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀Qwen-Image-Edit](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀Qwen-Image](https://github.com/vipshop/cache-dit/raw/main/examples)  
 - [🚀FLUX.1-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
@@ -141,7 +150,13 @@ Currently, **cache-dit** library supports almost **Any** Diffusion Transformers
 - [🚀Wan2.2-T2V](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀Wan2.1-T2V](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀Wan2.1-FLF2V](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀mochi-1-preview](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀HunyuanVideo](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀HunyuanDiT](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀HiDream-I1-Full](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀PixArt-Alpha](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀PixArt-Sigma](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀SD-3/3.5](https://github.com/vipshop/cache-dit/raw/main/examples)
 
 </details>
 

{cache_dit-0.2.30.dist-info → cache_dit-0.2.32.dist-info}/RECORD

@@ -1,30 +1,32 @@
 cache_dit/__init__.py,sha256=kX9V-FegZG4c8LMwI4PTmMqH794MEW0pzDArdhC0cJw,1241
-cache_dit/_version.py,sha256=6uKAYeE03adIcUS0SDwp52AaQx0KO8z_-07D_lPHrz8,706
+cache_dit/_version.py,sha256=J0YTFDgdG9rY1Xk5pUbWWGgbT2rbSasvUHcntxayVtA,706
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
 cache_dit/utils.py,sha256=WK7eqgH6gCYNHXNLmWyxBDU0XSHTPg7CfOcyXlGXBqE,10510
 cache_dit/cache_factory/.gitignore,sha256=5Cb-qT9wsTUoMJ7vACDF7ZcLpAXhi5v-xdcWSRit988,23
 cache_dit/cache_factory/__init__.py,sha256=Iw6-iJLFbdzCsIDZXXOw371L-HPmoeZO_P9a3sDjP5s,1103
-cache_dit/cache_factory/cache_adapters.py,sha256=TA_0mEHMdSQDrt4rYASeX4-BD8pJOznSJfMV3hkrGuk,17851
-cache_dit/cache_factory/cache_interface.py,sha256=y1nY6R3MucRmAnG2UJRI_tIKrRk27FktGWLbfckf3zE,8543
+cache_dit/cache_factory/cache_adapters.py,sha256=dmNX68nBD52HtQvHnNAuSn1zjDWrQdycD0qXy-w-mwc,18212
+cache_dit/cache_factory/cache_interface.py,sha256=LpyCy-tQ_GcTRAYLpMMf9hFVIktABHI6CObn5Ll8bMw,8548
 cache_dit/cache_factory/cache_types.py,sha256=ooukxQRG55uTLmaZ0SKw6gIeY6SQHhMxkbv55uj2Sqk,991
 cache_dit/cache_factory/forward_pattern.py,sha256=FumlCuZ-TSmSYH0hGBHctSJ-oGLCftdZjLygqhsmdR4,2258
 cache_dit/cache_factory/utils.py,sha256=XkVM9AXcB9zYq8-S8QKAsGz80r3tA6U3lBNGDGeHOe4,1871
-cache_dit/cache_factory/block_adapters/__init__.py,sha256=EA-4mEVy-JJ5vRDo6C3nJIOXu0ZDNc6FQ-ZLAKHDtB0,17251
-cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=OrKhuNdcGBCgSsPchdf4h32Ad-bQVUXNigMhPJ4cCvk,21069
+cache_dit/cache_factory/block_adapters/__init__.py,sha256=OZM5vJwmQIkoIwVmMxKXiHqKvs31NyAva1Z91C_ko3w,17547
+cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=EQBiJYyoInKU1ND69wTm7M0n5Ja4I8QW01SgRpBjSn8,21671
 cache_dit/cache_factory/block_adapters/block_registers.py,sha256=ZeN2wGPmuf2u3puSsBx8x-rl3wRo8-cWcuWNcrssVfA,2553
-cache_dit/cache_factory/cache_blocks/__init__.py,sha256=OWjnpJxA8EJVoRzuyb5miuiRphUFj831-bbtWsTDjnM,2750
+cache_dit/cache_factory/cache_blocks/__init__.py,sha256=08Ox7kD05lkRKCOsVTdEZeKAWBheqpxfrAT1Nz7eclI,2916
 cache_dit/cache_factory/cache_blocks/pattern_0_1_2.py,sha256=ElMps6_7uI74tSF9GDR_dEI0bZEhdzcepM29xFWnYo8,428
 cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py,sha256=nf2f5wdxp6tfq9AhFyMyBeKiZfxh63WG1g8q-c2BBSg,10182
-cache_dit/cache_factory/cache_blocks/pattern_base.py,sha256=XSDy3hsaKbAZPGZY92YgGA0qLgjQyIX8irQkb2R5T2c,20331
+cache_dit/cache_factory/cache_blocks/pattern_base.py,sha256=f1ojREQcDoBtDG3dzl8t1g_Vru8140LVDRPWlY-kAXw,21311
 cache_dit/cache_factory/cache_blocks/utils.py,sha256=dGOC1tMMOvcbvEgx44eTESKn_jsv-0RZ3tRHPa3wmQ4,1315
 cache_dit/cache_factory/cache_contexts/__init__.py,sha256=rqnJ5__zqnpVHK5A1OqWILpNh5Ss-0ZDTGgtxZMKGGo,250
 cache_dit/cache_factory/cache_contexts/cache_context.py,sha256=N88WLdd4KE9DuMWmpX8URcF55E2zWNwcKMxgVYkxMJY,13691
 cache_dit/cache_factory/cache_contexts/cache_manager.py,sha256=_NUXcMYYEIVfDHpc4HJr9RUjU5RUEkZmAgFGE8bh5Wc,34883
 cache_dit/cache_factory/cache_contexts/taylorseer.py,sha256=etSUIZzDvqW3ScKCbccTPcFaSmxV1T-xAXdk-p3e3wk,3802
-cache_dit/cache_factory/patch_functors/__init__.py,sha256=yK05iONMGILsTZ83ynrUUJtiJKJ_FDjxmVIzRLy416s,252
+cache_dit/cache_factory/patch_functors/__init__.py,sha256=06zdddrjvSCgBzJ0a8niRHd3ucF2qsbzlbL00d4aCvk,451
 cache_dit/cache_factory/patch_functors/functor_base.py,sha256=Ahk0fTfrHgNdEl-9JSkACvfyyv9G-Ei5OSz7XBIlX5o,357
-cache_dit/cache_factory/patch_functors/functor_chroma.py,sha256=N3UzszCM55g3GHeVdyXkid2_n72VJrfqBM2gdtD52gw,10042
-cache_dit/cache_factory/patch_functors/functor_flux.py,sha256=rJsbGEIxWPGnZyGI4ZwLLBdg8u6ZItsOeh0UoD_bVwk,9551
+cache_dit/cache_factory/patch_functors/functor_chroma.py,sha256=2iLxlsc-1dDHRveqCXaC07E9CeMNOuBNkvpJ1atpK7E,10048
+cache_dit/cache_factory/patch_functors/functor_flux.py,sha256=UMkyuEYjO7UO_zmXi9Djd-nD-XMgCUgE-qkYA3plWSM,9559
+cache_dit/cache_factory/patch_functors/functor_hidream.py,sha256=pi_vvpDy1lsgQHxu3eK3v93rdJL7oNwkt3WakRP8pbw,15375
+cache_dit/cache_factory/patch_functors/functor_hunyuan_dit.py,sha256=iSo5dD5uKnjQQeysDUIkKt0wdnK5bzXTc_F_lfHG70w,6401
 cache_dit/compile/__init__.py,sha256=FcTVzCeyypl-mxlc59_ehHL3lBNiDAFsXuRoJ-5Cfi0,56
 cache_dit/compile/utils.py,sha256=nN2OIrSdwRR5zGxJinKDqb07pXpvTNTF3g_OgLkeeBU,3858
 cache_dit/custom_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -39,9 +41,9 @@ cache_dit/quantize/__init__.py,sha256=kWYoMAyZgBXu9BJlZjTQ0dRffW9GqeeY9_iTkXrb70
 cache_dit/quantize/quantize_ao.py,sha256=mGspqYgQtenl3QnKPtsSYsSD7LbVX93f1M940bhXKLU,6066
 cache_dit/quantize/quantize_interface.py,sha256=2s_R7xPSKuJeFpEGeLwRxnq_CqJcBG3a3lzyW5wh-UM,1241
 cache_dit/quantize/quantize_svdq.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit-0.2.30.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.30.dist-info/METADATA,sha256=8Ln_X5fw14U3greCM7cSukrei1SRiMDpksFalg5ZBAU,22130
-cache_dit-0.2.30.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.30.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.2.30.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.30.dist-info/RECORD,,
+cache_dit-0.2.32.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.32.dist-info/METADATA,sha256=WQ9GP-Om05j3NBvtifkmbz5t20XBU_-KJQptrK7jQBs,24222
+cache_dit-0.2.32.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.32.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.32.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.32.dist-info/RECORD,,