cache-dit 0.2.34__py3-none-any.whl → 0.2.37__py3-none-any.whl

cache_dit/__init__.py

@@ -4,6 +4,10 @@ except ImportError:
     __version__ = "unknown version"
     version_tuple = (0, 0, "unknown version")
 
+from cache_dit.utils import summary
+from cache_dit.utils import strify
+from cache_dit.utils import disable_print
+from cache_dit.logger import init_logger
 from cache_dit.cache_factory import load_options
 from cache_dit.cache_factory import enable_cache
 from cache_dit.cache_factory import disable_cache
@@ -18,9 +22,7 @@ from cache_dit.cache_factory import supported_pipelines
 from cache_dit.cache_factory import get_adapter
 from cache_dit.compile import set_compile_configs
 from cache_dit.quantize import quantize
-from cache_dit.utils import summary
-from cache_dit.utils import strify
-from cache_dit.logger import init_logger
+
 
 NONE = CacheType.NONE
 DBCache = CacheType.DBCache

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.34'
-__version_tuple__ = version_tuple = (0, 2, 34)
+__version__ = version = '0.2.37'
+__version_tuple__ = version_tuple = (0, 2, 37)
 
 __commit_id__ = commit_id = None

cache_dit/metrics/clip_score.py

@@ -0,0 +1,135 @@
+import os
+import re
+import pathlib
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+import torch
+from transformers import CLIPProcessor, CLIPModel
+
+from typing import Tuple, Union
+from cache_dit.metrics.config import _IMAGE_EXTENSIONS
+from cache_dit.metrics.config import get_metrics_verbose
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+DISABLE_VERBOSE = not get_metrics_verbose()
+
+
+class CLIPScore:
+    def __init__(
+        self,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        clip_model_path: str = None,
+    ):
+        self.device = device
+        if clip_model_path is None:
+            clip_model_path = os.environ.get(
+                "CLIP_MODEL_DIR", "laion/CLIP-ViT-g-14-laion2B-s12B-b42K"
+            )
+
+        # Load models
+        self.clip_model = CLIPModel.from_pretrained(clip_model_path)
+        self.clip_model = self.clip_model.to(device)  # type: ignore
+        self.clip_processor = CLIPProcessor.from_pretrained(clip_model_path)
+
+    @torch.no_grad()
+    def compute_clip_score(
+        self,
+        img: Image.Image | np.ndarray,
+        prompt: str,
+    ) -> float:
+        if isinstance(img, Image.Image):
+            img_pil = img.convert("RGB")
+        elif isinstance(img, np.ndarray):
+            img_pil = Image.fromarray(img).convert("RGB")
+        else:
+            img_pil = Image.open(img).convert("RGB")
+        with torch.no_grad():
+            inputs = self.clip_processor(
+                text=prompt,
+                images=img_pil,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+            ).to(self.device)
+            outputs = self.clip_model(**inputs)
+        return outputs.logits_per_image.item()
+
+
+clip_score_instance: CLIPScore = None
+
+
+def compute_clip_score_img(
+    img: Image.Image | np.ndarray | str,
+    prompt: str,
+    clip_model_path: str = None,
+) -> float:
+    global clip_score_instance
+    if clip_score_instance is None:
+        clip_score_instance = CLIPScore(clip_model_path=clip_model_path)
+    assert clip_score_instance is not None
+    return clip_score_instance.compute_clip_score(img, prompt)
+
+
+def compute_clip_score(
+    img_dir: Image.Image | np.ndarray | str,
+    prompts: str | list[str],
+    clip_model_path: str = None,
+) -> Union[Tuple[float, int], Tuple[None, None]]:
+    if not os.path.isdir(img_dir) or (
+        not isinstance(prompts, list) and not os.path.isfile(prompts)
+    ):
+        return (
+            compute_clip_score_img(
+                img_dir,
+                prompts,
+                clip_model_path=clip_model_path,
+            ),
+            1,
+        )
+
+    # compute dir metric
+    def natural_sort_key(filename):
+        match = re.search(r"(\d+)\D*$", filename)
+        return int(match.group(1)) if match else filename
+
+    img_dir: pathlib.Path = pathlib.Path(img_dir)
+    img_files = [
+        file
+        for ext in _IMAGE_EXTENSIONS
+        for file in img_dir.rglob("*.{}".format(ext))
+    ]
+    img_files = [file.as_posix() for file in img_files]
+    img_files = sorted(img_files, key=natural_sort_key)
+
+    if os.path.isfile(prompts):
+        """Load prompts from file"""
+        with open(prompts, "r", encoding="utf-8") as f:
+            prompts_load = [line.strip() for line in f.readlines()]
+        prompts = prompts_load.copy()
+
+    vaild_len = min(len(img_files), len(prompts))
+    img_files = img_files[:vaild_len]
+    prompts = prompts[:vaild_len]
+
+    clip_scores = []
+
+    for img_file, prompt in tqdm(
+        zip(img_files, prompts),
+        total=vaild_len,
+        disable=not get_metrics_verbose(),
+    ):
+        clip_scores.append(
+            compute_clip_score_img(
+                img_file,
+                prompt,
+                clip_model_path=clip_model_path,
+            )
+        )
+
+    if vaild_len > 0:
+        return np.mean(clip_scores), vaild_len
+    return None, None

