model/PixDLM.py

from typing import List

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BitsAndBytesConfig, CLIPVisionModel
import copy
from utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
                         DEFAULT_IMAGE_PATCH_TOKEN)

from .llava.language_model.llava_llama import (LlavaLlamaForCausalLM,
                                               LlavaLlamaModel)
from .segment_anything import build_sam_vit_h
from .segment_anything.modeling import (
    MaskDecoder,
    PromptEncoder,
    TwoWayTransformer,
    LayerNorm2d,
    MaskDecoderMultiScale,
    Three_Level_Multi_Scale_Decoder,
)
from utils.matcher import match_pred
from typing import Any, Dict, List, Tuple


def dice_loss(
    inputs: torch.Tensor,
    targets: torch.Tensor,
    num_masks: float,
    scale=1000,             
    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.sum() / (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).sum() / (num_masks + 1e-8)
    return loss

def overlap_loss(inputs: torch.Tensor,
    targets: torch.Tensor,
    num_masks: float,
    batch_seg_token_count: int):
    if num_masks == 0:
        return inputs.sum() * 0
    batch_seg_token_count = batch_seg_token_count.cumsum(-1)  
    batch_seg_token_count = torch.cat(
            [torch.zeros(1).long().cuda(), batch_seg_token_count], dim=0
        )
    loss = 0

    for i in range(len(batch_seg_token_count) -1):
        start_i = batch_seg_token_count[i]
        end_i = batch_seg_token_count[i+1]
        assert end_i <= len(targets), (targets.shape, batch_seg_token_count)
        question_inputs = inputs[start_i:end_i]
        question_targets = targets[start_i:end_i]
        if len(question_targets) == 0:
            continue
        n, h, w = question_inputs.shape
        all_targets = torch.zeros_like(question_targets[0]).bool()
        for target in question_targets:
            all_targets = (all_targets | target.bool())
        bg_area = all_targets < 0
        bg_area = bg_area[None].repeat(n, 1, 1)

        overlap_area = (question_inputs > 0).sum(dim=0)
        overlap_area = overlap_area >= 2

        overlap_area = overlap_area[None].repeat(n, 1, 1)
        weight = torch.ones_like(question_inputs)
        weight[~overlap_area] = 0

        q_loss = F.binary_cross_entropy_with_logits(question_inputs, question_targets, weight=weight, reduction="none")
        q_loss = q_loss.flatten(1, 2).mean(1).sum() 
        loss = loss + q_loss
    loss = loss / (num_masks + 1e-8)
    return loss


class PixDLMMetaModel:
    def __init__(
        self,
        config,
        **kwargs,
    ):  
        super(PixDLMMetaModel, self).__init__(config)
        self.logger = kwargs.get("logger", None)
        self.local_rank = kwargs.get("local_rank", 1)
        self.config = config
           
        self.sam2_config = 'configs/sam2.1/sam2.1_hiera_l.yaml'
                                                                       
        if "three_level_multi_scale_decoder" in kwargs:
            self.config.three_level_multi_scale_decoder = kwargs["three_level_multi_scale_decoder"]
        
        if not hasattr(self.config, "train_mask_decoder"):
            self.config.train_mask_decoder = kwargs["train_mask_decoder"]
            self.config.out_dim = kwargs["out_dim"]
            self.vision_pretrained = kwargs.get("vision_pretrained", None)
        else:
            self.vision_pretrained = kwargs.get("vision_pretrained", None)
            self.initialize_pixdlm_modules(self.config)

    def initialize_pixdlm_modules(self, config):
             
        if self.config.vision_tower_for_mask:
            
            prompt_embed_dim = 256 
            image_size = config.resize_vision_tower_size
            mask_decoder_transformer_depth = 2
            if self.local_rank == 0 and self.logger is not None:
                self.logger.info('--------build_sam_decoder--------')
                self.logger.info('--------sam decoder image size {}--------'.format(image_size))
            vit_patch_size = 14
            image_embedding_size = image_size // vit_patch_size
            self.prompt_encoder=PromptEncoder(
                        embed_dim=prompt_embed_dim,
                        image_embedding_size=(image_embedding_size, image_embedding_size),
                        input_image_size=(image_size, image_size),
                        mask_in_chans=16,
                    )
             
                                                     
                                                                                         
            decoder_cls = (
                Three_Level_Multi_Scale_Decoder
                if getattr(config, "three_level_multi_scale_decoder", False)
                else MaskDecoderMultiScale
            )

            self.mask_decoder = decoder_cls(
                num_multimask_outputs=3,
                transformer=TwoWayTransformer(
                    depth=mask_decoder_transformer_depth,
                    embedding_dim=prompt_embed_dim,
                    mlp_dim=2048,
                    num_heads=8,
                ),
                transformer_dim=prompt_embed_dim,
                iou_head_depth=3,
                iou_head_hidden_dim=256,
                image_feature_scale_num=config.image_feature_scale_num
            ) 
            
            
            embed_dim = self.config.hidden_size 
            out_chans = prompt_embed_dim
            self.image_feature_neck = nn.Sequential(
                nn.Conv2d(
                    embed_dim,
                    out_chans,
                    kernel_size=1,
                    bias=False,
                ),
                LayerNorm2d(out_chans),
                nn.Conv2d(
                    out_chans,
                    out_chans,
                    kernel_size=3,
                    padding=1,
                    bias=False,
                ),
                LayerNorm2d(out_chans),
            )
            self.sam_to_embed_conv = nn.Sequential(  
            nn.Conv2d(256, embed_dim, kernel_size=1, bias=False),  
            LayerNorm2d(embed_dim),  
        )

        else:
            self.visual_model = build_sam_vit_h(self.vision_pretrained)
            for param in self.visual_model.parameters():
                param.requires_grad = False
            if config.train_mask_decoder:
                self.visual_model.mask_decoder.train()
                for param in self.visual_model.mask_decoder.parameters():
                    param.requires_grad = True

                          
        in_dim = config.hidden_size
        out_dim = config.out_dim
        text_fc = [
            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_fc)])
        self.text_hidden_fcs.train()
        for param in self.text_hidden_fcs.parameters():
            param.requires_grad = True


