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()