diffsynth/utils/state_dict_converters/stable_diffusion_vae.py

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