"""Single-image FastVLM inference using the native Transformers integration.

Setup: python -m pip install 'transformers==5.0.0' 'timm==1.0.29' torch pillow
Run: python fastvlm_transformers.py image.png --prompt 'What is in this image?'
Source: https://huggingface.co/docs/transformers/model_doc/fast_vlm
The default is a community conversion used by that guide, not Apple's original ZIP.
"""

import argparse
from pathlib import Path

import torch
from transformers import AutoConfig, AutoProcessor, FastVlmForConditionalGeneration


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('image', type=Path)
    parser.add_argument('--prompt', default='Describe this image briefly.')
    parser.add_argument('--model', default='KamilaMila/FastVLM-0.5B')
    args = parser.parse_args()
    if not args.image.is_file():
        parser.error('The image file does not exist.')

    device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
    dtype = torch.float16 if device == 'cuda' else torch.float32
    config = AutoConfig.from_pretrained(args.model)
    if args.model == 'KamilaMila/FastVLM-0.5B':
        # This older conversion records weight tying only in text_config.
        # Transformers 5 also needs it on the outer config to restore lm_head.
        config.tie_word_embeddings = config.text_config.tie_word_embeddings
    model = FastVlmForConditionalGeneration.from_pretrained(args.model, config=config, dtype=dtype).to(device).eval()
    processor = AutoProcessor.from_pretrained(args.model, use_fast=False)
    messages = [{'role': 'user', 'content': [
        {'type': 'image', 'path': str(args.image.resolve())},
        {'type': 'text', 'text': args.prompt},
    ]}]
    inputs = processor.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=True,
        return_dict=True, return_tensors='pt',
    ).to(device)
    if 'pixel_values' in inputs:
        inputs['pixel_values'] = inputs['pixel_values'].to(dtype=dtype)
    # Qwen chat turns may end with im_end instead of the legacy endoftext.
    stop_tokens = [processor.tokenizer.eos_token_id]
    chat_end = processor.tokenizer.get_vocab().get('<|im_end|>')
    if chat_end is not None:
        stop_tokens.append(chat_end)
    with torch.inference_mode():
        output = model.generate(
            **inputs, max_new_tokens=192, do_sample=False, repetition_penalty=1.2,
            eos_token_id=stop_tokens,
            pad_token_id=processor.tokenizer.pad_token_id,
        )
    answer = output[:, inputs['input_ids'].shape[1]:]
    print(processor.batch_decode(answer, skip_special_tokens=True)[0].strip())


if __name__ == '__main__':
    main()