checkpoint has error key

[xuchen-dev]

Hello, when I loaded the model trained with train_tdd.cpy, it reported an error with the incorrect key. When I manually removed the ’base_madel.model‘ prefix, the predicted image was noisy

[fastisrealslow]

@WangCunzheng 看起来是key的问题，存正帮忙转换下key

[WangCunzheng]

我这边因为学校毕设原因，暂时是离职的，然后手边也没有机器，所以可能需要小红书那边改一下。🙏

…

---- 回复的原邮件 ---- | 发件人 | ***@***.***> | | 日期 | 2025年02月10日 17:33 | | 收件人 | ***@***.***> | | 抄送至 | ***@***.***>***@***.***> | | 主题 | Re: [RedAIGC/Target-Driven-Distillation] checkpoint has error key (Issue #18) | @WangCunzheng 看起来是key的问题，存正帮忙转换下key — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you were mentioned.Message ID: ***@***.***>

[fastisrealslow]

import os
import cv2
import numpy as np
import torch
import argparse

from PIL import Image
from peft import LoraConfig, get_peft_model, get_peft_model_state_dict
from safetensors.torch import load_file, save_file
from accelerate import Accelerator


def get_module_kohya_state_dict(module, prefix: str, dtype: torch.dtype, adapter_name: str = "default"):
    kohya_ss_state_dict = {}
    for peft_key, weight in module.items():
        kohya_key = peft_key.replace("unet.base_model.model", prefix)
        kohya_key = kohya_key.replace("lora_A", "lora_down")
        kohya_key = kohya_key.replace("lora_B", "lora_up")
        kohya_key = kohya_key.replace(".", "_", kohya_key.count(".") - 2)
        kohya_ss_state_dict[kohya_key] = weight.to(dtype)
        # Set alpha parameter
        if "lora_down" in kohya_key:
            alpha_key = f'{kohya_key.split(".")[0]}.alpha'
            kohya_ss_state_dict[alpha_key] = torch.tensor(8).to(dtype)

    return kohya_ss_state_dict


# lora_weight_1 = load_file('TDD_uc0.2_etas0.3_ddim250_adv/checkpoint-20000/unet_lora/pytorch_lora_weights.safetensors')
# lora_state_dict_1 = get_module_kohya_state_dict(lora_weight_1, "lora_unet", torch.float16)
# save_file(lora_state_dict_1, os.path.join('tar_files/tdd_out', "pytorch_lora_weights.safetensors"))

@WangCunzheng @xuchen-dev