File size: 735 Bytes
7bfa997 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | """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()
|