update

Author: DN6

tests/lora/test_lora_layers_peft.py

@@ -46,6 +46,7 @@
     floats_tensor,
     load_image,
     nightly,
+    numpy_cosine_similarity_distance,
     require_peft_backend,
     require_torch_gpu,
     slow,
@@ -1753,7 +1754,8 @@ def test_sdxl_lcm_lora(self):
         image_np = pipe.image_processor.pil_to_numpy(image)
         expected_image_np = pipe.image_processor.pil_to_numpy(expected_image)
 
-        self.assertTrue(np.allclose(image_np, expected_image_np, atol=1e-2))
+        max_diff = numpy_cosine_similarity_distance(image_np.flatten(), expected_image_np.flatten())
+        assert max_diff < 1e-4
 
         pipe.unload_lora_weights()
 
@@ -1780,7 +1782,8 @@ def test_sdv1_5_lcm_lora(self):
         image_np = pipe.image_processor.pil_to_numpy(image)
         expected_image_np = pipe.image_processor.pil_to_numpy(expected_image)
 
-        self.assertTrue(np.allclose(image_np, expected_image_np, atol=1e-2))
+        max_diff = numpy_cosine_similarity_distance(image_np.flatten(), expected_image_np.flatten())
+        assert max_diff < 1e-4
 
         pipe.unload_lora_weights()
 
@@ -1816,7 +1819,8 @@ def test_sdv1_5_lcm_lora_img2img(self):
         image_np = pipe.image_processor.pil_to_numpy(image)
         expected_image_np = pipe.image_processor.pil_to_numpy(expected_image)
 
-        self.assertTrue(np.allclose(image_np, expected_image_np, atol=1e-2))
+        max_diff = numpy_cosine_similarity_distance(image_np.flatten(), expected_image_np.flatten())
+        assert max_diff < 1e-4
 
         pipe.unload_lora_weights()
 
@@ -1849,7 +1853,7 @@ def test_sdxl_1_0_lora_fusion(self):
         release_memory(pipe)
 
     def test_sdxl_1_0_lora_unfusion(self):
-        generator = torch.Generator().manual_seed(0)
+        generator = torch.Generator("cpu").manual_seed(0)
 
         pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
         lora_model_id = "hf-internal-testing/sdxl-1.0-lora"
@@ -1860,16 +1864,18 @@ def test_sdxl_1_0_lora_unfusion(self):
         pipe.enable_model_cpu_offload()
 
         images = pipe(
-            "masterpiece, best quality, mountain", output_type="np", generator=generator, num_inference_steps=2
+            "masterpiece, best quality, mountain", output_type="pil", generator=generator, num_inference_steps=3
         ).images
-        images_with_fusion = images[0, -3:, -3:, -1].flatten()
+        images[0].save("fused.png")
+        images_with_fusion = images.flatten()
 
         pipe.unfuse_lora()
-        generator = torch.Generator().manual_seed(0)
+        generator = torch.Generator("cpu").manual_seed(0)
         images = pipe(
-            "masterpiece, best quality, mountain", output_type="np", generator=generator, num_inference_steps=2
+            "masterpiece, best quality, mountain", output_type="pil", generator=generator, num_inference_steps=3
         ).images
-        images_without_fusion = images[0, -3:, -3:, -1].flatten()
+        images[0].save("unfused.png")
+        images_without_fusion = images.flatten()
 
         self.assertTrue(np.allclose(images_with_fusion, images_without_fusion, atol=1e-3))
         release_memory(pipe)
@@ -1913,10 +1919,8 @@ def test_sdxl_1_0_lora_fusion_efficiency(self):
         lora_model_id = "hf-internal-testing/sdxl-1.0-lora"
         lora_filename = "sd_xl_offset_example-lora_1.0.safetensors"
 
-        pipe = DiffusionPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
-        )
-        pipe.load_lora_weights(lora_model_id, weight_name=lora_filename, torch_dtype=torch.bfloat16)
+        pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
+        pipe.load_lora_weights(lora_model_id, weight_name=lora_filename, torch_dtype=torch.float16)
         pipe.enable_model_cpu_offload()
 
         start_time = time.time()
@@ -1929,10 +1933,8 @@ def test_sdxl_1_0_lora_fusion_efficiency(self):
 
         del pipe
 
-        pipe = DiffusionPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
-        )
-        pipe.load_lora_weights(lora_model_id, weight_name=lora_filename, torch_dtype=torch.bfloat16)
+        pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
+        pipe.load_lora_weights(lora_model_id, weight_name=lora_filename, torch_dtype=torch.float16)
         pipe.fuse_lora()
         # We need to unload the lora weights since in the previous API `fuse_lora` led to lora weights being
         # silently deleted - otherwise this will CPU OOM