from PIL import Image from itertools import zip_longest import os import requests from urllib.parse import urlparse import torch from torch import nn import torch.nn.functional as F import cv2 import numpy as np from model.llava.mm_utils import tokenizer_image_token from model.SAM.utils.transforms import ResizeLongestSide from dataset.utils.blip_processors import Blip2ImageEvalProcessor from tools.prompting import write_list as _write_list from tools.utils import ( DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX, ) def is_url(image): try: result = urlparse(image) return all([result.scheme, result.netloc]) except: return False class Chat: IMG_MEAN = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1) IMG_STD = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1) IMG_SIZE = 1024 IGNORE_LABEL = 255 WHITE_BAND_WIDTH = 10 def __init__( self, model, tokenizer, qformer_tokenizer, args, device="cuda:0", ): self.device = device self.model = model self.global_enc_processor = Blip2ImageEvalProcessor() self.tokenizer = tokenizer self.qformer_tokenizer = qformer_tokenizer image_size = args.image_size self.transform = ResizeLongestSide(image_size) self.use_mm_start_end = args.use_mm_start_end # Match each image stream to the corresponding encoder dtype/device. try: first_param = next(self.model.parameters()) self.runtime_dtype = first_param.dtype self.runtime_device = first_param.device except Exception: self.runtime_dtype = torch.bfloat16 self.runtime_device = torch.device(self.device) try: ge_param = next( self.base_model.model.grounding_encoder.image_encoder.parameters() ) self.grounding_dtype = ge_param.dtype self.grounding_device = ge_param.device except Exception: self.grounding_dtype = self.runtime_dtype self.grounding_device = self.runtime_device def _load_image_to_rgb(self, image): if is_url(image): response = requests.get(image) image_np = np.frombuffer(response.content, np.uint8) image_np = cv2.imdecode(image_np, cv2.IMREAD_COLOR) else: image_np = cv2.imread(image) if image_np is None: raise ValueError(f"Unable to load image: {image}") return cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB) def _global_enc_from_array(self, image_np: np.ndarray) -> torch.Tensor: return self.global_enc_processor(Image.fromarray(image_np)) def _prepare_encodings(self, image_np: np.ndarray, image_identifier: str): original_size = image_np.shape[:2] global_enc_image = self._global_enc_from_array(image_np) global_enc_image = global_enc_image.to( device=self.runtime_device, dtype=self.runtime_dtype ) grounding_image = self.transform.apply_image(image_np) resize = grounding_image.shape[:2] grounding_enc_image = self.grounding_enc_processor( torch.from_numpy(grounding_image).permute(2, 0, 1).contiguous() ).to(device=self.grounding_device, dtype=self.grounding_dtype) return ( global_enc_image, grounding_enc_image, original_size, resize, image_np, image_identifier, ) def grounding_enc_processor(self, x: torch.Tensor) -> torch.Tensor: x = (x - self.IMG_MEAN) / self.IMG_STD h, w = x.shape[-2:] x = F.pad(x, (0, self.IMG_SIZE - w, 0, self.IMG_SIZE - h)) return x @property def base_model(self): if isinstance(self.model, nn.DataParallel): return self.model.module else: return self.model def ask(self, text, conv, num_images=2): text = text.replace("<", "<").replace(">", ">") prompt = text if len(conv.messages) == 0: if num_images > 1: img_tokens = [ f"{DEFAULT_IMAGE_TOKEN} (IMAGE{i + 1})" for i in range(num_images) ] img_tokens_list = _write_list(img_tokens) begin_sentence = ( f"The {img_tokens_list} provide an overview of the pictures.\n" ) elif num_images == 1: begin_sentence = ( f"The {DEFAULT_IMAGE_TOKEN} (IMAGE1) provides an overview of the picture.\n" ) else: begin_sentence = "" prompt = begin_sentence + prompt if self.use_mm_start_end: replace_token = ( DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN ) prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token) conv.append_message(conv.roles[0], prompt) def answer( self, conv, img_list, question, answer_prefix: str = "", max_new_tokens=512, max_length=2300, ): conv.append_message(conv.roles[1], "") prompt = conv.get_prompt() input_ids = tokenizer_image_token(prompt, self.tokenizer, return_tensors="pt") if answer_prefix: if DEFAULT_IMAGE_TOKEN in answer_prefix: prefix_ids = tokenizer_image_token(answer_prefix, self.tokenizer) else: prefix_ids = self.tokenizer(answer_prefix).input_ids if len(prefix_ids) > 0 and prefix_ids[0] == self.tokenizer.bos_token_id: prefix_ids = prefix_ids[1:] if len(prefix_ids) > 0 and prefix_ids[-1] == self.tokenizer.eos_token_id: prefix_ids = prefix_ids[:-1] if len(prefix_ids) > 0: prefix_ids = torch.tensor(prefix_ids, dtype=input_ids.dtype) input_ids = torch.cat([input_ids, prefix_ids], dim=0) input_ids = input_ids.unsqueeze(0).to(self.device) current_max_len = input_ids.shape[1] + max_new_tokens begin_idx = max(0, current_max_len - max_length) input_ids = input_ids[:, begin_idx:] attention_masks = input_ids.ne(self.tokenizer.pad_token_id) attention_masks = attention_masks.to(self.device) conv_offset = torch.tensor([0, 1], device=self.device) questions_list = question if isinstance(question, list) else [question] if self.qformer_tokenizer is not None: qformer_inputs = self.qformer_tokenizer( questions_list, padding=True, truncation=True, return_tensors="pt" ) qformer_input_ids = qformer_inputs.input_ids.to(self.device) qformer_attention_masks = qformer_inputs.attention_mask.to(self.device) else: qformer_input_ids = None qformer_attention_masks = None if len(img_list) > 0: if len(img_list) > 1: ( global_enc_images, grounding_enc_images, original_sizes, resizes, _, image_paths, ) = list(zip(*img_list)) global_enc_images = torch.stack(global_enc_images, dim=0) grounding_enc_images = torch.stack(grounding_enc_images, dim=0) original_size_list = list(zip_longest([None, *original_sizes])) resize_list = list(zip_longest([None, *resizes])) else: ( global_enc_image, grounding_enc_image, original_size, resize, _, image_path, ) = img_list[0] global_enc_images = global_enc_image.unsqueeze(0) grounding_enc_images = grounding_enc_image.unsqueeze(0) image_paths = [image_path] resizes = [resize] original_sizes = [original_size] slots = [None, *original_sizes] resize_slots = [None, *resizes] original_size_list = list(zip_longest(slots)) resize_list = list(zip_longest(resize_slots)) image_paths = [image_paths] image_offset = torch.tensor( [0, global_enc_images.shape[0]], device=self.device ) grounding_offset = torch.tensor( [0, grounding_enc_images.shape[0]], device=self.device ) output_ids, pred_masks = self.model.evaluate( global_enc_images, grounding_enc_images, input_ids, attention_masks, resize_list, original_size_list, conv_offset, image_offset, grounding_offset, qformer_input_ids, qformer_attention_masks, max_new_tokens=max_new_tokens, image_paths=image_paths, eos_token_id=self.tokenizer.eos_token_id, ) output_ids = output_ids[0][output_ids[0] != IMAGE_TOKEN_INDEX] else: image_offset = torch.tensor([0, 0], device=self.device) grounding_offset = torch.tensor([0, 0], device=self.device) output_ids, _ = self.model.evaluate( None, None, input_ids, attention_masks, None, None, conv_offset, image_offset, grounding_offset, qformer_input_ids, qformer_attention_masks, max_new_tokens=max_new_tokens, eos_token_id=self.tokenizer.eos_token_id, ) output_ids = output_ids[0] pred_masks = None text_output = self.tokenizer.decode(output_ids, skip_special_tokens=False) text_output = " ".join(text_output.split()) chunks = text_output.split("") llm_message = chunks[0] if len(chunks) <= 2 else chunks[1] llm_message = llm_message.split("ASSISTANT: ")[-1].strip() conv.messages.pop() conv.append_message(conv.roles[1], llm_message) return llm_message, pred_masks, text_output def upload_img(self, image, conv, img_list): image_np = self._load_image_to_rgb(image) img_list.append(self._prepare_encodings(image_np, image)) msg = "Received." return msg