Upgrade SD output image verification with CLIP score (huggingface#920)

Co-authored-by: Sayantan Sarkar <sasarkar@habana.ai>
Co-authored-by: regisss <15324346+regisss@users.noreply.github.com>

Author: 3 people

tests/clip_coco_utils.py

@@ -1,7 +1,12 @@
 import os
+
+# Calculate CLIP score
+from functools import partial
 from pathlib import Path
 from urllib.request import urlretrieve
 
+import torch
+from torchmetrics.functional.multimodal import clip_score
 from transformers import AutoImageProcessor, AutoTokenizer, VisionTextDualEncoderModel, VisionTextDualEncoderProcessor
 
 
@@ -39,3 +44,12 @@ def create_clip_roberta_model():
     processor.save_pretrained("clip-roberta")
 
     print("Model generated.")
+
+
+clip_score_fn = partial(clip_score, model_name_or_path="openai/clip-vit-base-patch16")
+
+
+def calculate_clip_score(images, prompts):
+    images_int = (images * 255).astype("uint8")
+    clip_score = clip_score_fn(torch.from_numpy(images_int).permute(0, 3, 1, 2), prompts).detach()
+    return round(float(clip_score), 4)

tests/test_diffusers.py

@@ -51,7 +51,7 @@
 )
 from optimum.habana.utils import set_seed
 
-from .clip_coco_utils import download_files
+from .clip_coco_utils import calculate_clip_score, download_files
 
 
 if os.environ.get("GAUDI2_CI", "0") == "1":
@@ -608,6 +608,8 @@ def test_no_throughput_regression_autocast(self):
 
     @slow
     def test_no_generation_regression(self):
+        seed = 27
+        set_seed(seed)
         model_name = "CompVis/stable-diffusion-v1-4"
         # fp32
         scheduler = GaudiDDIMScheduler.from_pretrained(model_name, subfolder="scheduler")
@@ -619,39 +621,33 @@ def test_no_generation_regression(self):
             use_hpu_graphs=True,
             gaudi_config=GaudiConfig(use_torch_autocast=False),
         )
-        set_seed(27)
+
+        prompt = "An image of a squirrel in Picasso style"
+        generator = torch.manual_seed(seed)
         outputs = pipeline(
-            prompt="An image of a squirrel in Picasso style",
+            prompt=prompt,
+            generator=generator,
             output_type="np",
         )
 
         if os.environ.get("GAUDI2_CI", "0") == "1":
-            expected_slice = np.array(
-                [
-                    0.68306947,
-                    0.6812112,
-                    0.67309505,
-                    0.70057267,
-                    0.6582885,
-                    0.6325019,
-                    0.6708976,
-                    0.6226433,
-                    0.58038336,
-                ]
-            )
+            target_score = 29.8925
         else:
-            expected_slice = np.array(
-                [0.70760196, 0.7136303, 0.7000798, 0.714934, 0.6776865, 0.6800843, 0.6923707, 0.6653969, 0.6408076]
-            )
+            target_score = 36.774
+
         image = outputs.images[0]
         pil_image = numpy_to_pil(image)[0]
         pil_image.save("test_no_generation_regression_output.png")
 
+        clip_score = calculate_clip_score(np.expand_dims(image, axis=0), [prompt])
+
         self.assertEqual(image.shape, (512, 512, 3))
-        self.assertLess(np.abs(expected_slice - image[-3:, -3:, -1].flatten()).max(), 5e-3)
+        self.assertGreaterEqual(clip_score, target_score)
 
     @slow
     def test_no_generation_regression_ldm3d(self):
+        seed = 27
+        set_seed(seed)
         model_name = "Intel/ldm3d-4c"
         # fp32
         scheduler = GaudiDDIMScheduler.from_pretrained(model_name, subfolder="scheduler")
@@ -663,61 +659,28 @@ def test_no_generation_regression_ldm3d(self):
             use_hpu_graphs=True,
             gaudi_config=GaudiConfig(),
         )
-        set_seed(27)
+
+        prompt = "An image of a squirrel in Picasso style"
+        generator = torch.manual_seed(seed)
         outputs = pipeline(
-            prompt="An image of a squirrel in Picasso style",
+            prompt=prompt,
+            generator=generator,
             output_type="np",
         )
 
         if os.environ.get("GAUDI2_CI", "0") == "1":
-            expected_slice_rgb = np.array(
-                [
-                    0.2099357,
-                    0.16664368,
-                    0.08352646,
-                    0.20643419,
-                    0.16748399,
-                    0.08781305,
-                    0.21379063,
-                    0.19943115,
-                    0.04389626,
-                ]
-            )
-            expected_slice_depth = np.array(
-                [
-                    0.68369114,
-                    0.6827824,
-                    0.6852779,
-                    0.6836072,
-                    0.6888298,
-                    0.6895473,
-                    0.6853674,
-                    0.67561126,
-                    0.660434,
-                ]
-            )
+            target_score = 28.0894
         else:
-            expected_slice_rgb = np.array([0.7083766, 1.0, 1.0, 0.70610344, 0.9867363, 1.0, 0.7214538, 1.0, 1.0])
-            expected_slice_depth = np.array(
-                [
-                    0.919621,
-                    0.92072034,
-                    0.9184986,
-                    0.91994286,
-                    0.9242079,
-                    0.93387043,
-                    0.92345214,
-                    0.93558526,
-                    0.9223714,
-                ]
-            )
+            target_score = 35.81
+
         rgb = outputs.rgb[0]
         depth = outputs.depth[0]
 
+        rgb_clip_score = calculate_clip_score(np.expand_dims(rgb, axis=0), [prompt])
+
         self.assertEqual(rgb.shape, (512, 512, 3))
         self.assertEqual(depth.shape, (512, 512, 1))
-        self.assertLess(np.abs(expected_slice_rgb - rgb[-3:, -3:, -1].flatten()).max(), 5e-3)
-        self.assertLess(np.abs(expected_slice_depth - depth[-3:, -3:, -1].flatten()).max(), 5e-3)
+        self.assertGreaterEqual(rgb_clip_score, target_score)
 
     @slow
     def test_no_generation_regression_upscale(self):