| """Strict loader preserving the complete trained UniTok state, including semantic heads.""" | |
| import json | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| import torch | |
| from safetensors.torch import load_file | |
| from native.models.unitok import UniTok | |
| def load_model(path=None, dtype=torch.float32, device='cpu'): | |
| root = Path(path or Path(__file__).parent) | |
| cfg = json.loads((root / 'config.json').read_text()) | |
| cfg.update(device='cpu', grad_ckpt=False) | |
| # The native hybrid embedder runs a CPU shape probe in its constructor. | |
| model = UniTok(SimpleNamespace(**cfg)) | |
| model.load_state_dict(load_file(root / 'model.safetensors'), strict=True, assign=True) | |
| return model.to(device=device, dtype=dtype).eval() | |