test.py

"""Test inference from an exported HF-ready VLM model."""

import argparse
import os
import sys
import json

import torch
from PIL import Image
from transformers import AutoTokenizer, AutoImageProcessor


IMAGE_TOKEN = "<image>"
TOKENS_PER_TILE = 729

# Candidate grids for multi-crop (same as model.py)
_CANDIDATE_GRIDS = [
    (1, 1), (1, 2), (2, 1), (2, 2),
    (1, 3), (3, 1), (2, 3), (3, 2),
]


def _select_best_grid(width, height, max_tiles=6, tile_size=384):
    aspect = width / max(height, 1)
    best_grid, best_score = (1, 1), float("inf")
    for rows, cols in _CANDIDATE_GRIDS:
        if rows * cols > max_tiles:
            continue
        score = abs(cols / rows - aspect)
        if score < best_score:
            best_score = score
            best_grid = (rows, cols)
    return best_grid


def _multi_crop(image, grid, tile_size=384):
    rows, cols = grid
    overview = image.resize((tile_size, tile_size), Image.LANCZOS)
    if rows == 1 and cols == 1:
        return [overview]
    resized = image.resize((cols * tile_size, rows * tile_size), Image.LANCZOS)
    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


def load_exported_model(model_dir: str):
    """Load the exported HF-ready model."""
    sys.path.insert(0, model_dir)

    from configuration_llm2vlm import LLM2VLMConfig
    from modeling_llm2vlm import LLM2VLMForConditionalGeneration

    print(f"Loading exported model from {model_dir}...")
    model = LLM2VLMForConditionalGeneration.from_pretrained(model_dir)

    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    image_processor = AutoImageProcessor.from_pretrained(model_dir)

    with open(os.path.join(model_dir, "config.json")) as f:
        config = json.load(f)

    image_token_id = config.get("image_token_id")
    use_multi_crop = config.get("use_multi_crop", False)
    max_crop_tiles = config.get("max_crop_tiles", 6)
    tile_size = config.get("tile_size", 384)

    return model, tokenizer, image_processor, image_token_id, use_multi_crop, max_crop_tiles, tile_size


def run_inference(
    model,
    tokenizer,
    image_processor,
    image_token_id,
    device,
    image_path: str,
    prompt: str = "Describe this image in detail.",
    max_new_tokens: int = 256,
    temperature: float = 0.7,
    top_p: float = 0.9,
    use_multi_crop: bool = False,
    max_crop_tiles: int = 6,
    tile_size: int = 384,
):
    """Run inference on a single image with multi-crop support."""
    image = Image.open(image_path).convert("RGB")

    # Multi-crop processing
    if use_multi_crop:
        w, h = image.size
        grid = _select_best_grid(w, h, max_tiles=max_crop_tiles, tile_size=tile_size)
        tiles = _multi_crop(image, grid, tile_size=tile_size)
    else:
        tiles = [image]

    processed = []
    for tile in tiles:
        pv = image_processor(images=tile, return_tensors="pt")["pixel_values"][0]
        processed.append(pv)
    pixel_values = torch.stack(processed).unsqueeze(0)  # [1, T, 3, 384, 384]
    pixel_values = pixel_values.to(device, dtype=torch.bfloat16)

    num_image_tokens = len(tiles) * TOKENS_PER_TILE
    image_placeholder = IMAGE_TOKEN * num_image_tokens
    user_content = f"{image_placeholder}\n{prompt}"

    chat_text = (
        f"<|start|>system<|message|>You are a helpful vision-language assistant.<|end|>"
        f"<|start|>user<|message|>{user_content}<|end|>"
        f"<|start|>assistant<|channel|>final<|message|>"
    )

    input_ids = tokenizer.encode(chat_text, add_special_tokens=False, return_tensors="pt")
    input_ids = input_ids.to(device)
    attention_mask = torch.ones_like(input_ids)
    image_token_mask = (input_ids == image_token_id)

    print(f"Image: {image_path}")
    print(f"Prompt: {prompt}")
    print(f"Tiles: {len(tiles)} ({grid if use_multi_crop else '1x1'})")
    print(f"Input tokens: {input_ids.shape[1]}")
    print(f"Image tokens: {image_token_mask.sum().item()}")
    print()

    with torch.no_grad():
        embed_layer = model.language_model.model.embed_tokens
        inputs_embeds = embed_layer(input_ids)

        vision_features = model.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]

        end_token_id = tokenizer.convert_tokens_to_ids("<|end|>")
        stop_ids = [tokenizer.eos_token_id, end_token_id]

        output_ids = model.language_model.generate(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=stop_ids,
        )

    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    print(f"Response: {response}")
    return response


