cache-dit 0.2.9__py3-none-any.whl → 0.2.11__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.9'
-__version_tuple__ = version_tuple = (0, 2, 9)
+__version__ = version = '0.2.11'
+__version_tuple__ = version_tuple = (0, 2, 11)

cache_dit/cache_factory/__init__.py

@@ -1,168 +1,3 @@
-from enum import Enum
-
-from diffusers import DiffusionPipeline
-
-from cache_dit.cache_factory.dual_block_cache.diffusers_adapters import (
-    apply_db_cache_on_pipe,
-)
-from cache_dit.cache_factory.first_block_cache.diffusers_adapters import (
-    apply_fb_cache_on_pipe,
-)
-from cache_dit.cache_factory.dynamic_block_prune.diffusers_adapters import (
-    apply_db_prune_on_pipe,
-)
-from cache_dit.logger import init_logger
-
-
-logger = init_logger(__name__)
-
-
-class CacheType(Enum):
-    NONE = "NONE"
-    FBCache = "First_Block_Cache"
-    DBCache = "Dual_Block_Cache"
-    DBPrune = "Dynamic_Block_Prune"
-
-    @staticmethod
-    def type(cache_type: "CacheType | str") -> "CacheType":
-        if isinstance(cache_type, CacheType):
-            return cache_type
-        return CacheType.cache_type(cache_type)
-
-    @staticmethod
-    def cache_type(cache_type: "CacheType | str") -> "CacheType":
-        if cache_type is None:
-            return CacheType.NONE
-
-        if isinstance(cache_type, CacheType):
-            return cache_type
-        if cache_type.lower() in (
-            "first_block_cache",
-            "fb_cache",
-            "fbcache",
-            "fb",
-        ):
-            return CacheType.FBCache
-        elif cache_type.lower() in (
-            "dual_block_cache",
-            "db_cache",
-            "dbcache",
-            "db",
-        ):
-            return CacheType.DBCache
-        elif cache_type.lower() in (
-            "dynamic_block_prune",
-            "db_prune",
-            "dbprune",
-            "dbp",
-        ):
-            return CacheType.DBPrune
-        elif cache_type.lower() in (
-            "none_cache",
-            "nonecache",
-            "no_cache",
-            "nocache",
-            "none",
-            "no",
-        ):
-            return CacheType.NONE
-        else:
-            raise ValueError(f"Unknown cache type: {cache_type}")
-
-    @staticmethod
-    def range(start: int, end: int, step: int = 1) -> list[int]:
-        if start > end or end <= 0 or step <= 1:
-            return []
-        # Always compute 0 and end - 1 blocks for DB Cache
-        return list(
-            sorted(set([0] + list(range(start, end, step)) + [end - 1]))
-        )
-
-    @staticmethod
-    def default_options(cache_type: "CacheType | str") -> dict:
-        _no_options = {
-            "cache_type": CacheType.NONE,
-        }
-
-        _fb_options = {
-            "cache_type": CacheType.FBCache,
-            "residual_diff_threshold": 0.08,
-            "warmup_steps": 8,
-            "max_cached_steps": 8,
-        }
-
-        _Fn_compute_blocks = 8
-        _Bn_compute_blocks = 8
-
-        _db_options = {
-            "cache_type": CacheType.DBCache,
-            "residual_diff_threshold": 0.12,
-            "warmup_steps": 8,
-            "max_cached_steps": -1,  # -1 means no limit
-            # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
-            "Fn_compute_blocks": _Fn_compute_blocks,
-            "Bn_compute_blocks": _Bn_compute_blocks,
-            "max_Fn_compute_blocks": 16,
-            "max_Bn_compute_blocks": 16,
-            "Fn_compute_blocks_ids": [],  # 0, 1, 2, ..., 7, etc.
-            "Bn_compute_blocks_ids": [],  # 0, 1, 2, ..., 7, etc.
-        }
-
-        _dbp_options = {
-            "cache_type": CacheType.DBPrune,
-            "residual_diff_threshold": 0.08,
-            "Fn_compute_blocks": _Fn_compute_blocks,
-            "Bn_compute_blocks": _Bn_compute_blocks,
-            "warmup_steps": 8,
-            "max_pruned_steps": -1,  # -1 means no limit
-        }
-
-        if cache_type == CacheType.FBCache:
-            return _fb_options
-        elif cache_type == CacheType.DBCache:
-            return _db_options
-        elif cache_type == CacheType.DBPrune:
-            return _dbp_options
-        elif cache_type == CacheType.NONE:
-            return _no_options
-        else:
-            raise ValueError(f"Unknown cache type: {cache_type}")
-
-
-def apply_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
-    assert isinstance(pipe, DiffusionPipeline)
-
-    if hasattr(pipe, "_is_cached") and pipe._is_cached:
-        return pipe
-
-    if hasattr(pipe, "_is_pruned") and pipe._is_pruned:
-        return pipe
-
-    cache_type = kwargs.pop("cache_type", None)
-    if cache_type is None:
-        logger.warning(
-            "No cache type specified, we will use DBCache by default. "
-            "Please specify the cache_type explicitly if you want to "
-            "use a different cache type."
-        )
-        # Force to use DBCache with default cache options
-        return apply_db_cache_on_pipe(
-            pipe,
-            **CacheType.default_options(CacheType.DBCache),
-        )
-
-    cache_type = CacheType.type(cache_type)
-
-    if cache_type == CacheType.FBCache:
-        return apply_fb_cache_on_pipe(pipe, *args, **kwargs)
-    elif cache_type == CacheType.DBCache:
-        return apply_db_cache_on_pipe(pipe, *args, **kwargs)
-    elif cache_type == CacheType.DBPrune:
-        return apply_db_prune_on_pipe(pipe, *args, **kwargs)
-    elif cache_type == CacheType.NONE:
-        logger.warning(
-            f"Cache type is {cache_type}, no caching will be applied."
-        )
-        return pipe
-    else:
-        raise ValueError(f"Unknown cache type: {cache_type}")
+from cache_dit.cache_factory.adapters import CacheType
+from cache_dit.cache_factory.adapters import apply_cache_on_pipe
+from cache_dit.cache_factory.utils import load_cache_options_from_yaml

