Files
DiffSynth-Studio/diffsynth/utils/state_dict_converters/stable_diffusion_vae.py
mi804 82e482286c sd
2026-04-23 17:35:24 +08:00

19 lines
748 B
Python

def SDVAEStateDictConverter(state_dict):
new_state_dict = {}
for key in state_dict:
if ".query." in key:
new_key = key.replace(".query.", ".to_q.")
new_state_dict[new_key] = state_dict[key]
elif ".key." in key:
new_key = key.replace(".key.", ".to_k.")
new_state_dict[new_key] = state_dict[key]
elif ".value." in key:
new_key = key.replace(".value.", ".to_v.")
new_state_dict[new_key] = state_dict[key]
elif ".proj_attn." in key:
new_key = key.replace(".proj_attn.", ".to_out.0.")
new_state_dict[new_key] = state_dict[key]
else:
new_state_dict[key] = state_dict[key]
return new_state_dict