vision_language_model.py

import json
import os
import tempfile
from dataclasses import asdict
from typing import Optional

from .model_utils import top_k_top_p_filtering
from .vision_transformer import ViT
from .language_model import LanguageModel
from .modality_projector import ModalityProjector
from .vlm_config import VLMConfig

from .processors import get_tokenizer

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_model, save_model

def _drop_if_all_attend(attention_mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
    """Drop all-ones masks (avoid decode length mismatch with a fixed prefill mask)."""
    if attention_mask is None:
        return None
    if bool(attention_mask.all()):
        return None
    return attention_mask

class VisionLanguageModel(nn.Module):
    def __init__(self, cfg: VLMConfig, load_backbone=True):
        super().__init__()
        self.cfg = cfg
        if load_backbone:
            print("Loading from backbone weights")
            self.vision_encoder = ViT.from_pretrained(cfg)
            self.decoder = LanguageModel.from_pretrained(cfg)
        else:
            self.vision_encoder = ViT(cfg)
            self.decoder = LanguageModel(cfg)
        self.MP = ModalityProjector(cfg)
        self.load_backbone = load_backbone
        self.tokenizer = get_tokenizer(cfg.lm_tokenizer, cfg.vlm_extra_tokens, cfg.lm_chat_template)

        self._compiled_decoder = None
        self._compiled_vision_encoder = None

    def _vision_encoder_for_prefill(self):
        if not getattr(self.cfg, "compile_inference", False):
            return self.vision_encoder
        if self._compiled_vision_encoder is None:
            self._compiled_vision_encoder = torch.compile(
                self.vision_encoder,
                mode=getattr(self.cfg, "compile_inference_mode", "default"),
                dynamic=True,
                fullgraph=False,
            )
        return self._compiled_vision_encoder

    def _decoder_for_decode(self):

        if not getattr(self.cfg, "compile_inference", False):
            return self.decoder
        if self._compiled_decoder is None:
            mode = getattr(self.cfg, "compile_inference_mode", "default")
            self._compiled_decoder = torch.compile(
                self.decoder,
                mode=mode,
                dynamic=(False if mode == "reduce-overhead" else True),
                fullgraph=False,
            )
        return self._compiled_decoder

    def _replace_img_tokens_with_embd(self, input_ids, token_embd, image_embd):
        image_token_id = getattr(self.tokenizer, "image_token_id", None)
        if image_token_id is None and hasattr(self.tokenizer, "image_token"):
            image_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.image_token)
        if image_token_id is None:
            image_token_id = self.tokenizer.convert_tokens_to_ids("<|image|>")
        mask = (input_ids == image_token_id).unsqueeze(-1)
        flat = image_embd.view(-1, image_embd.size(-1)).to(token_embd.dtype)
        token_embd.masked_scatter_(mask, flat)
        return token_embd

    def _process_images(self, images, device):
        if isinstance(images, list):
            if images and isinstance(images[0], list):
                images = [img for sublist in images for img in sublist]

            if not images:
                return None
            else:
                return torch.cat(images, dim=0).to(device)
        return images

    def forward(self, input_ids, images, attention_mask=None, targets=None):
        images_tensor = self._process_images(images, input_ids.device)
        token_embd = self.decoder.token_embedding(input_ids)

        if images_tensor is not None:
            image_embd = self.vision_encoder(images_tensor)
            image_embd = self.MP(image_embd)
            token_embd = self._replace_img_tokens_with_embd(input_ids, token_embd, image_embd)

        attention_mask = _drop_if_all_attend(attention_mask)

        logits, _ = self.decoder(token_embd, attention_mask=attention_mask)

        loss = None
        if targets is not None:
            logits = self.decoder.head(logits)
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-100)

        return logits, loss

    @torch.inference_mode()
    def generate(self, input_ids, images, attention_mask=None, max_new_tokens=5, top_k=50, top_p=0.9, temperature=0.5, greedy=False):
        images_tensor = self._process_images(images, input_ids.device)
        token_embd = self.decoder.token_embedding(input_ids)

        if images_tensor is not None:
            vision_encoder_fn = self._vision_encoder_for_prefill()
            image_embd = vision_encoder_fn(images_tensor)
            image_embd = self.MP(image_embd)
            token_embd = self._replace_img_tokens_with_embd(input_ids, token_embd, image_embd)

        current_total_seq_len = token_embd.size(1)
        batch_size = input_ids.size(0)
        device = input_ids.device

        cache_max_length = current_total_seq_len + max(int(max_new_tokens), 1)
        head_dim = self.cfg.lm_hidden_dim // self.cfg.lm_n_heads
        use_cuda_graphs_pre = (
            getattr(self.cfg, "compile_inference", False)
            and getattr(self.cfg, "compile_inference_mode", "default") == "reduce-overhead"
            and device.type == "cuda"
        )
        if use_cuda_graphs_pre:
            quantum = int(getattr(self.cfg, "cuda_graphs_cache_quantum", 128))
            if quantum > 1:
                cache_max_length = ((cache_max_length + quantum - 1) // quantum) * quantum
        cache_shape = (batch_size, self.cfg.lm_n_kv_heads, cache_max_length, head_dim)
        cache_dtype = next(self.decoder.parameters()).dtype
        use_cuda_graphs = use_cuda_graphs_pre
        kv_cache_list = []
        for _ in range(self.cfg.lm_n_blocks):
            key_cache = torch.zeros(cache_shape, dtype=cache_dtype, device=device)
            value_cache = torch.zeros(cache_shape, dtype=cache_dtype, device=device)
            if use_cuda_graphs:
                torch._dynamo.mark_static_address(key_cache)
                torch._dynamo.mark_static_address(value_cache)
            kv_cache_list.append({
                "cache_max_length": cache_max_length,
                "key_cache": key_cache,
                "value_cache": value_cache,
            })

        attention_mask = _drop_if_all_attend(attention_mask)

        prefill_output, kv_cache_list = self.decoder(
            token_embd,
            attention_mask=attention_mask,
            kv_cache=kv_cache_list,
            start_pos=0,
        )

        last_token_output_from_prefill = prefill_output[:, -1, :]

        if not self.decoder.lm_use_tokens:
            current_logits = self.decoder.head(last_token_output_from_prefill)
        else:
            current_logits = last_token_output_from_prefill

        newly_generated_ids_list = []

        decode_decoder = self._decoder_for_decode()

        if use_cuda_graphs:
            start_pos_buf = torch.zeros((), dtype=torch.long, device=device)
            torch._dynamo.mark_static_address(start_pos_buf)
        else:
            start_pos_buf = None

        eos_id = self.tokenizer.eos_token_id
        eos_check_interval = int(getattr(self.cfg, "eos_check_interval", 16))
        if eos_id is None or eos_check_interval <= 0:
            eos_check_interval = 0

        for step_idx in range(max_new_tokens):
            if greedy:
                next_token_id = torch.argmax(current_logits, dim=-1, keepdim=True)
            else:
                filtered_logits = top_k_top_p_filtering(current_logits, top_k=top_k, top_p=top_p)
                probs = torch.softmax(filtered_logits / temperature, dim=-1)
                next_token_id = torch.multinomial(probs, num_samples=1)

            newly_generated_ids_list.append(next_token_id)

            if (
                eos_check_interval > 0
                and (step_idx + 1) % eos_check_interval == 0
                and bool((next_token_id == eos_id).all().item())
            ):
                break

            next_token_embed = self.decoder.token_embedding(next_token_id)

            current_token_start_pos = current_total_seq_len
            current_total_seq_len += 1

            if start_pos_buf is not None:
                start_pos_buf.fill_(current_token_start_pos)
                start_pos_to_pass = start_pos_buf
            else:
                start_pos_to_pass = current_token_start_pos

            decode_step_output, kv_cache_list = decode_decoder(
                next_token_embed,
                attention_mask=attention_mask,
                kv_cache=kv_cache_list,
                start_pos=start_pos_to_pass,
            )

            last_token_output = decode_step_output[:, -1, :]

            if not self.decoder.lm_use_tokens:
                current_logits = self.decoder.head(last_token_output)
            else:
                current_logits = last_token_output
        if not newly_generated_ids_list:
            return torch.empty((batch_size,0), dtype=torch.long, device=input_ids.device)

        generated_ids = torch.cat(newly_generated_ids_list, dim=1)

        if self.tokenizer.eos_token_id is not None and generated_ids.numel() > 0:
            seq_len = generated_ids.size(1)
            device = generated_ids.device

            eos_mask = (generated_ids == self.tokenizer.eos_token_id)

            col_indices_for_min = torch.arange(seq_len, device=device)
            masked_col_indices = torch.where(eos_mask, col_indices_for_min.unsqueeze(0).expand_as(generated_ids), seq_len + 1)

            first_eos_indices_values = torch.min(masked_col_indices, dim=1).values
            actual_first_eos_indices = torch.clamp(first_eos_indices_values, max=seq_len)

            col_indices_for_comparison = torch.arange(seq_len, device=device).unsqueeze(0).expand_as(generated_ids)
            replace_mask = col_indices_for_comparison > actual_first_eos_indices.unsqueeze(1)
            generated_ids[replace_mask] = self.tokenizer.eos_token_id
        return generated_ids

    @classmethod
    def from_pretrained(
        cls, repo_id_or_path: str, *, revision: Optional[str] = None
    ) -> "VisionLanguageModel":

        if os.path.exists(repo_id_or_path):
            config_path = os.path.join(repo_id_or_path, "config.json")
            weights_path = os.path.join(repo_id_or_path, "model.safetensors")

            if not os.path.exists(config_path):
                raise ValueError(
                    f"Config file not found at {config_path}. Please provide a valid path."
                )
            if not os.path.exists(weights_path):
                raise ValueError(
                    f"Weights file not found at {weights_path}. Please provide a valid path."
                )
        else:
            from huggingface_hub import hf_hub_download

            config_path = hf_hub_download(
                repo_id=repo_id_or_path, filename="config.json", revision=revision
            )
            weights_path = hf_hub_download(
                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
            )

        with open(config_path, "r") as f:
            raw_cfg = json.load(f)
        from dataclasses import fields as dc_fields
        valid_keys = {fld.name for fld in dc_fields(VLMConfig)}
        cfg = VLMConfig(**{k: v for k, v in raw_cfg.items() if k in valid_keys})

        model = cls(cfg, load_backbone=False)
        model._saved_config_keys = set(raw_cfg.keys())

        load_model(model, weights_path)

        return model

    def save_pretrained(self, save_directory: str) -> None:

        os.makedirs(save_directory, exist_ok=True)

        with open(os.path.join(save_directory, "config.json"), "w") as f:
            f.write(json.dumps(asdict(self.cfg), indent=4))

        save_model(self, os.path.join(save_directory, "model.safetensors"))

    def push_to_hub(self, repo_id: str, private: bool = False) -> None:

        from huggingface_hub import create_repo, upload_folder

        repo_url = create_repo(repo_id=repo_id, private=private, exist_ok=True)
        repo_id = repo_url.repo_id
        print("Created repo: ", repo_url)

        with tempfile.TemporaryDirectory() as save_path:
            self.save_pretrained(save_path)

            with open(os.path.join(save_path, "README.md"), "w") as f:
                f.write(MODEL_CARD_TEMPLATE.format(repo_id=repo_id))

            return upload_folder(
                repo_id=repo_id,
                repo_type="model",
                folder_path=save_path,
                commit_message="Upload VisionPsyNano using push_to_hub",
            )

MODEL_CARD_TEMPLATE = """
---
library_name: transformers
pipeline_tag: image-text-to-text
tags:
  - vision-language
  - multimodal
---

```python
from transformers import AutoModelForImageTextToText, AutoProcessor
from PIL import Image

repo = "{repo_id}"
model = AutoModelForImageTextToText.from_pretrained(
    repo, trust_remote_code=True, dtype="auto"
).cuda().eval()
model.apply_deploy_profile(model.device)
processor = AutoProcessor.from_pretrained(repo, trust_remote_code=True)

inputs = processor(images=Image.open("image.jpg"), text="Describe the image.", return_tensors="pt")
inputs = {{k: v.cuda() if hasattr(v, "cuda") else v for k, v in inputs.items() if v is not None}}
inputs.pop("pixel_values", None)
out = model.generate(**inputs, max_new_tokens=64, greedy=True)
print(processor.batch_decode(out, skip_special_tokens=True)[0])
```
"""