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")
        
        # Vision
        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)
        
        # Text
        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)
        
        # Fusion
        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)

# Son işlem: 2026-05-14 09:34:06