| """Full native VILA-U visual package and separate language-space projector.""" |
| import json |
| from pathlib import Path |
| import torch |
| from torch import nn |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| from native.configuration_rqvaesigliptransformer import RQVAESIGLIPTransformerConfig |
| from native.modeling_rqvaesigliptransformer import RQVAESIGLIPTransformer |
|
|
| class Projector(nn.Module): |
| def __init__(self): |
| super().__init__();self.layers=nn.Sequential(nn.Linear(1024,4096),nn.GELU(),nn.Linear(4096,4096)) |
| def forward(self,x):return self.layers(x) |
|
|
| def load_model(path=None,dtype=torch.float32,device='cpu'): |
| root=Path(path or Path(__file__).parent) |
| cfg=json.loads((root/'source_config.json').read_text()) |
| cfg['rqvaesiglip']['siglip_config']=json.loads((root/'siglip_config.json').read_text()) |
| config=RQVAESIGLIPTransformerConfig(**cfg) |
| with init_empty_weights():model=RQVAESIGLIPTransformer(config) |
| model.load_state_dict(load_file(root/'model.safetensors'),strict=True,assign=True) |
| return model.to(device=device,dtype=dtype).eval() |
|
|
| def load_projector(path=None,dtype=torch.float32,device='cpu'): |
| root=Path(path or Path(__file__).parent) |
| with init_empty_weights():model=Projector() |
| model.load_state_dict(load_file(root/'projector/model.safetensors'),strict=True,assign=True) |
| return model.to(device=device,dtype=dtype).eval() |
|
|
