nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-04-11 05:13:31 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

fix ckpt_utils.py (#580)

Browse source
This commit is contained in:
ZXMMD 2024-07-12 14:52:21 +08:00 committed by GitHub
parent a7b6aacc99
commit a29424c237
1 changed files with 6 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
opensora/utils/ckpt_utils.py
View file

@ -193,6 +193,12 @@ def load_checkpoint(model, ckpt_path, save_as_pt=False, model_name="model", stri
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=strict)
get_logger().info("Missing keys: %s", missing_keys)
get_logger().info("Unexpected keys: %s", unexpected_keys)
elif ckpt_path.endswith(".safetensors"):
from safetensors.torch import load_file
state_dict = load_file(ckpt_path)
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
print(f"Missing keys: {missing_keys}")
print(f"Unexpected keys: {unexpected_keys}")
elif os.path.isdir(ckpt_path):
load_from_sharded_state_dict(model, ckpt_path, model_name, strict=strict)
get_logger().info("Model checkpoint loaded from %s", ckpt_path)