cache-dit 0.3.2__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/caching/cache_blocks/__init__.py

@@ -0,0 +1,226 @@
+import torch
+
+from cache_dit.caching import ForwardPattern
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.caching.cache_contexts.cache_context import CachedContext
+from cache_dit.caching.cache_contexts.prune_context import PrunedContext
+from cache_dit.caching.cache_contexts.cache_manager import (
+    CachedContextManager,
+)
+from cache_dit.caching.cache_contexts.prune_manager import (
+    PrunedContextManager,
+)
+
+from cache_dit.caching.cache_blocks.pattern_0_1_2 import (
+    CachedBlocks_Pattern_0_1_2,
+    PrunedBlocks_Pattern_0_1_2,
+)
+from cache_dit.caching.cache_blocks.pattern_3_4_5 import (
+    CachedBlocks_Pattern_3_4_5,
+    PrunedBlocks_Pattern_3_4_5,
+)
+from cache_dit.caching.cache_blocks.pattern_utils import apply_stats
+from cache_dit.caching.cache_blocks.pattern_utils import remove_stats
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class UnifiedBlocks:
+    def __new__(
+        cls,
+        # 0. Transformer blocks configuration
+        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',
+        # 'layers', 'single_stream_blocks', 'double_stream_blocks'
+        cache_prefix: str = None,  # cache_prefix maybe un-need.
+        # Usually, blocks_name, etc.
+        cache_context: CachedContext | PrunedContext | str = None,
+        context_manager: CachedContextManager | PrunedContextManager = None,
+        cache_type: CacheType = CacheType.DBCache,
+        **kwargs,
+    ):
+        if cache_type == CacheType.DBCache:
+            return CachedBlocks(
+                # 0. Transformer blocks configuration
+                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_context=cache_context,
+                context_manager=context_manager,
+                cache_type=cache_type,
+                **kwargs,
+            )
+        elif cache_type == CacheType.DBPrune:
+            return PrunedBlocks(
+                # 0. Transformer blocks configuration
+                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_context=cache_context,
+                context_manager=context_manager,
+                cache_type=cache_type,
+                **kwargs,
+            )
+        else:
+            raise ValueError(f"Cache type {cache_type} is not supported now!")
+
+
+class CachedBlocks:
+    def __new__(
+        cls,
+        # 0. Transformer blocks configuration
+        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',
+        # 'layers', 'single_stream_blocks', 'double_stream_blocks'
+        cache_prefix: str = None,  # cache_prefix maybe un-need.
+        # Usually, blocks_name, etc.
+        cache_context: CachedContext | PrunedContext | str = None,
+        context_manager: CachedContextManager | PrunedContextManager = None,
+        cache_type: CacheType = CacheType.DBCache,
+        **kwargs,
+    ):
+        assert transformer is not None, "transformer can't be None."
+        assert forward_pattern is not None, "forward_pattern can't be None."
+        assert cache_context is not None, "cache_context can't be None."
+        assert context_manager is not None, "context_manager can't be None."
+        if forward_pattern in CachedBlocks_Pattern_0_1_2._supported_patterns:
+            if cache_type == CacheType.DBCache:
+                assert isinstance(
+                    context_manager, CachedContextManager
+                ), "context_manager must be CachedContextManager for DBCache."
+                return CachedBlocks_Pattern_0_1_2(
+                    # 0. Transformer blocks configuration
+                    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_context=cache_context,
+                    context_manager=context_manager,
+                    cache_type=cache_type,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Cache type {cache_type} is not supported now!"
+                )
+        elif forward_pattern in CachedBlocks_Pattern_3_4_5._supported_patterns:
+            if cache_type == CacheType.DBCache:
+                assert isinstance(
+                    context_manager, CachedContextManager
+                ), "context_manager must be CachedContextManager for DBCache."
+                return CachedBlocks_Pattern_3_4_5(
+                    # 0. Transformer blocks configuration
+                    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_context=cache_context,
+                    context_manager=context_manager,
+                    cache_type=cache_type,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Cache type {cache_type} is not supported now!"
+                )
+        else:
+            raise ValueError(f"Pattern {forward_pattern} is not supported now!")
+
+
+class PrunedBlocks:
+    def __new__(
+        cls,
+        # 0. Transformer blocks configuration
+        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',
+        # 'layers', 'single_stream_blocks', 'double_stream_blocks'
+        cache_prefix: str = None,  # cache_prefix maybe un-need.
+        # Usually, blocks_name, etc.
+        cache_context: CachedContext | PrunedContext | str = None,
+        context_manager: CachedContextManager | PrunedContextManager = None,
+        cache_type: CacheType = CacheType.DBCache,
+        **kwargs,
+    ):
+        assert transformer is not None, "transformer can't be None."
+        assert forward_pattern is not None, "forward_pattern can't be None."
+        assert cache_context is not None, "cache_context can't be None."
+        assert context_manager is not None, "context_manager can't be None."
+        if forward_pattern in PrunedBlocks_Pattern_0_1_2._supported_patterns:
+            if cache_type == CacheType.DBPrune:
+                assert isinstance(
+                    context_manager, PrunedContextManager
+                ), "context_manager must be PrunedContextManager for DBPrune."
+                return PrunedBlocks_Pattern_0_1_2(
+                    # 0. Transformer blocks configuration
+                    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_context=cache_context,
+                    context_manager=context_manager,
+                    cache_type=cache_type,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Cache type {cache_type} is not supported now!"
+                )
+        elif forward_pattern in PrunedBlocks_Pattern_3_4_5._supported_patterns:
+            if cache_type == CacheType.DBPrune:
+                assert isinstance(
+                    context_manager, PrunedContextManager
+                ), "context_manager must be PrunedContextManager for DBPrune."
+                return PrunedBlocks_Pattern_3_4_5(
+                    # 0. Transformer blocks configuration
+                    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_context=cache_context,
+                    context_manager=context_manager,
+                    cache_type=cache_type,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Cache type {cache_type} is not supported now!"
+                )
+        else:
+            raise ValueError(f"Pattern {forward_pattern} is not supported now!")

