ola_vlm/model/language_model/base_probe_vlm.py

from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation.utils import GenerateOutput
from transformers.generation.utils import GenerateOutput

from ola_vlm.model.aux_heads import GenHead, DepthHead, DAv2_Head
from ola_vlm.model.aux_heads.depth_anything_v2.dpt import DepthAnythingV2
from ola_vlm.model.aux_heads.oneformer_head import OneFormerHead, OneFormerSegHead

from transformers import OneFormerProcessor

from diffusers import (
    DPMSolverMultistepScheduler,
    StableUnCLIPImg2ImgPipeline,
)

import torch.distributed as dist
try:
    import wandb
except:
    pass
import os
import matplotlib
from .base_lm import BaseCausalLM
from tqdm import tqdm

from ola_vlm.ola_utils import *

class BaseProbe_VLM(BaseCausalLM):

    def __init__(self, config):
        super(BaseCausalLM, self).__init__(config)
        self.steps = 0
        self.config = config
        self.num_layers = config.num_hidden_layers

        # Initialize weights and apply final processing
        self.post_init()
        self.is_trained = False
        if hasattr(config, "probe_mode"):
            self.is_trained = True
            self.init_heads(config)

        try:
            if dist.get_rank() == 0:
                wandb.init(project=os.environ['WANDB_PROJECT'], name=f"{os.environ['WANDB_NAME']}")
        except:
            pass

    def get_model(self):
        return self.model

    def init_heads(self, config):
        self.mode = config.probe_mode
        
        if self.mode == "gen":            
            self.image_gen_heads = nn.ModuleList([
                GenHead(config.image_gen, llm_hidden_size=config.hidden_size)
                for _ in range(self.num_layers)
            ])
            
            if not self.is_trained:
                self.pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(config.image_generator, torch_dtype=torch.float16, variant="fp16")
                self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
                self.gen_encoder = self.pipe.image_encoder
                self.feature_extractor = self.pipe.feature_extractor
                for p in self.gen_encoder.parameters():
                    p.requires_grad = False

        elif self.mode == "seg":
            if not self.is_trained:
                self.oneformer_processor = OneFormerProcessor.from_pretrained(config.image_segmentor)
                self.oneformer = OneFormerHead.from_pretrained(config.image_segmentor)
                for p in self.oneformer.parameters():
                    p.requires_grad = False
                try:
                    self.oneformer = self.oneformer.to("cuda")
                except:
                    pass
            self.image_seg_heads = nn.ModuleList([
                OneFormerSegHead(config.image_seg, llm_hidden_size=config.hidden_size)
                for _ in range(self.num_layers)
            ])
            

        if self.mode == "depth":
            self.image_depth_heads = nn.ModuleList([
                DepthHead(proj_config=config.image_depth, llm_hidden_size=config.hidden_size, use_intermediate_depth=False)
                for _ in range(self.num_layers)
            ])

            dav2_cfg = {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}
            self.dav2_backbone = DepthAnythingV2(**dav2_cfg)
            self.dav2_backbone.load_state_dict(torch.load(config.depth_estimator, map_location='cpu'))
            for p in self.dav2_backbone.parameters():
                p.requires_grad = False

            self.da_v2_head = DAv2_Head()
            self.da_v2_head.load_state_dict(torch.load(config.depth_estimator), strict=False)
            for p in self.da_v2_head.parameters():
                p.requires_grad = False
    
    def _get_layer_loss_weight(self, config, prefix):
        layer_indices = config[f"{prefix}_layer_indices"]
        layer_indices = layer_indices.split("-")
        layer_indices = [int(i) - 1 for i in layer_indices]
        loss_weight = config[f"{prefix}_loss_weight"]
        return layer_indices, loss_weight
    
    def log_gen(self, img_embeds, pil_images, layer_idx, is_train=False):
        device = "cuda" if torch.cuda.is_available() else "hip"
        pipe = self.pipe.to(device)

        images = []

        if len(pil_images) > 2:
            pil_images = pil_images[:2]
            img_embeds = img_embeds[:2]

        for img_embed in img_embeds:
            image = pipe(image_embeds=img_embed.float().detach(),
                    num_inference_steps=25,
                    # guidance_scale=1,,
                ).images[0]
            images.append(image)
        
        if not is_train:
            return images

        n = len(images)
        c = min(n, 16)
        r = n // c
        images = images[:c*r]
        image_grid = make_grid(images, pil_images)
        
        wandb.log({
            f"val_gen_images/step_{self.steps}": wandb.Image(image_grid, caption=f"Layer-{layer_idx}")
        })
    
    def log_depth(self, depth_preds, layer_idx, depth_targets=None, is_train=False):
        cmap = matplotlib.colormaps.get_cmap('Spectral_r')
        depth_preds = depth_preds.float().detach()
        def _visualize_depth(depth):
            depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
            depth = depth.cpu().numpy().astype(np.uint8)
            colored_depth = (cmap(depth)[:, :, :3] * 255).astype(np.uint8)
            return Image.fromarray(colored_depth)

        pred_depths, gt_depths = [], []

        if depth_targets is None:
            depth_targets = [None] * len(depth_preds)

        from tqdm import tqdm
        for pred, target in tqdm(zip(depth_preds, depth_targets), desc="Visualizing Depth..."):
            if target is not None:
                gt = _visualize_depth(target.float())
                gt_depths.append(gt)

            pred = _visualize_depth(pred)
            pred_depths.append(pred)
        
        if not is_train:
            return pred_depths

        n = len(pred_depths)
        c = min(n, 16)
        r = n // c
        pred_depths = pred_depths[:c*r]
        gt_depths = gt_depths[:c*r]
        masks_grid = make_grid(pred_depths, gt_depths)
        
        wandb.log({
            f"val_depth_images/step_{self.steps}": wandb.Image(masks_grid, caption=f"Layer-{layer_idx}")
        })
    
    def log_seg(self, seg_embeds, pil_images, layer_idx, seg_targets=None, is_train=False):
        def _oneformer_prepare_panoptic_instance_prediction(
            segmentation: torch.Tensor, segments_info: dict
        ):
            masks = []
            classes = []

            for segment in segments_info:
                id = segment["id"]
                label_id = segment["label_id"]
                label = self.oneformer.config.id2label[label_id]
                mask = segmentation == id
                masks.append(mask.float())
                classes.append(label)

            return masks, classes
        
        pred_masks, gt_masks = [], []

        seg_embeds = seg_embeds.detach()

        if seg_targets is None:
            seg_targets = [None] * len(seg_embeds)

        if len(pil_images) > 2:
            pil_images = pil_images[:2]
            seg_embeds = seg_embeds[:2]
            seg_targets = seg_targets[:2]

        from tqdm import tqdm
        for emb, target, img in tqdm(zip(seg_embeds, seg_targets, pil_images), desc=f"Predicting Segmentation Map..."):
            with torch.no_grad():
                inputs = self.oneformer_processor(img, ["panoptic"], return_tensors="pt")
                inputs["pixel_values"] = inputs["pixel_values"].to(emb.device, emb.dtype)
                inputs["task_inputs"] = inputs["task_inputs"].to(emb.device, emb.dtype)
                gt = self.oneformer.get_masks(**inputs, backbone_last_feature=target.unsqueeze(0))
                gt = self.oneformer_processor.post_process_panoptic_segmentation(
                                        gt, target_sizes=[img.size[::-1]]
                                    )[0]
                gt_msk, gt_cls = _oneformer_prepare_panoptic_instance_prediction(**gt)
                gt = visualize_oneformer_masks_on_image(img, gt_msk, gt_cls)

                pred = self.oneformer.get_masks(**inputs, backbone_last_feature=emb.unsqueeze(0))
                pred = self.oneformer_processor.post_process_panoptic_segmentation(
                                        pred, target_sizes=[img.size[::-1]]
                                    )[0]
                pred_msk, pred_cls = _oneformer_prepare_panoptic_instance_prediction(**pred)
                pred = visualize_oneformer_masks_on_image(img, pred_msk, pred_cls)

            gt_masks.append(gt)
            pred_masks.append(pred)

        if not is_train:
            return pred_masks

        n = len(pred_masks)
        c = min(n, 16)
        r = n // c
        pred_masks = pred_masks[:c*r]
        gt_masks = gt_masks[:c*r]
        masks_grid = make_grid(pred_masks, gt_masks)
        
        wandb.log({
            f"val_seg_images/step_{self.steps}": wandb.Image(masks_grid, caption=f"Layer-{layer_idx}")
        })
    

    def _emb_loss(self, emb_preds, emb_targets):
        emb_targets = emb_targets.to(emb_preds.dtype).to(emb_preds.device)

        if emb_targets.shape[0] != emb_preds.shape[0]:
            repeat_factor = emb_preds.shape[0] // emb_targets.shape[0]
            emb_targets = emb_targets.repeat(repeat_factor, 1, 1)

            if emb_targets.shape[0] != emb_preds.shape[0]:
                emb_targets = emb_targets[:emb_preds.shape[0]]
                emb_mask = emb_mask[:emb_preds.shape[0]]

        emb_loss = F.smooth_l1_loss(
            emb_preds.float(), emb_targets.float(), reduction="none"
        ).mean()

        return emb_loss


    def _get_gen_feats(self, pil_images, device):
        gen_feats = []
        for img in pil_images:
            with torch.no_grad():
                clip_ims = self.pipe.feature_extractor(images=img, return_tensors="pt").pixel_values.to(device)
                feat = self.pipe.image_encoder(clip_ims).image_embeds
                gen_feats.append(feat)

        gen_feats = torch.stack(gen_feats, dim=0)
        return gen_feats
    
    def _forward_gen(self, gen_preds, layer_index, pil_images, gen_targets):        
        gen_loss = self._emb_loss(gen_preds, gen_targets)

        if dist.get_rank() == 0:
            if self.steps % 500 == 0:
                try:
                    self.log_gen(gen_preds.detach(), pil_images, layer_index, is_train=True)
                except:
                    pass
    
        return gen_loss
    

    def _get_dav2_feats(self, pil_images, device):
        dav2_gts = []
        depth_targets = [[]]
        for img in pil_images:
            img = img.resize((336, 336))
            img = np.array(img)
            with torch.no_grad():
                feat = self.dav2_backbone.infer_image(img, is_dsg=True)
                depth_gts = self.da_v2_head([feat[-1]] * 4)
                depth_targets[0].append(feat[-1][0])
            min_val = depth_gts.amin(dim=(1, 2), keepdim=True)
            max_val = depth_gts.amax(dim=(1, 2), keepdim=True)
            depth_gts = (depth_gts - min_val) / (max_val - min_val)
            dav2_gts.append(depth_gts.to(device))
        dav2_gts = torch.stack(dav2_gts, dim=0).squeeze(1)
        for i in range(len(depth_targets)):
            depth_targets[i] = (torch.stack(depth_targets[i], dim=0).squeeze(1), None)
        return depth_targets, dav2_gts
    
    def _forward_depth(self, all_depth_feats, layer_index, all_depth_targets, depth_pred_maps, depth_gts):                

        depth_feats, depth_targets = all_depth_feats[0][0], all_depth_targets[0][0]
        depth_loss = self._emb_loss(depth_feats, depth_targets)

        if dist.get_rank() == 0:
            if self.steps % 200 == 0:
                try:
                    self.log_depth(depth_pred_maps.detach(), layer_index, depth_gts, is_train=True)
                except:
                    pass

        return depth_loss


    def _get_seg_targets(self, pil_images, seg_preds):
        def _get_feats(img):
            img = img.resize((768, 768))
            inputs = self.oneformer_processor(img, ["panoptic"], return_tensors="pt")
            inputs["pixel_values"] = inputs["pixel_values"].to(seg_preds.device, seg_preds.dtype)
            with torch.no_grad():
                feats = self.oneformer.forward_features(**inputs)
            return feats

        seg_targets = []
        for img in pil_images:
            feat = _get_feats(img)
            seg_targets.append(feat)

        seg_targets = torch.stack(seg_targets, dim=0).squeeze(1)
        return seg_targets

    def _forward_seg(self, seg_preds, layer_index, pil_images, seg_targets):
        
        seg_loss = self._emb_loss(seg_preds, seg_targets)
        
        if dist.get_rank() == 0:
            if self.steps % 200 == 0:
                try:
                    self.log_seg(seg_preds.detach(), pil_images, layer_index, seg_targets, is_train=True)
                except:
                    pass
    
        return seg_loss


    def forward_emb_predictor(self, layer_states, idx, i, heads):
        inp_tokens = layer_states[idx]
        task_emb = heads[i](inp_tokens)
        return task_emb

    def depth_emb_forward(self, pil_images, layer_states):
        depth_preds = []
        depth_embs = []
        depth_loss = 0
        log_dict = {}
        if self.mode == "depth": 
            if pil_images is not None:
                depth_targets, depth_gts = self._get_dav2_feats(pil_images, layer_states[0].device)
            else:
                depth_targets, depth_gts = None, None
            
            for i, idx in enumerate(self.num_layers):
                
                depth_feats = self.forward_emb_predictor(layer_states, idx, i, self.image_depth_heads)
                depth_embs.append(depth_feats)

                with torch.no_grad():
                    depth_pred = self.da_v2_head([depth_feats[0]] * 4)
                    min_val = depth_pred.amin(dim=(1, 2), keepdim=True)
                    max_val = depth_pred.amax(dim=(1, 2), keepdim=True)
                    depth_pred = (depth_pred - min_val) / (max_val - min_val)
                    depth_preds.append(depth_pred)

                if depth_targets is not None:
                    layer_depth_loss = self._forward_depth(depth_feats, idx+1, depth_targets, depth_pred, depth_gts)
                    depth_loss += layer_depth_loss
                    if dist.get_rank() == 0:
                        log_dict = {
                            **log_dict,
                            f"{idx}_depth_loss": layer_depth_loss.item(),
                        }
    
    
        return depth_preds, depth_embs, depth_loss, log_dict
    
    def seg_emb_forward(self, pil_images, hidden_states, layer_states):
        seg_embs = []
        seg_loss = 0
        log_dict = {}
        if "seg" in self.mode:
            if pil_images is not None:
                seg_targets = self._get_seg_targets(pil_images, hidden_states)
            else:
                seg_targets = None
            for i, idx in enumerate(self.num_layers):

                seg_emb = self.forward_emb_predictor(layer_states, idx, i, "seg", self.image_seg_heads)
                seg_embs.append(seg_emb)

                if seg_targets is not None:
                    layer_seg_loss = self._forward_seg(seg_emb, idx+1, pil_images, seg_targets)
                    seg_loss += layer_seg_loss
                    if dist.get_rank() == 0:
                        log_dict = {
                            **log_dict,
                            f"{idx}_seg_loss": layer_seg_loss.item(),
                        }
        
        
        return seg_embs, seg_loss, log_dict
    
    def gen_emb_forward(self, pil_images, hidden_states, layer_states):
        img_embs = []
        gen_loss = 0
        log_dict = {}
        if "gen" in self.mode:
            if pil_images is not None:
                gen_targets = self._get_gen_feats(pil_images, hidden_states.device)
            else:
                gen_targets = None
            
            for i, idx in enumerate(self.num_layers):
                
                img_emb = self.forward_emb_predictor(layer_states, idx, i, "gen", self.image_gen_heads)
                img_embs.append(img_emb)

                if gen_targets is not None:
                    layer_gen_loss = self._forward_gen(img_emb, idx+1, pil_images, gen_targets)
                    gen_loss += layer_gen_loss
                    if dist.get_rank() == 0:
                        log_dict = {
                            **log_dict,
                            f"{idx}_gen_loss": layer_gen_loss.item(),
                        }

        return img_embs, gen_loss, log_dict

    @torch.no_grad()
    def get_visual_interpretations(
        self,
        inputs: Optional[torch.Tensor] = None,
        images: Optional[torch.Tensor] = None,
        image_sizes: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        **kwargs
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        
        position_ids = kwargs.pop("position_ids", None)
        attention_mask = kwargs.pop("attention_mask", None)
        if True:
            (
                inputs,
                position_ids,
                attention_mask,
                _,
                inputs_embeds,
                _
            ) = self.prepare_inputs_labels_for_multimodal(
                inputs,
                position_ids,
                attention_mask,
                None,
                None,
                images,
                image_sizes=image_sizes
            )

        
        return self.forward(
            input_ids=inputs,
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            attention_mask=attention_mask,
            return_dict=True,
            output_attentions=output_attentions,
            output_hidden_states=True,
        )

    @torch.no_grad()
    def generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        images: Optional[torch.Tensor] = None,
        image_sizes: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Union[GenerateOutput, torch.LongTensor]:
        position_ids = kwargs.pop("position_ids", None)
        attention_mask = kwargs.pop("attention_mask", None)
        if "inputs_embeds" in kwargs:
            raise NotImplementedError("`inputs_embeds` is not supported")

        if images is not None:
            (
                inputs,
                position_ids,
                attention_mask,
                _,
                inputs_embeds,
                _
            ) = self.prepare_inputs_labels_for_multimodal(
                inputs,
                position_ids,
                attention_mask,
                None,
                None,
                images,
                image_sizes=image_sizes
            )
        else:
            inputs_embeds = self.get_model().embed_tokens(inputs)

        return super().generate(
            position_ids=position_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            **kwargs
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
                                      inputs_embeds=None, **kwargs):
        images = kwargs.pop("images", None)
        image_sizes = kwargs.pop("image_sizes", None)
        pil_images = kwargs.pop("pil_images", None)
        
        inputs = super().prepare_inputs_for_generation(
            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
        )
        if images is not None:
            inputs['images'] = images
        if image_sizes is not None:
            inputs['image_sizes'] = image_sizes
        if pil_images is not None:
            inputs['pil_images'] = pil_images
        return inputs