| from pathlib import Path | |
| import torch | |
| from accelerate import init_empty_weights | |
| from safetensors.torch import load_file | |
| from native.vq_model import TokenFlow, ModelArgs | |
| def load_model(path=None, dtype=torch.float32, device='cpu'): | |
| root = Path(path or Path(__file__).parent) | |
| args = ModelArgs(codebook_size=32768, teacher='clipb_224', enhanced_decoder=True, infer_interpolate=True) | |
| with init_empty_weights(): | |
| model = TokenFlow(args) | |
| model.load_state_dict(load_file(root / 'model.safetensors'), strict=True, assign=True) | |
| return model.to(device=device, dtype=dtype).eval() | |