cache_dit/caching/cache_blocks/offload_utils.py

@@ -0,0 +1,115 @@
+import torch
+import asyncio
+import logging
+from contextlib import contextmanager
+from typing import Generator, Optional, List
+from diffusers.hooks.group_offloading import _is_group_offload_enabled
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@torch.compiler.disable
+@contextmanager
+def maybe_onload(
+    block: torch.nn.Module,
+    reference_tensor: torch.Tensor,
+    pending_tasks: List[asyncio.Task] = [],
+) -> Generator:
+
+    if not _is_group_offload_enabled(block):
+        yield block
+        return
+
+    original_devices: Optional[List[torch.device]] = None
+    if hasattr(block, "parameters"):
+        params = list(block.parameters())
+        if params:
+            original_devices = [param.data.device for param in params]
+
+    target_device: torch.device = reference_tensor.device
+    move_task: Optional[asyncio.Task] = None
+    need_restore: bool = False
+
+    try:
+        if original_devices is not None:
+            unique_devices = list(set(original_devices))
+            if len(unique_devices) > 1 or unique_devices[0] != target_device:
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(
+                        f"Onloading from {unique_devices} to {target_device}"
+                    )
+
+                has_meta_params = any(
+                    dev.type == "meta" for dev in original_devices
+                )
+                if has_meta_params:  # compatible with sequential cpu offload
+                    block = block.to_empty(device=target_device)
+                else:
+                    block = block.to(target_device, non_blocking=False)
+                need_restore = True
+        yield block
+    finally:
+        if need_restore and original_devices:
+
+            async def restore_device():
+                for param, original_device in zip(
+                    block.parameters(), original_devices
+                ):
+                    param.data = await asyncio.to_thread(
+                        lambda p, d: p.to(d, non_blocking=True),
+                        param.data,  # type: torch.Tensor
+                        original_device,  # type: torch.device
+                    )  # type: ignore[assignment]
+
+            loop = get_event_loop()
+            move_task = loop.create_task(restore_device())
+            if move_task:
+                pending_tasks.append(move_task)
+
+
+def get_event_loop() -> asyncio.AbstractEventLoop:
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+    if not loop.is_running():
+
+        def run_loop() -> None:
+            asyncio.set_event_loop(loop)
+            loop.run_forever()
+
+        import threading
+
+        if not any(t.name == "_my_loop" for t in threading.enumerate()):
+            threading.Thread(
+                target=run_loop, name="_my_loop", daemon=True
+            ).start()
+
+    return loop
+
+
+@torch.compiler.disable
+def maybe_offload(
+    pending_tasks: List[asyncio.Task],
+) -> None:
+    if not pending_tasks:
+        return
+
+    loop = get_event_loop()
+
+    async def gather_tasks():
+        return await asyncio.gather(*pending_tasks)
+
+    future = asyncio.run_coroutine_threadsafe(gather_tasks(), loop)
+    try:
+        future.result(timeout=30.0)
+    except Exception as e:
+        logger.error(f"May Offload Error: {e}")
+
+    pending_tasks.clear()

cache_dit/caching/cache_blocks/pattern_0_1_2.py

@@ -0,0 +1,26 @@
+from cache_dit.caching import ForwardPattern
+from cache_dit.caching.cache_blocks.pattern_base import (
+    CachedBlocks_Pattern_Base,
+    PrunedBlocks_Pattern_Base,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class CachedBlocks_Pattern_0_1_2(CachedBlocks_Pattern_Base):
+    _supported_patterns = [
+        ForwardPattern.Pattern_0,
+        ForwardPattern.Pattern_1,
+        ForwardPattern.Pattern_2,
+    ]
+    ...
+
+
+class PrunedBlocks_Pattern_0_1_2(PrunedBlocks_Pattern_Base):
+    _supported_patterns = [
+        ForwardPattern.Pattern_0,
+        ForwardPattern.Pattern_1,
+        ForwardPattern.Pattern_2,
+    ]
+    ...