code/model/PRIMA.py

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