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