cache-dit 0.2.6__py3-none-any.whl → 0.2.8__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.6'
-__version_tuple__ = version_tuple = (0, 2, 6)
+__version__ = version = '0.2.8'
+__version_tuple__ = version_tuple = (0, 2, 8)

cache_dit/cache_factory/dual_block_cache/cache_context.py

@@ -988,7 +988,7 @@ def get_Bn_encoder_buffer(prefix: str = "Bn"):
 @torch.compiler.disable
 def apply_hidden_states_residual(
     hidden_states: torch.Tensor,
-    encoder_hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor = None,
     prefix: str = "Bn",
     encoder_prefix: str = "Bn_encoder",
 ):
@@ -1006,25 +1006,27 @@ def apply_hidden_states_residual(
         # If cache is not residual, we use the hidden states directly
         hidden_states = hidden_states_prev
 
-    if "Bn" in encoder_prefix:
-        encoder_hidden_states_prev = get_Bn_encoder_buffer(encoder_prefix)
-    else:
-        encoder_hidden_states_prev = get_Fn_encoder_buffer(encoder_prefix)
+    hidden_states = hidden_states.contiguous()
 
-    assert (
-        encoder_hidden_states_prev is not None
-    ), f"{prefix}_encoder_buffer must be set before"
+    if encoder_hidden_states is not None:
+        if "Bn" in encoder_prefix:
+            encoder_hidden_states_prev = get_Bn_encoder_buffer(encoder_prefix)
+        else:
+            encoder_hidden_states_prev = get_Fn_encoder_buffer(encoder_prefix)
 
-    if is_encoder_cache_residual():
-        encoder_hidden_states = (
-            encoder_hidden_states_prev + encoder_hidden_states
-        )
-    else:
-        # If encoder cache is not residual, we use the encoder hidden states directly
-        encoder_hidden_states = encoder_hidden_states_prev
+        assert (
+            encoder_hidden_states_prev is not None
+        ), f"{prefix}_encoder_buffer must be set before"
 
-    hidden_states = hidden_states.contiguous()
-    encoder_hidden_states = encoder_hidden_states.contiguous()
+        if is_encoder_cache_residual():
+            encoder_hidden_states = (
+                encoder_hidden_states_prev + encoder_hidden_states
+            )
+        else:
+            # If encoder cache is not residual, we use the encoder hidden states directly
+            encoder_hidden_states = encoder_hidden_states_prev
+
+        encoder_hidden_states = encoder_hidden_states.contiguous()
 
     return hidden_states, encoder_hidden_states
 
@@ -1160,7 +1162,6 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
         torch._dynamo.graph_break()
         if can_use_cache:
-            torch._dynamo.graph_break()
             add_cached_step()
             del Fn_hidden_states_residual
             hidden_states, encoder_hidden_states = apply_hidden_states_residual(
@@ -1187,7 +1188,6 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 )
             )
         else:
-            torch._dynamo.graph_break()
             set_Fn_buffer(Fn_hidden_states_residual, prefix="Fn_residual")
             if is_l1_diff_enabled():
                 # for hidden states L1 diff
@@ -1795,7 +1795,6 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 f"the number of single transformer blocks {len(self.single_transformer_blocks)}"
             )
 
-            torch._dynamo.graph_break()
             hidden_states = torch.cat(
                 [encoder_hidden_states, hidden_states], dim=1
             )
@@ -1827,13 +1826,11 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 ],
                 dim=1,
             )
-            torch._dynamo.graph_break()
         else:
             assert Bn_compute_blocks() <= len(self.transformer_blocks), (
                 f"Bn_compute_blocks {Bn_compute_blocks()} must be less than "
                 f"the number of transformer blocks {len(self.transformer_blocks)}"
             )
-            torch._dynamo.graph_break()
             if len(Bn_compute_blocks_ids()) > 0:
                 for i, block in enumerate(self._Bn_transformer_blocks()):
                     hidden_states, encoder_hidden_states = (
@@ -1862,7 +1859,6 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                                 encoder_hidden_states,
                                 hidden_states,
                             )
