File size: 2,371 Bytes
3b57e8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | 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)
|