mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
5ff8e59ff9
Co-authored-by: Aznix07 <sruhilmodi13@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
206 lines
7.2 KiB
Python
206 lines
7.2 KiB
Python
# 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 logging
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
from safetensors.torch import save_file
|
|
from transformers import CLIPTextConfig, CLIPTextModel
|
|
|
|
from diffusers.configuration_utils import ConfigMixin
|
|
from diffusers.loaders import StableDiffusionLoraLoaderMixin, lora_base
|
|
from diffusers.loaders.lora_base import LoraBaseMixin
|
|
from diffusers.loaders.peft import PeftAdapterMixin
|
|
from diffusers.models.modeling_utils import ModelMixin
|
|
from diffusers.utils import is_transformers_version
|
|
from diffusers.utils.import_utils import is_peft_available
|
|
|
|
from ..testing_utils import CaptureLogger, require_peft_backend
|
|
|
|
|
|
if is_peft_available():
|
|
from peft import LoraConfig
|
|
from peft.tuners.tuners_utils import BaseTunerLayer
|
|
|
|
|
|
LORA_KEY = "unet.test.lora_A.weight"
|
|
|
|
|
|
def _write_lora_weights(path):
|
|
save_file({LORA_KEY: torch.ones(1)}, path)
|
|
|
|
|
|
@pytest.fixture
|
|
def lora_weight_path(tmp_path):
|
|
weight_path = tmp_path / "adapter.safetensors"
|
|
_write_lora_weights(weight_path)
|
|
return weight_path
|
|
|
|
|
|
@pytest.fixture
|
|
def model_info_mock(monkeypatch):
|
|
model_info_mock = Mock()
|
|
monkeypatch.setattr(lora_base, "model_info", model_info_mock)
|
|
return model_info_mock
|
|
|
|
|
|
def test_local_directory_in_offline_mode(lora_weight_path, monkeypatch, model_info_mock):
|
|
monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)
|
|
|
|
state_dict, _ = StableDiffusionLoraLoaderMixin.lora_state_dict(lora_weight_path.parent)
|
|
|
|
assert torch.equal(state_dict[LORA_KEY], torch.ones(1))
|
|
model_info_mock.assert_not_called()
|
|
|
|
|
|
def test_local_directory_with_local_files_only(lora_weight_path, model_info_mock):
|
|
state_dict, _ = StableDiffusionLoraLoaderMixin.lora_state_dict(lora_weight_path.parent, local_files_only=True)
|
|
|
|
assert torch.equal(state_dict[LORA_KEY], torch.ones(1))
|
|
model_info_mock.assert_not_called()
|
|
|
|
|
|
def test_local_file_in_offline_mode(lora_weight_path, monkeypatch, model_info_mock):
|
|
monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)
|
|
|
|
state_dict, _ = StableDiffusionLoraLoaderMixin.lora_state_dict(lora_weight_path)
|
|
|
|
assert torch.equal(state_dict[LORA_KEY], torch.ones(1))
|
|
model_info_mock.assert_not_called()
|
|
|
|
|
|
def test_remote_repository_in_offline_mode_requires_weight_name(monkeypatch, model_info_mock):
|
|
monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)
|
|
|
|
with pytest.raises(ValueError, match="offline mode.*weight_name"):
|
|
StableDiffusionLoraLoaderMixin.lora_state_dict("organization/repository")
|
|
|
|
model_info_mock.assert_not_called()
|
|
|
|
|
|
def test_local_directory_without_matching_files_returns_none(tmp_path, monkeypatch):
|
|
(tmp_path / "notes.txt").touch()
|
|
monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)
|
|
|
|
weight_name = lora_base._best_guess_weight_name(tmp_path)
|
|
|
|
assert weight_name is None
|
|
|
|
|
|
def test_local_directory_with_multiple_files_warns_and_uses_first(tmp_path, monkeypatch):
|
|
first_path = tmp_path / "first.safetensors"
|
|
second_path = tmp_path / "second.safetensors"
|
|
first_path.touch()
|
|
second_path.touch()
|
|
monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)
|
|
# `os.listdir` returns entries in arbitrary order; pin it so `first.safetensors` is the one picked.
|
|
monkeypatch.setattr(lora_base.os, "listdir", lambda _: [first_path.name, second_path.name])
|
|
|
|
# `caplog` cannot see this warning: `diffusers` sets `propagate = False` on its library root logger, so
|
|
# records never reach the root handler pytest installs. `CaptureLogger` attaches to the logger directly.
|
|
lora_base.logger.setLevel(logging.WARNING)
|
|
with CaptureLogger(lora_base.logger) as cap_logger:
|
|
weight_name = lora_base._best_guess_weight_name(tmp_path)
|
|
|
|
assert weight_name == first_path.name
|
|
assert "contains more than one weights file" in cap_logger.out
|
|
|
|
|
|
@require_peft_backend
|
|
def test_unfuse_lora_partial_components_keeps_merged_adapters_in_sync():
|
|
"""Regression test for gh-14214.
|
|
|
|
Unfusing only a subset of components must keep _merged_adapters in sync
|
|
with the adapters still physically fused in the remaining components.
|
|
"""
|
|
|
|
class TinyModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
|
|
config_name = "config.json"
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear = nn.Linear(8, 8)
|
|
|
|
class FakePipeline(LoraBaseMixin):
|
|
_lora_loadable_modules = ["unet", "text_encoder"]
|
|
|
|
def __init__(self, unet, text_encoder):
|
|
self._merged_adapters = set()
|
|
self.unet, self.text_encoder = unet, text_encoder
|
|
|
|
unet = TinyModel()
|
|
text_encoder = TinyModel()
|
|
config = LoraConfig(r=4, lora_alpha=4, target_modules=["linear"], init_lora_weights=False)
|
|
unet.add_adapter(config, adapter_name="adapter")
|
|
text_encoder.add_adapter(config, adapter_name="adapter")
|
|
|
|
pipe = FakePipeline(unet, text_encoder)
|
|
pipe.fuse_lora(components=["unet", "text_encoder"], adapter_names=["adapter"])
|
|
assert pipe.num_fused_loras == 1
|
|
|
|
pipe.unfuse_lora(components=["text_encoder"])
|
|
assert "adapter" in pipe.fused_loras, "adapter should remain tracked while unet is still fused"
|
|
assert pipe.num_fused_loras == 1
|
|
|
|
unet_still_merged = any(isinstance(m, BaseTunerLayer) and len(m.merged_adapters) > 0 for m in unet.modules())
|
|
assert unet_still_merged, "unet should still be physically merged at the PEFT level"
|
|
|
|
pipe.unfuse_lora(components=["unet"])
|
|
assert pipe.num_fused_loras == 0
|
|
|
|
|
|
@require_peft_backend
|
|
@pytest.mark.skipif(
|
|
not is_transformers_version(">=", "5.6.0"),
|
|
reason="CLIPTextModel uses a flattened module namespace starting with transformers 5.6",
|
|
)
|
|
def test_load_lora_into_flattened_clip_text_encoder():
|
|
config = CLIPTextConfig(
|
|
vocab_size=100,
|
|
hidden_size=32,
|
|
intermediate_size=64,
|
|
num_hidden_layers=1,
|
|
num_attention_heads=4,
|
|
max_position_embeddings=77,
|
|
bos_token_id=0,
|
|
eos_token_id=1,
|
|
pad_token_id=2,
|
|
)
|
|
text_encoder = CLIPTextModel(config)
|
|
|
|
lora_a = torch.randn(4, 32)
|
|
lora_b = torch.randn(32, 4)
|
|
|
|
state_dict = {
|
|
"text_encoder.text_model.encoder.layers.0.self_attn.q_proj.lora_A.weight": lora_a,
|
|
"text_encoder.text_model.encoder.layers.0.self_attn.q_proj.lora_B.weight": lora_b,
|
|
}
|
|
|
|
StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder(
|
|
state_dict=state_dict,
|
|
network_alphas=None,
|
|
text_encoder=text_encoder,
|
|
prefix="text_encoder",
|
|
adapter_name="test",
|
|
low_cpu_mem_usage=False,
|
|
)
|
|
|
|
q_proj = text_encoder.encoder.layers[0].self_attn.q_proj
|
|
|
|
torch.testing.assert_close(q_proj.lora_A["test"].weight, lora_a)
|
|
torch.testing.assert_close(q_proj.lora_B["test"].weight, lora_b)
|