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