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