def run_test(
    model_dir: str,
    image_path: str,
    prompt: str = "Describe this image in detail.",
    max_new_tokens: int = 256,
    temperature: float = 0.7,
    top_p: float = 0.9,
):
    """Load model and run inference on a single image."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model, tokenizer, image_processor, image_token_id, use_mc, max_tiles, tile_sz = \
        load_exported_model(model_dir)
    model.eval()

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    return run_inference(
        model, tokenizer, image_processor, image_token_id, device,
        image_path, prompt, max_new_tokens, temperature, top_p,
        use_multi_crop=use_mc, max_crop_tiles=max_tiles, tile_size=tile_sz,
    )


_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp", ".gif"}


def _collect_images(image_dir: str, num_images: int = 0):
    """Collect image files from a directory. If num_images > 0, limit the count."""
    files = []
    for fname in sorted(os.listdir(image_dir)):
        if os.path.splitext(fname)[1].lower() in _IMAGE_EXTENSIONS:
            files.append(os.path.join(image_dir, fname))
    if num_images > 0:
        files = files[:num_images]
    return files


def batch_test(model_dir: str, image_dir: str, num_images: int = 0, prompt: str = "Describe this image in detail."):
    """Run batch test on all images in a directory, loading the model only once."""
    image_files = _collect_images(image_dir, num_images)

    if not image_files:
        print(f"No images found in {image_dir}")
        return

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

    model, tokenizer, image_processor, image_token_id, use_mc, max_tiles, tile_sz = \
        load_exported_model(model_dir)
    model.eval()

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"Testing on {len(image_files)} images from {image_dir}")
    print("=" * 60)

    successes = 0
    failures = 0
    for i, image_path in enumerate(image_files):
        print(f"\n--- Image {i+1}/{len(image_files)}: {os.path.basename(image_path)} ---")
        try:
            run_inference(
                model, tokenizer, image_processor, image_token_id, device,
                image_path, prompt,
                use_multi_crop=use_mc, max_crop_tiles=max_tiles, tile_size=tile_sz,
            )
            successes += 1
        except Exception as e:
            print(f"Error: {e}")
            failures += 1
        print()

    print("=" * 60)
    print(f"Done: {successes} succeeded, {failures} failed, {len(image_files)} total")


def main():
    parser = argparse.ArgumentParser(description="Test LLM2VLM exported model")
    parser.add_argument("--model_dir", type=str, required=True, help="Path to exported HF-ready model")
    parser.add_argument("--image", type=str, default=None, help="Path to test image")
    parser.add_argument("--prompt", type=str, default="Describe this image in detail.")
    parser.add_argument("--max_new_tokens", type=int, default=256)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top_p", type=float, default=0.9)
    parser.add_argument("--batch_test", type=str, nargs="?", const="./datasets/coco2014/train2014",
                        default=None, metavar="IMAGE_DIR",
                        help="Run batch test on all images in a folder (default: ./datasets/coco2014/train2014)")
    parser.add_argument("--num_images", type=int, default=0,
                        help="Limit number of images for batch test (0 = all)")
    args = parser.parse_args()

    if args.batch_test is not None:
        batch_test(args.model_dir, args.batch_test, args.num_images, args.prompt)
    elif args.image:
        run_test(
            model_dir=args.model_dir,
            image_path=args.image,
            prompt=args.prompt,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_p=args.top_p,
        )
    else:
        print("Please specify --image or --batch_test")


if __name__ == "__main__":
    main()