Fix Kandinksy V22 tests (#7699)

update

Author: DN6

tests/pipelines/kandinsky/test_kandinsky.py

@@ -299,7 +299,7 @@ def test_kandinsky_text2img(self):
         pipe_prior.to(torch_device)
 
         pipeline = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16)
-        pipeline = pipeline.to(torch_device)
+        pipeline.to(torch_device)
         pipeline.set_progress_bar_config(disable=None)
 
         prompt = "red cat, 4k photo"

tests/pipelines/kandinsky2_2/test_kandinsky.py

@@ -25,11 +25,12 @@
     enable_full_determinism,
     floats_tensor,
     load_numpy,
+    numpy_cosine_similarity_distance,
     require_torch_gpu,
     slow,
 )
 
-from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
+from ..test_pipelines_common import PipelineTesterMixin
 
 
 enable_full_determinism()
@@ -248,30 +249,29 @@ def test_kandinsky_text2img(self):
         pipeline = KandinskyV22Pipeline.from_pretrained(
             "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
         )
-        pipeline = pipeline.enable_model_cpu_offload()
+        pipeline.enable_model_cpu_offload()
         pipeline.set_progress_bar_config(disable=None)
 
         prompt = "red cat, 4k photo"
 
-        generator = torch.Generator(device="cuda").manual_seed(0)
+        generator = torch.Generator(device="cpu").manual_seed(0)
         image_emb, zero_image_emb = pipe_prior(
             prompt,
             generator=generator,
             num_inference_steps=3,
             negative_prompt="",
         ).to_tuple()
 
-        generator = torch.Generator(device="cuda").manual_seed(0)
+        generator = torch.Generator(device="cpu").manual_seed(0)
         output = pipeline(
             image_embeds=image_emb,
             negative_image_embeds=zero_image_emb,
             generator=generator,
             num_inference_steps=3,
             output_type="np",
         )
-
         image = output.images[0]
-
         assert image.shape == (512, 512, 3)
 
-        assert_mean_pixel_difference(image, expected_image)
+        max_diff = numpy_cosine_similarity_distance(expected_image.flatten(), image.flatten())
+        assert max_diff < 1e-4

tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py

@@ -33,10 +33,11 @@
     load_image,
     load_numpy,
     nightly,
+    numpy_cosine_similarity_distance,
     require_torch_gpu,
 )
 
-from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
+from ..test_pipelines_common import PipelineTesterMixin
 
 
 enable_full_determinism()
@@ -260,20 +261,20 @@ def test_kandinsky_controlnet(self):
         pipeline = KandinskyV22ControlnetPipeline.from_pretrained(
             "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16
         )
-        pipeline = pipeline.enable_model_cpu_offload()
+        pipeline.enable_model_cpu_offload()
         pipeline.set_progress_bar_config(disable=None)
 
         prompt = "A robot, 4k photo"
 
-        generator = torch.Generator(device="cuda").manual_seed(0)
+        generator = torch.Generator(device="cpu").manual_seed(0)
         image_emb, zero_image_emb = pipe_prior(
             prompt,
             generator=generator,
             num_inference_steps=2,
             negative_prompt="",
         ).to_tuple()
 
-        generator = torch.Generator(device="cuda").manual_seed(0)
+        generator = torch.Generator(device="cpu").manual_seed(0)
         output = pipeline(
             image_embeds=image_emb,
             negative_image_embeds=zero_image_emb,
@@ -287,4 +288,5 @@ def test_kandinsky_controlnet(self):
 
         assert image.shape == (512, 512, 3)
 
-        assert_mean_pixel_difference(image, expected_image)
+        max_diff = numpy_cosine_similarity_distance(expected_image.flatten(), image.flatten())
+        assert max_diff < 1e-4

tests/pipelines/kandinsky2_2/test_kandinsky_controlnet_img2img.py

@@ -34,10 +34,11 @@
     load_image,
     load_numpy,
     nightly,
+    numpy_cosine_similarity_distance,
     require_torch_gpu,
 )
 
-from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
+from ..test_pipelines_common import PipelineTesterMixin
 
 
 enable_full_determinism()
@@ -274,7 +275,7 @@ def test_kandinsky_controlnet_img2img(self):
         pipeline = KandinskyV22ControlnetImg2ImgPipeline.from_pretrained(
             "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16
         )
-        pipeline = pipeline.enable_model_cpu_offload()
+        pipeline.enable_model_cpu_offload()
 
         pipeline.set_progress_bar_config(disable=None)
 
@@ -289,6 +290,7 @@ def test_kandinsky_controlnet_img2img(self):
             num_inference_steps=5,
         ).to_tuple()
 
+        generator = torch.Generator(device="cpu").manual_seed(0)
         output = pipeline(
             image=init_image,
             image_embeds=image_emb,
@@ -306,4 +308,5 @@ def test_kandinsky_controlnet_img2img(self):
 
         assert image.shape == (512, 512, 3)
 
-        assert_mean_pixel_difference(image, expected_image)
+        max_diff = numpy_cosine_similarity_distance(expected_image.flatten(), image.flatten())
+        assert max_diff < 1e-4

tests/pipelines/kandinsky2_2/test_kandinsky_img2img.py

@@ -33,11 +33,12 @@
     floats_tensor,
     load_image,
     load_numpy,
+    numpy_cosine_similarity_distance,
     require_torch_gpu,
     slow,
 )
 
-from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
+from ..test_pipelines_common import PipelineTesterMixin
 
 
 enable_full_determinism()
@@ -270,8 +271,7 @@ def test_kandinsky_img2img(self):
         pipeline = KandinskyV22Img2ImgPipeline.from_pretrained(
             "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
         )
-        pipeline = pipeline.enable_model_cpu_offload()
-
+        pipeline.enable_model_cpu_offload()
         pipeline.set_progress_bar_config(disable=None)
 
         generator = torch.Generator(device="cpu").manual_seed(0)
@@ -282,6 +282,7 @@ def test_kandinsky_img2img(self):
             negative_prompt="",
         ).to_tuple()
 
+        generator = torch.Generator(device="cpu").manual_seed(0)
         output = pipeline(
             image=init_image,
             image_embeds=image_emb,
@@ -298,4 +299,5 @@ def test_kandinsky_img2img(self):
 
         assert image.shape == (768, 768, 3)
 
-        assert_mean_pixel_difference(image, expected_image)
+        max_diff = numpy_cosine_similarity_distance(expected_image.flatten(), image.flatten())
+        assert max_diff < 1e-4

tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py

@@ -34,12 +34,13 @@
     is_flaky,
     load_image,
     load_numpy,
+    numpy_cosine_similarity_distance,
     require_torch_gpu,
     slow,
     torch_device,
 )
 
-from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
+from ..test_pipelines_common import PipelineTesterMixin
 
 
 enable_full_determinism()
@@ -338,6 +339,7 @@ def test_kandinsky_inpaint(self):
             negative_prompt="",
         ).to_tuple()
 
+        generator = torch.Generator(device="cpu").manual_seed(0)
         output = pipeline(
             image=init_image,
             mask_image=mask,
@@ -354,4 +356,5 @@ def test_kandinsky_inpaint(self):
 
         assert image.shape == (768, 768, 3)
 
-        assert_mean_pixel_difference(image, expected_image)
+        max_diff = numpy_cosine_similarity_distance(expected_image.flatten(), image.flatten())
+        assert max_diff < 1e-4