| | 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 = 'Describe this image and its style in a very detailed manner.'
|
| | 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,
|
| |
|
| | 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)
|
| |
|
