projects/llava_sam2/models/lisa.py

import torch
import torch.nn as nn

from mmengine.model import BaseModel

from xtuner.registry import BUILDER
from xtuner.model.utils import get_peft_model_state_dict


class LisaModel(BaseModel):
    def __init__(self,
                 mllm,
                 tokenizer,
                 grounding_encoder,
                 loss_mask=None,
                 loss_dice=None,):
        super(LisaModel, self).__init__()
        self.mllm = BUILDER.build(mllm)

        if self.mllm.use_llm_lora:
            self.mllm.model.language_model.base_model.model.lm_head.requires_grad_(True)
            self.mllm.model.language_model.base_model.model.model.embed_tokens.requires_grad_(True)

        self.tokenizer = BUILDER.build(tokenizer)
        self._add_special_tokens()
        self.grounding_encoder = BUILDER.build(grounding_encoder)
        self.grounding_encoder.requires_grad_(False)
        self.grounding_encoder.mask_decoder.requires_grad_(True)
        
        in_dim = self.mllm.model.config.llm_config.hidden_size
        out_dim = self.grounding_encoder.mask_decoder.transformer_dim
        self.text_hidden_fcs = nn.Sequential(
            nn.Linear(in_dim, in_dim), nn.ReLU(inplace=True),
            nn.Linear(in_dim, out_dim), nn.Dropout(0.0)
        )

        self.loss_mask = BUILDER.build(loss_mask)
        self.loss_dice = BUILDER.build(loss_dice)
        
    def _add_special_tokens(self):
        special_tokens = ['[SEG]']
        num_new_tokens = self.tokenizer.add_tokens(
            special_tokens, special_tokens=True)
        if num_new_tokens > 0:
            self.mllm.model.language_model.resize_token_embeddings(len(self.tokenizer))

        self.seg_token_idx = self.tokenizer("[SEG]", add_special_tokens=False).input_ids[0]

    def _generate_and_postprocess_masks(self, pred_embeddings, image_embeddings, resize_list=None, orig_size_list=None):
        pred_masks = []
        for i, pred_embedding in enumerate(pred_embeddings):
            sparse_embeddings, dense_embeddings = self.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.grounding_encoder.mask_decoder(
                image_embeddings=image_embeddings[i].unsqueeze(0),
                image_pe=self.grounding_encoder.prompt_encoder.get_dense_pe(),
                sparse_prompt_embeddings=sparse_embeddings, dense_prompt_embeddings=dense_embeddings,
                multimask_output=False, )
            
            pred_mask = self.grounding_encoder.postprocess_masks(
                low_res_masks, input_size=resize_list[i], original_size=orig_size_list[i], )
            pred_masks.append(pred_mask[:, 0])
        return pred_masks
    
    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
        return super().load_state_dict(state_dict, strict, assign)
    
    def state_dict(self, *args, **kwargs):
        state_dict = super().state_dict(*args, **kwargs)
        from collections import OrderedDict

        to_return = OrderedDict()
        # Step 1. visual_encoder
        if self.mllm.use_visual_encoder_lora:
            to_return.update(
                get_peft_model_state_dict(
                    self.mllm.model.vision_model, state_dict=state_dict))
        elif not self.mllm.freeze_visual_encoder:
            to_return.update({
                k: v
                for k, v in state_dict.items() if 'visual_encoder.' in k
            })
        # Step 2. LLM
        if self.mllm.use_llm_lora:
            to_return.update(
                get_peft_model_state_dict(self.mllm.model.language_model, state_dict=state_dict))
        elif not self.mllm.freeze_llm:
            to_return.update(
                {k: v
                 for k, v in state_dict.items() if 'llm.' in k})
        # Step 3. Projector
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'mlp1.' in k})
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'grounding_encoder.mask_decoder.' in k})
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'text_hidden_fcs.' in k})
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'lm_head.weight' in k})
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'embed_tokens.weight' in k})
        return to_return
    
    def forward(self, data, data_samples=None, mode='loss'):
        if mode == 'loss':
            return self.compute_loss(data)
        elif mode == 'predict':
            return self.predict(data)
        elif mode == 'tensor':
            return self._forward(data)
        else:
            raise NotImplementedError
        
    def compute_loss(self,data, data_samples=None, mode='loss'):
        g_pixel_values = data.pop('g_pixel_values', None)
        gt_masks = data.pop('masks', None)
        input_ids = data['input_ids']
        output = self.mllm(data, data_samples, mode)
        if gt_masks is None:
            g_pixel_values = [
                torch.randn(3, 512, 1024).to(output.hidden_states[-1]) 
                    for _ in range(len(input_ids))]
            ori_size_list = [(512, 1024) for _ in range(len(input_ids))]
            seg_token_mask = torch.zeros_like(input_ids).bool()
            seg_token_mask[:, -2] = True
        else:
            ori_size_list = [mask.shape[-2:] for mask in gt_masks]
            seg_token_mask = input_ids == self.seg_token_idx

        resize_list = [pixel.shape[-2:] for pixel in g_pixel_values]
        g_pixel_values = torch.stack([
            self.grounding_encoder.preprocess(pixel) for pixel in g_pixel_values
        ])
        image_embeddings = self.grounding_encoder.image_encoder(g_pixel_values)

        seg_token_mask = seg_token_mask[:, 1:]
        seg_token_mask = torch.cat([
            seg_token_mask,
            seg_token_mask.new_zeros(seg_token_mask.shape[0], 1)], dim=-1)

        hidden_states = output.hidden_states
        hidden_states = self.text_hidden_fcs(hidden_states[-1])
        pred_embeddings = hidden_states[seg_token_mask]

        seg_token_counts = seg_token_mask.int().sum(-1)
        pred_embeddings_list = torch.split(pred_embeddings, seg_token_counts.tolist(), dim=0)
        
        pred_masks = self._generate_and_postprocess_masks(
            pred_embeddings_list, image_embeddings, resize_list, ori_size_list)
        
        if gt_masks is None:
            return {
                'loss_mask': pred_masks[0].sum() * 0.0,
                'loss_dice': pred_masks[0].sum() * 0.0,
                'llm_loss': output.loss,
            }
        bs = len(pred_masks)
        loss_mask, loss_dice = 0, 0
        for i in range(bs):
            pred_mask = pred_masks[i]
            gt_mask = gt_masks[i]

            sam_loss_mask = self.loss_mask(pred_mask, gt_mask)
            sam_loss_dice = self.loss_dice(pred_mask, gt_mask)
            accuracy = torch.eq((pred_mask.sigmoid() > 0.5), gt_mask).to(pred_mask).mean()
            loss_mask += sam_loss_mask
            loss_dice += sam_loss_dice

        loss_dict = {
            'loss_mask': loss_mask / bs,
            'loss_dice': loss_dice / bs,
            'llm_loss': output.loss,
        }
        return loss_dict

    def predict(self, data):
        generation_config = dict(max_new_tokens=1024, do_sample=False)
        eos_token_id = self.tokenizer.convert_tokens_to_ids('<|end|>')
        generation_config['eos_token_id'] = eos_token_id
        pixel_values = data.pop('pixel_values')
        attention_mask = data.pop('attention_mask', None)
        input_ids = data['input_ids']
        generate_output = self.mllm.generate(
            pixel_values=pixel_values,
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict_in_generate=True,
            **generation_config,
        )
        device = self.mllm.model.device

        hidden_states = generate_output.hidden_states
        last_hidden_states = [item[-1] for item in hidden_states[1:]] # remove input_ids
        last_hidden_states = torch.cat(last_hidden_states, dim=1)
        last_hidden_states = last_hidden_states[0] # remove batch dim
        output_ids = generate_output.sequences[0][:-1] # remove batch dim and eos token
        output_text = self.tokenizer.decode(output_ids)
        seg_mask = output_ids == self.seg_token_idx
        if seg_mask.sum() == 0:
            return dict(
                pred_mask_logits=None,
                output_text=output_text,
            )
        seg_embeds = self.text_hidden_fcs(last_hidden_states[seg_mask])
        
        g_pixel_values = data.pop('g_pixel_values', None)
        gt_masks = data['masks']
        
        ori_size_list = [mask.shape[-2:] for mask in gt_masks]
        resize_list = [pixel.shape[-2:] for pixel in g_pixel_values]
        g_pixel_values = torch.stack([
            self.grounding_encoder.preprocess(pixel.to(device)) for pixel in g_pixel_values
        ])
        image_embeddings = self.grounding_encoder.image_encoder(g_pixel_values)
        pred_masks = self._generate_and_postprocess_masks(
            [seg_embeds], image_embeddings, resize_list, ori_size_list)
        
        return dict(
            pred_mask_logits=pred_masks[0], # remove batch dim
            output_text=output_text,
        )

    def gradient_checkpointing_enable(self):
        self.activation_checkpointing_enable()

    def activation_checkpointing_enable(self):
        self.mllm.model.language_model.gradient_checkpointing_enable()

    def gradient_checkpointing_disable(self):
        self.activation_checkpointing_disable()

    def activation_checkpointing_disable(self):
        self.mllm.model.language_model.gradient_checkpointing_disable()