File size: 1,056 Bytes
6b38b7c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | import json
from pathlib import Path
from functools import partial
import torch
from torch import nn
from accelerate import init_empty_weights
from safetensors.torch import load_file
from native.transformer import VisionTransformer,LayerNorm
def construct(root):
cfg=json.loads((Path(root)/'source_config.json').read_text())['model_cfg']
v=cfg['vision_cfg']
return VisionTransformer(image_size=v['image_size'],patch_size=v['patch_size'],width=v['width'],layers=v['layers'],heads=v['width']//v['head_width'],mlp_ratio=4.0,in_channels=v['in_channels'],no_ln_pre=v['no_ln_pre'],pool_type=v['pool_type'],final_ln_after_pool=v['final_ln_after_pool'],norm_layer=partial(LayerNorm,**v['norm_kwargs']),act_layer=nn.GELU,output_tokens=v['output_tokens'],output_dim=cfg['embed_dim'])
def load_model(root,device='cpu',dtype=torch.float32):
with init_empty_weights(include_buffers=False):m=construct(root)
m.load_state_dict(load_file(str(Path(root)/'model.safetensors')),strict=True,assign=True)
return m.to(device=device,dtype=dtype).eval()
|