Skip to content

Commit 21d8a8d

Browse files
committed
Merge branch 'main' into z-image-tests
2 parents 5613ff0 + d96cbac commit 21d8a8d

File tree

3 files changed

+53
-0
lines changed

3 files changed

+53
-0
lines changed

src/diffusers/loaders/peft.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import safetensors
2323
import torch
2424

25+
from ..hooks.group_offloading import _maybe_remove_and_reapply_group_offloading
2526
from ..utils import (
2627
MIN_PEFT_VERSION,
2728
USE_PEFT_BACKEND,
@@ -794,6 +795,8 @@ def delete_adapters(self, adapter_names: Union[List[str], str]):
794795
if hasattr(self, "peft_config"):
795796
self.peft_config.pop(adapter_name, None)
796797

798+
_maybe_remove_and_reapply_group_offloading(self)
799+
797800
def enable_lora_hotswap(
798801
self, target_rank: int = 128, check_compiled: Literal["error", "warn", "ignore"] = "error"
799802
) -> None:

tests/lora/utils.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
AutoencoderKL,
2929
UNet2DConditionModel,
3030
)
31+
from diffusers.hooks.group_offloading import _GROUP_OFFLOADING, apply_group_offloading
3132
from diffusers.utils import logging
3233
from diffusers.utils.import_utils import is_peft_available
3334

@@ -2367,3 +2368,51 @@ def test_lora_loading_model_cpu_offload(self):
23672368

23682369
output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0]
23692370
self.assertTrue(np.allclose(output_lora, output_lora_loaded, atol=1e-3, rtol=1e-3))
2371+
2372+
@require_torch_accelerator
2373+
def test_lora_group_offloading_delete_adapters(self):
2374+
components, _, denoiser_lora_config = self.get_dummy_components()
2375+
_, _, inputs = self.get_dummy_inputs(with_generator=False)
2376+
pipe = self.pipeline_class(**components)
2377+
pipe = pipe.to(torch_device)
2378+
pipe.set_progress_bar_config(disable=None)
2379+
2380+
denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
2381+
denoiser.add_adapter(denoiser_lora_config)
2382+
self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.")
2383+
2384+
try:
2385+
with tempfile.TemporaryDirectory() as tmpdirname:
2386+
modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
2387+
lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
2388+
self.pipeline_class.save_lora_weights(
2389+
save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts
2390+
)
2391+
2392+
components, _, _ = self.get_dummy_components()
2393+
pipe = self.pipeline_class(**components)
2394+
denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
2395+
pipe.to(torch_device)
2396+
2397+
# Enable Group Offloading (leaf_level for more granular testing)
2398+
apply_group_offloading(
2399+
denoiser,
2400+
onload_device=torch_device,
2401+
offload_device="cpu",
2402+
offload_type="leaf_level",
2403+
)
2404+
2405+
pipe.load_lora_weights(tmpdirname, adapter_name="default")
2406+
2407+
out_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
2408+
2409+
# Delete the adapter
2410+
pipe.delete_adapters("default")
2411+
2412+
out_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
2413+
2414+
self.assertFalse(np.allclose(out_lora, out_no_lora, atol=1e-3, rtol=1e-3))
2415+
finally:
2416+
# Clean up the hooks to prevent state leak
2417+
if hasattr(denoiser, "_diffusers_hook"):
2418+
denoiser._diffusers_hook.remove_hook(_GROUP_OFFLOADING, recurse=True)

tests/models/transformers/test_models_transformer_hunyuan_1_5.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ class HunyuanVideo15Transformer3DTests(ModelTesterMixin, unittest.TestCase):
2929
model_class = HunyuanVideo15Transformer3DModel
3030
main_input_name = "hidden_states"
3131
uses_custom_attn_processor = True
32+
model_split_percents = [0.99, 0.99, 0.99]
3233

3334
text_embed_dim = 16
3435
text_embed_2_dim = 8

0 commit comments

Comments
 (0)