Repository: huggingface/diffusers
File: src/diffusers/loaders/lora_conversion_utils.py (Python)

Code before:
        raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}")

    for key in list(converted_state_dict.keys()):
        converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key)

    return converted_state_dict


def _convert_non_diffusers_wan_lora_to_diffusers(state_dict):
    converted_state_dict = {}
    original_state_dict = {k[len("diffusion_model.") :]: v for k, v in state_dict.items()}

    num_blocks = len({k.split("blocks.")[1].split(".")[0] for k in original_state_dict if "blocks." in k})
    is_i2v_lora = any("k_img" in k for k in original_state_dict) and any("v_img" in k for k in original_state_dict)
    lora_down_key = "lora_A" if any("lora_A" in k for k in original_state_dict) else "lora_down"
    lora_up_key = "lora_B" if any("lora_B" in k for k in original_state_dict) else "lora_up"

    diff_keys = [k for k in original_state_dict if k.endswith((".diff_b", ".diff"))]
    if diff_keys:
        for diff_k in diff_keys:
            param = original_state_dict[diff_k]
            threshold = 1.6e-2
            absdiff = param.abs().max() - param.abs().min()
            all_zero = torch.all(param == 0).item()
            all_absdiff_lower_than_threshold = absdiff < threshold
            if all_zero or all_absdiff_lower_than_threshold:
                logger.debug(
                    f"Removed {diff_k} key from the state dict as it's all zeros, or values lower than hardcoded threshold."
                )
                original_state_dict.pop(diff_k)
            else:
                print(diff_k, absdiff)

    # For the `diff_b` keys, we treat them as lora_bias.
    # https://huggingface.co/docs/peft/main/en/package_reference/lora#peft.LoraConfig.lora_bias

    for i in range(num_blocks):
        # Self-attention
        for o, c in zip(["q", "k", "v", "o"], ["to_q", "to_k", "to_v", "to_out.0"]):
            converted_state_dict[f"blocks.{i}.attn1.{c}.lora_A.weight"] = original_state_dict.pop(
                f"blocks.{i}.self_attn.{o}.{lora_down_key}.weight"
            )
            converted_state_dict[f"blocks.{i}.attn1.{c}.lora_B.weight"] = original_state_dict.pop(
                f"blocks.{i}.self_attn.{o}.{lora_up_key}.weight"
            )
            if f"blocks.{i}.self_attn.{o}.diff_b" in original_state_dict:
                converted_state_dict[f"blocks.{i}.attn1.{c}.lora_B.bias"] = original_state_dict.pop(
                    f"blocks.{i}.self_attn.{o}.diff_b"
                )

        # Cross-attention

Review comment (suggestion):
The magnitudes of the `.diff`-ending weights are very low (most are below 1e-4, some are  upto 1e-3, and 2 of them are about 1.6e-2). The low magnitudes mostly correspond to norm layers. I think ignoring them is the best option at the moment

Code after:
        raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}")

    for key in list(converted_state_dict.keys()):
        converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key)

    return converted_state_dict


def _convert_non_diffusers_wan_lora_to_diffusers(state_dict):
    converted_state_dict = {}
    original_state_dict = {k[len("diffusion_model.") :]: v for k, v in state_dict.items()}

    num_blocks = len({k.split("blocks.")[1].split(".")[0] for k in original_state_dict if "blocks." in k})
    is_i2v_lora = any("k_img" in k for k in original_state_dict) and any("v_img" in k for k in original_state_dict)
    lora_down_key = "lora_A" if any("lora_A" in k for k in original_state_dict) else "lora_down"
    lora_up_key = "lora_B" if any("lora_B" in k for k in original_state_dict) else "lora_up"

    diff_keys = [k for k in original_state_dict if k.endswith((".diff_b", ".diff"))]
    if diff_keys:
        for diff_k in diff_keys:
            param = original_state_dict[diff_k]
            # The magnitudes of the .diff-ending weights are very low (most are below 1e-4, some are upto 1e-3,
            # and 2 of them are about 1.6e-2 [the case with AccVideo lora]). The low magnitudes mostly correspond
            # to norm layers. Ignoring them is the best option at the moment until a better solution is found. It
            # is okay to ignore because they do not affect the model output in a significant manner.
            threshold = 1.6e-2
            absdiff = param.abs().max() - param.abs().min()
            all_zero = torch.all(param == 0).item()
            all_absdiff_lower_than_threshold = absdiff < threshold
            if all_zero or all_absdiff_lower_than_threshold:
                logger.debug(
                    f"Removed {diff_k} key from the state dict as it's all zeros, or values lower than hardcoded threshold."
                )
                original_state_dict.pop(diff_k)

    # For the `diff_b` keys, we treat them as lora_bias.
    # https://huggingface.co/docs/peft/main/en/package_reference/lora#peft.LoraConfig.lora_bias

    for i in range(num_blocks):
        # Self-attention
        for o, c in zip(["q", "k", "v", "o"], ["to_q", "to_k", "to_v", "to_out.0"]):
            converted_state_dict[f"blocks.{i}.attn1.{c}.lora_A.weight"] = original_state_dict.pop(
                f"blocks.{i}.self_attn.{o}.{lora_down_key}.weight"
            )
            converted_state_dict[f"blocks.{i}.attn1.{c}.lora_B.weight"] = original_state_dict.pop(
                f"blocks.{i}.self_attn.{o}.{lora_up_key}.weight"
            )
            if f"blocks.{i}.self_attn.{o}.diff_b" in original_state_dict:
                converted_state_dict[f"blocks.{i}.attn1.{c}.lora_B.bias"] = original_state_dict.pop(
                    f"blocks.{i}.self_attn.{o}.diff_b"
                )