Repository: huggingface/diffusers
File: tests/lora/test_lora_layers_flux.py (Python)

Code before:

            # Modify the state dict to exclude "x_embedder" related LoRA params.
            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))
            lora_state_dict_without_xembedder = {k: v for k, v in lora_state_dict.items() if "x_embedder" not in k}

        pipe.load_lora_weights(lora_state_dict_without_xembedder, adapter_name="two")
        pipe.set_adapters(["one", "two"])
        self.assertTrue(check_if_lora_correctly_set(pipe.transformer), "Lora not correctly set in transformer")
        images_lora_with_absent_keys = pipe(**inputs, generator=torch.manual_seed(0)).images

        self.assertFalse(
            np.allclose(images_lora, images_lora_with_absent_keys, atol=1e-3, rtol=1e-3),
            "Different LoRAs should lead to different results.",
        )
        self.assertFalse(
            np.allclose(output_no_lora, images_lora_with_absent_keys, atol=1e-3, rtol=1e-3),
            "LoRA should lead to different results.",
        )

    def test_lora_expansion_works_for_extra_keys(self):
        components, _, denoiser_lora_config = self.get_dummy_components(FlowMatchEulerDiscreteScheduler)
        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 = self.cached_non_lora_outputs[self.scheduler_cls.__name__]

        # Modify the config to have a layer which won't be present in the first LoRA we will load.
        modified_denoiser_lora_config = copy.deepcopy(denoiser_lora_config)
        modified_denoiser_lora_config.target_modules.add("x_embedder")

        pipe.transformer.add_adapter(modified_denoiser_lora_config)
        self.assertTrue(check_if_lora_correctly_set(pipe.transformer), "Lora not correctly set in transformer")

        images_lora = pipe(**inputs, generator=torch.manual_seed(0)).images
        self.assertFalse(
            np.allclose(images_lora, output_no_lora, atol=1e-3, rtol=1e-3),
            "LoRA should lead to different results.",
        )

        with tempfile.TemporaryDirectory() as tmpdirname:
            denoiser_state_dict = get_peft_model_state_dict(pipe.transformer)
            self.pipeline_class.save_lora_weights(tmpdirname, transformer_lora_layers=denoiser_state_dict)

            self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")))
            pipe.unload_lora_weights()
            # Modify the state dict to exclude "x_embedder" related LoRA params.
            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))
            lora_state_dict_without_xembedder = {k: v for k, v in lora_state_dict.items() if "x_embedder" not in k}
            pipe.load_lora_weights(lora_state_dict_without_xembedder, adapter_name="one")


Review comment (performance):
More of a nit, but I would create a method in the Mixin called `get_base_pipeline_output` and return from there, rather than directly accessing the cache. 

Code after:

            # Modify the state dict to exclude "x_embedder" related LoRA params.
            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))
            lora_state_dict_without_xembedder = {k: v for k, v in lora_state_dict.items() if "x_embedder" not in k}

        pipe.load_lora_weights(lora_state_dict_without_xembedder, adapter_name="two")
        pipe.set_adapters(["one", "two"])
        self.assertTrue(check_if_lora_correctly_set(pipe.transformer), "Lora not correctly set in transformer")
        images_lora_with_absent_keys = pipe(**inputs, generator=torch.manual_seed(0)).images

        self.assertFalse(
            np.allclose(images_lora, images_lora_with_absent_keys, atol=1e-3, rtol=1e-3),
            "Different LoRAs should lead to different results.",
        )
        self.assertFalse(
            np.allclose(output_no_lora, images_lora_with_absent_keys, atol=1e-3, rtol=1e-3),
            "LoRA should lead to different results.",
        )

    def test_lora_expansion_works_for_extra_keys(self):
        components, _, denoiser_lora_config = self.get_dummy_components(FlowMatchEulerDiscreteScheduler)
        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 = self.get_base_pipe_output()

        # Modify the config to have a layer which won't be present in the first LoRA we will load.
        modified_denoiser_lora_config = copy.deepcopy(denoiser_lora_config)
        modified_denoiser_lora_config.target_modules.add("x_embedder")

        pipe.transformer.add_adapter(modified_denoiser_lora_config)
        self.assertTrue(check_if_lora_correctly_set(pipe.transformer), "Lora not correctly set in transformer")

        images_lora = pipe(**inputs, generator=torch.manual_seed(0)).images
        self.assertFalse(
            np.allclose(images_lora, output_no_lora, atol=1e-3, rtol=1e-3),
            "LoRA should lead to different results.",
        )

        with tempfile.TemporaryDirectory() as tmpdirname:
            denoiser_state_dict = get_peft_model_state_dict(pipe.transformer)
            self.pipeline_class.save_lora_weights(tmpdirname, transformer_lora_layers=denoiser_state_dict)

            self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")))
            pipe.unload_lora_weights()
            # Modify the state dict to exclude "x_embedder" related LoRA params.
            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))
            lora_state_dict_without_xembedder = {k: v for k, v in lora_state_dict.items() if "x_embedder" not in k}
            pipe.load_lora_weights(lora_state_dict_without_xembedder, adapter_name="one")