VisualSearch/model/VSM.py

from typing import List

import torch
import torch.nn as nn
import torch.nn.functional as F

from VisualSearch.model.llava.model.language_model.llava_llama import (LlavaLlamaForCausalLM,
													 LlavaLlamaModel)

from .segment_anything.modeling import PromptEncoder, MaskDecoder, TwoWayTransformer

from .owlvit.owlvit import OwlViT

def dice_loss(
	inputs: torch.Tensor,
	targets: torch.Tensor,
	num_masks: float,
	scale=1000,  # 100000.0,
	eps=1e-6,
):
	"""
	Compute the DICE loss, similar to generalized IOU for masks
	Args:
		inputs: A float tensor of arbitrary shape.
				The predictions for each example.
		targets: A float tensor with the same shape as inputs. Stores the binary
				 classification label for each element in inputs
				(0 for the negative class and 1 for the positive class).
	"""
	inputs = inputs.sigmoid()
	inputs = inputs.flatten(1, 2)
	targets = targets.flatten(1, 2)
	numerator = 2 * (inputs / scale * targets).sum(-1)
	denominator = (inputs / scale).sum(-1) + (targets / scale).sum(-1)
	loss = 1 - (numerator + eps) / (denominator + eps)
	loss = loss / (num_masks + 1e-8)
	return loss

def sigmoid_ce_loss(
	inputs: torch.Tensor,
	targets: torch.Tensor,
	num_masks: float,
):
	"""
	Args:
		inputs: A float tensor of arbitrary shape.
				The predictions for each example.
		targets: A float tensor with the same shape as inputs. Stores the binary
				 classification label for each element in inputs
				(0 for the negative class and 1 for the positive class).
	Returns:
		Loss tensor
	"""
	loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
	loss = loss.flatten(1, 2).mean(1) / (num_masks + 1e-8)
	return loss

class VSMMetaModel:
	def __init__(
		self,
		config,
		**kwargs,
	):
		super(VSMMetaModel, self).__init__(config)

		self.config = config
		if not hasattr(self.config, "train_mask_decoder"):
			self.config.train_mask_decoder = kwargs["train_mask_decoder"]
			self.config.out_dim = kwargs["out_dim"]
		else:
			is_eval = kwargs.get('is_eval', False)			
			self.initialize_lisa_modules(self.config, is_eval)

	def initialize_lisa_modules(self, config, is_eval=False):
		# OWL-ViT
		self.owlvit = OwlViT(1, is_eval)
		self.owlvit.train()
		for param in self.owlvit.parameters():
			param.requires_grad = True

		for param in self.owlvit.vision_model.parameters():
			param.requires_grad = False
		self.owlvit.vision_model.eval()

		for param in self.owlvit.box_head.parameters():
			param.requires_grad = False

		self.visual_projection = nn.Linear(self.owlvit.vision_model.config.hidden_size, 256, bias=False)
		for param in self.visual_projection.parameters():
			param.requires_grad = True

		self.prompt_encoder=PromptEncoder(
			embed_dim=256,
			image_embedding_size=(48, 48),
			input_image_size=(768, 768),
			mask_in_chans=16,
		)
		self.prompt_encoder.train()
		for param in self.prompt_encoder.parameters():
			param.requires_grad = True
		self.mask_decoder=MaskDecoder(
			num_multimask_outputs=3,
			transformer=TwoWayTransformer(
				depth=2,
				embedding_dim=256,
				mlp_dim=2048,
				num_heads=8,
			),
			transformer_dim=256,
			iou_head_depth=3,
			iou_head_hidden_dim=256,
		)
		self.mask_decoder.train()
		for param in self.mask_decoder.parameters():
			param.requires_grad = True

		# Projection layer
		in_dim = config.hidden_size
		out_dim = config.out_dim
		text_fc_det = [
			nn.Linear(in_dim, in_dim),
			nn.ReLU(inplace=True),
			nn.Linear(in_dim, out_dim),
			nn.Dropout(0.0),
		]
		self.text_hidden_fcs_det = nn.ModuleList([nn.Sequential(*text_fc_det)])
		self.text_hidden_fcs_det.train()
		for param in self.text_hidden_fcs_det.parameters():
			param.requires_grad = True

		text_fc_seg = [
			nn.Linear(in_dim, in_dim),
			nn.ReLU(inplace=True),
			nn.Linear(in_dim, 256),
			nn.Dropout(0.0),
		]
		self.text_hidden_fcs_seg = nn.ModuleList([nn.Sequential(*text_fc_seg)])
		self.text_hidden_fcs_seg.train()
		for param in self.text_hidden_fcs_seg.parameters():
			param.requires_grad = True


