| 789 | |
| 790 | |
| 791 | def load_checkpoints(weight_dir): |
| 792 | checkpoint_files = [ |
| 793 | os.path.join(weight_dir, f) |
| 794 | for f in os.listdir(weight_dir) |
| 795 | if (f.startswith("pytorch_model") and f.endswith(".bin")) |
| 796 | or (f.startswith("model") and f.endswith(".safetensors")) |
| 797 | ] |
| 798 | if len(checkpoint_files) == 0: |
| 799 | print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!") |
| 800 | print("!No model weight files found! Using fake weights!") |
| 801 | print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!") |
| 802 | |
| 803 | if len(checkpoint_files) == 0: |
| 804 | return None |
| 805 | |
| 806 | state_dict = {} |
| 807 | print("Loading weights from disk") |
| 808 | is_safetensors = checkpoint_files[0].endswith(".safetensors") |
| 809 | for i in range(len(checkpoint_files)): |
| 810 | if is_safetensors: |
| 811 | state_dict = {**state_dict, **load_file(checkpoint_files[i], device="cpu")} |
| 812 | else: |
| 813 | state_dict = { |
| 814 | **state_dict, |
| 815 | **torch.load( |
| 816 | checkpoint_files[i], map_location="cpu", weights_only=True |
| 817 | ), |
| 818 | } |
| 819 | |
| 820 | return state_dict |
| 821 | |
| 822 | |
| 823 | def resolve_model_classes( |