Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 31 additions & 3 deletions src/diffusers/loaders/lora_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -635,6 +635,10 @@ def unfuse_lora(self, components: list[str] | None = None, **kwargs):
unfuse_text_encoder (`bool`, defaults to `True`):
Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
LoRA parameters then it won't have any effect.

Note that `fused_loras` is tracked pipeline-wide, not per component. An adapter fused into multiple
components only drops out of `fused_loras` once it has been unfused from all of them; unfusing it from
only some components leaves it listed as fused.
"""
if components is None:
components = []
Expand Down Expand Up @@ -664,6 +668,13 @@ def unfuse_lora(self, components: list[str] | None = None, **kwargs):
if len(components) == 0:
raise ValueError("`components` cannot be an empty list.")

# self._merged_adapters is shared across the whole pipeline (all _lora_loadable_modules), while
# merge state is tracked per component by PEFT. The same adapter name can be fused into multiple
# components independently, so unmerging it in only some of them doesn't mean it's fully unfused.
# Track candidates here and only drop them from the pipeline-wide set once confirmed unmerged
# everywhere, below.
unmerge_candidates: set[str] = set()

for fuse_component in components:
if fuse_component not in self._lora_loadable_modules:
raise ValueError(f"{fuse_component} is not found in {self._lora_loadable_modules=}.")
Expand All @@ -673,11 +684,28 @@ def unfuse_lora(self, components: list[str] | None = None, **kwargs):
if issubclass(model.__class__, (ModelMixin, PreTrainedModel)):
for module in model.modules():
if isinstance(module, BaseTunerLayer):
for adapter in set(module.merged_adapters):
if adapter and adapter in self._merged_adapters:
self._merged_adapters = self._merged_adapters - {adapter}
unmerge_candidates.update(module.merged_adapters)
module.unmerge()

for adapter in unmerge_candidates:
if adapter not in self._merged_adapters:
continue
if not self._is_adapter_merged_in_any_component(adapter):
self._merged_adapters = self._merged_adapters - {adapter}

def _is_adapter_merged_in_any_component(self, adapter_name: str) -> bool:
"""Whether `adapter_name` is still merged into the base weights of any `_lora_loadable_modules`
component, used to keep the pipeline-wide `_merged_adapters` bookkeeping in sync with the real, per-component
PEFT merge state after a partial `unfuse_lora(components=...)` call."""
for component in self._lora_loadable_modules:
model = getattr(self, component, None)
if model is None or not issubclass(model.__class__, (ModelMixin, PreTrainedModel)):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why issubclass(model.__class__, ...) and not isinstance(model, ...)?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's copied from the existing pattern right above it in the same unfuse_lora loop (and used throughout this file, 16 occurrences), not a deliberate choice on my part - happy to switch to isinstance() here if you'd rather not perpetuate it, but left it matching the surrounding code for now.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see. IMO, isinstance makes more sense but maybe I'm missing something. I'll leave that decision up to the Diffusers maintainers.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds like the consensus is to leave it as-is for now and revisit library-wide later, so no change here.

continue
for module in model.modules():
if isinstance(module, BaseTunerLayer) and adapter_name in module.merged_adapters:
return True
return False

def set_adapters(
self,
adapter_names: list[str] | str,
Expand Down
40 changes: 40 additions & 0 deletions tests/lora/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1711,6 +1711,46 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(
pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules)
self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")

def test_fuse_unfuse_partial_components_keeps_merged_adapter_bookkeeping(self):
"""
`_merged_adapters` (backing `num_fused_loras`/`fused_loras`) is shared across the whole pipeline,
while actual merge state is tracked per component by PEFT. Fusing an adapter into every loadable
component and then unfusing only *some* of them should not report the adapter as fully unfused,
it's still merged into the untouched component(s).
"""
if not self.supports_text_encoder_loras:
pytest.skip("Skipping test as text encoder LoRAs are not currently supported.")

components, text_lora_config, denoiser_lora_config = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

pipe.text_encoder.add_adapter(text_lora_config, "adapter-1")
self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder")

denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
denoiser_component_name = "unet" if self.unet_kwargs is not None else "transformer"
denoiser.add_adapter(denoiser_lora_config, "adapter-1")
self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.")

pipe.fuse_lora(components=["text_encoder", denoiser_component_name], adapter_names=["adapter-1"])
self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")

# Only unfuse the text encoder; the denoiser still has "adapter-1" merged into its base weights.
pipe.unfuse_lora(components=["text_encoder"])
self.assertTrue(
pipe.num_fused_loras == 1,
f"adapter-1 is still merged into {denoiser_component_name}, but num_fused_loras dropped to "
f"{pipe.num_fused_loras=}, {pipe.fused_loras=}",
)
self.assertIn("adapter-1", pipe.fused_loras)

# Now unfuse the remaining component; bookkeeping should correctly drop to 0.
pipe.unfuse_lora(components=[denoiser_component_name])
self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For completeness, I would also check "adapter-1" is not in pipe.fused_loras.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in f6c54bc - now also asserts adapter-1 is out of fused_loras once it's unfused from every component.

self.assertNotIn("adapter-1", pipe.fused_loras)

def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3):
attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class)

Expand Down
Loading