export.py

"""Export trained VLM to a single HuggingFace-ready repo with .safetensors."""

import argparse
import json
import os
import shutil

import torch
from safetensors.torch import save_file
from transformers import AutoTokenizer, AutoImageProcessor

from config import VLMConfig
from model import LLM2VLM, VisionProjector, PseudoDeepStackExtractor


def export_model(checkpoint_dir: str, output_dir: str):
    """Export checkpoint to HF-ready format with merged LoRA weights."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Load model from checkpoint
    print(f"Loading model from {checkpoint_dir}...")
    model = LLM2VLM.load_from_checkpoint(checkpoint_dir, device=str(device))
    model.vision_extractor.to(device)
    model.projector.to(device)

    os.makedirs(output_dir, exist_ok=True)

    # 1. Merge LoRA into base LLM
    print("Merging LoRA weights into base model...")
    model.llm = model.llm.merge_and_unload()

    # 2. Collect all weights into a single state dict
    print("Collecting model weights...")
    state_dict = {}

    # Vision encoder weights
    for name, param in model.vision_extractor.vision_model.named_parameters():
        state_dict[f"vision_encoder.{name}"] = param.data.contiguous().clone()

    # Projector weights
    for name, param in model.projector.named_parameters():
        state_dict[f"projector.{name}"] = param.data.contiguous().clone()

    # LLM weights (merged)
    for name, param in model.llm.named_parameters():
        state_dict[f"language_model.{name}"] = param.data.contiguous().clone()

    # LLM buffers (e.g., rotary embeddings)
    for name, buf in model.llm.named_buffers():
        state_dict[f"language_model.{name}"] = buf.data.contiguous().clone()

    # 3. Save as safetensors (sharded if large)
    print("Saving safetensors...")
    total_size = sum(p.numel() * p.element_size() for p in state_dict.values())
    print(f"  Total model size: {total_size / 1024**3:.2f} GB")

    # Shard at 5GB boundaries
    shard_size = 5 * 1024 ** 3  # 5 GB per shard
    shards = {}
    current_shard = {}
    current_size = 0
    shard_idx = 0
    weight_map = {}

    for key, tensor in state_dict.items():
        tensor_size = tensor.numel() * tensor.element_size()
        if current_size + tensor_size > shard_size and current_shard:
            shard_name = f"model-{shard_idx:05d}-of-TOTAL.safetensors"
            shards[shard_name] = current_shard
            current_shard = {}
            current_size = 0
            shard_idx += 1

        current_shard[key] = tensor
        current_size += tensor_size

    if current_shard:
        shard_name = f"model-{shard_idx:05d}-of-TOTAL.safetensors"
        shards[shard_name] = current_shard

    total_shards = len(shards)
    final_shards = {}
    for old_name, shard_data in shards.items():
        new_name = old_name.replace("TOTAL", f"{total_shards:05d}")
        final_shards[new_name] = shard_data
        for key in shard_data:
            weight_map[key] = new_name

    for shard_name, shard_data in final_shards.items():
        shard_path = os.path.join(output_dir, shard_name)
        save_file(shard_data, shard_path)
        print(f"  Saved {shard_name}")

    # Save index
    index = {
        "metadata": {"total_size": total_size},
        "weight_map": weight_map,
    }
    with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f:
        json.dump(index, f, indent=2)

    # 4. Save tokenizer and image processor
    print("Saving tokenizer and image processor...")
    model.tokenizer.save_pretrained(output_dir)
    model.image_processor.save_pretrained(output_dir)

    # 5. Save component configs for self-contained loading
    print("Saving component configs...")
    from transformers import AutoConfig
    import transformers

    # Save LLM config (so we don't need original model dir at load time)
    llm_config = AutoConfig.from_pretrained(model.config.llm_model, trust_remote_code=True)
    if hasattr(llm_config, "quantization_config"):
        delattr(llm_config, "quantization_config")
    llm_config.save_pretrained(os.path.join(output_dir, "llm_config"))

    # Save vision encoder config
    vision_config_src = AutoConfig.from_pretrained(model.config.vision_model)
    vision_config_src.save_pretrained(os.path.join(output_dir, "vision_config"))

    # 6. Save VLM config
    vlm_config = {
        "architectures": ["LLM2VLMForConditionalGeneration"],
        "model_type": "llm2vlm",
        "transformers_version": transformers.__version__,
        "vision_config": {
            "model_name": model.config.vision_model,
            "hidden_size": model.config.vision_hidden_size,
            "num_hidden_layers": model.config.vision_num_layers,
            "num_image_tokens": model.config.num_image_tokens,
            "pseudo_deep_stack_layers": list(model.config.pseudo_deep_stack_layers),
        },
        "projector_config": {
            "input_dim": model.config.projector_input_dim,
            "hidden_dim": model.config.projector_hidden_dim,
            "output_dim": model.config.projector_output_dim,
        },
        "text_config": {
            "model_name": model.config.llm_model,
            "hidden_size": model.config.llm_hidden_size,
            "num_hidden_layers": model.config.llm_num_layers,
        },
        "image_token": model.config.image_token,
        "image_token_id": model.config.image_token_id,
        "max_seq_length": model.config.max_seq_length,
        "use_multi_crop": model.config.use_multi_crop,
        "max_crop_tiles": model.config.max_crop_tiles,
        "tile_size": model.config.tile_size,
        "auto_map": {
            "AutoConfig": "configuration_llm2vlm.LLM2VLMConfig",
            "AutoModelForCausalLM": "modeling_llm2vlm.LLM2VLMForConditionalGeneration",
        },
    }
    with open(os.path.join(output_dir, "config.json"), "w") as f:
        json.dump(vlm_config, f, indent=2)

    # 6. Write custom modeling code for trust_remote_code
    _write_configuration_file(output_dir)
    _write_modeling_file(output_dir)

    print(f"\nExport complete! HF-ready model saved to: {output_dir}")
    print(f"Load with: AutoModel.from_pretrained('{output_dir}', trust_remote_code=True)")


def _write_configuration_file(output_dir: str):
    """Write configuration_llm2vlm.py for HF auto_map."""
    code = '''\
"""LLM2VLM configuration for HuggingFace."""

from transformers import PretrainedConfig


class LLM2VLMConfig(PretrainedConfig):
    model_type = "llm2vlm"

    def __init__(
        self,
        vision_config=None,
        projector_config=None,
        text_config=None,
        image_token="<image>",
        image_token_id=None,
        max_seq_length=4096,
        use_multi_crop=True,
        max_crop_tiles=6,
        tile_size=384,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vision_config = vision_config or {}
        self.projector_config = projector_config or {}
        self.text_config = text_config or {}
        self.image_token = image_token
        self.image_token_id = image_token_id
        self.max_seq_length = max_seq_length
        self.use_multi_crop = use_multi_crop
        self.max_crop_tiles = max_crop_tiles
        self.tile_size = tile_size
'''
    with open(os.path.join(output_dir, "configuration_llm2vlm.py"), "w") as f:
        f.write(code)


def _write_modeling_file(output_dir: str):
    """Write modeling_llm2vlm.py for HF auto_map."""
    code = '''\
"""LLM2VLM model for HuggingFace inference."""

import torch
import torch.nn as nn
from transformers import (
    PreTrainedModel,
    SiglipVisionModel,
    SiglipVisionConfig,
    AutoModelForCausalLM,
    AutoConfig,
    AutoImageProcessor,
)
from safetensors.torch import load_file
try:
    from .configuration_llm2vlm import LLM2VLMConfig
except ImportError:
    from configuration_llm2vlm import LLM2VLMConfig
import os
import json
import time


class VisionProjector(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        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):
        return self.linear2(self.act(self.linear1(x)))


class LLM2VLMForConditionalGeneration(PreTrainedModel):
    config_class = LLM2VLMConfig

    def __init__(self, config):
        super().__init__(config)
        vc = config.vision_config
        pc = config.projector_config
        self.num_image_tokens = vc.get("num_image_tokens", 729)
        self.pseudo_deep_stack_layers = vc.get("pseudo_deep_stack_layers", [8, 17, 26])
        self.image_token_id = config.image_token_id

        self.projector = VisionProjector(
            input_dim=pc.get("input_dim", 3456),
            hidden_dim=pc.get("hidden_dim", 3456),
            output_dim=pc.get("output_dim", 2880),
        )
        # Vision model and LLM are loaded from state dict
        self.vision_encoder = None
        self.language_model = None
        self._hooks = []
        self._features = {}

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

    def _make_hook(self, layer_idx):
        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 get_vision_features(self, pixel_values):
        """Process single or multi-tile images.

        Args:
            pixel_values: [B, 3, H, W] or [B, T, 3, H, W] for multi-tile
        """
        if pixel_values.dim() == 5:
            B, T = pixel_values.shape[:2]
            flat = pixel_values.view(B * T, *pixel_values.shape[2:])
        else:
            B, T = pixel_values.shape[0], 1
            flat = pixel_values

        self._features.clear()
        with torch.no_grad():
            self.vision_encoder(flat)
        features = [self._features[idx] for idx in self.pseudo_deep_stack_layers]
        concatenated = torch.cat(features, dim=-1)  # [B*T, 729, 3456]
        projected = self.projector(concatenated.to(self.projector.linear1.weight.dtype))

        if T > 1:
            projected = projected.view(B, T * projected.shape[1], projected.shape[2])
        return projected

    def forward(self, input_ids, attention_mask=None, pixel_values=None,
                labels=None, image_token_mask=None, **kwargs):
        embed_layer = self.language_model.model.embed_tokens
        inputs_embeds = embed_layer(input_ids)

        if pixel_values is not None and image_token_mask is not None:
            vision_features = self.get_vision_features(pixel_values)
            vision_features = vision_features.to(inputs_embeds.dtype)
            for i in range(input_ids.shape[0]):
                img_positions = image_token_mask[i].nonzero(as_tuple=True)[0]
                if len(img_positions) > 0:
                    n = min(len(img_positions), vision_features.shape[1])
                    inputs_embeds[i, img_positions[:n]] = vision_features[i, :n]

        return self.language_model(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            labels=labels,
            return_dict=True,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        """Load the full VLM from an exported HF-ready directory.

        Optimized loading:
        - Creates model shells on meta device (zero memory, instant)
        - Loads safetensors shard-by-shard, assigning weights directly to GPU
        - Peak CPU memory: ~5 GB (one shard at a time)
        """
        t0 = time.time()
        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
        with open(config_path) as f:
            config_dict = json.load(f)

        config = LLM2VLMConfig(**config_dict)
        model = cls(config)

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # --- Create model shells on meta device (no memory allocation) ---
        # Vision encoder
        vision_cfg_dir = os.path.join(pretrained_model_name_or_path, "vision_config")
        vc = config.vision_config
        if os.path.isdir(vision_cfg_dir):
            vision_cfg = SiglipVisionConfig.from_pretrained(vision_cfg_dir)
        else:
            vision_cfg = SiglipVisionConfig.from_pretrained(vc["model_name"])
        with torch.device("meta"):
            model.vision_encoder = SiglipVisionModel(vision_cfg)

        # LLM
        llm_cfg_dir = os.path.join(pretrained_model_name_or_path, "llm_config")
        tc = config.text_config
        if os.path.isdir(llm_cfg_dir):
            llm_config = AutoConfig.from_pretrained(llm_cfg_dir, trust_remote_code=True)
        else:
            llm_config = AutoConfig.from_pretrained(tc["model_name"], trust_remote_code=True)
        if hasattr(llm_config, "quantization_config"):
            delattr(llm_config, "quantization_config")
        with torch.device("meta"):
            model.language_model = AutoModelForCausalLM.from_config(
                llm_config, torch_dtype=torch.bfloat16, trust_remote_code=True,
            )
        print(f"  Model shells created on meta device ({time.time()-t0:.1f}s)")

        # --- Load safetensors shard by shard directly to target device ---
        index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json")
        with open(index_path) as f:
            index = json.load(f)

        # Build reverse map: shard_file -> list of (exported_key, component, param_name)
        weight_map = index["weight_map"]
        shard_files = sorted(set(weight_map.values()))

        def _assign_weight(root_module, dotted_name, tensor):
            """Assign a tensor to a module's param/buffer by dotted name."""
            parts = dotted_name.split(".")
            mod = root_module
            for p in parts[:-1]:
                mod = getattr(mod, p)
            last = parts[-1]
            if last in mod._parameters:
                mod._parameters[last] = nn.Parameter(tensor, requires_grad=False)
                return True
            elif last in mod._buffers:
                mod._buffers[last] = tensor
                return True
            return False

        # Resize LLM embeddings if needed (check from weight map)
        embed_exported_key = "language_model.model.embed_tokens.weight"
        if embed_exported_key in weight_map:
            shard_for_embed = weight_map[embed_exported_key]
            shard_path = os.path.join(pretrained_model_name_or_path, shard_for_embed)
            embed_tensor = load_file(shard_path, device="cpu")[embed_exported_key]
            target_vocab = embed_tensor.shape[0]
            model.language_model.resize_token_embeddings(target_vocab)
            del embed_tensor

        t1 = time.time()
        loaded_count = 0
        for shard_file in shard_files:
            shard_path = os.path.join(pretrained_model_name_or_path, shard_file)
            shard_data = load_file(shard_path, device="cpu")

            for key, tensor in shard_data.items():
                tensor_gpu = tensor.to(device=device, dtype=torch.bfloat16, non_blocking=True)

                if key.startswith("vision_encoder."):
                    param_name = key[len("vision_encoder."):]
                    if _assign_weight(model.vision_encoder, param_name, tensor_gpu):
                        loaded_count += 1
                elif key.startswith("projector."):
                    param_name = key[len("projector."):]
                    if _assign_weight(model.projector, param_name, tensor_gpu):
                        loaded_count += 1
                elif key.startswith("language_model."):
                    param_name = key[len("language_model."):]
                    if _assign_weight(model.language_model, param_name, tensor_gpu):
                        loaded_count += 1

            del shard_data
            print(f"  Loaded {shard_file} ({loaded_count} tensors so far)")

        print(f"  All weights loaded to {device} ({time.time()-t1:.1f}s)")

        # Materialize any remaining meta buffers (e.g. position_ids)
        for name, buf in list(model.named_buffers()):
            if buf.device.type == "meta":
                parts = name.split(".")
                mod = model
                for p in parts[:-1]:
                    mod = getattr(mod, p)
                if parts[-1] == "position_ids":
                    mod._buffers[parts[-1]] = torch.arange(
                        buf.shape[-1], device=device
                    ).unsqueeze(0)
                else:
                    mod._buffers[parts[-1]] = torch.zeros_like(buf, device=device)

        model.vision_encoder.eval()
        for p in model.vision_encoder.parameters():
            p.requires_grad = False
        model._register_hooks()

        print(f"  Total load time: {time.time()-t0:.1f}s")
        return model
'''
    with open(os.path.join(output_dir, "modeling_llm2vlm.py"), "w") as f:
        f.write(code)
    # Also write __init__.py so relative imports work
    with open(os.path.join(output_dir, "__init__.py"), "w") as f:
        f.write("")


def main():
    parser = argparse.ArgumentParser(description="Export LLM2VLM to HF-ready format")
    parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint directory")
    parser.add_argument("--output_dir", type=str, required=True, help="Output directory for HF-ready model")
    args = parser.parse_args()

    export_model(args.checkpoint, args.output_dir)


if __name__ == "__main__":
    main()