cache-dit 0.1.0__py3-none-any.whl

cache_dit/cache_factory/first_block_cache/diffusers_adapters/__init__.py

@@ -0,0 +1,53 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/__init__.py
+
+import importlib
+from typing import Callable
+
+from diffusers import DiffusionPipeline
+
+
+def apply_fb_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"
+    elif transformer_cls_name.startswith("Wan"):
+        adapter_name = "wan"
+    else:
+        raise ValueError(
+            f"Unknown transformer class name: {transformer_cls_name}"
+        )
+
+    adapter_module = importlib.import_module(f".{adapter_name}", __package__)
+    apply_cache_on_transformer_fn = getattr(
+        adapter_module, "apply_cache_on_transformer"
+    )
+    return apply_cache_on_transformer_fn(transformer, *args, **kwargs)
+
+
+def apply_fb_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"
+    elif pipe_cls_name.startswith("Wan"):
+        adapter_name = "wan"
+    else:
+        raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}")
+
+    adapter_module = importlib.import_module(f".{adapter_name}", __package__)
+    apply_cache_on_pipe_fn = getattr(adapter_module, "apply_cache_on_pipe")
+    return apply_cache_on_pipe_fn(pipe, *args, **kwargs)
+
+
+# re-export functions for compatibility
+apply_cache_on_transformer: Callable = apply_fb_cache_on_transformer
+apply_cache_on_pipe: Callable = apply_fb_cache_on_pipe

cache_dit/cache_factory/first_block_cache/diffusers_adapters/cogvideox.py

