diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index dc01475ac8de..402f7beda285 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -664,6 +664,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=}.") @@ -673,11 +680,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)): + 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, diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 38aec8ce4807..ad5aaf1ab868 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -1711,6 +1711,45 @@ 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 "text_encoder" not in self.pipeline_class._lora_loadable_modules: + return + + 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=}") + 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)