import os from itertools import zip_longest from typing import List, Optional, Union, Tuple, Dict import torch import torch.nn as nn from model.SAM import build_sam_vit_h from model.llava.model.language_model.prima_llama import ( PRIMALlamaForCausalLM, PRIMALlamaModel, ) from transformers.models.llama.modeling_llama import LlamaAttention from transformers.modeling_outputs import BaseModelOutputWithPast from tools.utils import IMAGE_TOKEN_INDEX from model.configuration_prima import PRIMAConfig from model.utils import ( calculate_dice_loss, compute_sigmoid_cross_entropy, ) class PRIMABaseModel: def __init__(self, config, **kwargs): super(PRIMABaseModel, self).__init__(config) self.config = config self.vision_pretrained = kwargs.get("vision_pretrained", None) # Set config attributes if they don't exist self.config.train_mask_decoder = getattr( self.config, "train_mask_decoder", kwargs.get("train_mask_decoder", False) ) self.config.out_dim = getattr( self.config, "out_dim", kwargs.get("out_dim", 512) ) self.initialize_prima_model(self.config) def initialize_prima_model(self, config): # Initialize the visual model self.grounding_encoder = build_sam_vit_h(self.vision_pretrained) self._configure_grounding_encoder(config) # Initialize the text projection layer self._initialize_text_projection_layer() def _configure_grounding_encoder(self, config): # Freezing visual model parameters for param in self.grounding_encoder.parameters(): param.requires_grad = False # Training mask decoder if specified if config.train_mask_decoder: self._train_mask_decoder() def _train_mask_decoder(self): self.grounding_encoder.mask_decoder.train() for param in self.grounding_encoder.mask_decoder.parameters(): param.requires_grad = True def _initialize_text_projection_layer(self): in_dim, out_dim = self.config.hidden_size, self.config.out_dim text_projection_layers = [ nn.Linear(in_dim, in_dim), nn.ReLU(inplace=True), nn.Linear(in_dim, out_dim), nn.Dropout(0.0), ] self.text_hidden_fcs = nn.ModuleList([nn.Sequential(*text_projection_layers)]) self.text_hidden_fcs.train() class PRIMAModel(PRIMABaseModel, PRIMALlamaModel): def __init__(self, config, **kwargs): super(PRIMAModel, self).__init__(config, **kwargs) self._configure_model_settings() def _configure_model_settings(self): self.config.use_cache = False self.config.vision_module = self.config.mm_vision_module self.config.select_feature_type = "patch" self.config.image_aspect = "square" self.config.image_grid_points = None self.config.tune_mlp_adapter = False self.config.freeze_mlp_adapter = True self.config.pretrain_mm_mlp_adapter = None self.config.use_image_patch_token = False def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, image_token_slices: Optional[List[slice]] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: return super().forward( input_ids=input_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) class PRIMAForCausalLM(PRIMALlamaForCausalLM): config_class = PRIMAConfig _tied_weights_keys = [r'model\.vision_tower\..*'] def __init__(self, config, **kwargs): for name in ('seg_token_idx', 'img_emb_len', 'seg_image_tokens', 'max_num_images'): if name not in kwargs and hasattr(config, name): kwargs[name] = getattr(config, name) self._set_model_configurations(config, kwargs) super().__init__(config) self.model = PRIMAModel(config, **kwargs) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() if getattr(config, 'self_contained', False): self._initialize_packaged_vision(config) def _initialize_packaged_vision(self, config): """Construct every vision module locally; from_pretrained loads their weights.""" from model.qformer.eva_vit import create_eva_vit_g from model.qformer.Qformer import BertConfig, BertLMHeadModel from model.qformer.blip2 import LayerNorm, disabled_train from model.llava.prima_arch import CrossAttention core = self.model core.visual_encoder = create_eva_vit_g( config.global_image_size, drop_path_rate=0, precision='bf16', load_pretrained=False) core.ln_vision = LayerNorm(core.visual_encoder.num_features) core.Qformer = BertLMHeadModel(BertConfig(**config.qformer_config)) core.Qformer.cls = None core.query_tokens = nn.Parameter(torch.zeros( 1, config.num_query_tokens, config.qformer_config['hidden_size'])) core.residual_query_tokens = nn.Parameter(torch.zeros_like(core.query_tokens)) core.residual_extractor = CrossAttention( qdim=config.qformer_config['hidden_size'], kvdim=core.visual_encoder.num_features) core.mm_projector = nn.Linear(config.qformer_config['hidden_size'], config.hidden_size) core.vision_tower = core.visual_encoder for module in (core.visual_encoder, core.ln_vision, core.Qformer): module.requires_grad_(False) for module in (core.visual_encoder, core.ln_vision): module.eval() module.train = disabled_train @torch.inference_mode() def predict(self, images, question, tokenizer, qformer_tokenizer, max_new_tokens=1024): """Answer a question about ordered local images and return native-resolution masks.""" from types import SimpleNamespace from tools.chat import Chat from tools.outputs import image_mask_slots from model.llava.conversation import conv_templates if not 1 <= len(images) <= self.config.max_num_images: raise ValueError('The image count must be within max_num_images.') args = SimpleNamespace(image_size=self.config.grounding_image_size, use_mm_start_end=True) chat = Chat(self, tokenizer, qformer_tokenizer, args, device=str(self.device)) conv = conv_templates['llava_v1'].copy() encoded = [] for image in images: chat.upload_img(str(image), conv, encoded) chat.ask(question, conv, num_images=len(images)) response, masks, _ = chat.answer( conv, encoded, question, max_new_tokens=max_new_tokens, max_length=self.config.model_max_length + max_new_tokens) masks = [slot > 0 if slot is not None else None for slot in image_mask_slots(masks, len(images))] return {'response': response, 'masks': masks} def _set_model_configurations(self, config, kwargs): config.mm_use_image_start_end = kwargs.pop("use_mm_start_end", True) config.mm_vision_module = 'eva_clip_g' self._initialize_loss_weights(kwargs) self.seg_token_idx = kwargs.pop("seg_token_idx") self.img_emb_len = kwargs.pop("img_emb_len") self.seg_image_tokens = kwargs.pop("seg_image_tokens") self.max_num_images = kwargs.pop("max_num_images") def _initialize_loss_weights(self, kwargs): 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) def get_grounding_encoder_embs(self, pixel_values: torch.FloatTensor): with torch.no_grad(): return torch.cat( [self._encode_single_image(img) for img in pixel_values], dim=0 ) def _encode_single_image(self, image): torch.cuda.empty_cache() return self.model.grounding_encoder.image_encoder(image.unsqueeze(0)) def forward(self, **kwargs): return ( super().forward(**kwargs) if "past_key_values" in kwargs else self.model_forward(**kwargs) ) def model_forward( self, global_enc_images: torch.FloatTensor, grounding_enc_images: torch.FloatTensor, input_ids: torch.LongTensor, labels: torch.LongTensor, attention_masks: torch.LongTensor, qformer_input_ids: torch.LongTensor, qformer_attention_masks: torch.LongTensor, conv_offset: torch.LongTensor, image_offset: torch.LongTensor, grounding_offset: torch.LongTensor, masks_list: List[torch.FloatTensor], label_list: List[torch.Tensor], resize_list: List[tuple], image_paths: List[str], inference: bool = False, **kwargs, ): if grounding_enc_images is not None: image_embeddings = self.get_grounding_encoder_embs(grounding_enc_images) assert image_embeddings.shape[0] == grounding_offset[-1] else: image_embeddings = None num_grounding_images = max( [ (grounding_offset[i + 1] - grounding_offset[i]) for i in range(len(grounding_offset) - 1) ] ) seg_token_masks = self._create_seg_token_mask_multi( input_ids, num_grounding_images, infer=inference, ) if image_embeddings is not None: image_embeddings = self._get_image_embs_per_image( image_embeddings, grounding_offset, num_grounding_images ) # Handle inference or training paths if inference: output_hidden_states = self._inference_path( input_ids, global_enc_images, attention_masks, conv_offset, image_offset, qformer_input_ids, qformer_attention_masks, ) else: output, output_hidden_states = self._training_path( global_enc_images, input_ids, labels, attention_masks, conv_offset, image_offset, qformer_input_ids, qformer_attention_masks, ) # Process hidden states hidden_states, pred_embeddings = self._process_hidden_states_multi( output_hidden_states, seg_token_masks, conv_offset ) # Generate and post-process masks pred_masks = ( self._generate_and_postprocess_masks_multi( pred_embeddings, image_embeddings, resize_list, label_list, infer=inference, ) if image_embeddings is not None else [] ) if inference: return { "pred_masks": pred_masks, "gt_masks": masks_list, } # Calculate losses return self._calculate_losses(pred_masks, masks_list, output) def _get_image_embs_per_image(self, image_embeddings, grounding_offset, num_images): new_image_embeddings = [list() for _ in range(num_images + 1)] for i in range(len(grounding_offset) - 1): start_i, end_i = grounding_offset[i], grounding_offset[i + 1] if end_i - start_i > 0: new_image_embeddings[0].append(None) for idx in range(end_i - start_i): new_image_embeddings[idx + 1].append( image_embeddings[idx + start_i] ) for idx in range(end_i - start_i + 1, num_images + 1): new_image_embeddings[idx].append(None) else: for idx in range(num_images + 1): new_image_embeddings[idx].append(None) return new_image_embeddings def _create_shifted_mask(self, input_ids, shift, token_idx): bsz = input_ids.shape[0] padding = shift - 1 new_mask = input_ids[:, shift:] == token_idx if padding == 0: return new_mask return torch.cat( [ new_mask, torch.zeros((bsz, padding), dtype=torch.bool, device=input_ids.device), ], dim=1, ) def _create_seg_token_mask_multi( self, input_ids, num_grounding_images, infer=False ): image_idx = 0 all_masks = [] bsz = input_ids.shape[0] if self.seg_image_tokens: for image_idx in range(num_grounding_images): mask = torch.ones_like(input_ids[:, 1:], device=input_ids.device).bool() for idx, token_idx in enumerate(self.seg_image_tokens[image_idx]): shifted_mask = self._create_shifted_mask( input_ids, idx + 1, token_idx ) mask &= shifted_mask all_masks.append(mask) seg_only_mask = self._create_shifted_mask(input_ids, 1, self.seg_token_idx) for mask in all_masks: seg_only_mask &= ~mask all_masks = [seg_only_mask, *all_masks] img_count = (input_ids == IMAGE_TOKEN_INDEX).sum(dim=1) all_new_masks = [] for image_idx, masks in enumerate(all_masks): new_masks = [] for batch_idx in range(bsz): mask = masks[batch_idx] mask = torch.cat( [ torch.zeros( (self.img_emb_len - 1) * img_count[batch_idx], device=input_ids.device, ).bool(), mask, ], dim=0, ) if not infer: mask = torch.cat( (mask, torch.zeros(1, device=input_ids.device).bool()), dim=0 ) new_masks.append(mask) if any(x.shape != new_masks[0].shape for x in new_masks): max_len = max(x.shape[0] for x in new_masks) new_masks_align = [] for cur_new_mask in new_masks: cur_new_mask = torch.cat( ( cur_new_mask, torch.zeros( ( max_len - cur_new_mask.shape[0] # , ), dtype=cur_new_mask.dtype, device=cur_new_mask.device, ), ), dim=0, ) new_masks_align.append(cur_new_mask) new_masks = new_masks_align new_masks = torch.stack(new_masks, dim=0) all_new_masks.append(new_masks) return all_new_masks def _inference_path( self, input_ids, global_enc_images, attention_masks, conv_offset, image_offset, qformer_input_ids, qformer_attention_masks, ): # Process and return inference output output = super().forward( images=global_enc_images, attention_mask=attention_masks, input_ids=input_ids, output_hidden_states=True, conv_offset=conv_offset, image_offset=image_offset, qformer_input_ids=qformer_input_ids, qformer_attention_masks=qformer_attention_masks, ) output_hidden_states = output.hidden_states return output_hidden_states def _training_path( self, global_enc_images, input_ids, labels, attention_masks, conv_offset, image_offset, qformer_input_ids, qformer_attention_masks, ): output = super().forward( images=global_enc_images, attention_mask=attention_masks, input_ids=input_ids, labels=labels, output_hidden_states=True, conv_offset=conv_offset, image_offset=image_offset, qformer_input_ids=qformer_input_ids, qformer_attention_masks=qformer_attention_masks, ) output_hidden_states = output.hidden_states return output, output_hidden_states def _process_hidden_states_multi( self, output_hidden_states, seg_token_masks, conv_offset, infer=False ): if torch.is_tensor(output_hidden_states): output_hidden_states = (output_hidden_states,) hidden_states = [self.model.text_hidden_fcs[0](output_hidden_states[-1])] last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1) all_pred_embeddings_lists = [] for img_idx in range(len(seg_token_masks)): if torch.all(seg_token_masks[img_idx] == False): all_pred_embeddings_lists.append(None) continue pred_embeddings = last_hidden_state[seg_token_masks[img_idx]] seg_token_counts = seg_token_masks[img_idx].int().sum(-1) seg_token_offset = seg_token_counts.cumsum(-1) seg_token_offset = torch.cat( [ torch.zeros(1, device=seg_token_offset.device).long(), seg_token_offset, ], dim=0, ) if not infer: seg_token_offset = seg_token_offset[conv_offset] pred_embeddings_list = [] for i in range(len(seg_token_offset) - 1): start_i, end_i = seg_token_offset[i], seg_token_offset[i + 1] if start_i == end_i: pred_embeddings_list.append(None) continue pred_embeddings_list.append(pred_embeddings[start_i:end_i]) all_pred_embeddings_lists.append(pred_embeddings_list) return hidden_states, all_pred_embeddings_lists # [[], [seg, seg seg]] def _generate_and_postprocess_masks_multi( self, pred_embeddings, image_embeddings, resize_list, label_list, infer=False, ): all_pred_masks = [] for image_idx in range(len(pred_embeddings)): if pred_embeddings[image_idx] is None: all_pred_masks.append([]) continue pred_masks = [] for i, pred_embedding in enumerate(pred_embeddings[image_idx]): if pred_embedding is None: pred_masks.append(None) continue if image_embeddings[image_idx][i] is None: pred_masks.append(None) continue sparse_embeddings, dense_embeddings = ( self.model.grounding_encoder.prompt_encoder( points=None, boxes=None, masks=None, text_embeds=pred_embedding.unsqueeze(1), ) ) sparse_embeddings = sparse_embeddings.to(pred_embedding.dtype) low_res_masks, _ = self.model.grounding_encoder.mask_decoder( image_embeddings=image_embeddings[image_idx][i].unsqueeze(0), image_pe=self.model.grounding_encoder.prompt_encoder.get_dense_pe(), sparse_prompt_embeddings=sparse_embeddings, dense_prompt_embeddings=dense_embeddings, multimask_output=False, ) orig_size = ( label_list[image_idx][i].shape if not infer else label_list[image_idx][i] ) # During inference, we have original size list in place of label list pred_mask = self.model.grounding_encoder.postprocess_masks( low_res_masks, input_size=resize_list[image_idx][i], original_size=orig_size, ) pred_masks.append(pred_mask[:, 0]) all_pred_masks.append(pred_masks) return all_pred_masks def _calculate_losses(self, pred_masks, masks_list, output): return self._compute_loss_components_multi(pred_masks, masks_list, output) def _compute_loss_components_multi(self, pred_masks, masks_list, output): # Initialize loss components ce_loss = output.loss * self.ce_loss_weight image_losses = torch.tensor(0.0, device=ce_loss.device) all_mask_bce_loss = torch.tensor(0.0, device=ce_loss.device) all_mask_dice_loss = torch.tensor(0.0, device=ce_loss.device) # Iterate over batch and compute mask-related losses for img_idx in range(len(pred_masks)): if pred_masks[img_idx] is None: continue num_masks = 0 mask_bce_loss = torch.tensor(0.0, device=ce_loss.device) mask_dice_loss = torch.tensor(0.0, device=ce_loss.device) for batch_idx, pred_mask in enumerate(pred_masks[img_idx]): if pred_mask is None: assert ( masks_list[img_idx][batch_idx] is None ), "no prediction but ground truth mask available" continue if pred_mask.numel() > 0: # Ensure pred_mask is not empty gt_mask = masks_list[img_idx][batch_idx].to(pred_mask.device) # Resize gt_mask to match pred_mask if needed if gt_mask.shape[0] != pred_mask.shape[0]: gt_mask = gt_mask[: pred_mask.shape[0]] assert ( gt_mask.shape[0] == pred_mask.shape[0] ), f"Shape mismatch: gt_mask {gt_mask.shape}, pred_mask {pred_mask.shape}" # Compute Binary Cross-Entropy Loss mask_bce_loss += ( compute_sigmoid_cross_entropy( pred_mask, gt_mask, mask_count=gt_mask.shape[0] ) * gt_mask.shape[0] ) # Compute Dice Loss mask_dice_loss += ( calculate_dice_loss( pred_mask, gt_mask, mask_count=gt_mask.shape[0] ) * gt_mask.shape[0] ) num_masks += gt_mask.shape[0] # Normalize the losses 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) image_losses += mask_bce_loss + mask_dice_loss all_mask_bce_loss += mask_bce_loss all_mask_dice_loss += mask_dice_loss def _has_any_mask(pm): if not pm: return False return any(x is not None for x in pm) num_images_with_masks = sum(1 for pm in pred_masks if _has_any_mask(pm)) mask_loss = ( image_losses / num_images_with_masks if num_images_with_masks > 0 else image_losses ) # Aggregate all loss components total_loss = ce_loss + mask_loss return { "loss": total_loss, "ce_loss": ce_loss, "mask_bce_loss": all_mask_bce_loss, "mask_dice_loss": all_mask_dice_loss, "mask_loss": mask_loss, } @torch.no_grad() def evaluate( self, global_enc_images, grounding_enc_images, input_ids, attention_masks, resize_list, orig_sizes, conv_offset, image_offset, grounding_offset, qformer_input_ids, qformer_attention_masks, max_new_tokens=None, eos_token_id=None, **kwargs, ): if max_new_tokens is None: raise ValueError(f"max_new_tokens not set") generation_outputs = self.generate( images=global_enc_images, input_ids=input_ids, attention_mask=attention_masks, conv_offset=conv_offset, image_offset=image_offset, qformer_input_ids=qformer_input_ids, qformer_attention_masks=qformer_attention_masks, max_new_tokens=max_new_tokens, num_beams=1, output_hidden_states=True, return_dict_in_generate=True, **({"eos_token_id": eos_token_id} if eos_token_id is not None else {}), ) output_hidden_states = generation_outputs.hidden_states generated_output_ids = generation_outputs.sequences num_grounding_images = max( [ (grounding_offset[i + 1] - grounding_offset[i]) for i in range(len(grounding_offset) - 1) ] ) seg_token_masks = self._create_seg_token_mask_multi( generated_output_ids, num_grounding_images, infer=True ) if not seg_token_masks: return generated_output_ids, None image_embeddings = self.get_grounding_encoder_embs(grounding_enc_images) image_embeddings = self._get_image_embs_per_image( image_embeddings, grounding_offset, num_grounding_images ) # Process hidden states hidden_states, predicted_embeddings = self._process_hidden_states_multi( output_hidden_states, seg_token_masks, conv_offset, infer=True ) # Generate and post-process masks pred_masks = self._generate_and_postprocess_masks_multi( predicted_embeddings, image_embeddings, resize_list, orig_sizes, infer=True, ) return generated_output_ids, pred_masks