-            torch._dynamo.graph_break()
 
         hidden_states = (
             hidden_states.reshape(-1)

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py

@@ -82,6 +82,6 @@ def apply_db_cache_on_pipe(
         pipe.__class__._is_cached = True
 
     if not shallow_patch:
-        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+        apply_db_cache_on_transformer(pipe.transformer)
 
     return pipe

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py

@@ -93,6 +93,6 @@ def apply_db_cache_on_pipe(
         pipe.__class__._is_cached = True
 
     if not shallow_patch:
-        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+        apply_db_cache_on_transformer(pipe.transformer)
 
     return pipe

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py

@@ -289,6 +289,6 @@ def apply_db_cache_on_pipe(
         pipe.__class__._is_cached = True
 
     if not shallow_patch:
-        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+        apply_db_cache_on_transformer(pipe.transformer)
 
     return pipe

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py

@@ -82,6 +82,6 @@ def apply_db_cache_on_pipe(
         pipe.__class__._is_cached = True
 
     if not shallow_patch:
-        apply_db_cache_on_transformer(pipe.transformer, **kwargs)
+        apply_db_cache_on_transformer(pipe.transformer)
 
     return pipe

cache_dit/metrics/__init__.py

@@ -6,6 +6,8 @@ from cache_dit.metrics.metrics import compute_video_ssim
 from cache_dit.metrics.metrics import compute_video_mse
 from cache_dit.metrics.metrics import entrypoint
 from cache_dit.metrics.fid import FrechetInceptionDistance
+from cache_dit.metrics.config import set_metrics_verbose
+from cache_dit.metrics.config import get_metrics_verbose
 
 
 def main():

cache_dit/metrics/config.py

@@ -0,0 +1,34 @@
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+_metrics_progress_verbose = False
+
+
+def set_metrics_verbose(verbose: bool):
+    global _metrics_progress_verbose
+    _metrics_progress_verbose = verbose
+    logger.debug(f"Metrics verbose: {verbose}")
+
+
+def get_metrics_verbose() -> bool:
+    global _metrics_progress_verbose
+    return _metrics_progress_verbose
+
+
+_IMAGE_EXTENSIONS = [
+    "bmp",
+    "jpg",
+    "jpeg",
+    "pgm",
+    "png",
+    "ppm",
+    "tif",
+    "tiff",
+    "webp",
+]
+
+_VIDEO_EXTENSIONS = [
+    "mp4",
+]

cache_dit/metrics/fid.py

@@ -9,11 +9,14 @@ import torch
 import torchvision.transforms as TF
 from torch.nn.functional import adaptive_avg_pool2d
 from cache_dit.metrics.inception import InceptionV3
+from cache_dit.metrics.config import _IMAGE_EXTENSIONS
+from cache_dit.metrics.config import _VIDEO_EXTENSIONS
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
+# Adapted from: https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py
 class ImagePathDataset(torch.utils.data.Dataset):
     def __init__(self, files_or_imgs, transforms=None):
         self.files_or_imgs = files_or_imgs
@@ -219,22 +222,7 @@ def calculate_activation_statistics(
     return mu, sigma
 
 
-_IMAGE_EXTENSIONS = {
-    "bmp",
-    "jpg",
-    "jpeg",
-    "pgm",
-    "png",
-    "ppm",
-    "tif",
-    "tiff",
-    "webp",
-}
-
-
 class FrechetInceptionDistance:
-    IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS
-
     def __init__(
         self,
         device="cuda" if torch.cuda.is_available() else "cpu",
@@ -258,7 +246,8 @@ class FrechetInceptionDistance:
         image_true: np.ndarray | str,
         image_test: np.ndarray | str,
     ):
-        """Calculates the FID of two file paths
+        """
+        Calculates the FID of two file paths
         FID = FrechetInceptionDistance()
         img_fid = FID.compute_fid("img_true.png", "img_test.png")
         img_dir_fid = FID.compute_fid("img_true_dir", "img_test_dir")
@@ -267,8 +256,8 @@ class FrechetInceptionDistance:
             if os.path.isfile(image_true) or os.path.isfile(image_test):
                 assert os.path.exists(image_true)
                 assert os.path.exists(image_test)
-                assert image_true.split(".")[-1] in self.IMAGE_EXTENSIONS
-                assert image_test.split(".")[-1] in self.IMAGE_EXTENSIONS
+                assert image_true.split(".")[-1] in _IMAGE_EXTENSIONS
+                assert image_test.split(".")[-1] in _IMAGE_EXTENSIONS
                 image_true_files = [image_true]
                 image_test_files = [image_test]
             else:
@@ -279,7 +268,7 @@ class FrechetInceptionDistance:
                 image_true_files = sorted(
                     [
                         file
-                        for ext in self.IMAGE_EXTENSIONS
+                        for ext in _IMAGE_EXTENSIONS
                         for file in image_true_dir.rglob("*.{}".format(ext))
                     ]
                 )
@@ -287,7 +276,7 @@ class FrechetInceptionDistance:
                 image_test_files = sorted(
                     [
                         file
-                        for ext in self.IMAGE_EXTENSIONS
+                        for ext in _IMAGE_EXTENSIONS
                         for file in image_test_dir.rglob("*.{}".format(ext))
                     ]
                 )
@@ -297,15 +286,32 @@ class FrechetInceptionDistance:
                 image_test_files = [
                     file.as_posix() for file in image_test_files
                 ]
+
+                # select valid files
+                image_true_files_selected = []
+                image_test_files_selected = []
+                for i in range(
+                    min(len(image_true_files), len(image_test_files))
+                ):
+                    selected_image_true = image_true_files[i]
+                    selected_image_test = image_test_files[i]
+                    # Image pair must have the same basename
+                    if os.path.basename(
+                        selected_image_test
+                    ) == os.path.basename(selected_image_true):
+                        image_true_files_selected.append(selected_image_true)
+                        image_test_files_selected.append(selected_image_test)
+                image_true_files = image_true_files_selected.copy()
+                image_test_files = image_test_files_selected.copy()
+                if len(image_true_files) == 0:
+                    logger.error(
+                        "No valid Image pairs, please note that Image "
+                        "pairs must have the same basename."
+                    )
+                    return None, None
+
                 logger.debug(f"image_true_files: {image_true_files}")
                 logger.debug(f"image_test_files: {image_test_files}")
-                assert len(image_true_files) == len(image_test_files)
-                for image_true, image_test in zip(
-                    image_true_files, image_test_files
-                ):
-                    assert os.path.basename(image_true) == os.path.basename(
-                        image_test
-                    ), f"image_true:{image_true} != image_test: {image_test}"
         else:
             image_true_files = [image_true]
             image_test_files = [image_test]
@@ -340,6 +346,115 @@ class FrechetInceptionDistance:
         return fid_value, len(image_true_files)
 
     def compute_video_fid(
+        self,
+        # file or dir
+        video_true: str,
+        video_test: str,
+    ):
+        if os.path.isfile(video_true) and os.path.isfile(video_test):
+            video_true_frames, video_test_frames, valid_frames = (
+                self._fetch_video_frames(
+                    video_true=video_true,
+                    video_test=video_test,
+                )
+            )
+        elif os.path.isdir(video_true) and os.path.isdir(video_test):
+            # Glob videos
+            video_true_dir: pathlib.Path = pathlib.Path(video_true)
+            video_true_files = sorted(
+                [
+                    file
+                    for ext in _VIDEO_EXTENSIONS
+                    for file in video_true_dir.rglob("*.{}".format(ext))
+                ]
+            )
+            video_test_dir: pathlib.Path = pathlib.Path(video_test)
+            video_test_files = sorted(
+                [
+                    file
+                    for ext in _VIDEO_EXTENSIONS
+                    for file in video_test_dir.rglob("*.{}".format(ext))
+                ]
+            )
+            video_true_files = [file.as_posix() for file in video_true_files]
+            video_test_files = [file.as_posix() for file in video_test_files]
+
+            # select valid video files
+            video_true_files_selected = []
+            video_test_files_selected = []
+            for i in range(min(len(video_true_files), len(video_test_files))):
+                selected_video_true = video_true_files[i]
+                selected_video_test = video_test_files[i]
+                # Video pair must have the same basename
+                if os.path.basename(selected_video_test) == os.path.basename(
+                    selected_video_true
+                ):
+                    video_true_files_selected.append(selected_video_true)
+                    video_test_files_selected.append(selected_video_test)
+
+            video_true_files = video_true_files_selected.copy()
+            video_test_files = video_test_files_selected.copy()
+            if len(video_true_files) == 0:
+                logger.error(
+                    "No valid Video pairs, please note that Video "
+                    "pairs must have the same basename."
+                )
+                return None, None
+            logger.debug(f"video_true_files: {video_true_files}")
+            logger.debug(f"video_test_files: {video_test_files}")
+
+            # Fetch all frames
+            video_true_frames = []
+            video_test_frames = []
+            valid_frames = 0
+
+            for video_true_, video_test_ in zip(
+                video_true_files, video_test_files
+            ):
+                video_true_frames_, video_test_frames_, valid_frames_ = (
+                    self._fetch_video_frames(
+                        video_true=video_true_, video_test=video_test_
+                    )
+                )
+                video_true_frames.extend(video_true_frames_)
+                video_test_frames.extend(video_test_frames_)
+                valid_frames += valid_frames_
+        else:
+            raise ValueError("video_true and video_test must be files or dirs.")
+
+        if valid_frames <= 0:
+            logger.debug("No valid frames to compare")
+            return None, None
+
+        batch_size = min(16, self.batch_size)
+        m1, s1 = calculate_activation_statistics(
+            video_true_frames,
+            self.model,
+            batch_size,
+            self.dims,
+            self.device,
+            self.num_workers,
+            self.disable_tqdm,
+        )
+        m2, s2 = calculate_activation_statistics(
+            video_test_frames,
+            self.model,
+            batch_size,
+            self.dims,
+            self.device,
+            self.num_workers,
+            self.disable_tqdm,
+        )
+        fid_value = calculate_frechet_distance(
+            m1,
+            s1,
+            m2,
+            s2,
+        )
+
+        return fid_value, valid_frames
+
+    def _fetch_video_frames(
         self,
         video_true: str,
         video_test: str,
@@ -349,7 +464,7 @@ class FrechetInceptionDistance:
 
         if not cap1.isOpened() or not cap2.isOpened():
             logger.error("Could not open video files")
-            return None, None
+            return [], [], 0
 
         frame_count = min(
             int(cap1.get(cv2.CAP_PROP_FRAME_COUNT)),
@@ -378,32 +493,6 @@ class FrechetInceptionDistance:
         cap2.release()
 
         if valid_frames <= 0:
-            return None, None
+            return [], [], 0
 
-        batch_size = min(16, self.batch_size)
-        m1, s1 = calculate_activation_statistics(
-            video_true_frames,
-            self.model,
-            batch_size,
-            self.dims,
-            self.device,
-            self.num_workers,
-            self.disable_tqdm,
-        )
-        m2, s2 = calculate_activation_statistics(
-            video_test_frames,
-            self.model,
-            batch_size,
-            self.dims,
-            self.device,
-            self.num_workers,
-            self.disable_tqdm,
-        )
-        fid_value = calculate_frechet_distance(
-            m1,
-            s1,
-            m2,
-            s2,
-        )
-
-        return fid_value, valid_frames
+        return video_true_frames, video_test_frames, valid_frames

cache_dit/metrics/metrics.py

@@ -3,16 +3,24 @@ import cv2
 import pathlib
 import argparse
 import numpy as np
+from tqdm import tqdm
 from functools import partial
 from skimage.metrics import mean_squared_error
 from skimage.metrics import peak_signal_noise_ratio
 from skimage.metrics import structural_similarity
 from cache_dit.metrics.fid import FrechetInceptionDistance
+from cache_dit.metrics.config import set_metrics_verbose
+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
 
 logger = init_logger(__name__)
 
 
+DISABLE_VERBOSE = not get_metrics_verbose()
+
+
 def compute_psnr_file(
     image_true: np.ndarray | str,
     image_test: np.ndarray | str,
@@ -72,19 +80,6 @@ def compute_ssim_file(
     )
 
 
-_IMAGE_EXTENSIONS = {
-    "bmp",
-    "jpg",
-    "jpeg",
-    "pgm",
-    "png",
-    "ppm",
-    "tif",
-    "tiff",
-    "webp",
-}
-
-
 def compute_dir_metric(
     image_true_dir: np.ndarray | str,
     image_test_dir: np.ndarray | str,
@@ -117,17 +112,38 @@ def compute_dir_metric(
     )
     image_true_files = [file.as_posix() for file in image_true_files]
     image_test_files = [file.as_posix() for file in image_test_files]
+
+    # select valid files
+    image_true_files_selected = []
+    image_test_files_selected = []
+    for i in range(min(len(image_true_files), len(image_test_files))):
+        selected_image_true = image_true_files[i]
+        selected_image_test = image_test_files[i]
+        # Image pair must have the same basename
+        if os.path.basename(selected_image_test) == os.path.basename(
+            selected_image_true
+        ):
+            image_true_files_selected.append(selected_image_true)
+            image_test_files_selected.append(selected_image_test)
+    image_true_files = image_true_files_selected.copy()
+    image_test_files = image_test_files_selected.copy()
+    if len(image_true_files) == 0:
+        logger.error(
+            "No valid Image pairs, please note that Image "
+            "pairs must have the same basename."
+        )
+        return None, None
+
     logger.debug(f"image_true_files: {image_true_files}")
     logger.debug(f"image_test_files: {image_test_files}")
-    assert len(image_true_files) == len(image_test_files)
-    for image_true, image_test in zip(image_true_files, image_test_files):
-        assert os.path.basename(image_true) == os.path.basename(
-            image_test
-        ), f"image_true:{image_true} != image_test: {image_test}"
 
     total_metric = 0.0
     valid_files = 0
-    for image_true, image_test in zip(image_true_files, image_test_files):
+    for image_true, image_test in tqdm(
+        zip(image_true_files, image_test_files),
+        total=len(image_true_files),
+        disable=DISABLE_VERBOSE,
+    ):
         metric = compute_file_func(image_true, image_test)
         if metric != float("inf"):
             total_metric += metric
@@ -142,30 +158,25 @@ def compute_dir_metric(
         return None, None
 
 
-def compute_video_metric(
+def _fetch_video_frames(
     video_true: str,
     video_test: str,
-    compute_frame_func: callable = compute_psnr_file,
-) -> float:
-    """
-    video_true = "video_true.mp4"
-    video_test = "video_test.mp4"
-    PSNR = compute_video_psnr(video_true, video_test)
-    """
+):
     cap1 = cv2.VideoCapture(video_true)
     cap2 = cv2.VideoCapture(video_test)
 
     if not cap1.isOpened() or not cap2.isOpened():
         logger.error("Could not open video files")
-        return None
+        return [], [], 0
 
     frame_count = min(
         int(cap1.get(cv2.CAP_PROP_FRAME_COUNT)),
         int(cap2.get(cv2.CAP_PROP_FRAME_COUNT)),
     )
 
-    total_metric = 0.0
     valid_frames = 0
+    video_true_frames = []
+    video_test_frames = []
 
     logger.debug(f"Total frames: {frame_count}")
 
@@ -176,18 +187,115 @@ def compute_video_metric(
         if not ret1 or not ret2:
             break
 
-        metric = compute_frame_func(frame1, frame2)
+        video_true_frames.append(frame1)
+        video_test_frames.append(frame2)
 
-        if metric != float("inf"):
-            total_metric += metric
-            valid_frames += 1
-
-        if valid_frames % 10 == 0:
-            logger.debug(f"Processed {valid_frames}/{frame_count} frames")
+        valid_frames += 1
 
     cap1.release()
     cap2.release()
 
+    if valid_frames <= 0:
+        return [], [], 0
+
+    return video_true_frames, video_test_frames, valid_frames
+
+
+def compute_video_metric(
+    video_true: str,
+    video_test: str,
+    compute_frame_func: callable = compute_psnr_file,
+) -> float:
+    """
+    video_true = "video_true.mp4"
+    video_test = "video_test.mp4"
+    PSNR = compute_video_psnr(video_true, video_test)
+    """
+    if os.path.isfile(video_true) and os.path.isfile(video_test):
+        video_true_frames, video_test_frames, valid_frames = (
+            _fetch_video_frames(
+                video_true=video_true,
+                video_test=video_test,
+            )
+        )
+    elif os.path.isdir(video_true) and os.path.isdir(video_test):
+        # Glob videos
+        video_true_dir: pathlib.Path = pathlib.Path(video_true)
+        video_true_files = sorted(
+            [
+                file
+                for ext in _VIDEO_EXTENSIONS
+                for file in video_true_dir.rglob("*.{}".format(ext))
+            ]
+        )
+        video_test_dir: pathlib.Path = pathlib.Path(video_test)
+        video_test_files = sorted(
+            [
+                file
+                for ext in _VIDEO_EXTENSIONS
+                for file in video_test_dir.rglob("*.{}".format(ext))
+            ]
+        )
+        video_true_files = [file.as_posix() for file in video_true_files]
+        video_test_files = [file.as_posix() for file in video_test_files]
+
+        # select valid video files
+        video_true_files_selected = []
+        video_test_files_selected = []
+        for i in range(min(len(video_true_files), len(video_test_files))):
+            selected_video_true = video_true_files[i]
+            selected_video_test = video_test_files[i]
+            # Video pair must have the same basename
+            if os.path.basename(selected_video_test) == os.path.basename(
+                selected_video_true
+            ):
+                video_true_files_selected.append(selected_video_true)
+                video_test_files_selected.append(selected_video_test)
+
+        video_true_files = video_true_files_selected.copy()
+        video_test_files = video_test_files_selected.copy()
+        if len(video_true_files) == 0:
+            logger.error(
+                "No valid Video pairs, please note that Video "
+                "pairs must have the same basename."
+            )
+            return None, None
+        logger.debug(f"video_true_files: {video_true_files}")
+        logger.debug(f"video_test_files: {video_test_files}")
+
+        # Fetch all frames
+        video_true_frames = []
+        video_test_frames = []
+        valid_frames = 0
+
+        for video_true_, video_test_ in zip(video_true_files, video_test_files):
+            video_true_frames_, video_test_frames_, valid_frames_ = (
+                _fetch_video_frames(
+                    video_true=video_true_, video_test=video_test_
+                )
+            )
+            video_true_frames.extend(video_true_frames_)
+            video_test_frames.extend(video_test_frames_)
+            valid_frames += valid_frames_
+    else:
+        raise ValueError("video_true and video_test must be files or dirs.")
+
+    if valid_frames <= 0:
+        logger.debug("No valid frames to compare")
+        return None, None
+
+    total_metric = 0.0
+    valid_frames = 0  # reset
+    for frame1, frame2 in tqdm(
+        zip(video_true_frames, video_test_frames),
+        total=len(video_true_frames),
+        disable=DISABLE_VERBOSE,
+    ):
+        metric = compute_frame_func(frame1, frame2)
+        if metric != float("inf"):
+            total_metric += metric
+            valid_frames += 1
+
     if valid_frames > 0:
         average_metric = total_metric / valid_frames
         logger.debug(f"Average: {average_metric:.2f}")
@@ -265,14 +373,21 @@ def get_args():
         "-v1",
         type=str,
         default=None,
-        help="Path to ground truth video",
+        help="Path to ground truth video or Dir to ground truth videos",
     )
     parser.add_argument(
         "--video-test",
         "-v2",
         type=str,
         default=None,
-        help="Path to predicted video",
+        help="Path to predicted video or Dir to predicted videos",
+    )
+    parser.add_argument(
+        "--enable-verbose",
+        "-verbose",
+        action="store_true",
+        default=False,
+        help="Show metrics progress verbose",
     )
     return parser.parse_args()
 
@@ -281,6 +396,11 @@ def entrypoint():
     args = get_args()
     logger.debug(args)
 
+    if args.enable_verbose:
+        global DISABLE_VERBOSE
+        set_metrics_verbose(True)
+        DISABLE_VERBOSE = not get_metrics_verbose()
+
     if args.img_true is not None and args.img_test is not None:
         if any(
             (
@@ -306,7 +426,7 @@ def entrypoint():
                 f"{args.img_true} vs {args.img_test}, Num: {n},  MSE: {img_mse}"
             )
         if args.metric == "fid" or args.metric == "all":
-            FID = FrechetInceptionDistance()
+            FID = FrechetInceptionDistance(disable_tqdm=DISABLE_VERBOSE)
             img_fid, n = FID.compute_fid(args.img_true, args.img_test)
             logger.info(
                 f"{args.img_true} vs {args.img_test}, Num: {n},  FID: {img_fid}"
@@ -319,36 +439,29 @@ def entrypoint():
             )
         ):
             return
+        # video_true and video_test can be files or dirs
         if args.metric == "psnr" or args.metric == "all":
-            assert not os.path.isdir(args.video_true)
-            assert not os.path.isdir(args.video_test)
             video_psnr, n = compute_video_psnr(args.video_true, args.video_test)
             logger.info(
-                f"{args.video_true} vs {args.video_test}, Num: {n}, PSNR: {video_psnr}"
+                f"{args.video_true} vs {args.video_test}, Frames: {n}, PSNR: {video_psnr}"
             )
         if args.metric == "ssim" or args.metric == "all":
-            assert not os.path.isdir(args.video_true)
-            assert not os.path.isdir(args.video_test)
             video_ssim, n = compute_video_ssim(args.video_true, args.video_test)
             logger.info(
-                f"{args.video_true} vs {args.video_test}, Num: {n}, SSIM: {video_ssim}"
+                f"{args.video_true} vs {args.video_test}, Frames: {n}, SSIM: {video_ssim}"
             )
         if args.metric == "mse" or args.metric == "all":
-            assert not os.path.isdir(args.video_true)
-            assert not os.path.isdir(args.video_test)
             video_mse, n = compute_video_mse(args.video_true, args.video_test)
             logger.info(
-                f"{args.video_true} vs {args.video_test}, Num: {n},  MSE: {video_mse}"
+                f"{args.video_true} vs {args.video_test}, Frames: {n},  MSE: {video_mse}"
             )
         if args.metric == "fid" or args.metric == "all":
-            assert not os.path.isdir(args.video_true)
-            assert not os.path.isdir(args.video_test)
-            FID = FrechetInceptionDistance()
+            FID = FrechetInceptionDistance(disable_tqdm=DISABLE_VERBOSE)
             video_fid, n = FID.compute_video_fid(
                 args.video_true, args.video_test
             )
             logger.info(
-                f"{args.video_true} vs {args.video_test}, Num: {n},  FID: {video_fid}"
+                f"{args.video_true} vs {args.video_test}, Frames: {n},  FID: {video_fid}"
             )
 
 

{cache_dit-0.2.6.dist-info → cache_dit-0.2.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.6
+Version: 0.2.8
 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
@@ -519,23 +519,21 @@ from cache_dit.metrics import FrechetInceptionDistance  # FID
 
 FID = FrechetInceptionDistance()
 image_psnr, n = compute_psnr("true.png", "test.png") # Num: n
-image_fid,  n = FID.compute_fid("true.png", "test.png")
-video_psnr, n = compute_video_psnr("true.mp4", "test.mp4")
+image_fid,  n = FID.compute_fid("true_dir", "test_dir")
+video_psnr, n = compute_video_psnr("true.mp4", "test.mp4") # Frames: n
 ```
 
 Please check [test_metrics.py](./tests/test_metrics.py) for more details. Or, you can use `cache-dit-metrics-cli` tool. For examples: 
 
 ```bash
 cache-dit-metrics-cli -h  # show usage
-cache-dit-metrics-cli all -v1 true.mp4 -v2 test.mp4  # compare video
-cache-dit-metrics-cli all -i1 true.png -i2 test.png  # compare image
-cache-dit-metrics-cli all -i1 true_dir -i2 test_dir  # compare image dir
-cache-dit-metrics-cli all -i1 BASELINE -i2 OPTIMIZED # compare image dir
-
-INFO 07-09 20:59:40 [metrics.py:295] BASELINE vs OPTIMIZED, Num: 1000, PSNR: 38.742413478199005
-INFO 07-09 21:00:32 [metrics.py:300] BASELINE vs OPTIMIZED, Num: 1000, SSIM: 0.9863484896791567
-INFO 07-09 21:00:45 [metrics.py:305] BASELINE vs OPTIMIZED, Num: 1000,  MSE: 12.287594770695606
-INFO 07-09 21:01:04 [metrics.py:311] BASELINE vs OPTIMIZED, Num: 1000,  FID: 5.983550108647762
+# all: PSNR, FID, SSIM, MSE, ..., etc.
+cache-dit-metrics-cli all  -i1 true.png -i2 test.png  # image
+cache-dit-metrics-cli all  -i1 true_dir -i2 test_dir  # image dir
+cache-dit-metrics-cli all  -v1 true.mp4 -v2 test.mp4  # video
+cache-dit-metrics-cli all  -v1 true_dir -v2 test_dir  # video dir
+cache-dit-metrics-cli fid  -i1 true_dir -i2 test_dir  # FID
+cache-dit-metrics-cli psnr -i1 true_dir -i2 test_dir  # PSNR
 ```
 
 ## 👋Contribute 

{cache_dit-0.2.6.dist-info → cache_dit-0.2.8.dist-info}/RECORD

@@ -1,17 +1,17 @@
 cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/_version.py,sha256=nObnONsicQ3YX6SG5MVBxmIp5dmRacXDauSqZijWQbY,511
+cache_dit/_version.py,sha256=zkhRarrvPoGA1yWjS9_zVM80dWqpDesNn9DiHcF4JWM,511
 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=5RNuhWakvvqrOV4vkqrEBA7d-V1LwcNSsjtW14mkqK8,5255
 cache_dit/cache_factory/taylorseer.py,sha256=LKSNo2ode69EVo9xrxjxAMEjz0yDGiGADeDYnEqddA8,3987
 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=wE_xYp7DRbgB-fD8dpr75o4Cvvl2s-jnT2fRyqWm_RM,71286
+cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=itVEb6gT2eZuncAHUmP51ZS0r6v6cGtRvnPjyeXqKH8,71156
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py,sha256=krNAICf-aS3JLmSG8vOB9tpLa04uYRcABsC8PMbVUKY,1870
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py,sha256=fibkeU-FHa30BNT-uPV2Eqcd5IRli07EKb25tMDp23c,2270
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py,sha256=fddSpTHXU24COMGAY-Z21EmHHAEArZBv_-XLRFD6ADU,2625
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py,sha256=wcZdBhjUB8WSfz40A268BtSe3nr_hRsIi2BNlg1FHRU,9965
-cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py,sha256=Cmy0KHRDgwXqtmqfkrr7kw0CP6CmkSnuz29gDHcD6sQ,2262
+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/wan.py,sha256=EREHM5E1wxnL-uRXRAEege4HXraRp1oD_r1Zx4CsiKk,2596
 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=so1wGdb8W0ATwrjv7E5IEZLPcobybaY1HJa6hBYlOOQ,34698
@@ -33,13 +33,14 @@ cache_dit/compile/__init__.py,sha256=DfMdPleFFGADXLsr7zXui8BTz_y9futY6rNmNdh9y7k
 cache_dit/compile/utils.py,sha256=KU60xc474Anbj7Y_FLRFmNxEjVYLLXkhbtCLXO7o_Tc,3699
 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=5yavk_6b50EhaHij0C-7nQegOA0uD79-qX96dI1i_8s,461
-cache_dit/metrics/fid.py,sha256=RIC-wW0RFv6W9IW6nc6Ih4dAxunUTnzcugMbYZjdqRM,12891
+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/metrics.py,sha256=FUrpc58ofg0LKyM3Y_7kfUVbevMfCFFcaeaxeOrj9iY,10498
-cache_dit-0.2.6.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.6.dist-info/METADATA,sha256=aTTyx_2gM7Z-mcIYd3_LlGljiZOwIq9iadDXOFPDzwA,27848
-cache_dit-0.2.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.6.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.2.6.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.6.dist-info/RECORD,,
+cache_dit/metrics/metrics.py,sha256=tzAtG_-fM1xPIBfRVFIBupvOWYzIO3xDq29Vy5rOBWc,14730
+cache_dit-0.2.8.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.8.dist-info/METADATA,sha256=8E51DpSKDGqk3_cG9buahoXN-7fub6M8VCiPb_Idg64,27608
+cache_dit-0.2.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.8.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.8.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.8.dist-info/RECORD,,