File size: 1,546 Bytes
accba69 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | import json
from pathlib import Path
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from torch import nn
class Projector(nn.Module):
def __init__(self, config: dict):
super().__init__()
self.spatial_merge_size = 2
self.linear_1 = nn.Linear(config["input_size"], config["hidden_size"], bias=True)
self.linear_2 = nn.Linear(config["hidden_size"], config["output_size"], bias=True)
self.merge_linear_1 = nn.Linear(config["merged_hidden_size"], config["hidden_size"], bias=True)
self.merge_linear_2 = nn.Linear(config["hidden_size"], config["output_size"], bias=True)
def forward(self, image_features: torch.Tensor) -> torch.Tensor:
hidden = torch.nn.functional.gelu(self.linear_1(image_features))
hidden = self.linear_2(hidden)
hidden = hidden.reshape(hidden.shape[0] // (self.spatial_merge_size ** 2), -1)
return self.merge_linear_2(torch.nn.functional.gelu(self.merge_linear_1(hidden)))
def load_projector(model: str | Path) -> Projector:
path = Path(model)
if path.is_dir():
config_path = path / "projector_config.json"
weights_path = path / "projector.safetensors"
else:
config_path = Path(hf_hub_download(str(model), "projector_config.json"))
weights_path = Path(hf_hub_download(str(model), "projector.safetensors"))
projector = Projector(json.loads(config_path.read_text()))
projector.load_state_dict(load_file(weights_path))
return projector
|