Skip to content

Commit 58da14b

Browse files
committed
handle missing keys and unexpected keys better.
1 parent a15b9b3 commit 58da14b

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

src/diffusers/loaders/lora_pipeline.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1893,6 +1893,7 @@ def load_lora_into_transformer(
18931893
state_dict = {
18941894
k.replace(f"{cls.transformer_name}.", ""): v for k, v in state_dict.items() if k in transformer_keys
18951895
}
1896+
print(f'{any("transformer_blocks." in k for k in state_dict)=}')
18961897

18971898
if len(state_dict.keys()) > 0:
18981899
# check with first key if is not in peft format
@@ -1944,17 +1945,19 @@ def load_lora_into_transformer(
19441945
# Check only for unexpected keys.
19451946
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
19461947
if unexpected_keys:
1948+
lora_unexpected_keys = [k for k in unexpected_keys if "lora" in k and adapter_name in k]
19471949
logger.warning(
19481950
f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
1949-
f" {unexpected_keys}. "
1951+
f" {lora_unexpected_keys}. "
19501952
)
19511953

19521954
# Filter missing keys specific to the current adapter.
19531955
missing_keys = getattr(incompatible_keys, "missing_keys", None)
19541956
if missing_keys:
1957+
lora_missing_keys = [k for k in missing_keys if "lora" in k and adapter_name in k]
19551958
logger.warning(
19561959
f"Loading adapter weights from state_dict led to missing keys in the model: "
1957-
f" {missing_keys}. "
1960+
f" {lora_missing_keys}. "
19581961
)
19591962

19601963
# Offload back.

0 commit comments

Comments
 (0)