Fix: serialize low_cpu_mem_usage LoRA injection to prevent thread-safety leak

src/diffusers/loaders/lora_base.py

@@ -332,6 +332,9 @@ def _load_lora_into_text_encoder(
 ):
     from ..hooks.group_offloading import _maybe_remove_and_reapply_group_offloading
 
+    # Imported lazily to avoid a circular import (peft.py imports from this module).
+    from .peft import _LOW_CPU_MEM_USAGE_INJECTION_LOCK
+
     if not USE_PEFT_BACKEND:
         raise ValueError("PEFT backend is required for this method.")
 
@@ -398,12 +401,21 @@ def _load_lora_into_text_encoder(
         )
         # inject LoRA layers and load the state dict
         # in transformers we automatically check whether the adapter name is already in use or not
-        text_encoder.load_adapter(
-            adapter_name=adapter_name,
-            adapter_state_dict=state_dict,
-            peft_config=lora_config,
-            **peft_kwargs,
-        )
+        if low_cpu_mem_usage:
+            with _LOW_CPU_MEM_USAGE_INJECTION_LOCK:
+                text_encoder.load_adapter(
+                    adapter_name=adapter_name,
+                    adapter_state_dict=state_dict,
+                    peft_config=lora_config,
+                    **peft_kwargs,
+                )
+        else:
+            text_encoder.load_adapter(
+                adapter_name=adapter_name,
+                adapter_state_dict=state_dict,
+                peft_config=lora_config,
+                **peft_kwargs,
+            )
 
         # scale LoRA layers with `lora_scale`
         scale_lora_layers(text_encoder, weight=lora_scale)

src/diffusers/loaders/peft.py

@@ -15,6 +15,7 @@
 import inspect
 import json
 import os
+import threading
 from collections import defaultdict
 from functools import partial
 from pathlib import Path
@@ -45,6 +46,12 @@
 
 logger = logging.get_logger(__name__)
 
+# PEFT's `init_empty_weights()` context monkey-patches `torch.nn.Module.register_parameter`
+# process-wide on entry and restores the value captured at entry on exit, which is not
+# thread-safe. Concurrent adapter injection with `low_cpu_mem_usage=True` can interleave these
+# capture/restore operations and leak the patch, so the low-memory injection path is serialized.
+_LOW_CPU_MEM_USAGE_INJECTION_LOCK = threading.Lock()
+
 _SET_ADAPTER_SCALE_FN_MAPPING = defaultdict(
     lambda: (lambda model_cls, weights: weights),
     {
@@ -122,7 +129,8 @@ def load_lora_adapter(
                 link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
             low_cpu_mem_usage (`bool`, *optional*):
                 Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
-                weights.
+                weights. When enabled, the low-memory loading path is serialized with a process-wide lock, making it
+                safe to call concurrently across threads.
             hotswap : (`bool`, *optional*)
                 Defaults to `False`. Whether to substitute an existing (LoRA) adapter with the newly loaded adapter
                 in-place. This means that, instead of loading an additional adapter, this will take the existing
@@ -324,9 +332,15 @@ def map_state_dict_for_hotswap(sd):
                     # it to None
                     incompatible_keys = None
                 else:
-                    inject_adapter_in_model(
-                        lora_config, self, adapter_name=adapter_name, state_dict=state_dict, **peft_kwargs
-                    )
+                    if low_cpu_mem_usage:
+                        with _LOW_CPU_MEM_USAGE_INJECTION_LOCK:
+                            inject_adapter_in_model(
+                                lora_config, self, adapter_name=adapter_name, state_dict=state_dict, **peft_kwargs
+                            )
+                    else:
+                        inject_adapter_in_model(
+                            lora_config, self, adapter_name=adapter_name, state_dict=state_dict, **peft_kwargs
+                        )
                     incompatible_keys = set_peft_model_state_dict(self, state_dict, adapter_name, **peft_kwargs)
 
                     if self._prepare_lora_hotswap_kwargs is not None:

tests/lora/test_lora_low_cpu_mem_thread_safety.py

@@ -0,0 +1,140 @@
+# coding=utf-8
+# Copyright 2026 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import gc
+import threading
+import unittest
+
+import torch
+
+from diffusers import UNet2DConditionModel
+
+from ..testing_utils import require_peft_backend
+
+
+def _build_lora_state_dict(unet, rank=4):
+    """Build a minimal LoRA state dict for the attention linear layers of the given UNet."""
+    state_dict = {}
+    for name, module in unet.named_modules():
+        if isinstance(module, torch.nn.Linear) and (
+            name.endswith("to_q") or name.endswith("to_k") or name.endswith("to_v") or name.endswith("to_out.0")
+        ):
+            state_dict[f"{name}.lora_A.weight"] = torch.randn(rank, module.in_features)
+            state_dict[f"{name}.lora_B.weight"] = torch.randn(module.out_features, rank)
+    return state_dict
+
+
+@require_peft_backend
+class LowCpuMemUsageThreadSafetyTests(unittest.TestCase):
+    unet_kwargs = {
+        "block_out_channels": (32, 64),
+        "layers_per_block": 2,
+        "sample_size": 32,
+        "in_channels": 4,
+        "out_channels": 4,
+        "down_block_types": ("DownBlock2D", "CrossAttnDownBlock2D"),
+        "up_block_types": ("CrossAttnUpBlock2D", "UpBlock2D"),
+        "cross_attention_dim": 32,
+    }
+
+    num_threads = 4
+    num_rounds = 10
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        torch.manual_seed(0)
+        probe_unet = UNet2DConditionModel(**cls.unet_kwargs)
+        cls.lora_state_dict = _build_lora_state_dict(probe_unet)
+        del probe_unet
+        gc.collect()
+
+    def _run_concurrent_injections(self, low_cpu_mem_usage):
+        original = torch.nn.Module.register_parameter
+        errors = []
+        try:
+            barrier = threading.Barrier(self.num_threads)
+
+            def worker(thread_id):
+                for _ in range(self.num_rounds):
+                    try:
+                        barrier.wait()
+                    except threading.BrokenBarrierError:
+                        return
+                    try:
+                        unet = UNet2DConditionModel(**self.unet_kwargs)
+                        unet.load_lora_adapter(
+                            self.lora_state_dict,
+                            adapter_name=f"adapter-{thread_id}",
+                            low_cpu_mem_usage=low_cpu_mem_usage,
+                            prefix=None,
+                        )
+                    except Exception as e:
+                        errors.append((thread_id, type(e).__name__, str(e)))
+                        return
+
+            threads = [threading.Thread(target=worker, args=(i,)) for i in range(self.num_threads)]
+            for thread in threads:
+                thread.start()
+            for thread in threads:
+                thread.join()
+
+            return {
+                "register_parameter_patched": torch.nn.Module.register_parameter is not original,
+                "fresh_linear_device": torch.nn.Linear(2, 2).weight.device.type,
+                "thread_errors": errors,
+            }
+        finally:
+            torch.nn.Module.register_parameter = original
+
+    def test_concurrent_low_cpu_mem_usage_injection_does_not_leak_register_parameter(self):
+        """Concurrent `low_cpu_mem_usage=True` LoRA injection must not leak the global patch."""
+        observed = self._run_concurrent_injections(low_cpu_mem_usage=True)
+        self.assertFalse(
+            observed["register_parameter_patched"],
+            f"Concurrent `low_cpu_mem_usage=True` injection leaked the global "
+            f"`torch.nn.Module.register_parameter` monkey patch. Observed: {observed}",
+        )
+        self.assertNotEqual(
+            observed["fresh_linear_device"],
+            "meta",
+            f"Concurrent `low_cpu_mem_usage=True` injection left newly created modules on the meta "
+            f"device. Observed: {observed}",
+        )
+        self.assertEqual(
+            observed["thread_errors"],
+            [],
+            f"Worker threads raised during concurrent `low_cpu_mem_usage=True` injection: {observed['thread_errors']}",
+        )
+
+    def test_concurrent_injection_without_low_cpu_mem_usage_is_safe(self):
+        """Sanity control: concurrent `low_cpu_mem_usage=False` injection must be safe."""
+        observed = self._run_concurrent_injections(low_cpu_mem_usage=False)
+        self.assertFalse(
+            observed["register_parameter_patched"],
+            f"Concurrent `low_cpu_mem_usage=False` injection leaked the global "
+            f"`torch.nn.Module.register_parameter` monkey patch. Observed: {observed}",
+        )
+        self.assertNotEqual(
+            observed["fresh_linear_device"],
+            "meta",
+            f"Concurrent `low_cpu_mem_usage=False` injection left newly created modules on the meta "
+            f"device. Observed: {observed}",
+        )
+        self.assertEqual(
+            observed["thread_errors"],
+            [],
+            f"Worker threads raised during concurrent `low_cpu_mem_usage=False` injection: "
+            f"{observed['thread_errors']}",
+        )