updatye

Author: sayakpaul

tests/lora/test_lora_layers_sdxl.py

@@ -22,6 +22,7 @@
 import numpy as np
 import torch
 from packaging import version
+from parameterized import parameterized
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
 from diffusers import (
@@ -117,17 +118,28 @@ def tearDown(self):
     def test_multiple_wrong_adapter_name_raises_error(self):
         super().test_multiple_wrong_adapter_name_raises_error()
 
-    def test_simple_inference_with_text_denoiser_lora_unfused(self):
+    @parameterized.expand(
+        [
+            # Test actions on text_encoder LoRA only
+            ("fused", "text_encoder_only"),
+            ("unloaded", "text_encoder_only"),
+            ("save_load", "text_encoder_only"),
+            # Test actions on both text_encoder and denoiser LoRA
+            ("fused", "text_and_denoiser"),
+            ("unloaded", "text_and_denoiser"),
+            ("unfused", "text_and_denoiser"),
+            ("save_load", "text_and_denoiser"),
+        ]
+    )
+    def test_lora_actions(self, action, components_to_add):
         if torch.cuda.is_available():
             expected_atol = 9e-2
             expected_rtol = 9e-2
         else:
             expected_atol = 1e-3
             expected_rtol = 1e-3
 
-        super().test_simple_inference_with_text_denoiser_lora_unfused(
-            expected_atol=expected_atol, expected_rtol=expected_rtol
-        )
+        super().test_lora_actions(expected_atol=expected_atol, expected_rtol=expected_rtol)
 
     def test_simple_inference_with_text_lora_denoiser_fused_multi(self):
         if torch.cuda.is_available():

tests/lora/utils.py

@@ -292,7 +292,7 @@ def test_simple_inference_with_partial_text_lora(self):
                 "Removing adapters should change the output",
             )
 
-    def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3):
+    def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, expected_rtol=1e-3):
         """
         A unified test for various LoRA actions (fusing, unloading, saving/loading, etc.)
         on different combinations of model components.
@@ -321,7 +321,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3)
             )
 
             output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
-            self.assertTrue(not np.allclose(output_lora, output_no_lora, atol=expected_atol, rtol=1e-3))
+            self.assertTrue(not np.allclose(output_lora, output_no_lora, atol=expected_atol, rtol=expected_rtol))
 
             # 3. Perform the specified action and assert the outcome
             if action == "fused":
@@ -330,7 +330,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3)
                     self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
                 output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0]
                 self.assertTrue(
-                    not np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=1e-3),
+                    not np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=expected_rtol),
                     "Fused LoRA should produce a different output from the base model.",
                 )
 
@@ -342,7 +342,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3)
                     )
                 output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0]
                 self.assertTrue(
-                    np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=1e-3),
+                    np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=expected_rtol),
                     "Output after unloading LoRA should match the original output.",
                 )
 
@@ -358,7 +358,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3)
                 output_unfused = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
                 self.assertTrue(
-                    np.allclose(output_fused, output_unfused, atol=expected_atol, rtol=1e-3),
+                    np.allclose(output_fused, output_unfused, atol=expected_atol, rtol=expected_rtol),
                     "Output after unfusing should match the fused output.",
                 )
 
@@ -382,7 +382,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3)
 
                     output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0]
                     self.assertTrue(
-                        np.allclose(output_lora, output_after_action, atol=expected_atol, rtol=1e-3),
+                        np.allclose(output_lora, output_after_action, atol=expected_atol, rtol=expected_rtol),
                         "Loading from a saved checkpoint should yield the same result.",
                     )