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