cache_dit/cache_factory/adapters.py

@@ -0,0 +1,169 @@
+from enum import Enum
+
+from diffusers import DiffusionPipeline
+
+from cache_dit.cache_factory.dual_block_cache.diffusers_adapters import (
+    apply_db_cache_on_pipe,
+)
+from cache_dit.cache_factory.first_block_cache.diffusers_adapters import (
+    apply_fb_cache_on_pipe,
+)
+from cache_dit.cache_factory.dynamic_block_prune.diffusers_adapters import (
+    apply_db_prune_on_pipe,
+)
+
+from cache_dit.logger import init_logger
+
+
+logger = init_logger(__name__)
+
+
+class CacheType(Enum):
+    NONE = "NONE"
+    FBCache = "First_Block_Cache"
+    DBCache = "Dual_Block_Cache"
+    DBPrune = "Dynamic_Block_Prune"
+
+    @staticmethod
+    def type(cache_type: "CacheType | str") -> "CacheType":
+        if isinstance(cache_type, CacheType):
+            return cache_type
+        return CacheType.cache_type(cache_type)
+
+    @staticmethod
+    def cache_type(cache_type: "CacheType | str") -> "CacheType":
+        if cache_type is None:
+            return CacheType.NONE
+
+        if isinstance(cache_type, CacheType):
+            return cache_type
+        if cache_type.lower() in (
+            "first_block_cache",
+            "fb_cache",
+            "fbcache",
+            "fb",
+        ):
+            return CacheType.FBCache
+        elif cache_type.lower() in (
+            "dual_block_cache",
+            "db_cache",
+            "dbcache",
+            "db",
+        ):
+            return CacheType.DBCache
+        elif cache_type.lower() in (
+            "dynamic_block_prune",
+            "db_prune",
+            "dbprune",
+            "dbp",
+        ):
+            return CacheType.DBPrune
+        elif cache_type.lower() in (
+            "none_cache",
+            "nonecache",
+            "no_cache",
+            "nocache",
+            "none",
+            "no",
+        ):
+            return CacheType.NONE
+        else:
+            raise ValueError(f"Unknown cache type: {cache_type}")
+
+    @staticmethod
+    def range(start: int, end: int, step: int = 1) -> list[int]:
+        if start > end or end <= 0 or step <= 1:
+            return []
+        # Always compute 0 and end - 1 blocks for DB Cache
+        return list(
+            sorted(set([0] + list(range(start, end, step)) + [end - 1]))
+        )
+
+    @staticmethod
+    def default_options(cache_type: "CacheType | str") -> dict:
+        _no_options = {
+            "cache_type": CacheType.NONE,
+        }
+
+        _fb_options = {
+            "cache_type": CacheType.FBCache,
+            "residual_diff_threshold": 0.08,
+            "warmup_steps": 8,
+            "max_cached_steps": 8,
+        }
+
+        _Fn_compute_blocks = 8
+        _Bn_compute_blocks = 8
+
+        _db_options = {
+            "cache_type": CacheType.DBCache,
+            "residual_diff_threshold": 0.12,
+            "warmup_steps": 8,
+            "max_cached_steps": -1,  # -1 means no limit
+            # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
+            "Fn_compute_blocks": _Fn_compute_blocks,
+            "Bn_compute_blocks": _Bn_compute_blocks,
+            "max_Fn_compute_blocks": 16,
+            "max_Bn_compute_blocks": 16,
+            "Fn_compute_blocks_ids": [],  # 0, 1, 2, ..., 7, etc.
+            "Bn_compute_blocks_ids": [],  # 0, 1, 2, ..., 7, etc.
+        }
+
+        _dbp_options = {
+            "cache_type": CacheType.DBPrune,
+            "residual_diff_threshold": 0.08,
+            "Fn_compute_blocks": _Fn_compute_blocks,
+            "Bn_compute_blocks": _Bn_compute_blocks,
+            "warmup_steps": 8,
+            "max_pruned_steps": -1,  # -1 means no limit
+        }
+
+        if cache_type == CacheType.FBCache:
+            return _fb_options
+        elif cache_type == CacheType.DBCache:
+            return _db_options
+        elif cache_type == CacheType.DBPrune:
+            return _dbp_options
+        elif cache_type == CacheType.NONE:
+            return _no_options
+        else:
+            raise ValueError(f"Unknown cache type: {cache_type}")
+
+
+def apply_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
+    assert isinstance(pipe, DiffusionPipeline)
+
+    if hasattr(pipe, "_is_cached") and pipe._is_cached:
+        return pipe
+
+    if hasattr(pipe, "_is_pruned") and pipe._is_pruned:
+        return pipe
+
+    cache_type = kwargs.pop("cache_type", None)
+    if cache_type is None:
+        logger.warning(
+            "No cache type specified, we will use DBCache by default. "
+            "Please specify the cache_type explicitly if you want to "
+            "use a different cache type."
+        )
+        # Force to use DBCache with default cache options
+        return apply_db_cache_on_pipe(
+            pipe,
+            **CacheType.default_options(CacheType.DBCache),
+        )
+
+    cache_type = CacheType.type(cache_type)
+
+    if cache_type == CacheType.FBCache:
+        return apply_fb_cache_on_pipe(pipe, *args, **kwargs)
+    elif cache_type == CacheType.DBCache:
+        return apply_db_cache_on_pipe(pipe, *args, **kwargs)
+    elif cache_type == CacheType.DBPrune:
+        return apply_db_prune_on_pipe(pipe, *args, **kwargs)
+    elif cache_type == CacheType.NONE:
+        logger.warning(
+            f"Cache type is {cache_type}, no caching will be applied."
+        )
+        return pipe
+    else:
+        raise ValueError(f"Unknown cache type: {cache_type}")

