| import argparse |
| import torch |
| from PIL import Image |
| from transformers import AutoTokenizer |
| from earthdial.model.internvl_chat import InternVLChatModel |
| from earthdial.train.dataset import build_transform |
|
|
| def run_single_inference(args): |
| print(f"Loading model and tokenizer from Hugging Face: {args.checkpoint}") |
| tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, use_fast=False) |
| model = InternVLChatModel.from_pretrained( |
| args.checkpoint, |
| low_cpu_mem_usage=True, |
| torch_dtype=torch.bfloat16, |
| device_map="auto" if args.auto else None, |
| load_in_8bit=args.load_in_8bit, |
| load_in_4bit=args.load_in_4bit |
| ).eval() |
|
|
| if not args.load_in_8bit and not args.load_in_4bit and not args.auto: |
| model = model.cuda() |
|
|
| image = Image.open(args.image_path).convert("RGB") |
| image_size = model.config.force_image_size or model.config.vision_config.image_size |
| transform = build_transform(is_train=False, input_size=image_size, normalize_type='imagenet') |
| pixel_values = transform(image).unsqueeze(0).cuda().to(torch.bfloat16) |
|
|
| generation_config = { |
| "num_beams": args.num_beams, |
| "max_new_tokens": 100, |
| "min_new_tokens": 1, |
| "do_sample": args.temperature > 0, |
| "temperature": args.temperature, |
| } |
|
|
| answer = model.chat( |
| tokenizer=tokenizer, |
| pixel_values=pixel_values, |
| question=args.question, |
| generation_config=generation_config, |
| verbose=True |
| ) |
|
|
| print("\n=== Inference Result ===") |
| print(f"Question: {args.question}") |
| print(f"Answer: {answer}") |
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--checkpoint', type=str, required=True, help='Model repo ID on Hugging Face Hub') |
| parser.add_argument('--image-path', type=str, required=True, help='Path to input image') |
| parser.add_argument('--question', type=str, required=True, help='Visual question to ask') |
| parser.add_argument('--num-beams', type=int, default=5) |
| parser.add_argument('--temperature', type=float, default=0.0) |
| parser.add_argument('--load-in-8bit', action='store_true') |
| parser.add_argument('--load-in-4bit', action='store_true') |
| parser.add_argument('--auto', action='store_true') |
|
|
| args = parser.parse_args() |
| run_single_inference(args) |
|
|
