cache-dit 0.1.0__py3-none-any.whl

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py

@@ -0,0 +1,45 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/__init__.py
+
+import importlib
+
+from diffusers import DiffusionPipeline
+
+
+def apply_db_cache_on_transformer(transformer, *args, **kwargs):
+    transformer_cls_name: str = transformer.__class__.__name__
+    if transformer_cls_name.startswith("Flux"):
+        adapter_name = "flux"
+    elif transformer_cls_name.startswith("Mochi"):
+        adapter_name = "mochi"
+    elif transformer_cls_name.startswith("CogVideoX"):
+        adapter_name = "cogvideox"
+    else:
+        raise ValueError(
+            f"Unknown transformer class name: {transformer_cls_name}"
+        )
+
+    adapter_module = importlib.import_module(f".{adapter_name}", __package__)
+    apply_db_cache_on_transformer_fn = getattr(
+        adapter_module, "apply_db_cache_on_transformer"
+    )
+    return apply_db_cache_on_transformer_fn(transformer, *args, **kwargs)
+
+
+def apply_db_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
+    assert isinstance(pipe, DiffusionPipeline)
+
+    pipe_cls_name: str = pipe.__class__.__name__
+    if pipe_cls_name.startswith("Flux"):
+        adapter_name = "flux"
+    elif pipe_cls_name.startswith("Mochi"):
+        adapter_name = "mochi"
+    elif pipe_cls_name.startswith("CogVideoX"):
+        adapter_name = "cogvideox"
+    else:
+        raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}")
+
+    adapter_module = importlib.import_module(f".{adapter_name}", __package__)
+    apply_db_cache_on_pipe_fn = getattr(
+        adapter_module, "apply_db_cache_on_pipe"
+    )
+    return apply_db_cache_on_pipe_fn(pipe, *args, **kwargs)

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py

