Spaces:
Runtime error
Runtime error
| import argparse | |
| import re | |
| import requests | |
| from PIL import Image | |
| from io import BytesIO | |
| import torch | |
| from transformers import PreTrainedModel | |
| from tinyllava.utils import * | |
| from tinyllava.data import * | |
| from tinyllava.model import * | |
| def image_parser(args): | |
| out = args.image_file.split(args.sep) | |
| return out | |
| def load_image(image_file): | |
| if image_file.startswith("http") or image_file.startswith("https"): | |
| response = requests.get(image_file) | |
| image = Image.open(BytesIO(response.content)).convert("RGB") | |
| else: | |
| image = Image.open(image_file).convert("RGB") | |
| return image | |
| def load_images(image_files): | |
| out = [] | |
| for image_file in image_files: | |
| image = load_image(image_file) | |
| out.append(image) | |
| return out | |
| def eval_model(args): | |
| # Model | |
| disable_torch_init() | |
| if args.model_path is not None: | |
| model, tokenizer, image_processor, context_len = load_pretrained_model(args.model_path) | |
| else: | |
| assert args.model is not None, 'model_path or model must be provided' | |
| model = args.model | |
| if hasattr(model.config, "max_sequence_length"): | |
| context_len = model.config.max_sequence_length | |
| else: | |
| context_len = 2048 | |
| tokenizer = model.tokenizer | |
| image_processor = model.vision_tower._image_processor | |
| qs = args.query | |
| qs = DEFAULT_IMAGE_TOKEN + "\n" + qs | |
| text_processor = TextPreprocess(tokenizer, args.conv_mode) | |
| data_args = model.config | |
| image_processor = ImagePreprocess(image_processor, data_args) | |
| model.cuda() | |
| msg = Message() | |
| msg.add_message(qs) | |
| result = text_processor(msg.messages, mode='eval') | |
| input_ids = result['input_ids'] | |
| prompt = result['prompt'] | |
| input_ids = input_ids.unsqueeze(0).cuda() | |
| image_files = image_parser(args) | |
| images = load_images(image_files)[0] | |
| images_tensor = image_processor(images) | |
| images_tensor = images_tensor.unsqueeze(0).half().cuda() | |
| stop_str = text_processor.template.separator.apply()[1] | |
| keywords = [stop_str] | |
| stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids) | |
| with torch.inference_mode(): | |
| output_ids = model.generate( | |
| input_ids, | |
| images=images_tensor, | |
| do_sample=True if args.temperature > 0 else False, | |
| temperature=args.temperature, | |
| top_p=args.top_p, | |
| num_beams=args.num_beams, | |
| pad_token_id=tokenizer.pad_token_id, | |
| max_new_tokens=args.max_new_tokens, | |
| use_cache=True, | |
| stopping_criteria=[stopping_criteria], | |
| ) | |
| outputs = tokenizer.batch_decode( | |
| output_ids, skip_special_tokens=True | |
| )[0] | |
| outputs = outputs.strip() | |
| if outputs.endswith(stop_str): | |
| outputs = outputs[: -len(stop_str)] | |
| outputs = outputs.strip() | |
| print(outputs) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--model-path", type=str, default=None) | |
| parser.add_argument("--model", type=PreTrainedModel, default=None) | |
| parser.add_argument("--image-file", type=str, required=True) | |
| parser.add_argument("--query", type=str, required=True) | |
| parser.add_argument("--conv-mode", type=str, default=None) | |
| parser.add_argument("--sep", type=str, default=",") | |
| parser.add_argument("--temperature", type=float, default=0.2) | |
| parser.add_argument("--top_p", type=float, default=None) | |
| parser.add_argument("--num_beams", type=int, default=1) | |
| parser.add_argument("--max_new_tokens", type=int, default=512) | |
| args = parser.parse_args() | |
| eval_model(args) |