|
|
| import os, json, torch |
| import torch.nn as nn |
| from PIL import Image |
| from transformers import AutoModelForCausalLM, AutoModel, AutoProcessor, AutoTokenizer |
|
|
| class MLPProjector(nn.Module): |
| def __init__(self, input_dim=1152, hidden_dim=8192, output_dim=3072): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.LayerNorm(input_dim), |
| nn.Linear(input_dim, hidden_dim), |
| nn.GELU(), |
| nn.Linear(hidden_dim, output_dim), |
| nn.LayerNorm(output_dim), |
| ) |
| def forward(self, x): |
| return self.net(x) |
|
|
| class KumruVLM(nn.Module): |
| def __init__(self, model_path, device="cuda"): |
| super().__init__() |
| self.device = device |
| with open(os.path.join(model_path, "kumru_vlm_config.json"), "r") as f: |
| self.cfg = json.load(f) |
| |
| self.tokenizer = AutoTokenizer.from_pretrained(model_path) |
| self.llm = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).to(device).eval() |
| |
| vision_id = self.cfg["vision_model_id"] |
| self.processor = AutoProcessor.from_pretrained(vision_id) |
| vision_full = AutoModel.from_pretrained(vision_id, torch_dtype=torch.bfloat16).to(device).eval() |
| self.vision = vision_full.vision_model |
| |
| self.projector = MLPProjector(input_dim=1152, hidden_dim=8192, output_dim=self.llm.config.hidden_size).to(device, torch.bfloat16) |
| self.projector.load_state_dict(torch.load(os.path.join(model_path, "projector_final.pt"), map_location="cpu")) |
| |
| self.image_token_id = self.tokenizer.convert_tokens_to_ids("<image>") |
|
|
| @torch.no_grad() |
| def generate(self, image, prompt, **kwargs): |
| if isinstance(image, str): |
| image = Image.open(image).convert("RGB") |
| |
| |
| pv = self.processor(images=image, return_tensors="pt")["pixel_values"].to(self.device, torch.bfloat16) |
| v_output = self.vision(pixel_values=pv).last_hidden_state |
| patches = v_output[:, (0 if v_output.shape[1] == 256 else 1):, :] |
| v_embeds = self.projector(patches) |
| |
| |
| full_prompt = f"<image>\nUSER: {prompt}\nASSISTANT:" |
| ids = self.tokenizer(full_prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(self.device) |
| t_embeds = self.llm.get_input_embeddings()(ids) |
| |
| |
| p = (ids[0] == self.image_token_id).nonzero()[0].item() |
| inputs_embeds = torch.cat([t_embeds[:, :p], v_embeds, t_embeds[:, p+1:]], dim=1) |
| |
| out = self.llm.generate( |
| inputs_embeds=inputs_embeds, |
| eos_token_id=self.tokenizer.eos_token_id, |
| pad_token_id=self.tokenizer.pad_token_id, |
| **kwargs |
| ) |
| |
| decoded = self.tokenizer.decode(out[0], skip_special_tokens=True) |
| return decoded.split("ASSISTANT:")[-1].strip() |
|
|
| def load_kumru_vlm(model_path, device="cuda"): |
| return KumruVLM(model_path, device) |
|
|
| |
|
|