cache_dit/metrics/fid.py

@@ -1,6 +1,8 @@
 import os
 import cv2
 import pathlib
+import warnings
+
 import numpy as np
 from PIL import Image
 from tqdm import tqdm
@@ -8,13 +10,21 @@ from scipy import linalg
 import torch
 import torchvision.transforms as TF
 from torch.nn.functional import adaptive_avg_pool2d
+
+from typing import Tuple, Union
 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.metrics.config import get_metrics_verbose
+from cache_dit.utils import disable_print
 from cache_dit.logger import init_logger
 
+warnings.filterwarnings("ignore")
+
 logger = init_logger(__name__)
 
+DISABLE_VERBOSE = not get_metrics_verbose()
+
 
 # Adapted from: https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py
 class ImagePathDataset(torch.utils.data.Dataset):
@@ -496,3 +506,35 @@ class FrechetInceptionDistance:
             return [], [], 0
 
         return video_true_frames, video_test_frames, valid_frames
+
+
+fid_instance: FrechetInceptionDistance = None
+
+
+def compute_fid(
+    image_true: np.ndarray | str,
+    image_test: np.ndarray | str,
+) -> Union[Tuple[float, int], Tuple[None, None]]:
+    global fid_instance
+    if fid_instance is None:
+        with disable_print():
+            fid_instance = FrechetInceptionDistance(
+                disable_tqdm=not get_metrics_verbose(),
+            )
+    assert fid_instance is not None
+    return fid_instance.compute_fid(image_true, image_test)
+
+
+def compute_video_fid(
+    # file or dir
+    video_true: str,
+    video_test: str,
+) -> Union[Tuple[float, int], Tuple[None, None]]:
+    global fid_instance
+    if fid_instance is None:
+        with disable_print():
+            fid_instance = FrechetInceptionDistance(
+                disable_tqdm=not get_metrics_verbose(),
+            )
+    assert fid_instance is not None
+    return fid_instance.compute_fid(video_true, video_test)

cache_dit/metrics/image_reward.py

