File size: 599 Bytes
b9a3411 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | 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()
|