import torch import os import json from tqdm import tqdm from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN from llava.conversation import conv_templates, SeparatorStyle from llava.model.builder import load_pretrained_model from llava.utils import disable_torch_init from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria from PIL import Image import math import time import glob as gb class LLavaAgent: def __init__(self, model_path, device='cuda', conv_mode='vicuna_v1', load_8bit=False, load_4bit=False): self.device = device if torch.device(self.device).index is not None: device_map = {'model': torch.device(self.device).index, 'lm_head': torch.device(self.device).index} else: device_map = 'auto' model_path = os.path.expanduser(model_path) model_name = get_model_name_from_path(model_path) tokenizer, model, image_processor, context_len = load_pretrained_model( model_path, None, model_name, device=self.device, device_map=device_map, load_8bit=load_8bit, load_4bit=load_4bit) self.model = model self.image_processor = image_processor self.tokenizer = tokenizer self.context_len = context_len self.qs = '请详细描述该图片及其风格。' self.conv_mode = conv_mode if self.model.config.mm_use_im_start_end: self.qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + self.qs else: self.qs = DEFAULT_IMAGE_TOKEN + '\n' + self.qs self.conv = conv_templates[self.conv_mode].copy() self.conv.append_message(self.conv.roles[0], self.qs) self.conv.append_message(self.conv.roles[1], None) prompt = self.conv.get_prompt() self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze( 0).to(self.device) def update_qs(self, qs=None): if qs is None: qs = self.qs else: if self.model.config.mm_use_im_start_end: qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs else: qs = DEFAULT_IMAGE_TOKEN + '\n' + qs self.conv = conv_templates[self.conv_mode].copy() self.conv.append_message(self.conv.roles[0], qs) self.conv.append_message(self.conv.roles[1], None) prompt = self.conv.get_prompt() self.input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze( 0).to(self.device) def gen_image_caption(self, imgs, temperature=0.2, top_p=0.7, num_beams=1, qs=None): ''' [PIL.Image, ...] ''' self.update_qs(qs) bs = len(imgs) input_ids = self.input_ids.repeat(bs, 1) img_tensor_list = [] for image in imgs: _image_tensor = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0] img_tensor_list.append(_image_tensor) image_tensor = torch.stack(img_tensor_list, dim=0).half().to(self.device) stop_str = self.conv.sep if self.conv.sep_style != SeparatorStyle.TWO else self.conv.sep2 with torch.inference_mode(): output_ids = self.model.generate( input_ids, images=image_tensor, do_sample=True if temperature > 0 else False, temperature=temperature, top_p=top_p, num_beams=num_beams, # no_repeat_ngram_size=3, max_new_tokens=512, use_cache=True) input_token_len = input_ids.shape[1] outputs = self.tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True) img_captions = [] for output in outputs: output = output.strip() if output.endswith(stop_str): output = output[:-len(stop_str)] output = output.strip().replace('\n', ' ').replace('\r', ' ') img_captions.append(output) return img_captions if __name__ == '__main__': llava_agent = LLavaAgent("/opt/data/private/AIGC_pretrain/LLaVA1.5/llava-v1.5-13b") img = [Image.open('/opt/data/private/LV_Dataset/DiffGLV-Test-All/RealPhoto60/LQ/02.png')] caption = llava_agent.gen_image_caption(img)