cache-dit 0.2.12__py3-none-any.whl → 0.2.13__py3-none-any.whl

cache_dit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.12'
-__version_tuple__ = version_tuple = (0, 2, 12)
+__version__ = version = '0.2.13'
+__version_tuple__ = version_tuple = (0, 2, 13)

cache_dit/metrics/lpips.py

@@ -0,0 +1,43 @@
+import builtins as __builtin__
+import contextlib
+import warnings
+
+import lpips
+import torch
+
+warnings.filterwarnings("ignore")
+
+lpips_loss_fn_vgg = None
+lpips_loss_fn_alex = None
+
+
+def dummy_print(*args, **kwargs):
+    pass
+
+
+@contextlib.contextmanager
+def disable_print():
+    origin_print = __builtin__.print
+    __builtin__.print = dummy_print
+    yield
+    __builtin__.print = origin_print
+
+
+def compute_lpips_img(img0, img1, net: str = "alex"):
+    global lpips_loss_fn_vgg
+    global lpips_loss_fn_alex
+    if net.lower() == "alex":
+        if lpips_loss_fn_alex is None:
+            with disable_print():
+                lpips_loss_fn_alex = lpips.LPIPS(net="alex")
+        loss_fn = lpips_loss_fn_alex
+    elif net.lower() == "vgg":
+        if lpips_loss_fn_vgg is None:
+            with disable_print():
+                lpips_loss_fn_vgg = lpips.LPIPS(net="vgg")
+        loss_fn = lpips_loss_fn_vgg
+    else:
+        assert False, f"unsupport net {net}"
+
+    with torch.no_grad():
+        return loss_fn(img0, img1).item()

cache_dit/metrics/metrics.py

@@ -14,6 +14,7 @@ from cache_dit.metrics.config import get_metrics_verbose
 from cache_dit.metrics.config import _IMAGE_EXTENSIONS
 from cache_dit.metrics.config import _VIDEO_EXTENSIONS
 from cache_dit.logger import init_logger
+from cache_dit.metrics.lpips import compute_lpips_img
 
 logger = init_logger(__name__)
 
@@ -21,6 +22,35 @@ logger = init_logger(__name__)
 DISABLE_VERBOSE = not get_metrics_verbose()
 
 
+def compute_lpips_file(
+    image_true: np.ndarray | str,
+    image_test: np.ndarray | str,
+) -> float:
+    import torch
+    from PIL import Image
+    from torchvision.transforms.v2.functional import (
+        convert_image_dtype,
+        normalize,
+        pil_to_tensor,
+    )
+
+    def load_img_as_tensor(path):
+        pil = Image.open(path)
+        img = pil_to_tensor(pil)
+        img = convert_image_dtype(img, dtype=torch.float32)
+        img = normalize(img, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        return img
+
+    if isinstance(image_true, str):
+        image_true = load_img_as_tensor(image_true)
+    if isinstance(image_test, str):
+        image_test = load_img_as_tensor(image_test)
+    return compute_lpips_img(
+        image_true,
+        image_test,
+    )
+
+
 def compute_psnr_file(
     image_true: np.ndarray | str,
     image_test: np.ndarray | str,
@@ -305,6 +335,11 @@ def compute_video_metric(
         return None, None
 
 
+compute_lpips = partial(
+    compute_dir_metric,
+    compute_file_func=compute_lpips_file,
+)
+
 compute_psnr = partial(
     compute_dir_metric,
     compute_file_func=compute_psnr_file,
@@ -320,6 +355,10 @@ compute_mse = partial(
     compute_file_func=compute_mse_file,
 )
 
+compute_video_lpips = partial(
+    compute_video_metric,
+    compute_frame_func=compute_lpips_file,
+)
 compute_video_psnr = partial(
     compute_video_metric,
     compute_frame_func=compute_psnr_file,
@@ -335,6 +374,7 @@ compute_video_mse = partial(
 
 
 METRICS_CHOICES = [
+    "lpips",
     "psnr",
     "ssim",
     "mse",
@@ -522,6 +562,9 @@ def entrypoint():
                 METRICS_META[msg] = value
                 logger.info(msg)
 
+            if metric == "lpips" or metric == "all":
+                img_lpips, n = compute_lpips(img_true, img_test)
+                _logging_msg(img_lpips, "lpips", n)
             if metric == "psnr" or metric == "all":
                 img_psnr, n = compute_psnr(img_true, img_test)
                 _logging_msg(img_psnr, "psnr", n)
@@ -558,6 +601,9 @@ def entrypoint():
                 METRICS_META[msg] = value
                 logger.info(msg)
 
+            if metric == "lpips" or metric == "all":
+                video_lpips, n = compute_video_lpips(video_true, video_test)
+                _logging_msg(video_lpips, "lpips", n)
             if metric == "psnr" or metric == "all":
                 video_psnr, n = compute_video_psnr(video_true, video_test)
                 _logging_msg(video_psnr, "psnr", n)

{cache_dit-0.2.12.dist-info → cache_dit-0.2.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.12
+Version: 0.2.13
 Summary: 🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration Toolbox for Diffusion Transformers
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
@@ -16,6 +16,7 @@ Requires-Dist: transformers>=4.51.3
 Requires-Dist: diffusers>=0.33.1
 Requires-Dist: scikit-image
 Requires-Dist: scipy
+Requires-Dist: lpips==0.1.4
 Provides-Extra: all
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"

{cache_dit-0.2.12.dist-info → cache_dit-0.2.13.dist-info}/RECORD

@@ -1,5 +1,5 @@
 cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/_version.py,sha256=7CFcHKqzy7OglwTX58ipGg1TD8MpDORqWvEBO3W1dHI,513
+cache_dit/_version.py,sha256=2ECxD0Bipdh9vxnyteM0k9jxi9NOpPR7YxTi7Ad1ors,513
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
 cache_dit/primitives.py,sha256=A2iG9YLot3gOsZSPp-_gyjqjLgJvWQRx8aitD4JQ23Y,3877
 cache_dit/cache_factory/__init__.py,sha256=iYQwLwB_XLoYl0OB9unZGDbBtrYvZaLkOAmhGRwdW2E,191
@@ -38,10 +38,11 @@ cache_dit/metrics/__init__.py,sha256=RaUhl5dieF40RqnizGzR30qoJJ9dyMUEADwgwMaMQrE
 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/metrics.py,sha256=1TTbfaj_-vdUfxopLnc5kVrXs5rMpAoSi8D0ItYdPu8,26439
-cache_dit-0.2.12.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.12.dist-info/METADATA,sha256=-AIWGVOFsY-nhMkDeFErUFcELTWmza96-0IUN3od88A,28219
-cache_dit-0.2.12.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.12.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.2.12.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.12.dist-info/RECORD,,
+cache_dit/metrics/lpips.py,sha256=I2qCNi6qJh5TRsaIsdxO0WoRX1DN7U_H3zS0oCSahYM,1032
+cache_dit/metrics/metrics.py,sha256=8jvM1sF-nDxUuwCRy44QEoo4dYVLCQVh1QyAMs4eaQY,27840
+cache_dit-0.2.13.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.13.dist-info/METADATA,sha256=at8DNFeGI5aVnBTi7_6zJgAi_QdgsItpBMzSGl8HEME,28247
+cache_dit-0.2.13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.13.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.13.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.13.dist-info/RECORD,,