diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index e3e9359168c4..88fce58cd00a 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2323,19 +2323,22 @@ def convert_key(key: str) -> str: has_non_diffusers_lora_id = any(down_key in k or up_key in k for k in all_keys) has_diffusers_lora_id = any(a_key in k or b_key in k for k in all_keys) - if has_non_diffusers_lora_id: - - def get_alpha_scales(down_weight, alpha_key): - rank = down_weight.shape[0] - alpha = state_dict.pop(alpha_key).item() - scale = alpha / rank # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - return scale_down, scale_up + def get_alpha_scales(down_weight, alpha_key): + rank = down_weight.shape[0] + alpha = state_dict.pop(alpha_key, None) + if alpha is None: + return 1.0, 1.0 + scale = ( + alpha.item() / rank + ) # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here + scale_down = scale + scale_up = 1.0 + while scale_down * 2 < scale_up: + scale_down *= 2 + scale_up /= 2 + return scale_down, scale_up + if has_non_diffusers_lora_id: for k in all_keys: if k.endswith(down_key): diffusers_down_key = k.replace(down_key, ".lora_A.weight") @@ -2348,13 +2351,32 @@ def get_alpha_scales(down_weight, alpha_key): converted_state_dict[diffusers_down_key] = down_weight * scale_down converted_state_dict[diffusers_up_key] = up_weight * scale_up - # Already in diffusers format (lora_A/lora_B), just pop + # Already in diffusers format (lora_A/lora_B), apply alpha scaling and pop. elif has_diffusers_lora_id: for k in all_keys: - if a_key in k or b_key in k: - converted_state_dict[k] = state_dict.pop(k) - elif ".alpha" in k: - state_dict.pop(k) + if not k.endswith(a_key): + continue + + diffusers_up_key = k.replace(a_key, b_key) + if diffusers_up_key not in state_dict: + raise ValueError(f"Missing lora_B weight for {k.removesuffix(a_key)}") + + alpha_key = k.replace(a_key, ".alpha") + down_weight = state_dict.pop(k) + up_weight = state_dict.pop(diffusers_up_key) + scale_down, scale_up = get_alpha_scales(down_weight, alpha_key) + converted_state_dict[k] = down_weight * scale_down + converted_state_dict[diffusers_up_key] = up_weight * scale_up + + for k in all_keys: + if k.endswith(b_key) and k in state_dict: + raise ValueError(f"Missing lora_A weight for {k.removesuffix(b_key)}") + if ".alpha" in k: + state_dict.pop(k, None) + + for k in all_keys: + if ".alpha" in k: + state_dict.pop(k, None) if len(state_dict) > 0: raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}")