class VSMModel(VSMMetaModel, LlavaLlamaModel):
	def __init__(
		self,
		config,
		**kwargs,
	):
		super(VSMModel, self).__init__(config, **kwargs)

		self.config.use_cache = False
		self.config.vision_tower = self.config.mm_vision_tower
		self.config.mm_vision_select_feature = "patch"
		self.config.image_aspect_ratio = "square"
		self.config.image_grid_pinpoints = None
		self.config.tune_mm_mlp_adapter = False
		self.config.freeze_mm_mlp_adapter = True
		self.config.pretrain_mm_mlp_adapter = None
		self.config.mm_use_im_patch_token = False


class VSMForCausalLM(LlavaLlamaForCausalLM):
	def __init__(
		self,
		config,
		**kwargs,
	):
		if not hasattr(config, "train_mask_decoder"):
			config.mm_use_im_start_end = kwargs.pop("use_mm_start_end", True)
			config.mm_vision_tower = kwargs.get(
				"vision_tower", "openai/clip-vit-large-patch14"
			)
			self.ce_loss_weight = kwargs.pop("ce_loss_weight", None)
			self.dice_loss_weight = kwargs.pop("dice_loss_weight", None)
			self.bce_loss_weight = kwargs.pop("bce_loss_weight", None)
			self.det_loss_weight = kwargs.pop("det_loss_weight", None)
		else:
			config.mm_vision_tower = config.vision_tower

		self.loc_token_idx = kwargs.pop("loc_token_idx")

		super().__init__(config)

		self.model = VSMModel(config, **kwargs)

		self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

		# Initialize weights and apply final processing
		self.post_init()

	def get_visual_embs(self, pixel_values: torch.FloatTensor):
		with torch.no_grad():
			image_embeddings = self.model.owlvit.get_visual_embs(pixel_values)
		return image_embeddings

	def forward(self, **kwargs):
		if "past_key_values" in kwargs:
			return super().forward(**kwargs)
		return self.model_forward(**kwargs)

	def model_forward(
		self,
		images: torch.FloatTensor,
		images_clip: torch.FloatTensor,
		input_ids: torch.LongTensor,
		labels: torch.LongTensor,
		attention_masks: torch.LongTensor,
		offset: torch.LongTensor,
		masks_list: List[torch.FloatTensor],
		label_list: List[torch.Tensor],
		bboxes_labels_list: List[torch.FloatTensor],
		bboxes_valid_list: torch.Tensor,
		masks_valid_list: List[torch.Tensor],
		resize_list: List[tuple],
		inference: bool = False,
		**kwargs,
	):
		image_embeddings = self.get_visual_embs(images)
		batch_size = image_embeddings.shape[0]
		assert batch_size == len(offset) - 1

		loc_token_mask = input_ids[:, 1:] == self.loc_token_idx
		loc_token_mask = torch.cat(
			[
				loc_token_mask,
				torch.zeros((loc_token_mask.shape[0], 1)).bool().cuda(),
			],
			dim=1,
		)
		# hack for IMAGE_TOKEN_INDEX (we suppose that there is only one image, and it is in the front)
		loc_token_mask = torch.cat(
			[torch.zeros((loc_token_mask.shape[0], 255)).bool().cuda(), loc_token_mask],
			dim=1,
		)

		if inference:
			n_batch = 1
			length = input_ids.shape[0]
			assert images_clip.shape[0] == 1
			images_clip_extend = images_clip.expand(length, -1, -1, -1).contiguous()

			output_hidden_states = []
			for i in range(n_batch):
				start_i, end_i = i * length, min((i + 1) * length, input_ids.shape[0])
				output_i = super().forward(
					images=images_clip_extend[: end_i - start_i],
					attention_mask=attention_masks[start_i:end_i],
					input_ids=input_ids[start_i:end_i],
					output_hidden_states=True,
				)
				output_hidden_states.append(output_i.hidden_states)
				torch.cuda.empty_cache()

			output_hidden_states_list = []
			output_hidden_states_level = torch.cat(output_hidden_states, dim=0)
			output_hidden_states_list.append(output_hidden_states_level)
			output_hidden_states = output_hidden_states_list
			output = None

		else:
			images_clip_list = []
			for i in range(len(offset) - 1):
				start_i, end_i = offset[i], offset[i + 1]
				images_clip_i = (
					images_clip[i]
					.unsqueeze(0)
					.expand(end_i - start_i, -1, -1, -1)
					.contiguous()
				)
				images_clip_list.append(images_clip_i)
			images_clip = torch.cat(images_clip_list, dim=0)

			output = super().forward(
				images=images_clip,
				attention_mask=attention_masks,
				input_ids=input_ids,
				labels=labels,
				output_hidden_states=True,
			)
			output_hidden_states = output.hidden_states

		# seg
		hidden_states_seg = []
		assert len(self.model.text_hidden_fcs_seg) == 1
		hidden_states_seg.append(self.model.text_hidden_fcs_seg[0](output_hidden_states[-1]))

		last_hidden_state_seg = torch.stack(hidden_states_seg, dim=-1).sum(dim=-1)

		# det
		hidden_states_det = []

		assert len(self.model.text_hidden_fcs_det) == 1
		hidden_states_det.append(self.model.text_hidden_fcs_det[0](output_hidden_states[-1]))
		last_hidden_state_det = torch.stack(hidden_states_det, dim=-1).sum(dim=-1)

		pred_embeddings_seg = last_hidden_state_seg[loc_token_mask]
		pred_embeddings_det = last_hidden_state_det[loc_token_mask]
		loc_token_counts = loc_token_mask.int().sum(-1)  # [bs, ]

		loc_token_offset = loc_token_counts.cumsum(-1)
		loc_token_offset = torch.cat(
			[torch.zeros(1).long().cuda(), loc_token_offset], dim=0
		)

		loc_token_offset = loc_token_offset[offset]

		pred_embeddings_seg_ = []
		for i in range(len(loc_token_offset) - 1):
			start_i, end_i = loc_token_offset[i], loc_token_offset[i + 1]
			pred_embeddings_seg_.append(pred_embeddings_seg[start_i:end_i])
		pred_embeddings_seg = pred_embeddings_seg_

		pred_embeddings_det_ = []
		for i in range(len(loc_token_offset) - 1):
			start_i, end_i = loc_token_offset[i], loc_token_offset[i + 1]
			pred_embeddings_det_.append(pred_embeddings_det[start_i:end_i])
		pred_embeddings_det = pred_embeddings_det_

		# seg branch 
		multimask_output = False
		pred_masks = []
		for i in range(len(pred_embeddings_seg)):
			(
				sparse_embeddings,
				dense_embeddings,
			) = self.model.prompt_encoder(
				points=None,
				boxes=None,
				masks=None,
				text_embeds=pred_embeddings_seg[i].unsqueeze(1),
			)
			sparse_embeddings = sparse_embeddings.to(pred_embeddings_seg[i].dtype)
			low_res_masks, iou_predictions = self.model.mask_decoder(
				image_embeddings=self.model.visual_projection(image_embeddings[i].unsqueeze(0)).permute(0, 3, 1, 2),
				image_pe=self.model.prompt_encoder.get_dense_pe(),
				sparse_prompt_embeddings=sparse_embeddings,
				dense_prompt_embeddings=dense_embeddings,
				multimask_output=multimask_output,
			)
			pred_mask = F.interpolate(
			low_res_masks, label_list[i].shape, mode="bilinear", align_corners=False
		)
			pred_masks.append(pred_mask[:, 0])

		gt_masks = masks_list

		# det branch
		detection_result_batch = []
		for i in range(len(pred_embeddings_det)):
			bs = pred_embeddings_det[i].shape[0]
			detection_result = self.model.owlvit(image_embeddings[i].unsqueeze(0).repeat(bs, 1, 1, 1), pred_embeddings_det[i].unsqueeze(1))
			detection_result_batch.append(detection_result)


		pred_logits = torch.cat([detection_result['pred_logits'] for detection_result in detection_result_batch], 0)
		pred_boxes = torch.cat([detection_result['pred_boxes'] for detection_result in detection_result_batch], 0)
		if inference:
			return {
				"pred_masks": pred_masks,
				"gt_masks": gt_masks,
				"pred_logits": pred_logits,
				"pred_boxes": pred_boxes,
				"gt_bboxes": bboxes_labels_list
			}
		
		num_boxes = 0
		for bboxes_labels, bboxes_valid in zip(bboxes_labels_list, bboxes_valid_list):
			if bboxes_valid:
				num_boxes += len(bboxes_labels)
		num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=image_embeddings.device)
		num_boxes = torch.clamp(num_boxes, min=1).item()
		
		detection_result_batch = {'pred_logits':pred_logits, 'pred_boxes':pred_boxes}

		target_det = []
		all_bboxes_valid = []
		for bboxes_label, bboxes_valid in zip(bboxes_labels_list, bboxes_valid_list):
			target_det.append({"labels":torch.zeros(len(bboxes_label)).to(bboxes_label.device, torch.long), "boxes":bboxes_label})
			if bboxes_valid:
				all_bboxes_valid.append(torch.ones((min(24*24, len(bboxes_label)), 1)).to(bboxes_label.device, torch.long))
			else:
				all_bboxes_valid.append(torch.zeros((min(24*24, len(bboxes_label)), 1)).to(bboxes_label.device, torch.long))
		all_bboxes_valid = torch.cat(all_bboxes_valid, 0)
		
		loss_dict = self.model.owlvit.criterion(detection_result_batch, target_det, num_boxes)

		for loss_k, loss_v in loss_dict.items():
			if "loss_ce" in loss_k:
				loss_dict[loss_k] = (loss_v*bboxes_valid_list.unsqueeze(-1)).mean()
			else:
				loss_dict[loss_k] = (loss_v*all_bboxes_valid).sum()

		weight_dict = self.model.owlvit.criterion.weight_dict
		detection_loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
		detection_loss = detection_loss*self.det_loss_weight

		model_output = output
		output = model_output.logits

		ce_loss = model_output.loss
		ce_loss = ce_loss * self.ce_loss_weight
		mask_bce_loss = 0
		mask_dice_loss = 0
		num_masks = 0
		for batch_idx in range(len(pred_masks)):
			gt_mask = gt_masks[batch_idx]
			pred_mask = pred_masks[batch_idx]
			masks_valid = masks_valid_list[batch_idx]

			mask_bce_loss += (
				sigmoid_ce_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
				* gt_mask.shape[0] * masks_valid
			).sum()
			mask_dice_loss += (
				dice_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
				* gt_mask.shape[0] * masks_valid
			).sum()
			num_masks += masks_valid.sum()

		mask_bce_loss = self.bce_loss_weight * mask_bce_loss / (num_masks + 1e-8)
		mask_dice_loss = self.dice_loss_weight * mask_dice_loss / (num_masks + 1e-8)
		mask_loss = mask_bce_loss + mask_dice_loss

		loss = ce_loss + mask_loss + detection_loss

		return {
			"loss": loss,
			"ce_loss": ce_loss,
			"mask_bce_loss": mask_bce_loss,
			"mask_dice_loss": mask_dice_loss,
			"mask_loss": mask_loss,
			"detection_loss": detection_loss,
			"detection_loss_ce": loss_dict['loss_ce'],
			"detection_loss_bbox": loss_dict['loss_bbox'],
			"detection_loss_giou": loss_dict['loss_giou'],
		}

	def inference(
		self,
		images_clip,
		images,
		input_ids,
		resize_list,
		original_size_list,
		max_new_tokens=32,
		tokenizer=None,
		mode = 'vqa'
	):
		assert mode in ['vqa', 'segmentation', 'detection']
		with torch.no_grad():
			outputs = self.generate(
				images=images_clip,
				input_ids=input_ids,
				max_new_tokens=max_new_tokens,
				num_beams=1,
				output_hidden_states=True,
				return_dict_in_generate=True,
			)
			output_hidden_states = outputs.hidden_states[-1]
			output_ids = outputs.sequences

			if mode == 'vqa':
				return output_ids, None, None

			loc_token_mask = output_ids[:, 1:] == self.loc_token_idx
			# hack for IMAGE_TOKEN_INDEX (we suppose that there is only one image, and it is in the front)
			loc_token_mask = torch.cat(
				[
					torch.zeros((loc_token_mask.shape[0], 255)).bool().cuda(),
					loc_token_mask,
				],
				dim=1,
			)

			# seg
			hidden_states_seg = []
			assert len(self.model.text_hidden_fcs_seg) == 1
			hidden_states_seg.append(self.model.text_hidden_fcs_seg[0](output_hidden_states))

			last_hidden_state_seg = torch.stack(hidden_states_seg, dim=-1).sum(dim=-1)

			# det
			hidden_states_det = []

			assert len(self.model.text_hidden_fcs_det) == 1
			hidden_states_det.append(self.model.text_hidden_fcs_det[0](output_hidden_states))
			last_hidden_state_det = torch.stack(hidden_states_det, dim=-1).sum(dim=-1)

			pred_embeddings_seg = last_hidden_state_seg[loc_token_mask]
			pred_embeddings_det = last_hidden_state_det[loc_token_mask]
			loc_token_counts = loc_token_mask.int().sum(-1)  # [bs, ]

			loc_token_offset = loc_token_counts.cumsum(-1)
			loc_token_offset = torch.cat(
				[torch.zeros(1).long().cuda(), loc_token_offset], dim=0
			)


			pred_embeddings_seg_ = []
			for i in range(len(loc_token_offset) - 1):
				start_i, end_i = loc_token_offset[i], loc_token_offset[i + 1]
				pred_embeddings_seg_.append(pred_embeddings_seg[start_i:end_i])
			pred_embeddings_seg = pred_embeddings_seg_

			pred_embeddings_det_ = []
			for i in range(len(loc_token_offset) - 1):
				start_i, end_i = loc_token_offset[i], loc_token_offset[i + 1]
				pred_embeddings_det_.append(pred_embeddings_det[start_i:end_i])
			pred_embeddings_det = pred_embeddings_det_

			image_embeddings = self.get_visual_embs(images)

			multimask_output = False
			pred_masks = []
			for i in range(len(pred_embeddings_seg)):
				(
					sparse_embeddings,
					dense_embeddings,
				) = self.model.prompt_encoder(
					points=None,
					boxes=None,
					masks=None,
					text_embeds=pred_embeddings_seg[i].unsqueeze(1),
				)

				sparse_embeddings = sparse_embeddings.to(pred_embeddings_seg[i].dtype)
				low_res_masks, iou_predictions = self.model.mask_decoder(
					image_embeddings=self.model.visual_projection(image_embeddings[i].unsqueeze(0)).permute(0, 3, 1, 2),
					image_pe=self.model.prompt_encoder.get_dense_pe(),
					sparse_prompt_embeddings=sparse_embeddings,
					dense_prompt_embeddings=dense_embeddings,
					multimask_output=multimask_output,
				)
				pred_mask = F.interpolate(
				low_res_masks.float(), original_size_list[i], mode="bilinear", align_corners=False
			)
				pred_masks.append(pred_mask[:, 0])

			if mode == 'segmentation':
				return None, pred_masks, None

			# detection model
			detection_result_batch = []
			for i in range(len(pred_embeddings_det)):
				bs = pred_embeddings_det[i].shape[0]
				detection_result = self.model.owlvit(image_embeddings[i].unsqueeze(0).repeat(bs, 1, 1, 1), pred_embeddings_det[i].unsqueeze(1))
				detection_result_batch.append(detection_result)


			pred_logits = torch.cat([detection_result['pred_logits'] for detection_result in detection_result_batch], 0)
			pred_boxes = torch.cat([detection_result['pred_boxes'] for detection_result in detection_result_batch], 0)
			detection_result_batch = {'pred_logits':pred_logits, 'pred_boxes':pred_boxes}

		return None, pred_masks, detection_result_batch