diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index d1692bd61ba5..6f584a5f0e0a 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2129,6 +2129,10 @@ def _convert_non_diffusers_ltxv_lora_to_diffusers(state_dict, non_diffusers_pref def _convert_non_diffusers_qwen_lora_to_diffusers(state_dict): + has_diffusion_model = any(k.startswith("diffusion_model.") for k in state_dict) + if has_diffusion_model: + state_dict = {k.removeprefix("diffusion_model."): v for k, v in state_dict.items()} + has_lora_unet = any(k.startswith("lora_unet_") for k in state_dict) if has_lora_unet: state_dict = {k.removeprefix("lora_unet_"): v for k, v in state_dict.items()} @@ -2201,29 +2205,44 @@ def convert_key(key: str) -> str: all_keys = list(state_dict.keys()) down_key = ".lora_down.weight" up_key = ".lora_up.weight" + a_key = ".lora_A.weight" + b_key = ".lora_B.weight" - 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 + 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) - for k in all_keys: - if k.endswith(down_key): - diffusers_down_key = k.replace(down_key, ".lora_A.weight") - diffusers_up_key = k.replace(down_key, up_key).replace(up_key, ".lora_B.weight") - alpha_key = k.replace(down_key, ".alpha") - - down_weight = state_dict.pop(k) - up_weight = state_dict.pop(k.replace(down_key, up_key)) - scale_down, scale_up = 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 + 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 + + for k in all_keys: + if k.endswith(down_key): + diffusers_down_key = k.replace(down_key, ".lora_A.weight") + diffusers_up_key = k.replace(down_key, up_key).replace(up_key, ".lora_B.weight") + alpha_key = k.replace(down_key, ".alpha") + + down_weight = state_dict.pop(k) + up_weight = state_dict.pop(k.replace(down_key, up_key)) + scale_down, scale_up = 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 + 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 len(state_dict) > 0: raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}") diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 572ace472f1d..7e89066f1fc4 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -6684,7 +6684,8 @@ def lora_state_dict( has_alphas_in_sd = any(k.endswith(".alpha") for k in state_dict) has_lora_unet = any(k.startswith("lora_unet_") for k in state_dict) - if has_alphas_in_sd or has_lora_unet: + has_diffusion_model = any(k.startswith("diffusion_model.") for k in state_dict) + if has_alphas_in_sd or has_lora_unet or has_diffusion_model: state_dict = _convert_non_diffusers_qwen_lora_to_diffusers(state_dict) out = (state_dict, metadata) if return_lora_metadata else state_dict