kohya-ss / kohya-ss/sd-scripts
"Missing keys in state dict" solved when training lora on some anima ft checkpoints
- Dominant language
- Python
- Stars
- 7.2k
- Forks
- 1.2k
- Avg merge
- 11m
- Merged PRs (30d)
- 2
Description
The actual reason for this issue seems that ComfyUI modified the key prefixes of the models during export, causing the training script to be unable to find the corresponding keys. Use the following script to restore the modified key prefixes.
I tested that the repaired image generation had no impact, and it passed the model validation during training lora.
```
import os
import torch
from safetensors.torch import load_file, save_file
# ==================== 请在此处修改您的文件路径 ====================
base_model_path = r"D:\anima\anima-base-v1.0.safetensors" # 官方原版底模路径
ft_model_path = r"D:\anima\ft.safetensors" # 您的微调/合并模型路径
output_model_path = r"D:\anima\ft_repaired.safetensors" # 对齐重构后的输出路径
# ==================================================================
def get_suffix(key):
# 智能剥离常见前缀,提取核心键名
prefixes = ["model.diffusion_model.", "diffusion_model.", "net."]
for p in prefixes:
if key.startswith(p):
return key[len(p):]
return key
def main():
if not os.path.exists(base_model_path):
print(f"错误: 找不到原版底模,请检查路径: {base_model_path}")
return
if not os.path.exists(ft_model_path):
print(f"错误: 找不到微调模型,请检查路径: {ft_model_path}")
return
print("1. 正在加载官方原版底模权重...")
base_state = load_file(base_model_path)
print("2. 正在加载微调/合并模型权重...")
ft_state = load_file(ft_model_path)
# 3. 建立原版底模的 [核心键名(suffix) -> 官方原版键名(net.xxx)] 映射表
print("3. 正在构建原版键名映射表...")
suffix_to_official_key = {}
for official_key in base_state.keys():
suffix = get_suffix(official_key)
suffix_to_official_key[suffix] = official_key
# 4. 建立微调模型的 [核心键名(suffix) -> 微调权重张量] 映射表
print("4. 正在提取微调模型权重...")
suffix_to_ft_tensor = {}
for ft_key, tensor in ft_state.items():
suffix = get_suffix(ft_key)
suffix_to_ft_tensor[suffix] = tensor
# 5. 全面重构并对齐键名
print("5. 正在对齐并重构键名结构...")
aligned_state = {}
use_ft_count = 0
fallback_base_count = 0
for suffix, official_key in suffix_to_official_key.items():
if suffix in suffix_to_ft_tensor:
# 采用微调模型中的权重,并将其键名强行对齐为官方的 "net.xxx" 格式
aligned_state[official_key] = suffix_to_ft_tensor[suffix]
use_ft_count += 1
else:
# 如果微调模型里确实彻底丢失了该键(例如部分组件被精简),则使用原版底模的权重兜底
aligned_state[official_key] = base_state[official_key]
fallback_base_count += 1
print(f"\n--- 重构结果统计 ---")
print(f"成功从您的微调模型中对齐并迁移了 {use_ft_count} 个权重参数(键名已被强行对齐为官方 net. 格式)")
if fallback_base_count > 0:
print(f"微调模型中缺失了 {fallback_base_count} 个参数,已自动使用原版底模权重完成补全")
print("--------------------\n")
print("6. 正在保存重构后的完整模型...")
metadata = {"format": "pt"}
save_file(aligned_state, output_model_path, metadata=metadata)
print(f"保存成功!输出路径: {output_model_path}")
print("该模型现在的结构与官方原版完全一致,您可以直接用它作为底模进行训练了。")
if __name__ == "__main__":
main()
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reproducing LoRA training validation with the affected Anima fine-tuned checkpoints and inspect the state-dict prefixes used by the exported models. Compare those keys with the official checkpoint and the provided repair script; the work is done when affected checkpoints are accepted without manual key rewriting and model validation still passes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100