cache-dit 0.1.1.dev2__py3-none-any.whl

cache_dit/primitives.py

@@ -0,0 +1,152 @@
+# Adapted from: https://github.com/chengzeyi/ParaAttention/blob/main/src/para_attn/primitives.py
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+
+if dist.is_available():
+    import torch.distributed._functional_collectives as ft_c
+    import torch.distributed.distributed_c10d as c10d
+else:
+    ft_c = None
+    c10d = None
+
+
+def get_group(group=None):
+    if group is None:
+        group = c10d._get_default_group()
+
+    if isinstance(group, dist.ProcessGroup):
+        pg: Union[dist.ProcessGroup, List[dist.ProcessGroup]] = group
+    else:
+        pg = group.get_group()
+
+    return pg
+
+
+def get_world_size(group=None):
+    pg = get_group(group)
+    return dist.get_world_size(pg)
+
+
+def get_rank(group=None):
+    pg = get_group(group)
+    return dist.get_rank(pg)
+
+
+def _maybe_wait(tensor: torch.Tensor) -> torch.Tensor:
+    """
+    When tracing the code, the result tensor is not an AsyncCollectiveTensor,
+    so we cannot call ``wait()``.
+    """
+    if isinstance(tensor, ft_c.AsyncCollectiveTensor):
+        return tensor.wait()
+    return tensor
+
+
+def all_gather_tensor_sync(x, *args, group=None, **kwargs):
+    group = get_group(group)
+    x_shape = x.shape
+    x = x.flatten()
+    x_numel = x.numel()
+    x = ft_c.all_gather_tensor(x, *args, group=group, **kwargs)
+    x = _maybe_wait(x)
+    x_shape = list(x_shape)
+    x_shape[0] *= x.numel() // x_numel
+    x = x.reshape(x_shape)
+    return x
+
+
+def all_gather_tensor_autograd_sync(x, *args, group=None, **kwargs):
+    group = get_group(group)
+    x_shape = x.shape
+    x = x.flatten()
+    x_numel = x.numel()
+    x = ft_c.all_gather_tensor_autograd(x, *args, group=group, **kwargs)
+    x = _maybe_wait(x)
+    x_shape = list(x_shape)
+    x_shape[0] *= x.numel() // x_numel
+    x = x.reshape(x_shape)
+    return x
+
+
+def all_to_all_single_sync(x, *args, **kwargs):
+    x_shape = x.shape
+    x = x.flatten()
+    x = ft_c.all_to_all_single(x, *args, **kwargs)
+    x = _maybe_wait(x)
+    x = x.reshape(x_shape)
+    return x
+
+
+def all_to_all_single_autograd_sync(x, *args, **kwargs):
+    x_shape = x.shape
+    x = x.flatten()
+    x = ft_c.all_to_all_single_autograd(x, *args, **kwargs)
+    x = _maybe_wait(x)
+    x = x.reshape(x_shape)
+    return x
+
+
+def all_reduce_sync(x, *args, group=None, **kwargs):
+    group = get_group(group)
+    x = ft_c.all_reduce(x, *args, group=group, **kwargs)
+    x = _maybe_wait(x)
+    return x
+
+
+def get_buffer(
+    shape_or_tensor: Union[Tuple[int], torch.Tensor],
+    *,
+    repeats: int = 1,
+    dim: int = 0,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[torch.device] = None,
+    group=None,
+) -> torch.Tensor:
+    if repeats is None:
+        repeats = get_world_size(group)
+
+    if isinstance(shape_or_tensor, torch.Tensor):
+        shape = shape_or_tensor.shape
+        dtype = shape_or_tensor.dtype
+        device = shape_or_tensor.device
+
+    assert dtype is not None
+    assert device is not None
+
+    shape = list(shape)
+    if repeats > 1:
+        shape[dim] *= repeats
+
+    buffer = torch.empty(shape, dtype=dtype, device=device)
+    return buffer
+
+
+def get_assigned_chunk(
+    tensor: torch.Tensor,
+    dim: int = 0,
+    idx: Optional[int] = None,
+    group=None,
+) -> torch.Tensor:
+    if idx is None:
+        idx = get_rank(group)
+    world_size = get_world_size(group)
+    total_size = tensor.shape[dim]
+    assert (
+        total_size % world_size == 0
+    ), f"tensor.shape[{dim}]={total_size} is not divisible by world_size={world_size}"
+    return tensor.chunk(world_size, dim=dim)[idx]
+
+
+def get_complete_tensor(
+    tensor: torch.Tensor,
+    *,
+    dim: int = 0,
+    group=None,
+) -> torch.Tensor:
+    tensor = tensor.transpose(0, dim).contiguous()
+    output_tensor = all_gather_tensor_sync(tensor, gather_dim=0, group=group)
+    output_tensor = output_tensor.transpose(0, dim)
+    return output_tensor

cache_dit-0.1.1.dev2.dist-info/METADATA

@@ -0,0 +1,31 @@
+Metadata-Version: 2.4
+Name: cache_dit
+Version: 0.1.1.dev2
+Summary: ⚡️DBCache: A Training-free UNet-style Cache Acceleration for Diffusion Transformers
+Author: DefTruth, vipshop.com, etc.
+Maintainer: DefTruth, vipshop.com, etc
+Project-URL: Repository, https://github.com/vipshop/DBCache.git
+Project-URL: Homepage, https://github.com/vipshop/DBCache.git
+Requires-Python: >=3.10
+Requires-Dist: packaging
+Requires-Dist: torch
+Requires-Dist: transformers
+Requires-Dist: diffusers
+Provides-Extra: all
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest<8.0.0,>=7.0.0; extra == "dev"
+Requires-Dist: pytest-html; extra == "dev"
+Requires-Dist: expecttest; extra == "dev"
+Requires-Dist: hypothesis; extra == "dev"
+Requires-Dist: transformers; extra == "dev"
+Requires-Dist: diffusers; extra == "dev"
+Requires-Dist: accelerate; extra == "dev"
+Requires-Dist: peft; extra == "dev"
+Requires-Dist: protobuf; extra == "dev"
+Requires-Dist: sentencepiece; extra == "dev"
+Requires-Dist: opencv-python-headless; extra == "dev"
+Requires-Dist: ftfy; extra == "dev"
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python

cache_dit-0.1.1.dev2.dist-info/RECORD

@@ -0,0 +1,30 @@
+cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cache_dit/_version.py,sha256=lIErBp1sZ_uGq2rboUGas8Ch-hnYN8OtqGh3G0mtds0,524
+cache_dit/logger.py,sha256=dKfNe_RRk9HJwfgHGeRR1f0LbskJpKdGmISCbL9roQs,3443
+cache_dit/primitives.py,sha256=A2iG9YLot3gOsZSPp-_gyjqjLgJvWQRx8aitD4JQ23Y,3877
+cache_dit/cache_factory/__init__.py,sha256=plAOUMsne-dTYA-cq1RLbE7dlH-kFA_Hst9MzbWPqiI,5224
+cache_dit/cache_factory/taylorseer.py,sha256=0W29ykJg3MnyLAB2KFicsl11Xe41cDYPgI60bquG_NY,2495
+cache_dit/cache_factory/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cache_dit/cache_factory/dual_block_cache/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=EJ-uhA2-sWMW1jNDhcBtjHDqSn8lUzfKbYoPfZDQhZU,49665
+cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py,sha256=C6tfXHpdY8YFV3gk74dr_IpYH4bO4ItbPCQYud3NgAM,1667
+cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py,sha256=g3ua-hmTpVeTOQNVYjUX2gsHuG2NV0B81iKHGa51wwk,2401
+cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py,sha256=UbE6nIF-EtA92QxIZVMzIssdZKQSPAVX1hchF9R8drU,2754
+cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py,sha256=qxMu1L3ycT8F-uxpGsmFQBY_BH1vDiGIOXgS_Qbb7dM,2391
+cache_dit/cache_factory/dynamic_block_prune/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cache_dit/cache_factory/dynamic_block_prune/prune_context.py,sha256=cE27f5NPgQ_COmTnF__e85Uz5pWyXID0ut-tmtSQfVQ,34597
+cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/__init__.py,sha256=8IjJjZOs5XRzsj7Ni2MXpR2Z1PUyRSONIhmfAn1G0eM,1667
+cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/cogvideox.py,sha256=ORJpdkXkgziDUo-rpebC6pUemgYaDCoeu0cwwLz175U,2407
+cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/flux.py,sha256=KbEkLSsHtS6xwLWNh3jlOlXRyGRdrI2pWV1zyQxMTj4,2757
+cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/mochi.py,sha256=rgeXfww-7WX6URSDg7mF1HuxSmYmoJVjMVoNGuxjwxc,2395
+cache_dit/cache_factory/first_block_cache/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cache_dit/cache_factory/first_block_cache/cache_context.py,sha256=DpDhtK095PlrvACf7sbjOt2-QpVkV1arr1qGEKJqgaQ,23502
+cache_dit/cache_factory/first_block_cache/diffusers_adapters/__init__.py,sha256=h80pgoZ3l8qC4rbm9KY0jSN8hOsmGgyvvFxD-xznHdw,1959
+cache_dit/cache_factory/first_block_cache/diffusers_adapters/cogvideox.py,sha256=qO5CWyurtwW30mvOe6cxeQPTSXLDlPJcezm72zEjDq8,2375
+cache_dit/cache_factory/first_block_cache/diffusers_adapters/flux.py,sha256=Dcd4OzABCtyQCZNX2KNnUTdVoO1E1ApM7P8gcVYzcK0,2733
+cache_dit/cache_factory/first_block_cache/diffusers_adapters/mochi.py,sha256=lQTClo52OwPbNEE4jiBZQhfC7hbtYqnYIABp_vbm_dk,2363
+cache_dit/cache_factory/first_block_cache/diffusers_adapters/wan.py,sha256=IVH-lroOzvYb4XKLk9MOw54EtijBtuzVaKcVGz0KlBA,2656
+cache_dit-0.1.1.dev2.dist-info/METADATA,sha256=pQr1yJwVuMWqB9b-IfPUq5x9UrpzvUauKqAzJWmQIZ0,1150
+cache_dit-0.1.1.dev2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.1.1.dev2.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.1.1.dev2.dist-info/RECORD,,

cache_dit-0.1.1.dev2.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

cache_dit-0.1.1.dev2.dist-info/top_level.txt

@@ -0,0 +1 @@
+cache_dit