Test: Add regression test for group offloading + delete_adapters

Author: adi776borate

tests/lora/utils.py

@@ -30,6 +30,7 @@
 )
 from diffusers.utils import logging
 from diffusers.utils.import_utils import is_peft_available
+from diffusers.hooks.group_offloading import apply_group_offloading
 
 from ..testing_utils import (
     CaptureLogger,
@@ -2367,3 +2368,42 @@ def test_lora_loading_model_cpu_offload(self):
 
         output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0]
         self.assertTrue(np.allclose(output_lora, output_lora_loaded, atol=1e-3, rtol=1e-3))
+
+    @require_torch_accelerator
+    def test_lora_group_offloading_delete_adapters(self):
+        components, _, denoiser_lora_config = self.get_dummy_components()
+        _, _, inputs = self.get_dummy_inputs(with_generator=False)
+        pipe = self.pipeline_class(**components)
+        pipe = pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
+        denoiser.add_adapter(denoiser_lora_config)
+        self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.")
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
+            lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
+            self.pipeline_class.save_lora_weights(
+                save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts
+            )
+
+            components, _, _ = self.get_dummy_components()
+            pipe = self.pipeline_class(**components)
+            pipe.to(torch_device)
+
+            denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
+
+            # Enable Group Offloading (leaf_level)
+            apply_group_offloading(
+                denoiser,
+                onload_device=torch_device,
+                offload_device="cpu",
+                offload_type="leaf_level",
+            )
+
+            pipe.load_lora_weights(tmpdirname, adapter_name="default")
+            pipe(**inputs, generator=torch.manual_seed(0))
+            # Delete the adapter
+            pipe.delete_adapters("default")
+            pipe(**inputs, generator=torch.manual_seed(0))