update

Author: sayakpaul

tests/lora/test_lora_layers_wan.py

@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import sys
 import tempfile
 import unittest
 
-import numpy as np
+import safetensors.torch
 import torch
 from transformers import AutoTokenizer, T5EncoderModel
 
@@ -26,12 +27,18 @@
     WanPipeline,
     WanTransformer3DModel,
 )
-from diffusers.utils.testing_utils import floats_tensor, require_peft_backend, skip_mps, torch_device
+from diffusers.utils.testing_utils import (
+    floats_tensor,
+    require_peft_backend,
+    require_peft_version_greater,
+    skip_mps,
+    torch_device,
+)
 
 
 sys.path.append(".")
 
-from utils import PeftLoraLoaderMixinTests  # noqa: E402
+from utils import PeftLoraLoaderMixinTests, check_module_lora_metadata  # noqa: E402
 
 
 @require_peft_backend
@@ -142,35 +149,40 @@ def test_simple_inference_with_text_lora_save_load(self):
 
     # Refer to
     # https://github.com/huggingface/diffusers/pull/11806 for more details.
+    @require_peft_version_greater("0.13.2")
     def test_lora_exclude_modules_for_wan(self):
         scheduler_cls = self.scheduler_classes[0]
         components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls)
         pipe = self.pipeline_class(**components).to(torch_device)
         _, _, inputs = self.get_dummy_inputs(with_generator=False)
 
-        output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
-        self.assertTrue(output_no_lora.shape == self.output_shape)
-
-        pipe, _ = self.check_if_adapters_added_correctly(
+        # Only denoiser for now.
+        denoiser_lora_config.target_modules = ["to_q", "to_k", "to_v", "out"]
+        denoiser_lora_config.exclude_modules = ["proj_out"]
+        pipe, _ = self.add_adapters_to_pipeline(
             pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config
         )
-        output_lora_exclude_modules = pipe(**inputs, generator=torch.manual_seed(0))[0]
+        # Inference works.
+        _ = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
         with tempfile.TemporaryDirectory() as tmpdir:
             modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
             lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
             lora_metadatas = self._get_lora_adapter_metadata(modules_to_save)
             self.pipeline_class.save_lora_weights(save_directory=tmpdir, **lora_state_dicts, **lora_metadatas)
             pipe.unload_lora_weights()
-            pipe.load_lora_weights(tmpdir)
-
-            output_lora_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
-            self.assertTrue(
-                not np.allclose(output_no_lora, output_lora_exclude_modules, atol=1e-3, rtol=1e-3),
-                "LoRA should change outputs.",
-            )
-            self.assertTrue(
-                np.allclose(output_lora_exclude_modules, output_lora_pretrained, atol=1e-3, rtol=1e-3),
-                "Lora outputs should match.",
+            # Check the state dict. It should not have any `proj_out` related modules.
+            state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
+            # There should not be any `proj_out` modules, but there should still be some modules for `out`.
+            self.assertTrue(not any("proj_out" in k for k in state_dict))
+            self.assertTrue("out" in k for k in state_dict)
+
+            # Check if the metadata matches.
+            out = pipe.lora_state_dict(tmpdir, return_lora_metadata=True)
+            _, parsed_metadata = out
+            check_module_lora_metadata(
+                parsed_metadata=parsed_metadata, lora_metadatas=lora_metadatas, module_key="transformer"
             )
+
+            # Inference matching is already tested in `test_lora_exclude_modules`.

tests/lora/utils.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import copy
 import inspect
 import os
 import re
@@ -2340,12 +2341,14 @@ def test_lora_unload_add_adapter(self):
         )
         _ = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
+    @require_peft_version_greater("0.13.2")
     def test_lora_exclude_modules(self):
         """
         Test to check if `exclude_modules` works or not. It works in the following way:
         we first create a pipeline and insert LoRA config into it. We then derive a `set`
         of modules to exclude by investigating its denoiser state dict and denoiser LoRA
         state dict.
+
         We then create a new LoRA config to include the `exclude_modules` and perform tests.
         """
         scheduler_cls = self.scheduler_classes[0]
@@ -2356,6 +2359,16 @@ def test_lora_exclude_modules(self):
         output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
         self.assertTrue(output_no_lora.shape == self.output_shape)
 
+        # only supported for `denoiser` now
+        pipe_cp = copy.deepcopy(pipe)
+        pipe_cp, _ = self.add_adapters_to_pipeline(
+            pipe_cp, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config
+        )
+        denoiser_exclude_modules = self._get_exclude_modules(pipe_cp)
+        pipe_cp.to("cpu")
+        del pipe_cp
+
+        denoiser_lora_config.exclude_modules = denoiser_exclude_modules
         pipe, _ = self.add_adapters_to_pipeline(
             pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config
         )