Fix _merged_adapters bookkeeping desync on partial unfuse_lora(components=...)

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,

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)