@@ -0,0 +1,177 @@
+import os
+import re
+import pathlib
+import warnings
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+import torch
+import ImageReward as RM
+import torchvision.transforms.v2.functional as TF
+import torchvision.transforms.v2 as T
+
+from typing import Tuple, Union
+from cache_dit.metrics.config import _IMAGE_EXTENSIONS
+from cache_dit.metrics.config import get_metrics_verbose
+from cache_dit.utils import disable_print
+from cache_dit.logger import init_logger
+
+warnings.filterwarnings("ignore")
+
+
+logger = init_logger(__name__)
+
+
+DISABLE_VERBOSE = not get_metrics_verbose()
+
+
+class ImageRewardScore:
+    def __init__(
+        self,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        imagereward_model_path: str = None,
+    ):
+        self.device = device
+        if imagereward_model_path is None:
+            imagereward_model_path = os.environ.get(
+                "IMAGEREWARD_MODEL_DIR", None
+            )
+
+        # Load ImageReward model
+        self.med_config = os.path.join(
+            imagereward_model_path, "med_config.json"
+        )
+        self.imagereward_path = os.path.join(
+            imagereward_model_path, "ImageReward.pt"
+        )
+        if imagereward_model_path is not None:
+            self.imagereward_model = RM.load(
+                self.imagereward_path,
+                download_root=imagereward_model_path,
+                med_config=self.med_config,
+            ).to(self.device)
+        else:
+            self.imagereward_model = RM.load(
+                "ImageReward-v1.0",  # download from huggingface
+            ).to(self.device)
+
+        # ImageReward transform
+        self.reward_transform = T.Compose(
+            [
+                T.Resize(224, interpolation=T.InterpolationMode.BICUBIC),
+                T.CenterCrop(224),
+                T.ToImage(),
+                T.ToDtype(torch.float32, scale=True),
+                T.Normalize(
+                    (0.48145466, 0.4578275, 0.40821073),
+                    (0.26862954, 0.26130258, 0.27577711),
+                ),
+            ]
+        )
+
+    @torch.no_grad()
+    def compute_reward_score(
+        self,
+        img: Image.Image | np.ndarray,
+        prompt: str,
+    ) -> float:
+        if isinstance(img, Image.Image):
+            img_pil = img.convert("RGB")
+        elif isinstance(img, np.ndarray):
+            img_pil = Image.fromarray(img).convert("RGB")
+        else:
+            img_pil = Image.open(img).convert("RGB")
+        with torch.no_grad():
+            img_tensor = TF.pil_to_tensor(img_pil).unsqueeze(0).to(self.device)
+            img_reward = self.reward_transform(img_tensor)
+            inputs = self.imagereward_model.blip.tokenizer(
+                [prompt],
+                padding="max_length",
+                truncation=True,
+                max_length=512,
+                return_tensors="pt",
+            ).to(self.device)
+            score = self.imagereward_model.score_gard(
+                inputs.input_ids, inputs.attention_mask, img_reward
+            )
+        return score.item()
+
+
+image_reward_score_instance: ImageRewardScore = None
+
+
+def compute_reward_score_img(
+    img: Image.Image | np.ndarray | str,
+    prompt: str,
+    imagereward_model_path: str = None,
+) -> float:
+    global image_reward_score_instance
+    if image_reward_score_instance is None:
+        with disable_print():
+            image_reward_score_instance = ImageRewardScore(
+                imagereward_model_path=imagereward_model_path
+            )
+    assert image_reward_score_instance is not None
+    return image_reward_score_instance.compute_reward_score(img, prompt)
+
+
+def compute_reward_score(
+    img_dir: Image.Image | np.ndarray | str,
+    prompts: str | list[str],
+    imagereward_model_path: str = None,
+) -> Union[Tuple[float, int], Tuple[None, None]]:
+    if not os.path.isdir(img_dir) or (
+        not isinstance(prompts, list) and not os.path.isfile(prompts)
+    ):
+        return (
+            compute_reward_score_img(
+                img_dir,
+                prompts,
+                imagereward_model_path=imagereward_model_path,
+            ),
+            1,
+        )
+
+    # compute dir metric
+    def natural_sort_key(filename):
+        match = re.search(r"(\d+)\D*$", filename)
+        return int(match.group(1)) if match else filename
+
+    img_dir: pathlib.Path = pathlib.Path(img_dir)
+    img_files = [
+        file
+        for ext in _IMAGE_EXTENSIONS
+        for file in img_dir.rglob("*.{}".format(ext))
+    ]
+    img_files = [file.as_posix() for file in img_files]
+    img_files = sorted(img_files, key=natural_sort_key)
+
+    if os.path.isfile(prompts):
+        """Load prompts from file"""
+        with open(prompts, "r", encoding="utf-8") as f:
+            prompts_load = [line.strip() for line in f.readlines()]
+        prompts = prompts_load.copy()
+
+    vaild_len = min(len(img_files), len(prompts))
+    img_files = img_files[:vaild_len]
+    prompts = prompts[:vaild_len]
+
+    reward_scores = []
+
+    for img_file, prompt in tqdm(
+        zip(img_files, prompts),
+        total=vaild_len,
+        disable=not get_metrics_verbose(),
+    ):
+        reward_scores.append(
+            compute_reward_score_img(
+                img_file,
+                prompt,
+                imagereward_model_path=imagereward_model_path,
+            )
+        )
+
+    if vaild_len > 0:
+        return np.mean(reward_scores), vaild_len
+    return None, None

cache_dit/metrics/lpips.py

@@ -1,9 +1,9 @@
-import builtins as __builtin__
-import contextlib
 import warnings
 
 import lpips
 import torch
+from cache_dit.utils import disable_print
+
 
 warnings.filterwarnings("ignore")
 
@@ -11,18 +11,6 @@ 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