"""Strict loader for the complete official Emu3.5 image VQ codec."""
import json
from pathlib import Path
import torch
from safetensors.torch import load_file
from native.configuration_emu3p5visionvq import Emu3p5VisionVQConfig
from native.modeling_emu3p5visionvq import Emu3p5VisionVQModel


def load_model(path=None, dtype=torch.float32, device='cpu'):
    root = Path(path or Path(__file__).parent)
    config = Emu3p5VisionVQConfig(**json.loads((root / 'config.json').read_text()))
    with torch.device('meta'):
        model = Emu3p5VisionVQModel(config)
    model.load_state_dict(load_file(root / 'model.safetensors'), strict=True, assign=True)
    return model.to(device=device, dtype=dtype).eval()