cache_dit/cache_factory/utils.py

@@ -0,0 +1,53 @@
+import yaml
+from cache_dit.cache_factory.adapters import CacheType
+
+
+def load_cache_options_from_yaml(yaml_file_path):
+    try:
+        with open(yaml_file_path, "r") as f:
+            config = yaml.safe_load(f)
+
+        required_keys = [
+            "cache_type",
+            "warmup_steps",
+            "max_cached_steps",
+            "Fn_compute_blocks",
+            "Bn_compute_blocks",
+            "residual_diff_threshold",
+        ]
+        for key in required_keys:
+            if key not in config:
+                raise ValueError(
+                    f"Configuration file missing required item: {key}"
+                )
+
+        # Convert cache_type to CacheType enum
+        if isinstance(config["cache_type"], str):
+            try:
+                config["cache_type"] = CacheType[config["cache_type"]]
+            except KeyError:
+                valid_types = [ct.name for ct in CacheType]
+                raise ValueError(
+                    f"Invalid cache_type value: {config['cache_type']}, "
+                    f"valid values are: {valid_types}"
+                )
+        elif not isinstance(config["cache_type"], CacheType):
+            raise ValueError(
+                f"cache_type must be a string or CacheType enum, "
+                f"got: {type(config['cache_type'])}"
+            )
+
+        # Handle default value for taylorseer_kwargs
+        if "taylorseer_kwargs" not in config and config.get(
+            "enable_taylorseer", False
+        ):
+            config["taylorseer_kwargs"] = {"n_derivatives": 2}
+
+        return config
+
+    except FileNotFoundError:
+        raise FileNotFoundError(
+            f"Configuration file not found: {yaml_file_path}"
+        )
+    except yaml.YAMLError as e:
+        raise yaml.YAMLError(f"YAML file parsing error: {str(e)}")

