| """Load Ace-Mini safetensors and evaluate prepared casting-frame query points.""" |
| import json |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from safetensors.torch import load_file |
| from stream_surrogate import StreamNet,features |
| from warp_surrogate import WarpNet |
|
|
| class AceMini: |
| def __init__(self,directory=None,device='cpu'): |
| root=Path(directory or Path(__file__).parent) |
| self.config=json.loads((root/'config.json').read_text()) |
| self.device=torch.device(device) |
| self.model=StreamNet(self.config['query_dim']).to(self.device).eval() |
| self.model.load_state_dict(load_file(str(root/'model.safetensors'),device=str(self.device))) |
| self.warp=WarpNet(self.config['query_dim']).to(self.device).eval() |
| self.warp.load_state_dict(load_file(str(root/'deformation.safetensors'),device=str(self.device))) |
|
|
| @torch.inference_mode() |
| def predict(self,spec,data,time_s=0.,deformation_phase=1.): |
| """All geometry must already be in a frame with gravity along -Z. |
| |
| Returns de-normalized SI fields (defect percentages as percentage points). |
| Deformation phase is normalized progress, not physical time. |
| """ |
| if not np.isfinite(time_s) or time_s<0:raise ValueError('time_s must be finite and nonnegative') |
| if not 0<=deformation_phase<=1:raise ValueError('deformation_phase must be in [0,1]') |
| cloud,q=features(spec,data) |
| if not np.isfinite(cloud).all() or not np.isfinite(q).all():raise ValueError('Inputs must be finite') |
| cloud=torch.from_numpy(cloud).to(self.device);q=torch.from_numpy(q).to(self.device) |
| code=self.model.encode(cloud);result={} |
| for task,stats in self.config['stats'].items(): |
| t=0. if task in ('maps','mechanical') else float(time_s) |
| y=self.model(cloud,q,torch.full((len(q),),t,device=self.device),task,code).cpu().numpy() |
| result[task]=y*np.array(stats['std'])+np.array(stats['mean']) |
| final_mm=torch.tensor(result['mechanical'][:,:3]*1000,dtype=torch.float32,device=self.device) |
| phase=torch.full((len(q),),float(deformation_phase),device=self.device) |
| result['deformation_m']=self.warp(q,code,phase,final_mm).cpu().numpy()/1000 |
| return result
|
|
|
