wkea/blockdiffusion-api

1from __future__ import annotations
2
3import os
4from typing import Any, Dict, Optional
5
6import torch
7
8
9def save_checkpoint(path: str, payload: Dict[str, Any]) -> None:
10    """
11    保存 checkpoint（会自动创建目录）。
12    """
13    os.makedirs(os.path.dirname(path), exist_ok=True)
14    torch.save(payload, path)
15
16
17def load_checkpoint(path: str, map_location: Optional[str] = None) -> Dict[str, Any]:
18    """
19    读取 checkpoint。
20    """
21    return torch.load(path, map_location=map_location)
22
23
24