diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index e1bc530eb29d..7055ceaa28c5 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2411,6 +2411,39 @@ def _convert_non_diffusers_anima_lora_to_diffusers(state_dict): converted_state_dict[new_key] = value + for down_key_suffix, up_key_suffix in ( + (".lora_down.weight", ".lora_up.weight"), + (".lora_A.weight", ".lora_B.weight"), + ): + all_keys = list(converted_state_dict.keys()) + for key in all_keys: + if not key.endswith(down_key_suffix) or key not in converted_state_dict: + continue + + up_key = key[: -len(down_key_suffix)] + up_key_suffix + alpha_key = key[: -len(down_key_suffix)] + ".alpha" + if up_key not in converted_state_dict: + continue + + down_weight = converted_state_dict.pop(key) + up_weight = converted_state_dict.pop(up_key) + alpha = converted_state_dict.pop(alpha_key, None) + + if alpha is None: + scale_down, scale_up = 1.0, 1.0 + else: + scale = alpha.item() / down_weight.shape[0] + scale_down, scale_up = scale, 1.0 + while scale_down * 2 < scale_up: + scale_down *= 2 + scale_up /= 2 + + converted_down_key = key[: -len(down_key_suffix)] + ".lora_A.weight" + converted_up_key = key[: -len(down_key_suffix)] + ".lora_B.weight" + converted_state_dict[converted_down_key] = down_weight * scale_down + converted_state_dict[converted_up_key] = up_weight * scale_up + + converted_state_dict = {k: v for k, v in converted_state_dict.items() if not k.endswith(".alpha")} return converted_state_dict diff --git a/tests/loaders/test_lora_conversion_utils.py b/tests/loaders/test_lora_conversion_utils.py new file mode 100644 index 000000000000..045841dda259 --- /dev/null +++ b/tests/loaders/test_lora_conversion_utils.py @@ -0,0 +1,55 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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 unittest + +import torch + +from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_anima_lora_to_diffusers + + +class AnimaLoraConversionUtilsTests(unittest.TestCase): + def test_comfy_lora_alpha_is_folded_into_converted_weights(self): + state_dict = { + "diffusion_model.blocks.0.cross_attn.k_proj.lora_down.weight": torch.ones(4, 8), + "diffusion_model.blocks.0.cross_attn.k_proj.lora_up.weight": torch.ones(8, 4), + "diffusion_model.blocks.0.cross_attn.k_proj.alpha": torch.tensor(2.0), + } + + converted_state_dict = _convert_non_diffusers_anima_lora_to_diffusers(state_dict) + + down_key = "transformer.transformer_blocks.0.attn2.to_k.lora_A.weight" + up_key = "transformer.transformer_blocks.0.attn2.to_k.lora_B.weight" + self.assertEqual(set(converted_state_dict), {down_key, up_key}) + self.assertTrue(torch.allclose(converted_state_dict[down_key], torch.full((4, 8), 0.5))) + self.assertTrue(torch.allclose(converted_state_dict[up_key], torch.ones(8, 4))) + + def test_comfy_diffusers_style_lora_alpha_is_removed(self): + state_dict = { + "diffusion_model.blocks.0.cross_attn.k_proj.lora_A.weight": torch.ones(4, 8), + "diffusion_model.blocks.0.cross_attn.k_proj.lora_B.weight": torch.ones(8, 4), + "diffusion_model.blocks.0.cross_attn.k_proj.alpha": torch.tensor(2.0), + } + + converted_state_dict = _convert_non_diffusers_anima_lora_to_diffusers(state_dict) + + down_key = "transformer.transformer_blocks.0.attn2.to_k.lora_A.weight" + up_key = "transformer.transformer_blocks.0.attn2.to_k.lora_B.weight" + self.assertEqual(set(converted_state_dict), {down_key, up_key}) + self.assertTrue(torch.allclose(converted_state_dict[down_key], torch.full((4, 8), 0.5))) + self.assertTrue(torch.allclose(converted_state_dict[up_key], torch.ones(8, 4))) + + +if __name__ == "__main__": + unittest.main()