@@ -0,0 +1,89 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/cogvideox.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import CogVideoXTransformer3DModel, DiffusionPipeline
+
+from cache_dit.cache_factory.dual_block_cache import cache_context
+
+
+def apply_db_cache_on_transformer(
+    transformer: CogVideoXTransformer3DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            cache_context.DBCachedTransformerBlocks(
+                transformer.transformer_blocks,
+                transformer=transformer,
+            )
+        ]
+    )
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with unittest.mock.patch.object(
+            self,
+            "transformer_blocks",
+            cached_transformer_blocks,
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_cached = True
+
+    return transformer
+
+
+def apply_db_cache_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.04,
+    downsample_factor=1,
+    warmup_steps=0,
+    max_cached_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = cache_context.collect_cache_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            "warmup_steps": warmup_steps,
+            "max_cached_steps": max_cached_steps,
+        },
+        **kwargs,
+    )
+    if not getattr(pipe, "_is_cached", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with cache_context.cache_context(
+                cache_context.create_cache_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_cached = True
+
+    if not shallow_patch:
+        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py

@@ -0,0 +1,100 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/flux.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import DiffusionPipeline, FluxTransformer2DModel
+
+from cache_dit.cache_factory.dual_block_cache import cache_context
+
+
+def apply_db_cache_on_transformer(
+    transformer: FluxTransformer2DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            cache_context.DBCachedTransformerBlocks(
+                transformer.transformer_blocks,
+                transformer.single_transformer_blocks,
+                transformer=transformer,
+                return_hidden_states_first=False,
+            )
+        ]
+    )
+    dummy_single_transformer_blocks = torch.nn.ModuleList()
+
+    original_forward = transformer.forward
+
+    @functools.wraps(original_forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with (
+            unittest.mock.patch.object(
+                self,
+                "transformer_blocks",
+                cached_transformer_blocks,
+            ),
+            unittest.mock.patch.object(
+                self,
+                "single_transformer_blocks",
+                dummy_single_transformer_blocks,
+            ),
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_cached = True
+
+    return transformer
+
+
+def apply_db_cache_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.05,
+    downsample_factor=1,
+    warmup_steps=0,
+    max_cached_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = cache_context.collect_cache_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            "warmup_steps": warmup_steps,
+            "max_cached_steps": max_cached_steps,
+        },
+        **kwargs,
+    )
+
+    if not getattr(pipe, "_is_cached", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with cache_context.cache_context(
+                cache_context.create_cache_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_cached = True
+
+    if not shallow_patch:
+        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py

@@ -0,0 +1,88 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/mochi.py
+import functools
+import unittest
+
+import torch
+from diffusers import DiffusionPipeline, MochiTransformer3DModel
+
+from cache_dit.cache_factory.dual_block_cache import cache_context
+
+
+def apply_db_cache_on_transformer(
+    transformer: MochiTransformer3DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            cache_context.DBCachedTransformerBlocks(
+                transformer.transformer_blocks,
+                transformer=transformer,
+            )
+        ]
+    )
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with unittest.mock.patch.object(
+            self,
+            "transformer_blocks",
+            cached_transformer_blocks,
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_cached = True
+
+    return transformer
+
+
+def apply_db_cache_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.06,
+    downsample_factor=1,
+    warmup_steps=0,
+    max_cached_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = cache_context.collect_cache_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            "warmup_steps": warmup_steps,
+            "max_cached_steps": max_cached_steps,
+        },
+        **kwargs,
+    )
+    if not getattr(pipe, "_is_cached", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with cache_context.cache_context(
+                cache_context.create_cache_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_cached = True
+
+    if not shallow_patch:
+        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/__init__.py

@@ -0,0 +1,45 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/__init__.py
+
+import importlib
+
+from diffusers import DiffusionPipeline
+
+
+def apply_db_prune_on_transformer(transformer, *args, **kwargs):
+    transformer_cls_name: str = transformer.__class__.__name__
+    if transformer_cls_name.startswith("Flux"):
+        adapter_name = "flux"
+    elif transformer_cls_name.startswith("Mochi"):
+        adapter_name = "mochi"
+    elif transformer_cls_name.startswith("CogVideoX"):
+        adapter_name = "cogvideox"
+    else:
+        raise ValueError(
+            f"Unknown transformer class name: {transformer_cls_name}"
+        )
+
+    adapter_module = importlib.import_module(f".{adapter_name}", __package__)
+    apply_db_cache_on_transformer_fn = getattr(
+        adapter_module, "apply_db_prune_on_transformer"
+    )
+    return apply_db_cache_on_transformer_fn(transformer, *args, **kwargs)
+
+
+def apply_db_prune_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
+    assert isinstance(pipe, DiffusionPipeline)
+
+    pipe_cls_name: str = pipe.__class__.__name__
+    if pipe_cls_name.startswith("Flux"):
+        adapter_name = "flux"
+    elif pipe_cls_name.startswith("Mochi"):
+        adapter_name = "mochi"
+    elif pipe_cls_name.startswith("CogVideoX"):
+        adapter_name = "cogvideox"
+    else:
+        raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}")
+
+    adapter_module = importlib.import_module(f".{adapter_name}", __package__)
+    apply_db_cache_on_pipe_fn = getattr(
+        adapter_module, "apply_db_prune_on_pipe"
+    )
+    return apply_db_cache_on_pipe_fn(pipe, *args, **kwargs)

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/cogvideox.py

@@ -0,0 +1,89 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/cogvideox.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import CogVideoXTransformer3DModel, DiffusionPipeline
+
+from cache_dit.cache_factory.dynamic_block_prune import prune_context
+
+
+def apply_db_prune_on_transformer(
+    transformer: CogVideoXTransformer3DModel,
+):
+    if getattr(transformer, "_is_pruned", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            prune_context.DBPrunedTransformerBlocks(
+                transformer.transformer_blocks,
+                transformer=transformer,
+            )
+        ]
+    )
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with unittest.mock.patch.object(
+            self,
+            "transformer_blocks",
+            cached_transformer_blocks,
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_pruned = True
+
+    return transformer
+
+
+def apply_db_prune_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.04,
+    downsample_factor=1,
+    warmup_steps=0,
+    max_pruned_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = prune_context.collect_prune_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            "warmup_steps": warmup_steps,
+            "max_pruned_steps": max_pruned_steps,
+        },
+        **kwargs,
+    )
+    if not getattr(pipe, "_is_pruned", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with prune_context.prune_context(
+                prune_context.create_prune_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_pruned = True
+
+    if not shallow_patch:
+        apply_db_prune_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/flux.py

@@ -0,0 +1,100 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/flux.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import DiffusionPipeline, FluxTransformer2DModel
+
+from cache_dit.cache_factory.dynamic_block_prune import prune_context
+
+
+def apply_db_prune_on_transformer(
+    transformer: FluxTransformer2DModel,
+):
+    if getattr(transformer, "_is_pruned", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            prune_context.DBPrunedTransformerBlocks(
+                transformer.transformer_blocks,
+                transformer.single_transformer_blocks,
+                transformer=transformer,
+                return_hidden_states_first=False,
+            )
+        ]
+    )
+    dummy_single_transformer_blocks = torch.nn.ModuleList()
+
+    original_forward = transformer.forward
+
+    @functools.wraps(original_forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with (
+            unittest.mock.patch.object(
+                self,
+                "transformer_blocks",
+                cached_transformer_blocks,
+            ),
+            unittest.mock.patch.object(
+                self,
+                "single_transformer_blocks",
+                dummy_single_transformer_blocks,
+            ),
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_pruned = True
+
+    return transformer
+
+
+def apply_db_prune_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.05,
+    downsample_factor=1,
+    warmup_steps=0,
+    max_pruned_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = prune_context.collect_prune_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            "warmup_steps": warmup_steps,
+            "max_pruned_steps": max_pruned_steps,
+        },
+        **kwargs,
+    )
+
+    if not getattr(pipe, "_is_pruned", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with prune_context.prune_context(
+                prune_context.create_prune_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_pruned = True
+
+    if not shallow_patch:
+        apply_db_prune_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/mochi.py

@@ -0,0 +1,89 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/diffusers_adapters/mochi.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import DiffusionPipeline, MochiTransformer3DModel
+
+from cache_dit.cache_factory.dynamic_block_prune import prune_context
+
+
+def apply_db_prune_on_transformer(
+    transformer: MochiTransformer3DModel,
+):
+    if getattr(transformer, "_is_pruned", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            prune_context.DBPrunedTransformerBlocks(
+                transformer.transformer_blocks,
+                transformer=transformer,
+            )
+        ]
+    )
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with unittest.mock.patch.object(
+            self,
+            "transformer_blocks",
+            cached_transformer_blocks,
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_pruned = True
+
+    return transformer
+
+
+def apply_db_prune_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.06,
+    downsample_factor=1,
+    warmup_steps=0,
+    max_pruned_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = prune_context.collect_prune_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            "warmup_steps": warmup_steps,
+            "max_pruned_steps": max_pruned_steps,
+        },
+        **kwargs,
+    )
+    if not getattr(pipe, "_is_pruned", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with prune_context.prune_context(
+                prune_context.create_prune_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_pruned = True
+
+    if not shallow_patch:
+        apply_db_prune_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe