cache-dit 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

cache_dit/cache_factory/first_block_cache/diffusers_adapters/wan.py

@@ -1,98 +0,0 @@
-# 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, WanTransformer3DModel
-
-from cache_dit.cache_factory.first_block_cache import cache_context
-
-
-def apply_cache_on_transformer(
-    transformer: WanTransformer3DModel,
-):
-    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-0.2.15.dist-info/RECORD

@@ -1,50 +0,0 @@
-cache_dit/__init__.py,sha256=0-B173-fLi3IA8nJXoS71zK0zD33Xplysd9skmLfEOY,171
-cache_dit/_version.py,sha256=TspyaLI34cMH7HNr-lrUKFnY4CF39yGfG_421vC3fhg,513
-cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
-cache_dit/primitives.py,sha256=A2iG9YLot3gOsZSPp-_gyjqjLgJvWQRx8aitD4JQ23Y,3877
-cache_dit/utils.py,sha256=4cFNh0asch6Zgsixq0bS1ElfwBu_6BG5ZSmaa1khjyg,144
-cache_dit/cache_factory/__init__.py,sha256=iYQwLwB_XLoYl0OB9unZGDbBtrYvZaLkOAmhGRwdW2E,191
-cache_dit/cache_factory/adapters.py,sha256=QMCaXnmqM7NT7sx4bCF1mMLn-QcXX9h1RmgLAypDedg,5256
-cache_dit/cache_factory/taylorseer.py,sha256=LKSNo2ode69EVo9xrxjxAMEjz0yDGiGADeDYnEqddA8,3987
-cache_dit/cache_factory/utils.py,sha256=V-Mb5Jn07geEUUWo4QAfh6pmSzkL-2OGDn0VAXbG6hQ,1799
-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=sJ9yxQlcrX4qkPln94FrL0WDe2WIn3_UD2-Mk8YtjSw,73301
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py,sha256=uSqF5aD2-feHB25vEbx1STBQVjWVAOn_wYTdAEmS4NU,2045
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py,sha256=3xUjvDzor9AkBkDUc0N7kZqM86MIdajuigesnicNzXE,2260
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py,sha256=cIsov6Pf0dRyddqkzTA2CU-jSDotof8LQr-HIoY9T9M,2615
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py,sha256=SO4q39PQuQ5QVHy5Z-ubiKdstzvQPedONN2J5oiGUh0,9955
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py,sha256=8W9m-WeEVE2ytYi9udKEA8Wtb0EnvP3eT2A1Tu-d29k,2252
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/qwen_image.py,sha256=ZIjB4GFIL0_xY8FkXdOJJ9-Xcft54rnBCrz43VWZLi0,2296
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py,sha256=EH2PpYFJ76KQLb35za6T-nM8Q10owLNatv6cd480ydE,2584
-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=1qarKAsEFiaaN2_ghko2dqGz_R7BTQSOyGtb_eQq38Y,35716
-cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/__init__.py,sha256=hVBTXj9MMGFGVezT3j8MntFRBiphSaUL4YhSOd8JtuY,1870
-cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/cogvideox.py,sha256=KP8NxtHAKzzBOoX0lhvlMgY_5dmP4Z3T5TOfwl4SSyg,2273
-cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/flux.py,sha256=kCB7lL4OIq8TZn-baMIF8D_PVPTFW60omCMVQCb8ebs,2628
-cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/hunyuan_video.py,sha256=xAkd40BGsfuCKdW3Abrx35VwgZQg4CZFz13P4VY71eY,9968
-cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/mochi.py,sha256=zXgoRDDjus3a2WSjtNh4ERtQp20ceb6nzohHMDlo2zY,2265
-cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/wan.py,sha256=PA7nuLgfAelnaI8usQx0Kxi8XATzMapyR1WndEdFoZA,2604
-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=qn4zWJ_eEMIPYzrxXoslunxbzK0WueuNtC54Pp5Q57k,23241
-cache_dit/cache_factory/first_block_cache/diffusers_adapters/__init__.py,sha256=-FFgA2MoudEo7uDacg4aWgm1KwfLZFsEDTVxatgbq9M,2146
-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/hunyuan_video.py,sha256=OL7W4ukYlZz0IDmBR1zVV6XT3Mgciglj9Hqzv1wUAkQ,10092
-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=dBNzHBECAuTTA1a7kLdvZL20YzaKTAS3iciVLzKKEWA,2638
-cache_dit/compile/__init__.py,sha256=DfMdPleFFGADXLsr7zXui8BTz_y9futY6rNmNdh9y7k,63
-cache_dit/compile/utils.py,sha256=N4A55_8uIbEd-S4xyJPcrdKceI2MGM9BTIhJE63jyL4,3786
-cache_dit/custom_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/custom_ops/triton_taylorseer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/metrics/__init__.py,sha256=RaUhl5dieF40RqnizGzR30qoJJ9dyMUEADwgwMaMQrE,575
-cache_dit/metrics/config.py,sha256=ieOgD9ayz722RjVzk24bSIqS2D6o7TZjGk8KeXV-OLQ,551
-cache_dit/metrics/fid.py,sha256=9Ivtazl6mW0Bon2VXa-Ia5Xj2ewxRD3V1Qkd69zYM3Y,17066
-cache_dit/metrics/inception.py,sha256=pBVe2X6ylLPIXTG4-GWDM9DWnCviMJbJ45R3ulhktR0,12759
-cache_dit/metrics/lpips.py,sha256=I2qCNi6qJh5TRsaIsdxO0WoRX1DN7U_H3zS0oCSahYM,1032
-cache_dit/metrics/metrics.py,sha256=8jvM1sF-nDxUuwCRy44QEoo4dYVLCQVh1QyAMs4eaQY,27840
-cache_dit-0.2.15.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.15.dist-info/METADATA,sha256=7cxw9ZaYpCFbtA1iDt-CWhk07mDXuRPECxfuO-wB0IE,25153
-cache_dit-0.2.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.15.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.2.15.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.15.dist-info/RECORD,,