model.py

"""VLM model: SigLIP2 vision encoder + PseudoDeepStack + MLP projector + GPT-OSS-20B with LoRA."""

import os
import json
import glob

import torch
import torch.nn as nn
from typing import Optional, Tuple
from transformers import (
    SiglipVisionModel,
    AutoModelForCausalLM,
    AutoConfig,
    AutoTokenizer,
    AutoImageProcessor,
    BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model, PeftModel
from safetensors import safe_open

from config import VLMConfig


# ---------------------------------------------------------------------------
# MXFP4 dequantization for ROCm compatibility
# ---------------------------------------------------------------------------

def dequantize_mxfp4(blocks: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
    """Dequantize MXFP4 (E2M1 data + E8M0 scales) to bfloat16.

    Args:
        blocks: [*, num_blocks, 16] uint8 — 32 FP4 values packed into 16 bytes
        scales: [*, num_blocks] uint8 — E8M0 shared exponent per block of 32

    Returns:
        Dense tensor in bfloat16: [*, num_blocks * 32]
    """
    orig_shape = blocks.shape[:-2]
    num_blocks = blocks.shape[-2]
    blocks = blocks.reshape(-1, num_blocks, 16)
    scales = scales.reshape(-1, num_blocks)

    # Unpack uint8 -> two 4-bit values (low nibble first, then high nibble)
    low = (blocks & 0x0F).to(torch.int8)
    high = (blocks >> 4).to(torch.int8)
    unpacked = torch.stack([low, high], dim=-1).reshape(-1, num_blocks, 32)

    # Decode FP4 E2M1: bit3=sign, bit2-1=exponent, bit0=mantissa
    sign = ((unpacked >> 3) & 1).float()
    exp = ((unpacked >> 1) & 0x3).int()
    mantissa = (unpacked & 1).float()

    sign_mult = 1.0 - 2.0 * sign
    normal_val = (2.0 ** (exp.float() - 1.0)) * (1.0 + mantissa * 0.5)
    subnormal_val = mantissa * 0.5
    fp4_val = torch.where(exp > 0, normal_val, subnormal_val) * sign_mult

    # Apply E8M0 scale: 2^(scale - 127)
    scale_mult = (2.0 ** (scales.float() - 127.0)).unsqueeze(-1)
    result = fp4_val * scale_mult

    in_dim = num_blocks * 32
    return result.reshape(*orig_shape, in_dim).to(torch.bfloat16)


def _load_gptoss_dequantized(model_path: str, device_map="auto") -> AutoModelForCausalLM:
    """Load GPT-OSS-20B with MXFP4 expert weights dequantized to bf16.

    The MXFP4 triton kernels don't support autograd backward pass, so we
    dequantize to bf16 for training (gradient flow through expert MLPs).
    """
    print("  Loading GPT-OSS-20B with MXFP4->bf16 dequantization...")

    # Load config without quantization
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    if hasattr(config, "quantization_config"):
        delattr(config, "quantization_config")

    # Instantiate model shell (expert weights will be random/empty)
    model = AutoModelForCausalLM.from_config(
        config, torch_dtype=torch.bfloat16, trust_remote_code=True,
    )

    # Find safetensors files
    st_files = sorted(glob.glob(os.path.join(model_path, "model-*.safetensors")))
    if not st_files:
        st_files = sorted(glob.glob(os.path.join(model_path, "model.safetensors")))

    # Pre-load all blocks and scales from all safetensors files
    all_tensors = {}
    for st_file in st_files:
        with safe_open(st_file, framework="pt") as f:
            for key in f.keys():
                all_tensors[key] = f.get_tensor(key)

    # Build state dict from safetensors, dequantizing MXFP4 experts
    state_dict = model.state_dict()
    loaded_keys = set()
    processed_bases = set()

    for key, tensor in all_tensors.items():
        if key.endswith("_blocks"):
            base_name = key[:-7]  # strip "_blocks"
            if base_name in processed_bases:
                continue
            processed_bases.add(base_name)

            scales_key = base_name + "_scales"
            if scales_key not in all_tensors:
                continue

            # Dequantize: result shape [experts, out_dim, in_dim]
            dense = dequantize_mxfp4(tensor, all_tensors[scales_key])
            # MXFP4 stores as [E, out, in] but model expects [E, in, out]
            dense = dense.transpose(-1, -2).contiguous()

            if base_name in state_dict:
                state_dict[base_name] = dense
                loaded_keys.add(base_name)
            continue

        if key.endswith("_scales"):
            continue

        # Non-quantized weight — load directly
        if key in state_dict:
            state_dict[key] = tensor.to(torch.bfloat16) if tensor.is_floating_point() else tensor
            loaded_keys.add(key)

    missing = set(state_dict.keys()) - loaded_keys
    if missing:
        print(f"  Warning: {len(missing)} keys not loaded from checkpoint: {list(missing)[:5]}...")

    model.load_state_dict(state_dict, strict=False)

    if device_map == "auto" or device_map == "cuda":
        model = model.to("cuda")
    elif device_map != "cpu":
        model = model.to(device_map)

    print(f"  Model loaded: {sum(p.numel() for p in model.parameters()) / 1e9:.1f}B params")
    return model


# ---------------------------------------------------------------------------
# Multi-crop image processing
# ---------------------------------------------------------------------------

# Candidate grids: (rows, cols) for splitting high-res images
_CANDIDATE_GRIDS = [
    (1, 1), (1, 2), (2, 1), (2, 2),
    (1, 3), (3, 1), (2, 3), (3, 2),
]


def select_best_grid(width: int, height: int, max_tiles: int = 6, tile_size: int = 384):
    """Pick the grid layout that best matches image aspect ratio within tile budget.

    Only considers grids where the image is large enough to benefit from tiling.
    An image smaller than 1.5x the tile size in both dimensions stays at 1x1.
    """
    # Don't tile small images — not enough pixels to fill multiple tiles
    if width < tile_size * 1.5 and height < tile_size * 1.5:
        return (1, 1)

    aspect = width / max(height, 1)
    best_grid = (1, 1)
    best_score = (float("inf"), 0)  # (aspect_error, -num_tiles)
    for rows, cols in _CANDIDATE_GRIDS:
        num_tiles = rows * cols
        if num_tiles > max_tiles:
            continue
        # Skip grids that upscale the image too much (target < 2x upscale)
        if cols * tile_size > width * 2 or rows * tile_size > height * 2:
            continue
        grid_aspect = cols / rows
        score = (abs(grid_aspect - aspect), -num_tiles)  # prefer more tiles on tie
        if score < best_score:
            best_score = score
            best_grid = (rows, cols)
    return best_grid


def multi_crop_image(image, grid, tile_size: int = 384):
    """Split image into grid tiles + one overview tile.

    Args:
        image: PIL Image
        grid: (rows, cols)
        tile_size: pixel size of each square tile

    Returns:
        List of PIL Images: [overview, tile_0_0, tile_0_1, ..., tile_r_c]
    """
    from PIL import Image as _PILImage
    rows, cols = grid
    # Overview: resize full image to tile_size x tile_size
    overview = image.resize((tile_size, tile_size), _PILImage.LANCZOS)
    if rows == 1 and cols == 1:
        return [overview]
    # Resize image to fill the grid
    grid_w, grid_h = cols * tile_size, rows * tile_size
    resized = image.resize((grid_w, grid_h), _PILImage.LANCZOS)
    # Split into tiles
    tiles = [overview]
    for r in range(rows):
        for c in range(cols):
            box = (c * tile_size, r * tile_size, (c + 1) * tile_size, (r + 1) * tile_size)
            tiles.append(resized.crop(box))
    return tiles


# ---------------------------------------------------------------------------
# Vision components
# ---------------------------------------------------------------------------

class PseudoDeepStackExtractor(nn.Module):
    """Extract multi-scale features from SigLIP2 vision encoder layers."""

    def __init__(self, vision_model: SiglipVisionModel, extract_layers: Tuple[int, ...] = (8, 17, 26)):
        super().__init__()
        self.vision_model = vision_model
        self.extract_layers = extract_layers
        self._features = {}
        self._hooks = []
        self._register_hooks()

    def _register_hooks(self):
        encoder_layers = self.vision_model.vision_model.encoder.layers
        for layer_idx in self.extract_layers:
            hook = encoder_layers[layer_idx].register_forward_hook(
                self._make_hook(layer_idx)
            )
            self._hooks.append(hook)

    def _make_hook(self, layer_idx: int):
        def hook_fn(module, input, output):
            if isinstance(output, tuple):
                self._features[layer_idx] = output[0]
            else:
                self._features[layer_idx] = output
        return hook_fn

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """[B, 3, 384, 384] -> [B, 729, 3456]"""
        self._features.clear()
        with torch.no_grad():
            self.vision_model(pixel_values)
        features = [self._features[idx] for idx in self.extract_layers]
        return torch.cat(features, dim=-1)

    def remove_hooks(self):
        for hook in self._hooks:
            hook.remove()
        self._hooks.clear()


class VisionProjector(nn.Module):
    """2-layer MLP projector: maps vision features to LLM embedding space."""

    def __init__(self, input_dim: int = 3456, hidden_dim: int = 3456, output_dim: int = 2880):
        super().__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim)
        self.act = nn.GELU()
        self.linear2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear2(self.act(self.linear1(x)))


# ---------------------------------------------------------------------------
# Embedding accessor helper
# ---------------------------------------------------------------------------

def _get_embed_layer(llm):
    """Navigate through PeftModel/model wrappers to find embed_tokens."""
    m = llm
    for attr in ("model", "model", "model"):
        if hasattr(m, "embed_tokens"):
            return m.embed_tokens
        if hasattr(m, attr):
            m = getattr(m, attr)
        else:
            break
    if hasattr(m, "embed_tokens"):
        return m.embed_tokens
    raise AttributeError("Cannot find embed_tokens in LLM model hierarchy")


# ---------------------------------------------------------------------------
# Main VLM model
# ---------------------------------------------------------------------------

class LLM2VLM(nn.Module):
    """Full VLM model combining vision encoder, projector, and LLM."""

    def __init__(self, config: VLMConfig, device: torch.device = None):
        super().__init__()
        self.config = config
        self._target_device = device  # specific GPU for DDP, None for auto

        # Vision components
        self.image_processor = AutoImageProcessor.from_pretrained(config.vision_model)
        vision_model = SiglipVisionModel.from_pretrained(
            config.vision_model, torch_dtype=torch.bfloat16
        )
        vision_model.eval()
        for p in vision_model.parameters():
            p.requires_grad = False
        self.vision_extractor = PseudoDeepStackExtractor(
            vision_model, config.pseudo_deep_stack_layers
        )

        # Projector (trainable)
        self.projector = VisionProjector(
            input_dim=config.projector_input_dim,
            hidden_dim=config.projector_hidden_dim,
            output_dim=config.projector_output_dim,
        )

        # LLM
        self.tokenizer = AutoTokenizer.from_pretrained(
            config.llm_model, trust_remote_code=True
        )
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        self._load_llm(config)
        self._setup_image_token(config)

    def _load_llm(self, config: VLMConfig):
        if config.use_unsloth:
            self._load_llm_unsloth(config)
        else:
            self._load_llm_standard(config)

    def _load_llm_standard(self, config: VLMConfig):
        """Standard loading: dequantize MXFP4 for ROCm, then apply LoRA."""
        # Resolve device_map: use specific device for DDP, "auto" for single-GPU
        device_map = str(self._target_device) if self._target_device is not None else "auto"

        if config.use_qlora:
            # QLoRA: load dequantized then re-quantize with bitsandbytes NF4
            base_model = _load_gptoss_dequantized(config.llm_model, device_map="cpu")
            from transformers import AutoModelForCausalLM as _Auto
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.bfloat16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
            )
            import tempfile
            with tempfile.TemporaryDirectory() as tmp:
                base_model.save_pretrained(tmp, safe_serialization=True)
                cfg_path = os.path.join(tmp, "config.json")
                with open(cfg_path) as f:
                    cfg = json.load(f)
                cfg.pop("quantization_config", None)
                with open(cfg_path, "w") as f:
                    json.dump(cfg, f, indent=2)
                self.llm = _Auto.from_pretrained(
                    tmp,
                    quantization_config=bnb_config,
                    torch_dtype=torch.bfloat16,
                    trust_remote_code=True,
                    device_map=device_map,
                )
            del base_model
        else:
            self.llm = _load_gptoss_dequantized(config.llm_model, device_map=device_map)

        # Freeze LLM base
        for p in self.llm.parameters():
            p.requires_grad = False

        # Apply LoRA
        lora_config = LoraConfig(
            r=config.lora_r,
            lora_alpha=config.lora_alpha,
            lora_dropout=config.lora_dropout,
            target_modules=config.lora_target_modules,
            bias="none",
            task_type="CAUSAL_LM",
        )
        self.llm = get_peft_model(self.llm, lora_config)
        self.llm.print_trainable_parameters()

        # Gradient checkpointing: recompute activations during backward to save memory.
        # Reduces activation memory from O(layers) to O(sqrt(layers)) at ~30% compute cost.
        if config.gradient_checkpointing:
            self.llm.gradient_checkpointing_enable(
                gradient_checkpointing_kwargs={"use_reentrant": False}
            )
            print("  Gradient checkpointing enabled")

    def _load_llm_unsloth(self, config: VLMConfig):
        """Unsloth loading path: faster LoRA finetuning."""
        from unsloth import FastLanguageModel

        self.llm, self.tokenizer = FastLanguageModel.from_pretrained(
            model_name=config.llm_model,
            max_seq_length=config.max_seq_length,
            dtype=torch.bfloat16,
            load_in_4bit=config.use_qlora,
            trust_remote_code=True,
        )

        self.llm = FastLanguageModel.get_peft_model(
            self.llm,
            r=config.lora_r,
            lora_alpha=config.lora_alpha,
            lora_dropout=config.lora_dropout,
            target_modules=config.lora_target_modules,
            bias="none",
            max_seq_length=config.max_seq_length,
        )
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

    def _setup_image_token(self, config: VLMConfig):
        if config.image_token not in self.tokenizer.get_vocab():
            self.tokenizer.add_special_tokens(
                {"additional_special_tokens": [config.image_token]}
            )
            self.llm.resize_token_embeddings(len(self.tokenizer))
        config.image_token_id = self.tokenizer.convert_tokens_to_ids(config.image_token)

    def get_vision_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Process single or multi-tile images through vision encoder + projector.

        Args:
            pixel_values: [B, 3, 384, 384] for single-tile
                          OR [B, T, 3, 384, 384] for multi-tile (T tiles per image)

        Returns:
            [B, num_tokens, 2880] where num_tokens = T * 729
        """
        if pixel_values.dim() == 5:
            # Multi-tile: [B, T, 3, 384, 384]
            B, T = pixel_values.shape[:2]
            flat = pixel_values.view(B * T, *pixel_values.shape[2:])
            features = self.vision_extractor(flat)  # [B*T, 729, 3456]
            features = features.view(B, T * features.shape[1], features.shape[2])
        else:
            # Single tile: [B, 3, 384, 384]
            features = self.vision_extractor(pixel_values)  # [B, 729, 3456]
        return self.projector(features.to(self.projector.linear1.weight.dtype))

    def _inject_vision(self, inputs_embeds, pixel_values, image_token_mask):
        """Replace <image> token embeddings with projected vision features."""
        if pixel_values is None or image_token_mask is None:
            return inputs_embeds
        # Clone to avoid in-place ops on embedding weight views
        # (required for gradient checkpointing with use_reentrant=False)
        inputs_embeds = inputs_embeds.clone()
        vision_features = self.get_vision_features(pixel_values).to(inputs_embeds.dtype)
        for i in range(inputs_embeds.shape[0]):
            pos = image_token_mask[i].nonzero(as_tuple=True)[0]
            if len(pos) > 0:
                n = min(len(pos), vision_features.shape[1])
                inputs_embeds[i, pos[:n]] = vision_features[i, :n]
        return inputs_embeds

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.LongTensor,
        pixel_values: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        image_token_mask: Optional[torch.BoolTensor] = None,
    ) -> dict:
        embed_layer = _get_embed_layer(self.llm)
        inputs_embeds = embed_layer(input_ids)
        inputs_embeds = self._inject_vision(inputs_embeds, pixel_values, image_token_mask)

        outputs = self.llm(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            labels=labels,
            return_dict=True,
        )
        return {"loss": outputs.loss, "logits": outputs.logits}

    @torch.no_grad()
    def generate(
        self,
        pixel_values: torch.Tensor,
        input_ids: torch.LongTensor,
        attention_mask: torch.LongTensor,
        image_token_mask: torch.BoolTensor,
        max_new_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
        do_sample: bool = True,
    ) -> torch.LongTensor:
        embed_layer = _get_embed_layer(self.llm)
        inputs_embeds = embed_layer(input_ids)
        inputs_embeds = self._inject_vision(inputs_embeds, pixel_values, image_token_mask)

        # Stop at <|end|> (end-of-turn) or <|return|> (end-of-generation)
        end_token_id = self.tokenizer.convert_tokens_to_ids("<|end|>")
        stop_ids = [self.tokenizer.eos_token_id, end_token_id]

        return self.llm.generate(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=do_sample,
            pad_token_id=self.tokenizer.pad_token_id,
            eos_token_id=stop_ids,
        )

    def save_trainable(self, save_dir: str):
        """Save only the trainable components (projector + LoRA)."""
        os.makedirs(save_dir, exist_ok=True)

        torch.save(self.projector.state_dict(), os.path.join(save_dir, "projector.pt"))

        lora_dir = os.path.join(save_dir, "lora_adapter")
        self.llm.save_pretrained(lora_dir)

        config_dict = {
            "vision_model": self.config.vision_model,
            "llm_model": self.config.llm_model,
            "projector_input_dim": self.config.projector_input_dim,
            "projector_hidden_dim": self.config.projector_hidden_dim,
            "projector_output_dim": self.config.projector_output_dim,
            "pseudo_deep_stack_layers": list(self.config.pseudo_deep_stack_layers),
            "num_image_tokens": self.config.num_image_tokens,
            "image_token": self.config.image_token,
            "image_token_id": self.config.image_token_id,
            "lora_r": self.config.lora_r,
            "lora_alpha": self.config.lora_alpha,
            "lora_target_modules": self.config.lora_target_modules,
            "use_qlora": self.config.use_qlora,
            "use_unsloth": self.config.use_unsloth,
            "max_seq_length": self.config.max_seq_length,
            "use_multi_crop": self.config.use_multi_crop,
            "max_crop_tiles": self.config.max_crop_tiles,
            "tile_size": self.config.tile_size,
        }
        with open(os.path.join(save_dir, "vlm_config.json"), "w") as f:
            json.dump(config_dict, f, indent=2)

        self.tokenizer.save_pretrained(save_dir)
        print(f"Saved trainable weights to {save_dir}")

    @classmethod
    def load_from_checkpoint(cls, checkpoint_dir: str, device: str = "cuda") -> "LLM2VLM":
        """Load model from a saved checkpoint.

        Args:
            checkpoint_dir: path to checkpoint directory
            device: target device string (e.g. "cuda", "cuda:0", "cpu")
        """
        with open(os.path.join(checkpoint_dir, "vlm_config.json")) as f:
            config_dict = json.load(f)

        config = VLMConfig()
        for key, value in config_dict.items():
            if key == "pseudo_deep_stack_layers":
                setattr(config, key, tuple(value))
            elif hasattr(config, key):
                setattr(config, key, value)

        model = cls.__new__(cls)
        nn.Module.__init__(model)
        model.config = config
        model._target_device = torch.device(device) if device != "auto" else None

        # Vision
        model.image_processor = AutoImageProcessor.from_pretrained(config.vision_model)
        vision_model = SiglipVisionModel.from_pretrained(
            config.vision_model, torch_dtype=torch.bfloat16
        )
        vision_model.eval()
        for p in vision_model.parameters():
            p.requires_grad = False
        model.vision_extractor = PseudoDeepStackExtractor(
            vision_model, config.pseudo_deep_stack_layers
        )

        # Projector
        model.projector = VisionProjector(
            input_dim=config.projector_input_dim,
            hidden_dim=config.projector_hidden_dim,
            output_dim=config.projector_output_dim,
        )
        model.projector.load_state_dict(
            torch.load(os.path.join(checkpoint_dir, "projector.pt"),
                       map_location="cpu", weights_only=True)
        )

        # Tokenizer
        model.tokenizer = AutoTokenizer.from_pretrained(
            checkpoint_dir, trust_remote_code=True
        )
        if model.tokenizer.pad_token is None:
            model.tokenizer.pad_token = model.tokenizer.eos_token

        # LLM (dequantized for ROCm)
        device_map = device if device not in ("auto",) else "auto"
        model.llm = _load_gptoss_dequantized(config.llm_model, device_map=device_map)
        model.llm.resize_token_embeddings(len(model.tokenizer))

        # Load LoRA adapter
        lora_dir = os.path.join(checkpoint_dir, "lora_adapter")
        model.llm = PeftModel.from_pretrained(model.llm, lora_dir)

        config.image_token_id = model.tokenizer.convert_tokens_to_ids(config.image_token)
        model.config = config

        return model