class PixDLMModel(PixDLMMetaModel, LlavaLlamaModel):
    def __init__(
        self,
        config,
        **kwargs,
    ):
        super(PixDLMModel, 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 PixDLMForCausalLM(LlavaLlamaForCausalLM):
    def __init__(
        self,
        config,
        **kwargs,
    ):
                                    
                                 
                                                      
                                     
        kwargs.setdefault("image_feature_scale_num", 3)
        kwargs.setdefault("pad_train_clip_images", True)
        kwargs.setdefault("resize_vision_tower", True)
        kwargs.setdefault("resize_vision_tower_size", 448)
        kwargs.setdefault("vision_tower_for_mask", True)
        kwargs.setdefault("separate_mm_projector", False)
        self.logger = kwargs.get("logger", None)
        config.resize_vision_tower = kwargs.get("resize_vision_tower", False)
        config.resize_vision_tower_size = kwargs.get("resize_vision_tower_size", 224)
        config.pad_train_clip_images = kwargs.get("pad_train_clip_images", False)
        config.vision_tower_for_mask = kwargs.get("vision_tower_for_mask", False)
        config.separate_mm_projector = kwargs.get("separate_mm_projector", False)
        config.mm_projector_hidden_dim = 2
        config.mm_projector_out_dim = 1
        self.image_feature_scale_num = kwargs.get("image_feature_scale_num", 1)
        config.image_feature_scale_num = kwargs.get("image_feature_scale_num", 1)
        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)
        else:
            config.mm_vision_tower = config.vision_tower
                                                                  
            self.ce_loss_weight = kwargs.pop("ce_loss_weight", getattr(config, "ce_loss_weight", 1.0))
            self.dice_loss_weight = kwargs.pop("dice_loss_weight", getattr(config, "dice_loss_weight", 1.0))
            self.bce_loss_weight = kwargs.pop("bce_loss_weight", getattr(config, "bce_loss_weight", 1.0))
        
        self.vision_tower_for_mask = kwargs.get("vision_tower_for_mask", False)
        self.seg_token_idx = kwargs.pop("seg_token_idx")
        self.seg_token_num = kwargs.get("seg_token_num", 1)
        self.tokenizer = kwargs.get("tokenizer", None)
        self.local_rank = kwargs.get("local_rank", 1)
        self.pad_train_clip_images = kwargs.get("pad_train_clip_images", False)
        self.masks_process_with_clip = kwargs.get("masks_process_with_clip", False)
          
        self.is_multipath_encoder = kwargs.get("is_multipath_encoder", True)
        config.is_multipath_encoder=self.is_multipath_encoder
        self.freeze_vision=kwargs.get("freeze_vision", False)
        config.freeze_vision=self.freeze_vision
                                                          
                                                          
        kwargs_value = kwargs.get("three_level_multi_scale_decoder", None)
        config_value = getattr(config, "three_level_multi_scale_decoder", False)
        if kwargs_value is not None:
            self.three_level_multi_scale_decoder = kwargs_value
        else:
            self.three_level_multi_scale_decoder = config_value
        config.three_level_multi_scale_decoder = self.three_level_multi_scale_decoder
        
                       
        logger = kwargs.get("logger", None)
        if isinstance(self.seg_token_idx, list):
            if self.local_rank == 0 and logger is not None:
                logger.info("Initialize multi-seg scalar")
            seg_token_num = len(self.seg_token_idx)
            scalar = 1 / seg_token_num
            self.multiseg_scalar = [torch.nn.Parameter(torch.ones([]) * scalar) for _ in range(seg_token_num)]
        if self.image_feature_scale_num > 1:
            scalar = 1 / self.image_feature_scale_num
            self.multiscale_scalar = [torch.nn.Parameter(torch.ones([]) * scalar) for _ in range(self.image_feature_scale_num)]
        super().__init__(config)
        self.model = PixDLMModel(config, **kwargs)
           
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
                                                       
        self.post_init()
        self.iter = 0
        self.iter1 = 0
         
        if config.resize_vision_tower_size != 224:
            if self.local_rank == 0 and self.logger is not None:
                self.logger.info('--------mm_projector requires grad--------')
            for n, p in self.model.named_parameters():
                if any([x in n for x in ["mm_projector"]]):
                    p.requires_grad = True
                    
    def get_visual_embs(self, pixel_values: torch.FloatTensor):
        with torch.no_grad():
            image_embeddings_list = []
            for i in range(pixel_values.shape[0]):
                torch.cuda.empty_cache()
                image_embeddings = self.model.visual_model.image_encoder(
                    pixel_values[i].unsqueeze(0)
                )
                image_embeddings_list.append(image_embeddings)
            torch.cuda.empty_cache()
            image_embeddings = torch.cat(image_embeddings_list, 0)
        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],
        resize_list: List[tuple],
        inference: bool = False,
        clip_resize_list = None,
        txt_feat=None,
        **kwargs,
    ):
                                                                 
        multi_reason_list = kwargs.get('multi_reason_list', None)
        if not self.vision_tower_for_mask:
            image_embeddings = self.get_visual_embs(images)
        batch_size = images.shape[0]
        assert batch_size == len(offset) - 1
        

        if isinstance(self.seg_token_idx, list):
            seg_token_num = self.seg_token_num
            seg_token_mask = torch.zeros_like(input_ids[:, 1:]).bool()
            for seg_token_idx in self.seg_token_idx:
                seg_token_mask = seg_token_mask | (input_ids[:, 1:] == seg_token_idx)  
        else:
            seg_token_num = self.seg_token_num
            seg_token_mask = input_ids[:, 1:] == self.seg_token_idx

                                  
                  
       
        seg_token_mask = torch.cat(
            [
                seg_token_mask,
                torch.zeros((seg_token_mask.shape[0], 1)).bool().cuda(),
            ],
            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()
            extend_clip_resize_list = [clip_resize_list[0]] * length
            output_hidden_states = []
            output_image_features = []
            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,
                    clip_resize_list=extend_clip_resize_list,
                    txt_feat=txt_feat
                )
                num_image_tokens = self._last_visual_token_num
                seg_token_mask = torch.cat(
                    [torch.zeros((seg_token_mask.shape[0], num_image_tokens)).bool().cuda(), seg_token_mask],
                    dim=1,
                )

                output_image_feature_i = torch.stack(output_i.image_features, dim=0)
                                                                              
                output_hidden_states.append(output_i.hidden_states)
                output_image_features.append(output_image_feature_i)
                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_image_features = torch.cat(output_image_features, dim=1)
            output = None

        else:
            images_clip_list = []
            extend_clip_resize_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()
                )
                extend_clip_resize_list.extend([clip_resize_list[i]] * (end_i - start_i))
                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,
                txt_feat=txt_feat,
                clip_resize_list=extend_clip_resize_list
            )
            num_image_tokens = self._last_visual_token_num
            seg_token_mask = torch.cat(
                [torch.zeros((seg_token_mask.shape[0], num_image_tokens)).bool().cuda(), seg_token_mask],
                dim=1,
            )

            output_hidden_states = output.hidden_states
            output_image_features = output.image_features

        hidden_states = []

         
        
        assert len(self.model.text_hidden_fcs) == 1
        hidden_states.append(self.model.text_hidden_fcs[0](output_hidden_states[-1]))
                                                        
        last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1)
                                                                
                                                       
        pred_embeddings = last_hidden_state[seg_token_mask]
        seg_token_counts = seg_token_mask.int().sum(-1)          

        seg_token_offset = seg_token_counts.cumsum(-1)
        seg_token_offset = torch.cat(
            [torch.zeros(1).long().cuda(), seg_token_offset], dim=0
        )
        
        seg_token_offset = seg_token_offset[offset]
        feat_scale_num = self.image_feature_scale_num
        pred_embeddings_ = []
        batch_seg_token_counts = []
        for i in range(len(seg_token_offset) - 1):
            start_i, end_i = seg_token_offset[i], seg_token_offset[i + 1]
            batch_pred_embeddings = pred_embeddings[start_i:end_i]
            batch_seg_token_counts.append(seg_token_counts[offset[i]:offset[i+1]] // seg_token_num*feat_scale_num)
            assert len(batch_pred_embeddings) % seg_token_num == 0
            batch_pred_embeddings = batch_pred_embeddings.view(len(batch_pred_embeddings) // (seg_token_num*feat_scale_num), feat_scale_num, seg_token_num, batch_pred_embeddings.shape[-1])                            
            if seg_token_num > 1:
                fused_batch_pred_embeddings = batch_pred_embeddings[:, :, 0] * 0                   
                for i in range(seg_token_num):
                    fused_batch_pred_embeddings = fused_batch_pred_embeddings + self.multiseg_scalar[i] * batch_pred_embeddings[:, :, i]
                batch_pred_embeddings = fused_batch_pred_embeddings
            else:
                batch_pred_embeddings = batch_pred_embeddings[:, :, 0]
            
            pred_embeddings_.append(batch_pred_embeddings)
       
        pred_embeddings = pred_embeddings_                                                 
        multi_scale_num = len(output_image_features)


        if not inference:
            output_image_features = torch.stack(output_image_features, dim=0)                     
        img_embeddings = output_image_features.flatten(1, 2)                     
        img_token_mask = torch.ones(output_image_features.shape[1], output_image_features.shape[2]).to(seg_token_mask)
        img_token_counts = img_token_mask.int().sum(-1) 
        patch_count = int(img_token_counts[0])
        
        patch_size = int(patch_count**0.5)
        img_token_offset = img_token_counts.cumsum(-1)                 
        img_token_offset = torch.cat(
            [torch.zeros(1).long().cuda(), img_token_offset], dim=0                    
        )
        img_token_offset = img_token_offset[offset]                 
        img_embeddings_ = []
        single_img_embeddings = []                  
        for i in range(len(img_token_offset) - 1):
            start_i, end_i = img_token_offset[i], img_token_offset[i + 1]
            question_num = pred_embeddings_[i].shape[0] 
            img_num = img_embeddings[:, start_i:end_i].shape[1] // patch_count
            single_img_embeddings.append(img_embeddings[:, start_i:end_i].view(multi_scale_num, img_num, patch_count, img_embeddings.shape[-1]).permute(0, 1, 3, 2).view(multi_scale_num, img_num, img_embeddings.shape[-1], patch_size, patch_size)[:, 0])
            if question_num == 0:
                batch_img_embeddings = torch.zeros(multi_scale_num, 0, 4096, patch_size, patch_size).to(img_embeddings)
            else:
                batch_img_embeddings = img_embeddings[:, start_i:end_i].view(multi_scale_num, img_num, patch_count, img_embeddings.shape[-1])
                batch_img_embeddings = batch_img_embeddings.permute(0, 1, 3, 2).view(multi_scale_num, img_num, img_embeddings.shape[-1], patch_size, patch_size)
            img_embeddings_.append(batch_img_embeddings)

        img_embeddings = img_embeddings_

        
        multimask_output = False
        pred_masks = []
        mask_scores = []
        pred_depths = []
         
        for i in range(len(pred_embeddings)):
            if self.vision_tower_for_mask:
                sparse_embeddings, dense_embeddings = self.model.prompt_encoder(points=None, boxes=None,masks=None,text_embeds=pred_embeddings[i],)                               
                sparse_embeddings = sparse_embeddings.to(pred_embeddings[i].dtype)
                                                                                            
                _img_embeddings = None
                use_sam_multilayer = False
                if (
                    getattr(self, "three_level_multi_scale_decoder", False)
                    and hasattr(self.model, "vision_tower")
                    and hasattr(self.model.vision_tower, "forward_sam_multilayer_features")
                ):
                                        
                    if i < len(images_clip):
                        img_i = images_clip[i].unsqueeze(0)                
                        sam_multilayer_feats = self.model.vision_tower.forward_sam_multilayer_features(img_i)
                                                                    
                        if isinstance(sam_multilayer_feats, list) and len(sam_multilayer_feats) > 0:
                            use_sam_multilayer = True
                                                           
                            feats_to_use = sam_multilayer_feats[: self.image_feature_scale_num]
                            sam_feats_list = []
                            for feat in feats_to_use:
                                                               
                                feat = feat.squeeze(0)
                                if hasattr(self.model, "image_feature_neck"):
                                    embed_dim = self.model.image_feature_neck[0].in_channels
                                    target_dtype = self.model.image_feature_neck[0].weight.dtype
                                    target_device = feat.device

                                                                            
                                    if feat.shape[0] != embed_dim:  
                                        feat = feat.to(target_dtype)  
                                        feat = self.model.sam_to_embed_conv(feat.unsqueeze(0)).squeeze(0)  
                                    else:  
                                        feat = feat.to(target_dtype)

                                                                                         
                                    feat = self.model.image_feature_neck(feat.unsqueeze(0)).squeeze(0)
                                sam_feats_list.append(feat)

                            if len(sam_feats_list) > 0:
                                _img_embeddings = torch.stack(sam_feats_list, dim=0)

                                                   
                if _img_embeddings is None:
                    print("hh")
                    _img_embeddings = self.model.image_feature_neck(single_img_embeddings[i])                     

                out_size = 128
                low_res_masks = torch.zeros([sparse_embeddings.shape[0], 1, out_size, out_size]).to(_img_embeddings)
                if self.image_feature_scale_num > 1:  
                    for l in range(self.image_feature_scale_num):  
                                       
                        feat_h, feat_w = _img_embeddings[l].shape[1], _img_embeddings[l].shape[2]  
                        
                                                         
                        dense_embeddings_adjusted = dense_embeddings  
                        if dense_embeddings.shape[-1] != feat_w or dense_embeddings.shape[-2] != feat_h:  
                            dense_embeddings_adjusted = F.interpolate(  
                                dense_embeddings.float(),  
                                size=(feat_h, feat_w),  
                                mode='bilinear',  
                                align_corners=False  
                            ).to(dense_embeddings.dtype)  
                        
                                  
                        default_embedding_size = self.model.prompt_encoder.image_embedding_size[0]  
                        if feat_h != default_embedding_size or feat_w != default_embedding_size:  
                            default_pe = self.model.prompt_encoder.get_dense_pe()  
                            image_pe = F.interpolate(  
                                default_pe.float(),  
                                size=(feat_h, feat_w),  
                                mode='bilinear',  
                                align_corners=False  
                            ).to(default_pe.dtype)  
                        else:  
                            image_pe = self.model.prompt_encoder.get_dense_pe()  
                        
                                                   
                        l_low_res_masks, iou_predictions = self.model.mask_decoder(  
                            image_embeddings=_img_embeddings[l].unsqueeze(0),  
                            image_pe=image_pe,  
                            sparse_prompt_embeddings=sparse_embeddings[:, l].unsqueeze(1),  
                            dense_prompt_embeddings=dense_embeddings_adjusted,  
                            multimask_output=multimask_output,  
                            previous_masks=l_low_res_masks if l>0 else None,  
                            level_num=l  
                        )  
                        low_res_masks = low_res_masks + self.multiscale_scalar[l] * F.interpolate(l_low_res_masks.float(), (out_size, out_size),mode="bilinear",align_corners=False,).to(l_low_res_masks)  
                else:  
                                    
                    feat_h, feat_w = _img_embeddings[0].shape[1], _img_embeddings[0].shape[2]  
                    
                    dense_embeddings_adjusted = dense_embeddings  
                    if dense_embeddings.shape[-1] != feat_w or dense_embeddings.shape[-2] != feat_h:  
                        dense_embeddings_adjusted = F.interpolate(  
                            dense_embeddings.float(),  
                            size=(feat_h, feat_w),  
                            mode='bilinear',  
                            align_corners=False  
                        ).to(dense_embeddings.dtype)  
                    
                    default_embedding_size = self.model.prompt_encoder.image_embedding_size[0]  
                    if feat_h != default_embedding_size or feat_w != default_embedding_size:  
                        default_pe = self.model.prompt_encoder.get_dense_pe()  
                        image_pe = F.interpolate(  
                            default_pe.float(),  
                            size=(feat_h, feat_w),  
                            mode='bilinear',  
                            align_corners=False  
                        ).to(default_pe.dtype)  
                    else:  
                        image_pe = self.model.prompt_encoder.get_dense_pe()  
                    
                    low_res_masks, iou_predictions = self.model.mask_decoder(  
                        image_embeddings=_img_embeddings[0].unsqueeze(0),  
                        image_pe=image_pe,  
                        sparse_prompt_embeddings=sparse_embeddings[:, 0].unsqueeze(1),  
                        dense_prompt_embeddings=dense_embeddings_adjusted,  
                        multimask_output=multimask_output,  
                    )
                pred_mask = self.postprocess_masks(
                    low_res_masks,
                    input_size=clip_resize_list[i],
                    original_size=label_list[i].shape,
                )

            

            else:

                (
                    sparse_embeddings,
                    dense_embeddings,
                ) = self.model.visual_model.prompt_encoder(
                    points=None,
                    boxes=None,
                    masks=None,
                    text_embeds=pred_embeddings[i],
                )
                sparse_embeddings = sparse_embeddings.to(pred_embeddings[i].dtype)
                low_res_masks, iou_predictions = self.model.visual_model.mask_decoder(
                    image_embeddings=image_embeddings[i].unsqueeze(0),
                    image_pe=self.model.visual_model.prompt_encoder.get_dense_pe(),
                    sparse_prompt_embeddings=sparse_embeddings,
                    dense_prompt_embeddings=dense_embeddings,
                    multimask_output=multimask_output,
                )

                
                pred_mask = self.model.visual_model.postprocess_masks(
                    low_res_masks,
                    input_size=resize_list[i],
                    original_size=label_list[i].shape,
                )
                pred_depths.append([])
            

            
            pred_masks.append(pred_mask[:, 0])
            mask_score = (pred_mask[:, 0].sigmoid().flatten(1) * (pred_mask[:, 0] > 0).flatten(1)).sum(1) / ((pred_mask[:, 0] > 0).flatten(1).sum(1) + 1e-6)
            mask_scores.append(mask_score)


        model_output = output
        gt_masks = masks_list


        if inference:
            return {
                "pred_masks": pred_masks,
                "gt_masks": gt_masks,
                "batch_seg_token_counts": batch_seg_token_counts,
                "mask_scores": mask_scores,
            }

        output = model_output.logits
        ce_loss = model_output.loss
        ce_loss = ce_loss * self.ce_loss_weight
        loss = ce_loss
        mask_bce_loss = pred_masks[0].sum() * 0
        mask_dice_loss = pred_masks[0].sum() * 0
        mask_overlap_loss = pred_masks[0].sum() * 0
        num_masks = 0

        for batch_idx in range(len(pred_masks)):
            gt_mask = gt_masks[batch_idx]
            pred_mask = pred_masks[batch_idx]
            batch_seg_token_count = batch_seg_token_counts[batch_idx]
            assert (
                gt_mask.shape[0] == pred_mask.shape[0]
            ), "gt_mask.shape: {}, pred_mask.shape: {}".format(
                gt_mask.shape, pred_mask.shape
            )
             
      
            mask_bce_loss += (
                sigmoid_ce_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
                * gt_mask.shape[0]
            )
            mask_dice_loss += (
                dice_loss(pred_mask, gt_mask, num_masks=gt_mask.shape[0])
                * gt_mask.shape[0]
            )
                                    
                                                                                               
                                        
                       
            num_masks += gt_mask.shape[0]
             
            
        
        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 = loss + mask_loss


        
        return {
            "loss": loss,
            "ce_loss": ce_loss,
            "mask_bce_loss": mask_bce_loss,
            "mask_dice_loss": mask_dice_loss,
            "mask_loss": mask_loss,
        }

    def evaluate(
        self,
        images_clip,
        images,
        input_ids,
        resize_list,
        clip_resize_list,
        original_size_list,
        max_new_tokens=32,
        tokenizer=None,
        txt_feat=None
                            
    ):
        import time  

        all_pred_embeddings = []
        all_output_ids = []
        batch_seg_token_counts = []
          
        all_pred_embeddings = []  
        all_output_ids = []  
        batch_seg_token_counts = []  
        
                  
        total_start_time = time.time()  
        
                  
        image_encoding_time = 0  
        text_gen_time = 0  
        embedding_process_time = 0  
        mask_gen_time = 0  
        num_images = len(input_ids)  

        with torch.no_grad():
            encoding_start = time.time()  
                                                                                                                                   
            output_image_features, _, _, _, _, _ = self.prepare_inputs_labels_for_multimodal(
            input_ids=torch.ones(images_clip.shape[0], 2).long().cuda() * IMAGE_TOKEN_INDEX,
            attention_mask=None,
            past_key_values=None,
            labels=None,
            images=images_clip,
            clip_resize_list=clip_resize_list,
            txt_feat=txt_feat,
        )
            multi_scale_num = self.image_feature_scale_num
            output_image_features = torch.stack(output_image_features, dim=0)
            image_encoding_time = time.time() - encoding_start  
            for idx, input_id in enumerate(input_ids):
                if 0 in input_id:
                    unk_start = torch.where(input_id==0)[0].min()
                    _input_id = input_id[:unk_start]
                else:
                    _input_id = input_id
                gen_start = time.time()  
                outputs = self.generate(
                    images=images_clip,
                    input_ids=_input_id[None],
                    max_new_tokens=max_new_tokens,
                    num_beams=1,
                    output_hidden_states=True,
                    return_dict_in_generate=True,
                    clip_resize_list=clip_resize_list
                )
                text_gen_time += (time.time() - gen_start)  
                embed_start = time.time()  
                output_hidden_states = outputs.hidden_states[-1]
                                                                          
                output_ids = outputs.sequences
                all_output_ids.append(output_ids)
       
                if isinstance(self.seg_token_idx, list):
                    seg_token_num = self.seg_token_num
                    seg_token_mask = torch.zeros_like(output_ids[:, 1:]).bool()
                             
                      
                    for seg_token_idx in self.seg_token_idx:
                        seg_token_mask = seg_token_mask | (output_ids[:, 1:] == seg_token_idx)  
                
                else:
                    seg_token_num = self.seg_token_num
                    seg_token_mask = output_ids[:, 1:] == self.seg_token_idx
                                                                                                              
                num_image_tokens = self._last_visual_token_num
                seg_token_mask = torch.cat(
                    [
                        torch.zeros((seg_token_mask.shape[0], num_image_tokens)).bool().cuda(),
                        seg_token_mask,
                    ],
                    dim=1,
                )

                hidden_states = []
            
                assert len(self.model.text_hidden_fcs) == 1
                hidden_states.append(self.model.text_hidden_fcs[0](output_hidden_states))
                feat_scale_num = self.image_feature_scale_num
                last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1)
                                                               
                                                                     
                pred_embeddings = last_hidden_state[seg_token_mask]

                if len(pred_embeddings) % (seg_token_num*feat_scale_num) != 0:
                    seg_token_mask = (seg_token_mask*0).bool()
                    pred_embeddings = last_hidden_state[seg_token_mask]

                seg_token_counts = seg_token_mask.int().sum(-1)          
                seg_token_offset = seg_token_counts.cumsum(-1)
                seg_token_offset = torch.cat(
                    [torch.zeros(1).long().cuda(), seg_token_offset], dim=0
                )
                seg_token_offset = seg_token_offset[[0, len(seg_token_offset)-1]]
                pred_embeddings_ = []
                
                
                for i in range(len(seg_token_offset) - 1):
                    start_i, end_i = seg_token_offset[i], seg_token_offset[i + 1]
                    batch_pred_embeddings = pred_embeddings[start_i:end_i]
                                                          
                    assert len(batch_pred_embeddings) % (seg_token_num*feat_scale_num) == 0
                    batch_pred_embeddings = batch_pred_embeddings.view(len(batch_pred_embeddings) // (seg_token_num*feat_scale_num), feat_scale_num, seg_token_num, batch_pred_embeddings.shape[-1])                            
                    if seg_token_num > 1:
                        fused_batch_pred_embeddings = batch_pred_embeddings[:, :, 0] * 0                   
                        for i in range(seg_token_num):
                            fused_batch_pred_embeddings = fused_batch_pred_embeddings + self.multiseg_scalar[i] * batch_pred_embeddings[:, :, i]
                        batch_pred_embeddings = fused_batch_pred_embeddings
                    else:
                        batch_pred_embeddings = batch_pred_embeddings[:, :, 0]
                    pred_embeddings_.append(batch_pred_embeddings)
                batch_seg_token_counts.append(len(batch_pred_embeddings))
                pred_embeddings = pred_embeddings_                                                 
                all_pred_embeddings.extend(pred_embeddings)
                embedding_process_time += (time.time() - embed_start)  
            
            batch_seg_token_counts = [torch.tensor(batch_seg_token_counts).to(seg_token_counts)]
            pred_embeddings = [torch.cat(all_pred_embeddings)]
            mask_start = time.time() 
            multimask_output = False
            pred_masks = []
            mask_scores = []
            if not self.vision_tower_for_mask:
                image_embeddings = self.get_visual_embs(images)
            else:
                
                img_embeddings = output_image_features.flatten(1, 2)                                       
                img_embeddings = [img_embeddings.view(multi_scale_num, 1024, img_embeddings.shape[-1]).permute(0, 2, 1).view(multi_scale_num, img_embeddings.shape[-1], 32, 32)]
             
            for i in range(len(pred_embeddings)):
                if self.vision_tower_for_mask:
                    sparse_embeddings, dense_embeddings = self.model.prompt_encoder(points=None, boxes=None,masks=None,text_embeds=pred_embeddings[i],)
                    sparse_embeddings = sparse_embeddings.to(pred_embeddings[i].dtype)
                    _img_embeddings = self.model.image_feature_neck(img_embeddings[i]) 
                    out_size = 128
                    low_res_masks = torch.zeros([sparse_embeddings.shape[0], 1, out_size, out_size]).to(_img_embeddings)
                    if self.image_feature_scale_num > 1:
                        for l in range(self.image_feature_scale_num):
                            l_low_res_masks, iou_predictions = self.model.mask_decoder(image_embeddings=_img_embeddings[l].unsqueeze(0), image_pe=self.model.prompt_encoder.get_dense_pe(), sparse_prompt_embeddings=sparse_embeddings[:, l].unsqueeze(1), dense_prompt_embeddings=dense_embeddings, multimask_output=multimask_output, previous_masks=l_low_res_masks if l>0 else None, level_num=l)
                            low_res_masks = low_res_masks + self.multiscale_scalar[l] * F.interpolate(l_low_res_masks.float(), (out_size, out_size),mode="bilinear",align_corners=False,).to(l_low_res_masks)
                    else:
                        low_res_masks, iou_predictions = self.model.mask_decoder(image_embeddings=_img_embeddings[0].unsqueeze(0), image_pe=self.model.prompt_encoder.get_dense_pe(), sparse_prompt_embeddings=sparse_embeddings[:, 0].unsqueeze(1), dense_prompt_embeddings=dense_embeddings, multimask_output=multimask_output, )


                    pred_mask = self.postprocess_masks(
                        low_res_masks,
                        input_size=clip_resize_list[i],
                        original_size=original_size_list[i],
                    )
                else:
                    (
                        sparse_embeddings,
                        dense_embeddings,
                    ) = self.model.visual_model.prompt_encoder(
                        points=None,
                        boxes=None,
                        masks=None,
                        text_embeds=pred_embeddings[i],
                    )

                    sparse_embeddings = sparse_embeddings.to(pred_embeddings[i].dtype)
                    low_res_masks, iou_predictions = self.model.visual_model.mask_decoder(
                        image_embeddings=image_embeddings[i].unsqueeze(0),
                        image_pe=self.model.visual_model.prompt_encoder.get_dense_pe(),
                        sparse_prompt_embeddings=sparse_embeddings,
                        dense_prompt_embeddings=dense_embeddings,
                        multimask_output=multimask_output,
                    )
                    pred_mask = self.model.visual_model.postprocess_masks(
                        low_res_masks,
                        input_size=resize_list[i],
                        original_size=original_size_list[i],
                    )
                
                pred_masks.append(pred_mask[:, 0])
                mask_score = (pred_mask[:, 0].sigmoid().flatten(1) * (pred_mask[:, 0] > 0).flatten(1)).sum(1) / ((pred_mask[:, 0] > 0).flatten(1).sum(1) + 1e-6)
                mask_scores.append(mask_score)
            mask_gen_time = time.time() - mask_start 
        total_time = time.time() - total_start_time 
        avg_time = total_time / num_images  
        fps = num_images / total_time  
        
                  
        print(f"\n{'='*50}")  
        print(f"Inference Speed Statistics")  
        print(f"{'='*50}")  
        print(f"Total images: {num_images}")  
        print(f"Total time: {total_time:.4f}s")  
        print(f"  - Image encoding: {image_encoding_time:.4f}s ({image_encoding_time/total_time*100:.1f}%)")  
        print(f"  - Text generation: {text_gen_time:.4f}s ({text_gen_time/total_time*100:.1f}%)")  
        print(f"  - Embedding processing: {embedding_process_time:.4f}s ({embedding_process_time/total_time*100:.1f}%)")  
        print(f"  - Mask generation: {mask_gen_time:.4f}s ({mask_gen_time/total_time*100:.1f}%)")  
        print(f"Average time per image: {avg_time:.4f}s")  
        print(f"FPS: {fps:.2f}")  
        print(f"{'='*50}\n") 
        return all_output_ids, pred_masks, batch_seg_token_counts, mask_scores


    def postprocess_masks(
            self,
            masks: torch.Tensor,
            input_size: Tuple[int, ...],
            original_size: Tuple[int, ...],
        ) -> torch.Tensor:
            """
            Remove padding and upscale masks to the original image size.

            Arguments:
            masks (torch.Tensor): Batched masks from the mask_decoder,
                in BxCxHxW format.
            input_size (tuple(int, int)): The size of the image input to the
                model, in (H, W) format. Used to remove padding.
            original_size (tuple(int, int)): The original size of the image
                before resizing for input to the model, in (H, W) format.

            Returns:
            (torch.Tensor): Batched masks in BxCxHxW format, where (H, W)
                is given by original_size.
            """
           
            target_size = max(input_size)
            dtype = masks.dtype
            if self.vision_tower_for_mask:
                masks = F.interpolate(
                    masks.float(),
                    (target_size, target_size),
                    mode="bilinear",
                    align_corners=False,
                )
            
            if not self.masks_process_with_clip:
                assert input_size[0] <= target_size
                assert input_size[1] <= target_size
                masks = masks[..., : input_size[0], : input_size[1]]
                masks = F.interpolate(
                    masks, original_size, mode="bilinear", align_corners=False
                )
            
            masks = masks.to(dtype)
              
            return masks