File size: 3,387 Bytes
3e77c56 237a3b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """CPU inference helpers for the demo (loads a trained checkpoint; handles 1- and 3-channel heads).
Kept dependency-light so it also works inside a Hugging Face Space. The model architecture is
imported from the installed ``stress_operator`` package; the HF export bundles a self-contained
copy (see scripts/05_export_to_hf.sh).
"""
from __future__ import annotations
import os
import sys
import numpy as np
import torch
# Make `stress_operator` importable in BOTH layouts: the local repo (package under ../src) and a
# bundled Hugging Face Space (package is a sibling of this file).
_HERE = os.path.dirname(os.path.abspath(__file__))
for _p in (_HERE, os.path.join(_HERE, "..", "src")):
if _p not in sys.path:
sys.path.insert(0, _p)
from stress_operator.models.transolver import build_model # noqa: E402
def load_checkpoint(ckpt_path: str, device: str = "cpu"):
"""Load a trained model. Supports two formats:
- ``*.safetensors`` (HF deployment): weights from safetensors + ``config.json`` (sibling)
holding the model config, normalizer mean/std, and scale_S.
- ``*.pt`` (local training checkpoint): a dict with state_dict / normalizer / config.
"""
if ckpt_path.endswith(".safetensors"):
import json
from safetensors.torch import load_file
cfg_path = os.path.join(os.path.dirname(ckpt_path) or ".", "config.json")
with open(cfg_path) as f:
cfg = json.load(f)
model_cfg = cfg["model"]
state_dict = load_file(ckpt_path, device=device)
mean = torch.tensor(float(cfg["normalizer"]["mean"]), device=device).reshape(1, 1, 1)
std = torch.tensor(float(cfg["normalizer"]["std"]), device=device).reshape(1, 1, 1)
scale_S = cfg.get("scale_S", None)
metrics = cfg.get("metrics", {})
else:
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
model_cfg = ckpt["config"]["model"]
state_dict = ckpt["state_dict"]
mean = ckpt["normalizer"]["mean"].to(device)
std = ckpt["normalizer"]["std"].to(device)
scale_S = ckpt.get("scale_S", None)
metrics = ckpt.get("metrics", {})
model = build_model(model_cfg).to(device)
model.load_state_dict(state_dict)
model.eval()
info = {
"out_dim": model_cfg.get("out_dim", 1),
"scale_S": scale_S,
"attention": model_cfg.get("attention", "physics"),
"metrics": metrics,
}
return model, (mean, std), info
@torch.no_grad()
def predict_stress(model, coords: np.ndarray, norm, info, device: str = "cpu") -> np.ndarray:
"""coords: (N, 2) -> per-node von Mises stress (N,) in physical units."""
mean, std = norm
x = torch.as_tensor(coords, dtype=torch.float32, device=device).unsqueeze(0) # (1,N,2)
out = model(x, None)[0] # (N, out_dim)
if info["out_dim"] == 3:
from stress_operator.losses.equilibrium import von_mises
S = info["scale_S"] or 1.0
stress = von_mises(out) * S # 3-channel tensor -> von Mises, undo target scale
else:
# de-normalize scalar prediction; reshape so the broadcast of a (1,1[,1]) normalizer
# against out[:,0] (N,) cannot leak a leading axis (audit bug 1). Always returns (N,).
stress = (out[:, 0] * std.reshape(()) + mean.reshape(()))
return stress.reshape(-1).detach().cpu().numpy()
|