allow big lora tests to run on the CI.

Author: sayakpaul

tests/lora/test_lora_layers_flux.py

@@ -31,9 +31,12 @@
     numpy_cosine_similarity_distance,
     require_peft_backend,
     require_torch_gpu,
+    require_big_gpu_with_torch_cuda,
     slow,
     torch_device,
+    print_tensor_test
 )
+import pytest
 
 
 if is_peft_available():
@@ -169,8 +172,8 @@ def test_modify_padding_mode(self):
 @nightly
 @require_torch_gpu
 @require_peft_backend
-@unittest.skip("We cannot run inference on this model with the current CI hardware")
-# TODO (DN6, sayakpaul): move these tests to a beefier GPU
+@require_big_gpu_with_torch_cuda
+@pytest.mark.big_gpu_with_torch_cuda
 class FluxLoRAIntegrationTests(unittest.TestCase):
     """internal note: The integration slices were obtained on audace.
 
@@ -211,6 +214,7 @@ def test_flux_the_last_ben(self):
             generator=torch.manual_seed(self.seed),
         ).images
         out_slice = out[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.1855, 0.1855, 0.1836, 0.1855, 0.1836, 0.1875, 0.1777, 0.1758, 0.2246])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), out_slice)
@@ -233,6 +237,7 @@ def test_flux_kohya(self):
         ).images
 
         out_slice = out[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.6367, 0.6367, 0.6328, 0.6367, 0.6328, 0.6289, 0.6367, 0.6328, 0.6484])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), out_slice)
@@ -255,6 +260,7 @@ def test_flux_kohya_with_text_encoder(self):
         ).images
 
         out_slice = out[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.4023, 0.4023, 0.4023, 0.3965, 0.3984, 0.3965, 0.3926, 0.3906, 0.4219])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), out_slice)
@@ -277,6 +283,7 @@ def test_flux_xlabs(self):
             generator=torch.manual_seed(self.seed),
         ).images
         out_slice = out[0, -3:, -3:, -1].flatten()
+        print_tensor_test(out_slice)
         expected_slice = np.array([0.3965, 0.4180, 0.4434, 0.4082, 0.4375, 0.4590, 0.4141, 0.4375, 0.4980])
 
         max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), out_slice)

tests/lora/test_lora_layers_sd3.py

@@ -34,7 +34,12 @@
     require_peft_backend,
     require_torch_gpu,
     torch_device,
+    slow,
+    nightly,
+    require_big_gpu_with_torch_cuda,
+    print_tensor_test
 )
+import pytest
 
 
 if is_peft_available():
@@ -130,9 +135,13 @@ def test_modify_padding_mode(self):
         pass
 
 
+@slow
+@nightly
 @require_torch_gpu
 @require_peft_backend
-class LoraSD3IntegrationTests(unittest.TestCase):
+@require_big_gpu_with_torch_cuda
+@pytest.mark.big_gpu_with_torch_cuda
+class SD3LoraIntegrationTests(unittest.TestCase):
     pipeline_class = StableDiffusion3Img2ImgPipeline
     repo_id = "stabilityai/stable-diffusion-3-medium-diffusers"
 
@@ -173,6 +182,7 @@ def test_sd3_img2img_lora(self):
 
         image = pipe(**inputs).images[0]
         image_slice = image[0, :10, :10]
+        print_tensor_test(image[0, -3:, -3:, -1].flatten())
         expected_slice = np.array(
             [
                 0.47827148,