[WWB]: Add ImageText-to-Image pipeline validation

tools/who_what_benchmark/tests/test_cli_image.py

@@ -20,6 +20,8 @@ def run_wwb(args):
 @pytest.mark.parametrize(
     ("model_id", "model_type", "backend"),
     [
+        ("hf-internal-testing/tiny-stable-diffusion-torch", "image-to-image", "hf"),
+        ("hf-internal-testing/tiny-stable-diffusion-xl-pipe", "image-to-image", "hf"),
         ("hf-internal-testing/tiny-stable-diffusion-torch", "text-to-image", "hf"),
         ("hf-internal-testing/tiny-stable-diffusion-torch", "text-to-image", "openvino"),
         ("hf-internal-testing/tiny-stable-diffusion-xl-pipe", "text-to-image", "hf"),
@@ -65,6 +67,7 @@ def test_image_model_types(model_id, model_type, backend):
 @pytest.mark.parametrize(
     ("model_id", "model_type"),
     [
+        ("echarlaix/tiny-random-stable-diffusion-xl", "image-to-image"),
         ("echarlaix/tiny-random-stable-diffusion-xl", "text-to-image"),
     ],
 )
@@ -81,7 +84,7 @@ def test_image_model_genai(model_id, model_type):
 
         wwb_args = [
             "--base-model",
-            MODEL_PATH,
+            model_id,
             "--num-samples",
             "1",
             "--gt-data",
@@ -90,6 +93,7 @@ def test_image_model_genai(model_id, model_type):
             "CPU",
             "--model-type",
             model_type,
+            "--hf",
         ]
         result = run_wwb(wwb_args)
         assert result.returncode == 0
@@ -131,6 +135,7 @@ def test_image_model_genai(model_id, model_type):
             model_type,
             "--output",
             output_dir,
+            "--genai",
         ]
         result = run_wwb(wwb_args)
         assert result.returncode == 0

tools/who_what_benchmark/whowhatbench/__init__.py

@@ -3,6 +3,7 @@
 from .text_evaluator import TextEvaluator as Evaluator
 from .text2image_evaluator import Text2ImageEvaluator
 from .visualtext_evaluator import VisualTextEvaluator
+from .image2image import Image2ImageEvaluator
 
 
 __all__ = [
@@ -11,5 +12,6 @@
     "TextEvaluator",
     "Text2ImageEvaluator",
     "VisualTextEvaluator",
+    "Image2ImageEvaluator",
     "EVALUATOR_REGISTRY",
 ]

tools/who_what_benchmark/whowhatbench/image2image.py

@@ -0,0 +1,143 @@
+import os
+from typing import Any, Union
+
+import datasets
+import pandas as pd
+from tqdm import tqdm
+from transformers import set_seed
+import torch
+import openvino_genai
+
+from .registry import register_evaluator
+from .text2image_evaluator import Text2ImageEvaluator
+
+from .whowhat_metrics import ImageSimilarity
+
+
+class Generator(openvino_genai.Generator):
+    def __init__(self, seed, rng, mu=0.0, sigma=1.0):
+        openvino_genai.Generator.__init__(self)
+        self.mu = mu
+        self.sigma = sigma
+        self.rng = rng
+
+    def next(self):
+        return torch.randn(1, generator=self.rng, dtype=torch.float32).item()
+
+
+def preprocess_fn(example):
+    return {
+        "prompts": example["Instruction_VLM-LLM"],
+        "images": example["source_img"],
+    }
+
+
+def prepare_default_data(num_samples=None):
+    DATASET_NAME = "paint-by-inpaint/PIPE"
+    NUM_SAMPLES = 10 if num_samples is None else num_samples
+    set_seed(42)
+    default_dataset = datasets.load_dataset(
+        DATASET_NAME, split="test", streaming=True
+    ).filter(lambda example: example["Instruction_VLM-LLM"] != "").take(NUM_SAMPLES)
+    return default_dataset.map(
+        lambda x: preprocess_fn(x), remove_columns=default_dataset.column_names
+    )
+
+
+@register_evaluator("image-to-image")
+class Image2ImageEvaluator(Text2ImageEvaluator):
+    def __init__(
+        self,
+        base_model: Any = None,
+        gt_data: str = None,
+        test_data: Union[str, list] = None,
+        metrics="similarity",
+        similarity_model_id: str = "openai/clip-vit-large-patch14",
+        resolution=(512, 512),
+        num_inference_steps=4,
+        crop_prompts=True,
+        num_samples=None,
+        gen_image_fn=None,
+        seed=42,
+        is_genai=False,
+    ) -> None:
+        assert (
+            base_model is not None or gt_data is not None
+        ), "Text generation pipeline for evaluation or ground trush data must be defined"
+
+        self.test_data = test_data
+        self.metrics = metrics
+        self.resolution = resolution
+        self.crop_prompt = crop_prompts
+        self.num_samples = num_samples
+        self.num_inference_steps = num_inference_steps
+        self.seed = seed
+        self.similarity = None
+        self.similarity = ImageSimilarity(similarity_model_id)
+        self.last_cmp = None
+        self.gt_dir = os.path.dirname(gt_data)
+        self.generation_fn = gen_image_fn
+        self.is_genai = is_genai
+
+        if base_model:
+            base_model.resolution = self.resolution
+            self.gt_data = self._generate_data(
+                base_model, gen_image_fn, os.path.join(self.gt_dir, "reference")
+            )
+        else:
+            self.gt_data = pd.read_csv(gt_data, keep_default_na=False)
+
+    def _generate_data(self, model, gen_image_fn=None, image_dir="reference"):
+        def default_gen_image_fn(model, prompt, image, num_inference_steps, generator=None):
+            with torch.no_grad():
+                output = model(
+                    prompt,
+                    image=image,
+                    num_inference_steps=num_inference_steps,
+                    output_type="pil",
+                    width=self.resolution[0],
+                    height=self.resolution[0],
+                    generator=generator,
+                )
+            return output.images[0]
+
+        generation_fn = gen_image_fn or default_gen_image_fn
+
+        if self.test_data:
+            if isinstance(self.test_data, str):
+                data = pd.read_csv(self.test_data)
+            else:
+                if isinstance(self.test_data, dict):
+                    assert "prompts" in self.test_data
+                    assert "images" in self.test_data
+                    data = dict(self.test_data)
+                data = pd.DataFrame.from_dict(data)
+        else:
+            data = pd.DataFrame.from_dict(prepare_default_data(self.num_samples))
+
+        prompts = data["prompts"]
+        images = data["images"]
+        output_images = []
+        rng = torch.Generator(device="cpu")
+
+        if not os.path.exists(image_dir):
+            os.makedirs(image_dir)
+
+        for i, (prompt, image) in tqdm(enumerate(zip(prompts, images)), desc="Evaluate pipeline"):
+            set_seed(self.seed)
+            rng = rng.manual_seed(self.seed)
+            output = generation_fn(
+                model,
+                prompt,
+                image=image,
+                num_inference_steps=self.num_inference_steps,
+                generator=Generator(self.seed, rng) if self.is_genai else rng
+            )
+            image_path = os.path.join(image_dir, f"{i}.png")
+            output.save(image_path)
+            output_images.append(image_path)
+
+        res_data = {"prompts": list(prompts), "images": output_images}
+        df = pd.DataFrame(res_data)
+
+        return df