cache_dit/compile/utils.py

@@ -39,6 +39,14 @@ def set_custom_compile_configs(
     # https://github.com/pytorch/pytorch/issues/153791
     torch._inductor.config.autotune_local_cache = False
 
+    if dist.is_initialized():
+        # Enable compute comm overlap
+        torch._inductor.config.reorder_for_compute_comm_overlap = True
+        # L20 64 GB/s, PCIe; A100/A800 NVLink 300 GB/s.
+        torch._inductor.config.intra_node_bw = (
+            64 if "L20" in torch.cuda.get_device_name() else 300
+        )
+
     FORCE_DISABLE_CUSTOM_COMPILE_CONFIG = (
         os.environ.get("CACHE_DIT_FORCE_DISABLE_CUSTOM_COMPILE_CONFIG", "0")
         == "1"
@@ -51,14 +59,6 @@ def set_custom_compile_configs(
         )
         return
 
-    if dist.is_initialized():
-        # Enable compute comm overlap
-        torch._inductor.config.reorder_for_compute_comm_overlap = True
-        # L20 64 GB/s, PCIe; A100/A800 NVLink 300 GB/s.
-        torch._inductor.config.intra_node_bw = (
-            64 if "L20" in torch.cuda.get_device_name() else 300
-        )
-
     # Below are default settings for torch.compile, you can change
     # them to your needs and test the performance
     torch._inductor.config.max_fusion_size = 64

cache_dit/metrics/metrics.py

@@ -334,22 +334,26 @@ compute_video_mse = partial(
 )
 
 
+METRICS_CHOICES = [
+    "psnr",
+    "ssim",
+    "mse",
+    "fid",
+    "all",
+]
+
+
 # Entrypoints
 def get_args():
+    global METRICS_CHOICES
     parser = argparse.ArgumentParser(
         description="CacheDiT's Metrics CLI",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    METRICS_CHOICES = [
-        "psnr",
-        "ssim",
-        "mse",
-        "fid",
-        "all",
-    ]
     parser.add_argument(
-        "metric",
+        "metrics",
         type=str,
+        nargs="+",
         default="psnr",
         choices=METRICS_CHOICES,
         help=f"Metric choices: {METRICS_CHOICES}",
@@ -382,6 +386,49 @@ def get_args():
         default=None,
         help="Path to predicted video or Dir to predicted videos",
     )
+
+    # Image 1 vs N pattern
+    parser.add_argument(
+        "--img-source-dir",
+        "-d",
+        type=str,
+        default=None,
+        help="Path to dir that contains dirs of images",
+    )
+    parser.add_argument(
+        "--ref-img-dir",
+        "-r",
+        type=str,
+        default=None,
+        help="Path to ref dir that contains ground truth images",
+    )
+
+    # Video 1 vs N pattern
+    parser.add_argument(
+        "--video-source-dir",
+        "-vd",
+        type=str,
+        default=None,
+        help="Path to dir that contains many videos",
+    )
+    parser.add_argument(
+        "--ref-video",
+        "-rv",
+        type=str,
+        default=None,
+        help="Path to ground truth video",
+    )
+
+    # FID batch size
+    parser.add_argument(
+        "--fid-batch-size",
+        "-b",
+        type=int,
+        default=1,
+        help="Batch size for FID compute",
+    )
+
+    # Verbose
     parser.add_argument(
         "--enable-verbose",
         "-verbose",
@@ -389,10 +436,20 @@ def get_args():
         default=False,
         help="Show metrics progress verbose",
     )
+
+    # Format output
+    parser.add_argument(
+        "--sort-output",
+        "-sort",
+        action="store_true",
+        default=False,
+        help="Sort the outupt metrics results",
+    )
     return parser.parse_args()
 
 
 def entrypoint():
+    global METRICS_CHOICES
     args = get_args()
     logger.debug(args)
 
@@ -401,68 +458,270 @@ def entrypoint():
         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(
-            (
-                not os.path.exists(args.img_true),
-                not os.path.exists(args.img_test),
-            )
-        ):
+    if "all" in args.metrics or "fid" in args.metrics:
+        FID = FrechetInceptionDistance(
+            disable_tqdm=DISABLE_VERBOSE,
+            batch_size=args.fid_batch_size,
+        )
+
+    METRICS_META: dict[str, float] = {}
+
+    # run one metric
+    def _run_metric(
+        metric: str,
+        img_true: str = None,
+        img_test: str = None,
+        video_true: str = None,
+        video_test: str = None,
+    ) -> None:
+        nonlocal FID
+        nonlocal METRICS_META
+        metric = metric.lower()
+        if img_true is not None and img_test is not None:
+            if any(
+                (
+                    not os.path.exists(img_true),
+                    not os.path.exists(img_test),
+                )
+            ):
+                return
+            # img_true and img_test can be files or dirs
+            img_true_info = os.path.basename(img_true)
+            img_test_info = os.path.basename(img_test)
+
+            def _logging_msg(value: float, n: int):
+                if value is None or n is None:
+                    return
+                msg = (
+                    f"{img_true_info} vs {img_test_info}, "
+                    f"Num: {n}, {metric.upper()}: {value:.5f}"
+                )
+                METRICS_META[msg] = value
+                logger.info(msg)
+
+            if metric == "psnr" or metric == "all":
+                img_psnr, n = compute_psnr(img_true, img_test)
+                _logging_msg(img_psnr, n)
+            if metric == "ssim" or metric == "all":
+                img_ssim, n = compute_ssim(img_true, img_test)
+                _logging_msg(img_ssim, n)
+            if metric == "mse" or metric == "all":
+                img_mse, n = compute_mse(img_true, img_test)
+                _logging_msg(img_mse, n)
+            if metric == "fid" or metric == "all":
+                img_fid, n = FID.compute_fid(img_true, img_test)
+                _logging_msg(img_fid, n)
+
+        if video_true is not None and video_test is not None:
+            if any(
+                (
+                    not os.path.exists(video_true),
+                    not os.path.exists(video_test),
+                )
+            ):
+                return
+
+            # video_true and video_test can be files or dirs
+            video_true_info = os.path.basename(video_true)
+            video_test_info = os.path.basename(video_test)
+
+            def _logging_msg(value: float, n: int):
+                if value is None or n is None:
+                    return
+                msg = (
+                    f"{video_true_info} vs {video_test_info}, "
+                    f"Frames: {n}, {metric.upper()}: {value:.5f}"
+                )
+                METRICS_META[msg] = value
+                logger.info(msg)
+
+            if metric == "psnr" or metric == "all":
+                video_psnr, n = compute_video_psnr(video_true, video_test)
+                _logging_msg(video_psnr, n)
+            if metric == "ssim" or metric == "all":
+                video_ssim, n = compute_video_ssim(video_true, video_test)
+                _logging_msg(video_ssim, n)
+            if metric == "mse" or metric == "all":
+                video_mse, n = compute_video_mse(video_true, video_test)
+                _logging_msg(video_mse, n)
+            if metric == "fid" or metric == "all":
+                video_fid, n = FID.compute_video_fid(video_true, video_test)
+                _logging_msg(video_fid, n)
+
+    # run selected metrics
+    if not DISABLE_VERBOSE:
+        logger.info(f"Selected metrics: {args.metrics}")
+
+    def _is_image_1vsN_pattern() -> bool:
+        return args.img_source_dir is not None and args.ref_img_dir is not None
+
+    def _is_video_1vsN_pattern() -> bool:
+        return args.video_source_dir is not None and args.ref_video is not None
+
+    assert not all((_is_image_1vsN_pattern(), _is_video_1vsN_pattern()))
+
+    if _is_image_1vsN_pattern():
+        # Glob Image dirs
+        if not os.path.exists(args.img_source_dir):
+            logger.error(f"{args.img_source_dir} not exist!")
             return
-        # img_true and img_test can be files or dirs
-        if args.metric == "psnr" or args.metric == "all":
-            img_psnr, n = compute_psnr(args.img_true, args.img_test)
-            logger.info(
-                f"{args.img_true} vs {args.img_test}, Num: {n}, PSNR: {img_psnr}"
-            )
-        if args.metric == "ssim" or args.metric == "all":
-            img_ssim, n = compute_ssim(args.img_true, args.img_test)
-            logger.info(
-                f"{args.img_true} vs {args.img_test}, Num: {n}, SSIM: {img_ssim}"
-            )
-        if args.metric == "mse" or args.metric == "all":
-            img_mse, n = compute_mse(args.img_true, args.img_test)
-            logger.info(
-                f"{args.img_true} vs {args.img_test}, Num: {n},  MSE: {img_mse}"
-            )
-        if args.metric == "fid" or args.metric == "all":
-            FID = FrechetInceptionDistance(disable_tqdm=DISABLE_VERBOSE)
-            img_fid, n = FID.compute_fid(args.img_true, args.img_test)
+        if not os.path.exists(args.ref_img_dir):
+            logger.error(f"{args.ref_img_dir} not exist!")
+            return
+
+        directories = []
+        for item in os.listdir(args.img_source_dir):
+            item_path = os.path.join(args.img_source_dir, item)
+            if os.path.isdir(item_path):
+                if os.path.basename(item_path) == os.path.basename(
+                    args.ref_img_dir
+                ):
+                    continue
+                directories.append(item_path)
+
+        if len(directories) == 0:
+            return
+
+        directories = sorted(directories)
+        if not DISABLE_VERBOSE:
             logger.info(
-                f"{args.img_true} vs {args.img_test}, Num: {n},  FID: {img_fid}"
+                f"Compare {args.ref_img_dir} vs {directories}, "
+                f"Num compares: {len(directories)}"
             )
-    if args.video_true is not None and args.video_test is not None:
-        if any(
-            (
-                not os.path.exists(args.video_true),
-                not os.path.exists(args.video_test),
-            )
-        ):
+
+        for metric in args.metrics:
+            for img_test_dir in directories:
+                _run_metric(
+                    metric=metric,
+                    img_true=args.ref_img_dir,
+                    img_test=img_test_dir,
+                )
+
+    elif _is_video_1vsN_pattern():
+        # Glob videos
+        if not os.path.exists(args.video_source_dir):
+            logger.error(f"{args.video_source_dir} not exist!")
+            return
+        if not os.path.exists(args.ref_video):
+            logger.error(f"{args.ref_video} not exist!")
+            return
+
+        video_source_dir: pathlib.Path = pathlib.Path(args.video_source_dir)
+        video_source_files = sorted(
+            [
+                file
+                for ext in _VIDEO_EXTENSIONS
+                for file in video_source_dir.rglob("*.{}".format(ext))
+            ]
+        )
+        video_source_files = [file.as_posix() for file in video_source_files]
+
+        video_source_selected = []
+        for video_source_file in video_source_files:
+            if os.path.basename(video_source_file) == os.path.basename(
+                args.ref_video
+            ):
+                continue
+            video_source_selected.append(video_source_file)
+
+        if len(video_source_selected) == 0:
             return
-        # video_true and video_test can be files or dirs
-        if args.metric == "psnr" or args.metric == "all":
-            video_psnr, n = compute_video_psnr(args.video_true, args.video_test)
+
+        video_source_selected = sorted(video_source_selected)
+        if not DISABLE_VERBOSE:
             logger.info(
-                f"{args.video_true} vs {args.video_test}, Frames: {n}, PSNR: {video_psnr}"
+                f"Compare {args.ref_video} vs {video_source_selected}, "
+                f"Num compares: {len(video_source_selected)}"
             )
-        if args.metric == "ssim" or args.metric == "all":
-            video_ssim, n = compute_video_ssim(args.video_true, args.video_test)
-            logger.info(
-                f"{args.video_true} vs {args.video_test}, Frames: {n}, SSIM: {video_ssim}"
+
+        for metric in args.metrics:
+            for video_test in video_source_selected:
+                _run_metric(
+                    metric=metric,
+                    video_true=args.ref_video,
+                    video_test=video_test,
+                )
+
+    else:
+        for metric in args.metrics:
+            _run_metric(
+                metric=metric,
+                img_true=args.img_true,
+                img_test=args.img_test,
+                video_true=args.video_true,
+                video_test=args.video_test,
             )
-        if args.metric == "mse" or args.metric == "all":
-            video_mse, n = compute_video_mse(args.video_true, args.video_test)
-            logger.info(
-                f"{args.video_true} vs {args.video_test}, Frames: {n},  MSE: {video_mse}"
+
+    if args.sort_output:
+
+        def _parse_value(
+            text: str,
+            tag: str = "Num",
+        ) -> float:
+            import re
+
+            pattern = re.compile(
+                rf"{re.escape(tag)}:\s*(\d+\.?\d*)", re.IGNORECASE
             )
-        if args.metric == "fid" or args.metric == "all":
-            FID = FrechetInceptionDistance(disable_tqdm=DISABLE_VERBOSE)
-            video_fid, n = FID.compute_video_fid(
-                args.video_true, args.video_test
+
+            match = pattern.search(text)
+
+            if not match:
+                return None
+
+            if tag.lower() in METRICS_CHOICES:
+                return float(match.group(1))
+            return int(match.group(1))
+
+        def _format_item(
+            key: str,
+            metric: str,
+            value: float,
+            max_key_len: int,
+        ):
+            # U1-Q0-C0-NONE vs U4-Q1-C1-NONE
+            header = key.split(",")[0].strip()
+            # Num / Frames
+            if n := _parse_value(key, "Num"):
+                print(
+                    f"{header:<{max_key_len}}  Num: {n}  "
+                    f"{metric.upper()}: {value:<.4f}"
+                )
+            elif n := _parse_value(key, "Frames"):
+                print(
+                    f"{header:<{max_key_len}}  Frames: {n}  "
+                    f"{metric.upper()}: {value:<.4f}"
+                )
+            else:
+                raise ValueError("Num or Frames can not be NoneType.")
+
+        for metric in args.metrics:
+            selected_items = {}
+            for key in METRICS_META.keys():
+                if metric.upper() in key or metric.lower() in key:
+                    selected_items[key] = METRICS_META[key]
+
+            reverse = True if metric.lower() in ["psnr", "ssim"] else False
+            sorted_items = sorted(
+                selected_items.items(), key=lambda x: x[1], reverse=reverse
             )
-            logger.info(
-                f"{args.video_true} vs {args.video_test}, Frames: {n},  FID: {video_fid}"
+            selected_keys = [
+                key.split(",")[0].strip() for key in selected_items.keys()
+            ]
+            max_key_len = max(len(key) for key in selected_keys)
+
+            format_len = int(max_key_len * 1.5)
+            res_len = format_len - len(f"Summary: {metric.upper()}")
+            left_len = res_len // 2
+            right_len = res_len - left_len
+            print("-" * format_len)
+            print(
+                " " * left_len + f"Summary: {metric.upper()}" + " " * right_len
             )
+            print("-" * format_len)
+            for key, value in sorted_items:
+                _format_item(key, metric, value, max_key_len)
+            print("-" * format_len)
 
 
 if __name__ == "__main__":

{cache_dit-0.2.9.dist-info → cache_dit-0.2.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.9
+Version: 0.2.11
 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
@@ -10,6 +10,7 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: packaging
+Requires-Dist: pyyaml
 Requires-Dist: torch>=2.5.1
 Requires-Dist: transformers>=4.51.3
 Requires-Dist: diffusers>=0.33.1
@@ -62,9 +63,8 @@ Dynamic: requires-python
 </div> 
 
 ## 🔥News🔥
-
-- [2025-07-13] An end2end speedup example for FLUX using cache-dit is released! **[🤗flux-faster](https://github.com/xlite-dev/flux-faster)**: A forked version of [huggingface/flux-fast](https://github.com/huggingface/flux-fast) that **makes flux-fast even faster** with **[cache-dit](https://github.com/vipshop/cache-dit)**, **3.3x** speedup on NVIDIA L20 while still maintaining **high precision**.
-  
+- [2025-07-18] 🎉First caching mechanism in **[🤗huggingface/flux-fast](https://github.com/huggingface/flux-fast)** with **[cache-dit](https://github.com/vipshop/cache-dit)**, please check [PR](https://github.com/huggingface/flux-fast/pull/13). 
+- [2025-07-13] **[🤗flux-faster](https://github.com/xlite-dev/flux-faster)** is released! A forked version of **[🤗huggingface/flux-fast](https://github.com/huggingface/flux-fast)** that **makes flux-fast even faster** with **[cache-dit](https://github.com/vipshop/cache-dit)**, **3.3x** speedup on NVIDIA L20.
 
 ## 🤗 Introduction 
 

{cache_dit-0.2.9.dist-info → cache_dit-0.2.11.dist-info}/RECORD

@@ -1,10 +1,11 @@
 cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/_version.py,sha256=Iq6CyehddPOWDVsW9Hnb65BEkCEkAnt4bl0MAuqXKLA,511
+cache_dit/_version.py,sha256=Y72g1mojWf0yRnnMW5zEUr6skXUSsqAdPjRJUrxXSYc,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=5RNuhWakvvqrOV4vkqrEBA7d-V1LwcNSsjtW14mkqK8,5255
+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=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+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=itVEb6gT2eZuncAHUmP51ZS0r6v6cGtRvnPjyeXqKH8,71156
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py,sha256=krNAICf-aS3JLmSG8vOB9tpLa04uYRcABsC8PMbVUKY,1870
@@ -30,17 +31,17 @@ cache_dit/cache_factory/first_block_cache/diffusers_adapters/hunyuan_video.py,sh
 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=OTvkwcezSrApZ2M1IMkYtkEmFbkfpTknhHMgoBApd6U,3786
+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/metrics.py,sha256=tzAtG_-fM1xPIBfRVFIBupvOWYzIO3xDq29Vy5rOBWc,14730
-cache_dit-0.2.9.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.9.dist-info/METADATA,sha256=TdvKAftNWwijdCW8K-8iO7fITEcfllWX3FJdZ-qcRqA,28032
-cache_dit-0.2.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.9.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.2.9.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.9.dist-info/RECORD,,
+cache_dit/metrics/metrics.py,sha256=PAzyhJawos1UeMnHsxcu4edkwCSYMBmjDGRR_--I104,22410
+cache_dit-0.2.11.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.11.dist-info/METADATA,sha256=HExb-ldgaYzZRSCy6Tg6syONjd0uDoB8_8AShvbfA-0,28213
+cache_dit-0.2.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.11.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.11.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.11.dist-info/RECORD,,