[LoRA] add tests for partial text encoders LoRAs

Author: elias-gaeros

tests/lora/utils.py

@@ -395,6 +395,65 @@ def test_simple_inference_with_text_lora_save_load(self):
                 "Loading from saved checkpoints should give same results.",
             )
 
+    def test_simple_inference_with_partial_text_lora(self):
+        """
+        Tests a simple inference with lora attached on the text encoder
+        with different ranks and some adapters removed
+        and makes sure it works as expected
+        """
+        for scheduler_cls in [DDIMScheduler, LCMScheduler]:
+            components, _, _ = self.get_dummy_components(scheduler_cls)
+            text_lora_config = LoraConfig(
+                r=4,
+                rank_pattern={"q_proj": 1, "k_proj": 2, "v_proj": 3},
+                lora_alpha=4,
+                target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
+                init_lora_weights=False,
+                use_dora=False,
+            )
+            pipe = self.pipeline_class(**components)
+            pipe = pipe.to(torch_device)
+            pipe.set_progress_bar_config(disable=None)
+            _, _, inputs = self.get_dummy_inputs(with_generator=False)
+
+            output_no_lora = pipe(**inputs, generator=torch.manual_seed(0)).images
+            self.assertTrue(output_no_lora.shape == (1, 64, 64, 3))
+
+            pipe.text_encoder.add_adapter(text_lora_config)
+            self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder")
+            state_dict = {
+                f"text_encoder.{module_name}": param
+                for module_name, param in get_peft_model_state_dict(pipe.text_encoder).items()
+            }
+
+            if self.has_two_text_encoders:
+                pipe.text_encoder_2.add_adapter(text_lora_config)
+                self.assertTrue(
+                    check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2"
+                )
+                state_dict.update(
+                    {
+                        f"text_encoder.{module_name}": param
+                        for module_name, param in get_peft_model_state_dict(pipe.text_encoder_2).items()
+                    }
+                )
+
+            # Discard half of the adapters.
+            rng = np.random.default_rng(0)
+            key2adapters = {k: k.rsplit(".", 2)[0] for k in state_dict.keys()}
+            adapters = list(set(key2adapters.values()))
+            adapters = set(rng.choice(adapters, size=len(adapters) // 2, replace=False))
+            state_dict = {k: state_dict[k] for k, adapter in key2adapters.items() if adapter in adapters}
+
+            # Unload lora and load it back using the pipe.load_lora_weights machinery
+            pipe.unload_lora_weights()
+            pipe.load_lora_weights(state_dict)
+
+            output_lora = pipe(**inputs, generator=torch.manual_seed(0)).images
+            self.assertTrue(
+                not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output"
+            )
+
     def test_simple_inference_save_pretrained(self):
         """
         Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained