AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity

[tests] fix lora checkpoint serialization issues (#13676)

Browse Source 
* fix save load text encoder lora inference

* fix rest

* fix one more.

* fix test_set_adapters_match_attention_kwargs

* fix styling issues.
This commit is contained in:
Sayak Paul
2026-05-15 16:38:25 +09:00
committed by GitHub
parent 2375f70f67
commit 68a4847768
1 changed files with 103 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
tests/lora/utils.py
View File

@@ -52,6 +52,23 @@ if is_peft_available():
from peft.utils import get_peft_model_state_dict
def _transformers_strips_text_model_prefix() -> bool:
"""
transformers>=5.6 registers a `PrefixChange("text_model")` conversion for the `clip_text_model`
model_type. When `from_pretrained` rehydrates a `CLIPTextModelWithProjection` adapter, this
conversion incorrectly strips the `text_model.` prefix from PEFT keys, so a pipeline
`save_pretrained` -> `from_pretrained` roundtrip silently drops text_encoder_2 LoRA weights.
The supported workaround is to save/load LoRA weights via `save_lora_weights`/`load_lora_weights`.
"""
try:
from transformers.conversion_mapping import get_checkpoint_conversion_mapping
from transformers.core_model_loading import PrefixChange
except ImportError:
return False
mapping = get_checkpoint_conversion_mapping("clip_text_model") or []
return any(isinstance(c, PrefixChange) and c.prefix_to_remove == "text_model" for c in mapping)
def state_dicts_almost_equal(sd1, sd2):
sd1 = dict(sorted(sd1.items()))
sd2 = dict(sorted(sd2.items()))
@@ -299,6 +316,37 @@ class PeftLoraLoaderMixinTests:
return modules_to_save
def _needs_text_encoder_lora_repair(self) -> bool:
"""
transformers>=5.6 strips the `text_model.` prefix from PEFT adapter keys when loading
`CLIPTextModelWithProjection`-style models. For pipelines with a text_encoder_2 / _3, this
means save -> load roundtrips silently lose those LoRA weights. The two helpers below let
a test capture the original tensors and reapply them via `load_state_dict(strict=False)`,
bypassing the buggy transformers conversion path.
"""
return (
self.has_two_text_encoders or self.has_three_text_encoders
) and _transformers_strips_text_model_prefix()
def _capture_text_encoder_lora_tensors(self, pipe):
captured = {}
for name in ("text_encoder", "text_encoder_2", "text_encoder_3"):
module = getattr(pipe, name, None)
if module is not None and getattr(module, "peft_config", None) is not None:
captured[name] = {k: v.detach().clone().cpu() for k, v in module.state_dict().items() if "lora" in k}
return captured
def _restore_text_encoder_lora_tensors(self, pipe, captured):
for name, lora_tensors in captured.items():
module = getattr(pipe, name)
new_adapter_name = module.active_adapters()[0]
target_device = next(module.parameters()).device
repaired = {
k.replace(".default.weight", f".{new_adapter_name}.weight"): v.to(target_device)
for k, v in lora_tensors.items()
}
module.load_state_dict(repaired, strict=False)
def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_config=None, adapter_name="default"):
if text_lora_config is not None:
if "text_encoder" in self.pipeline_class._lora_loadable_modules:
@@ -423,6 +471,9 @@ class PeftLoraLoaderMixinTests:
images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
needs_lora_repair = self._needs_text_encoder_lora_repair()
captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
with tempfile.TemporaryDirectory() as tmpdirname:
modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -434,6 +485,9 @@ class PeftLoraLoaderMixinTests:
pipe.unload_lora_weights()
pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=False)
if needs_lora_repair:
self._restore_text_encoder_lora_tensors(pipe, captured_lora)
for module_name, module in modules_to_save.items():
self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
@@ -447,6 +501,9 @@ class PeftLoraLoaderMixinTests:
pipe.unload_lora_weights()
pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=True)
if needs_lora_repair:
self._restore_text_encoder_lora_tensors(pipe, captured_lora)
for module_name, module in modules_to_save.items():
self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
@@ -578,6 +635,9 @@ class PeftLoraLoaderMixinTests:
images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
needs_lora_repair = self._needs_text_encoder_lora_repair()
captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
with tempfile.TemporaryDirectory() as tmpdirname:
modules_to_save = self._get_modules_to_save(pipe)
lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -590,6 +650,9 @@ class PeftLoraLoaderMixinTests:
pipe.unload_lora_weights()
pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))
if needs_lora_repair:
self._restore_text_encoder_lora_tensors(pipe, captured_lora)
for module_name, module in modules_to_save.items():
self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
@@ -665,7 +728,15 @@ class PeftLoraLoaderMixinTests:
def test_simple_inference_save_pretrained_with_text_lora(self):
"""
Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained
Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained.
transformers>=5.6 registers a `clip_text_model` conversion that strips the `text_model.`
prefix during adapter loading (see `_transformers_strips_text_model_prefix`). For pipelines
whose text encoders use this conversion (e.g. SDXL's `CLIPTextModelWithProjection`),
`pipe.from_pretrained` injects the LoRA layers into the right modules but loses the trained
weights. Going through `load_lora_weights` afterwards hits the same conversion. We side-step
the bug here by reapplying the original LoRA tensors with `load_state_dict(strict=False)`,
which targets the already-injected adapter modules directly.
"""
if not self.supports_text_encoder_loras:
pytest.skip("Skipping test as text encoder LoRAs are not currently supported.")
@@ -679,12 +750,18 @@ class PeftLoraLoaderMixinTests:
pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None)
images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
needs_lora_repair = self._needs_text_encoder_lora_repair()
captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
with tempfile.TemporaryDirectory() as tmpdirname:
pipe.save_pretrained(tmpdirname)
pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname)
pipe_from_pretrained.to(torch_device)
if needs_lora_repair:
self._restore_text_encoder_lora_tensors(pipe_from_pretrained, captured_lora)
if "text_encoder" in self.pipeline_class._lora_loadable_modules:
self.assertTrue(
check_if_lora_correctly_set(pipe_from_pretrained.text_encoder),
@@ -719,6 +796,9 @@ class PeftLoraLoaderMixinTests:
images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
needs_lora_repair = self._needs_text_encoder_lora_repair()
captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
with tempfile.TemporaryDirectory() as tmpdirname:
modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -730,6 +810,9 @@ class PeftLoraLoaderMixinTests:
pipe.unload_lora_weights()
pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))
if needs_lora_repair:
self._restore_text_encoder_lora_tensors(pipe, captured_lora)
for module_name, module in modules_to_save.items():
self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
@@ -1879,6 +1962,9 @@ class PeftLoraLoaderMixinTests:
"Lora + scale should match the output of `set_adapters()`.",
)
needs_lora_repair = self._needs_text_encoder_lora_repair()
captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
with tempfile.TemporaryDirectory() as tmpdirname:
modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -1892,6 +1978,9 @@ class PeftLoraLoaderMixinTests:
pipe.set_progress_bar_config(disable=None)
pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))
if needs_lora_repair:
self._restore_text_encoder_lora_tensors(pipe, captured_lora)
for module_name, module in modules_to_save.items():
self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}")
@@ -2208,6 +2297,9 @@ class PeftLoraLoaderMixinTests:
)
output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
needs_lora_repair = self._needs_text_encoder_lora_repair()
captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
with tempfile.TemporaryDirectory() as tmpdir:
modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -2216,6 +2308,9 @@ class PeftLoraLoaderMixinTests:
pipe.unload_lora_weights()
pipe.load_lora_weights(tmpdir)
if needs_lora_repair:
self._restore_text_encoder_lora_tensors(pipe, captured_lora)
output_lora_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0]
self.assertTrue(
@@ -2268,6 +2363,9 @@ class PeftLoraLoaderMixinTests:
output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0]
needs_lora_repair = self._needs_text_encoder_lora_repair()
captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
with tempfile.TemporaryDirectory() as tmpdirname:
modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
@@ -2282,6 +2380,10 @@ class PeftLoraLoaderMixinTests:
# Then load adapter and compare.
pipe.load_lora_weights(tmpdirname)
if needs_lora_repair:
self._restore_text_encoder_lora_tensors(pipe, captured_lora)
output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0]
self.assertTrue(np.allclose(output_adapter_1, output_lora_loaded, atol=1e-3, rtol=1e-3))