@@ -0,0 +1,89 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/cogvideox.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import CogVideoXTransformer3DModel, DiffusionPipeline
+
+from cache_dit.cache_factory.first_block_cache import cache_context
+
+
+def apply_cache_on_transformer(
+    transformer: CogVideoXTransformer3DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            cache_context.CachedTransformerBlocks(
+                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_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_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/first_block_cache/diffusers_adapters/flux.py

@@ -0,0 +1,100 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/flux.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import DiffusionPipeline, FluxTransformer2DModel
+
+from cache_dit.cache_factory.first_block_cache import cache_context
+
+
+def apply_cache_on_transformer(
+    transformer: FluxTransformer2DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            cache_context.CachedTransformerBlocks(
+                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_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,
+    )  # noqa
+
+    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_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/first_block_cache/diffusers_adapters/mochi.py

@@ -0,0 +1,89 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/mochi.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import DiffusionPipeline, MochiTransformer3DModel
+
+from cache_dit.cache_factory.first_block_cache import cache_context
+
+
+def apply_cache_on_transformer(
+    transformer: MochiTransformer3DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    cached_transformer_blocks = torch.nn.ModuleList(
+        [
+            cache_context.CachedTransformerBlocks(
+                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_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_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/first_block_cache/diffusers_adapters/wan.py

@@ -0,0 +1,98 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache/wan.py
+
+import functools
+import unittest
+
+import torch
+from diffusers import DiffusionPipeline, HunyuanVideoTransformer3DModel
+
+from cache_dit.cache_factory.first_block_cache import cache_context
+
+
+def apply_cache_on_transformer(
+    transformer: HunyuanVideoTransformer3DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    blocks = torch.nn.ModuleList(
+        [
+            cache_context.CachedTransformerBlocks(
+                transformer.blocks,
+                transformer=transformer,
+                return_hidden_states_only=True,
+            )
+        ]
+    )
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with unittest.mock.patch.object(
+            self,
+            "blocks",
+            blocks,
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_cached = True
+
+    return transformer
+
+
+def apply_cache_on_pipe(
+    pipe: DiffusionPipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.03,
+    downsample_factor=1,
+    slg_layers=None,
+    slg_start: float = 0.0,
+    slg_end: float = 0.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,
+            "enable_alter_cache": True,
+            "slg_layers": slg_layers,
+            "slg_start": slg_start,
+            "slg_end": slg_end,
+            "num_inference_steps": kwargs.get("num_inference_steps", 50),
+            "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_cache_on_transformer(pipe.transformer, **kwargs)
+
+    return pipe

cache_dit/cache_factory/taylorseer.py

@@ -0,0 +1,76 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/first_block_cache/taylorseer.py
+import math
+
+
+class TaylorSeer:
+    def __init__(
+        self,
+        n_derivatives=2,
+        warmup_steps=1,
+        skip_interval_steps=1,
+        compute_step_map=None,
+    ):
+        self.n_derivatives = n_derivatives
+        self.ORDER = n_derivatives + 1
+        self.warmup_steps = warmup_steps
+        self.skip_interval_steps = skip_interval_steps
+        self.compute_step_map = compute_step_map
+        self.reset_cache()
+
+    def reset_cache(self):
+        self.state = {
+            "dY_prev": [None] * self.ORDER,
+            "dY_current": [None] * self.ORDER,
+        }
+        self.current_step = -1
+        self.last_non_approximated_step = -1
+
+    def should_compute_full(self, step=None):
+        step = self.current_step if step is None else step
+        if self.compute_step_map is not None:
+            return self.compute_step_map[step]
+        if (
+            step < self.warmup_steps
+            or (step - self.warmup_steps + 1) % self.skip_interval_steps == 0
+        ):
+            return True
+        return False
+
+    def approximate_derivative(self, Y):
+        dY_current = [None] * self.ORDER
+        dY_current[0] = Y
+        window = self.current_step - self.last_non_approximated_step
+        for i in range(self.n_derivatives):
+            if self.state["dY_prev"][i] is not None and self.current_step > 1:
+                dY_current[i + 1] = (
+                    dY_current[i] - self.state["dY_prev"][i]
+                ) / window
+            else:
+                break
+        return dY_current
+
+    def approximate_value(self):
+        elapsed = self.current_step - self.last_non_approximated_step
+        output = 0
+        for i, derivative in enumerate(self.state["dY_current"]):
+            if derivative is not None:
+                output += (1 / math.factorial(i)) * derivative * (elapsed**i)
+            else:
+                break
+        return output
+
+    def mark_step_begin(self):
+        self.current_step += 1
+
+    def update(self, Y):
+        self.state["dY_prev"] = self.state["dY_current"]
+        self.state["dY_current"] = self.approximate_derivative(Y)
+        self.last_non_approximated_step = self.current_step
+
+    def step(self, Y):
+        self.mark_step_begin()
+        if self.should_compute_full():
+            self.update(Y)
+            return Y
+        else:
+            return self.approximate_value()

cache_dit/logger.py

@@ -0,0 +1,97 @@
+import logging
+import os
+import sys
+
+_FORMAT = "%(levelname)s %(asctime)s [%(filename)s:%(lineno)d] %(message)s"
+_DATE_FORMAT = "%m-%d %H:%M:%S"
+
+_LOG_LEVEL = os.environ.get("CACHE_DIT_LOG_LEVEL", "info")
+_LOG_LEVEL = getattr(logging, _LOG_LEVEL.upper(), 0)
+_LOG_DIR = os.environ.get("CACHE_DIT_LOG_DIR", None)
+
+
+class NewLineFormatter(logging.Formatter):
+    """Adds logging prefix to newlines to align multi-line messages."""
+
+    def __init__(self, fmt, datefmt=None):
+        logging.Formatter.__init__(self, fmt, datefmt)
+
+    def format(self, record):
+        msg = logging.Formatter.format(self, record)
+        if record.message != "":
+            parts = msg.split(record.message)
+            msg = msg.replace("\n", "\r\n" + parts[0])
+        return msg
+
+
+_root_logger = logging.getLogger("CACHE_DIT")
+_default_handler = None
+_default_file_handler = None
+_inference_log_file_handler = {}
+
+
+def _setup_logger():
+    _root_logger.setLevel(_LOG_LEVEL)
+    global _default_handler
+    global _default_file_handler
+    fmt = NewLineFormatter(_FORMAT, datefmt=_DATE_FORMAT)
+
+    if _default_handler is None:
+        _default_handler = logging.StreamHandler(sys.stdout)
+        _default_handler.flush = sys.stdout.flush  # type: ignore
+        _default_handler.setLevel(_LOG_LEVEL)
+        _root_logger.addHandler(_default_handler)
+
+    if _default_file_handler is None and _LOG_DIR is not None:
+        if not os.path.exists(_LOG_DIR):
+            try:
+                os.makedirs(_LOG_DIR)
+            except OSError as e:
+                _root_logger.warn(f"Error creating directory {_LOG_DIR} : {e}")
+        _default_file_handler = logging.FileHandler(_LOG_DIR + "/default.log")
+        _default_file_handler.setLevel(_LOG_LEVEL)
+        _default_file_handler.setFormatter(fmt)
+        _root_logger.addHandler(_default_file_handler)
+
+    _default_handler.setFormatter(fmt)
+    # Setting this will avoid the message
+    # being propagated to the parent logger.
+    _root_logger.propagate = False
+
+
+# The logger is initialized when the module is imported.
+# This is thread-safe as the module is only imported once,
+# guaranteed by the Python GIL.
+_setup_logger()
+
+
+def init_logger(name: str):
+    pid = os.getpid()
+    # Use the same settings as above for root logger
+    logger = logging.getLogger(name)
+    logger.setLevel(_LOG_LEVEL)
+    logger.addHandler(_default_handler)
+    if _LOG_DIR is not None and pid is None:
+        logger.addHandler(_default_file_handler)
+    elif _LOG_DIR is not None:
+        if _inference_log_file_handler.get(pid, None) is not None:
+            logger.addHandler(_inference_log_file_handler[pid])
+        else:
+            if not os.path.exists(_LOG_DIR):
+                try:
+                    os.makedirs(_LOG_DIR)
+                except OSError as e:
+                    _root_logger.warn(
+                        f"Error creating directory {_LOG_DIR} : {e}"
+                    )
+            _inference_log_file_handler[pid] = logging.FileHandler(
+                _LOG_DIR + f"/process.{pid}.log"
+            )
+            _inference_log_file_handler[pid].setLevel(_LOG_LEVEL)
+            _inference_log_file_handler[pid].setFormatter(
+                NewLineFormatter(_FORMAT, datefmt=_DATE_FORMAT)
+            )
+            _root_logger.addHandler(_inference_log_file_handler[pid])
+            logger.addHandler(_inference_log_file_handler[pid])
+    logger.propagate = False
+    return logger