Safetensors export and infer.

infer.py

@@ -3,14 +3,13 @@ import os
 import pprint
 import time
 from typing import List, Tuple, Optional, Dict, Any
-import yaml
-
 import numpy as np
 import cv2
 import torch
 import torch.nn.functional as F
 from torch.amp import autocast
 from tqdm import tqdm
+from safetensors.torch import load_file as safe_load_file
 
 from src.wireseghr.model import WireSegHR
 from pathlib import Path
@@ -256,7 +255,7 @@ def main():
         "--ckpt",
         type=str,
         default="",
-        help="Optional checkpoint (.pt) with model state",
+        help="Optional checkpoint (.pt with {'model': state_dict} or .safetensors with pure state_dict)",
     )
     parser.add_argument(
         "--save_prob", action="store_true", help="Also save probability .npy"
@@ -356,8 +355,22 @@ def main():
     if ckpt_path:
         assert Path(ckpt_path).is_file(), f"Checkpoint not found: {ckpt_path}"
         print(f"[WireSegHR][infer] Loading checkpoint: {ckpt_path}")
-        state = torch.load(ckpt_path, map_location=device)
-        model.load_state_dict(state["model"])
+        suffix = Path(ckpt_path).suffix.lower()
+        if suffix == ".safetensors":
+            # Safetensors exports contain a pure state_dict
+            state_dict = safe_load_file(ckpt_path)
+            model.load_state_dict(state_dict)
+        else:
+            print(
+                "[WireSegHR][infer][WARN] Loading a PyTorch checkpoint. Prefer .safetensors for inference-only weights."
+            )
+            # PyTorch .pt/.pth checkpoints expected to have {'model': state_dict}
+            state = torch.load(ckpt_path, map_location=device)
+            assert "model" in state, (
+                "Expected a dict with key 'model' for PyTorch checkpoint. "
+                "Use scripts/strip_checkpoint.py or provide a .safetensors file."
+            )
+            model.load_state_dict(state["model"])
     model.eval()
 
     # Benchmark mode

requirements.txt

@@ -8,3 +8,4 @@ PyYAML>=6.0.1
 tqdm>=4.65.0
 gdown>=5.1.0
 pydrive2
+safetensors

scripts/strip_checkpoint.py

@@ -3,6 +3,7 @@
 import argparse
 from pathlib import Path
 import torch
+from safetensors.torch import save_file as safetensors_save_file
 
 
 def main():
@@ -10,7 +11,8 @@ def main():
         description="Strip training checkpoint to inference-only weights (FP32)."
     )
     parser.add_argument("--in", dest="inp", type=str, required=True, help="Path to training checkpoint .pt")
-    parser.add_argument("--out", dest="out", type=str, required=True, help="Path to save weights-only .pt")
+    parser.add_argument("--out", dest="out", type=str, required=True, help="Path to save weights-only .pt or .safetensors")
+    # Output format is inferred from --out extension
     args = parser.parse_args()
 
     in_path = Path(args.inp)
@@ -38,9 +40,14 @@ def main():
     #in the future, can cast to bfloat if necessary.
     # state_dict = {k: (v.float() if torch.is_floating_point(v) else v) for k, v in state_dict.items()}
 
-    to_save = {"model": state_dict}
-    torch.save(to_save, str(out_path))
-    print(f"[strip_checkpoint] Saved dict with only 'model' to: {out_path}")
+    suffix = out_path.suffix.lower()
+    if suffix == ".safetensors":
+        safetensors_save_file(state_dict, str(out_path))
+        print(f"[strip_checkpoint] Saved safetensors (pure state_dict) to: {out_path}")
+    else:
+        to_save = {"model": state_dict}
+        torch.save(to_save, str(out_path))
+        print(f"[strip_checkpoint] Saved dict with only 'model' to: {out_path}")
 
 
 if __name__ == "__main__":