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