Commit 0376da18 authored by AUTOMATIC's avatar AUTOMATIC
Browse files

make it possible to save nai model using safetensors

parent bb11bee2

modules/sd_models.py

+2 −2
Original line number Diff line number Diff line
@@ -144,8 +144,8 @@ def transform_checkpoint_dict_key(k): 




def get_state_dict_from_checkpoint(pl_sd):

    if "state_dict" in pl_sd:

        pl_sd = pl_sd["state_dict"]

    pl_sd = pl_sd.pop("state_dict", pl_sd)

    pl_sd.pop("state_dict", None)



    sd = {}

    for k, v in pl_sd.items():

CRA Git | Maintained and supported by SUSTech CRA and CCSE