Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
"""Checkpoint containers and logical-file discovery, independent of adapter layout."""
|
|
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
|
|
TORCH_EXTENSIONS = (".pt", ".pth", ".bin", ".ckpt")
|
|
EXTENSIONS = (".safetensors", *TORCH_EXTENSIONS)
|
|
|
|
|
|
def find_checkpoint(directory, stem, required=True):
|
|
for extension in EXTENSIONS:
|
|
path = directory / (stem + extension)
|
|
if path.is_file():
|
|
return path
|
|
if required:
|
|
raise FileNotFoundError(f"No {stem} checkpoint in {directory}; expected {', '.join(EXTENSIONS)}")
|
|
return None
|
|
|
|
|
|
def load_checkpoint(path):
|
|
path = Path(path)
|
|
extension = path.suffix.lower()
|
|
if extension == ".safetensors":
|
|
state = load_file(path, device="cpu")
|
|
elif extension in TORCH_EXTENSIONS:
|
|
state = torch.load(path, map_location="cpu", weights_only=True)
|
|
if isinstance(state, dict) and "state_dict" in state:
|
|
state = state["state_dict"]
|
|
else:
|
|
raise ValueError(f"Unsupported checkpoint container: {path}")
|
|
if not isinstance(state, dict) or not all(isinstance(value, torch.Tensor) for value in state.values()):
|
|
raise ValueError(f"Expected a tensor state dictionary: {path}")
|
|
return state
|