cache-dit 0.2.33__py3-none-any.whl → 0.2.36__py3-none-any.whl

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