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, TaskTokenGenHead, TaskTokenDepthHead from ola_vlm.model.aux_heads.depth_anything_v2.dpt import DepthAnythingV2 from ola_vlm.model.aux_heads.oneformer_head import OneFormerHead, OneFormerSegHead, OneFormerTaskTokenSegHead 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 BaseOLA_VLM(BaseCausalLM): def __init__(self, config): super(BaseCausalLM, self).__init__(config) self.steps = 0 self.config = config if hasattr(config, "image_gen"): 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_target_models(self, config): if hasattr(config, "image_gen") and "gen" in self.mode: if not os.path.exists(config.image_generator): config.image_generator = "stabilityai/stable-diffusion-2-1-unclip" self.pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(config.image_generator, torch_dtype=torch.float16, variant="fp16") self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config) for p in self.pipe.image_encoder.parameters(): p.requires_grad = False try: self.pipe = self.pipe.to("cuda") except: pass if hasattr(config, "image_depth") and "depth" in self.mode: dav2_cfg = {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]} self.dav2_backbone = DepthAnythingV2(**dav2_cfg) if not os.path.exists(config.depth_estimator): url = "https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/depth_anything_v2_vitl.pth?download=true" local_model_path = "depth_anything_v2_vitl.pth" if not os.path.exists(local_model_path): os.system(f"wget -O {local_model_path} {url}") config.depth_estimator = local_model_path config.depth_estimator = local_model_path 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 if hasattr(config, "image_seg") and "seg" in self.mode: if not os.path.exists(config.image_segmentor): config.image_segmentor = "oneformer/oneformer_coco_swin_large" 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 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 init_heads(self, config): self.mode = getattr(config, "aux_mode", "gen-depth-seg") self.pass_text_to_aux_head = getattr(config, "pass_text_to_aux", True) self.use_ce = getattr(config, "use_ce", False) self.contrastive_loss_weight = config.contrastive_loss_weight num_task_tokens = config.num_task_tokens if hasattr(config, "image_gen") and "gen" in self.mode: self.img_layer_indices, self.img_gen_loss_weight = self._get_layer_loss_weight(config.image_gen, "img") if getattr(config, "use_contrastive", True): self.gen_logit_scale = nn.Parameter(torch.tensor(2.0)) else: self.gen_logit_scale = None self.image_gen_heads = nn.ModuleList([ TaskTokenGenHead(config.image_gen, llm_hidden_size=config.hidden_size) if num_task_tokens > 0 else GenHead(proj_config=config.image_gen, llm_hidden_size=config.hidden_size) for _ in self.img_layer_indices ]) if hasattr(config, "image_depth") and "depth" in self.mode: self.depth_layer_indices, self.img_depth_loss_weight = self._get_layer_loss_weight(config.image_depth, "depth") self.img_depth_loss_weight = config.image_depth["depth_loss_weight"] if getattr(config, "use_contrastive", True): self.depth_logit_scale = nn.Parameter(torch.tensor(2.0)) else: self.depth_logit_scale = None self.use_intermediate_depth = config.image_depth.get("use_intermediate_depth", True) self.image_depth_heads = nn.ModuleList([ TaskTokenDepthHead(proj_config=config.image_depth, llm_hidden_size=config.hidden_size, use_intermediate_depth=self.use_intermediate_depth) if num_task_tokens > 0 else DepthHead(proj_config=config.image_depth, llm_hidden_size=config.hidden_size, use_intermediate_depth=self.use_intermediate_depth) for _ in self.depth_layer_indices ]) self.da_v2_head = DAv2_Head() if not os.path.exists(config.depth_estimator): url = "https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/depth_anything_v2_vitl.pth?download=true" local_model_path = "depth_anything_v2_vitl.pth" if not os.path.exists(local_model_path): os.system(f"wget -O {local_model_path} {url}") config.depth_estimator = local_model_path 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 if hasattr(config, "image_seg") and "seg" in self.mode: self.seg_layer_indices, self.img_seg_loss_weight = self._get_layer_loss_weight(config.image_seg, "seg") self.seg_teacher = config.image_seg.get("seg_teacher", "sam") assert self.seg_teacher in ["sam", "oneformer"] if getattr(config, "use_contrastive", True): self.seg_logit_scale = nn.Parameter(torch.tensor(2.0)) else: self.seg_logit_scale = None self.image_seg_heads = nn.ModuleList([ OneFormerTaskTokenSegHead(config.image_seg, llm_hidden_size=config.hidden_size) if num_task_tokens > 0 else OneFormerSegHead(config.image_seg, llm_hidden_size=config.hidden_size) for _ in self.seg_layer_indices ]) def log_gen(self, img_embeds, pil_images, layer_idx, is_train=False): pipe = self.pipe.to("cuda") images = [] for img_embed in img_embeds: image = pipe(image_embeds=img_embed.float().detach(), num_inference_steps=25, ).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) 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) 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) 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_mask, emb_targets, logit_scale): 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) emb_mask = emb_mask.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]] if emb_preds.ndim == 3: emb_mask = emb_mask.view(emb_preds.shape[0], 1, 1) else: emb_mask = emb_mask.view(emb_preds.shape[0], 1, 1, 1) sl1_loss = F.smooth_l1_loss( emb_preds.float(), emb_targets.float(), reduction="none" ) if logit_scale is not None: contrastive_loss = calculate_contrastive_loss(emb_preds, emb_targets, logit_scale) else: contrastive_loss = 0 sl1_loss = (sl1_loss * emb_mask.float()).mean() contrastive_loss = (self.contrastive_loss_weight * contrastive_loss * emb_mask.float()).mean() emb_loss = sl1_loss + contrastive_loss return emb_loss, sl1_loss, contrastive_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_mask, gen_targets): gen_loss, gen_sl1_loss, gen_cont_loss = self._emb_loss(gen_preds, gen_mask, gen_targets, self.gen_logit_scale) if dist.get_rank() == 0: if self.steps % 4000 == 0: try: self.log_gen(gen_preds.detach(), pil_images, layer_index, is_train=True) except: pass return gen_loss, gen_cont_loss, gen_sl1_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) ft_gt = (feat[0][0] + feat[1][0] + feat[2][0] + feat[3][0]) / 4 depth_gts = self.da_v2_head([(ft_gt, None)] * 4) depth_targets[0].append(ft_gt) 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, depth_mask, all_depth_targets, depth_pred_maps, depth_gts): depth_feats, depth_targets = all_depth_feats[0][0], all_depth_targets[0][0] depth_loss, sl1_loss, cont_loss = self._emb_loss(depth_feats, depth_mask, depth_targets, self.depth_logit_scale) if dist.get_rank() == 0: if self.steps % 1000 == 0: try: self.log_depth(depth_pred_maps.detach(), layer_index, depth_gts, is_train=True) except: pass return depth_loss, sl1_loss, cont_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_mask): seg_loss, sl1_loss, cont_loss = self._emb_loss(seg_preds, seg_mask, seg_targets, self.seg_logit_scale) if dist.get_rank() == 0: if self.steps % 1000 == 0: try: self.log_seg(seg_preds.detach(), pil_images, layer_index, seg_targets, is_train=True) except: pass return seg_loss, sl1_loss, cont_loss def forward_emb_predictor(self, layer_states, idx, i, task, heads, special_tokens): task_idx = self.token_order.index(task) task_start_idx = self.NUM_SYS_TOKENS + 576 + (self.num_task_tokens * task_idx) task_end_idx = task_start_idx + self.num_task_tokens end_idx = self.NUM_SYS_TOKENS + 576 + (self.num_task_tokens * len(self.token_order)) inp_tokens = layer_states[idx][:, :self.NUM_SYS_TOKENS+576] if self.num_task_tokens == 0 or layer_states[idx].shape[1] < 600: if self.pass_text_to_aux_head: inp_tokens = layer_states[idx] else: inp_tokens = torch.cat([inp_tokens, layer_states[idx][:, task_start_idx:task_end_idx]], dim=1) if self.pass_text_to_aux_head: inp_tokens = torch.cat([inp_tokens, layer_states[idx][:, end_idx:]], dim=1) if self.num_task_tokens == 0: task_emb = heads[i](inp_tokens) else: task_tokens = special_tokens if task != "gen": task_tokens = task_tokens.repeat(inp_tokens.shape[0], 1, 1) else: if not self.pass_text_to_aux_head: task_tokens = inp_tokens[:, -self.num_task_tokens:] else: task_tokens = inp_tokens[:, self.NUM_SYS_TOKENS+576:self.NUM_SYS_TOKENS+576+self.num_task_tokens] task_emb = heads[i](inp_tokens, task_tokens) return task_emb def depth_emb_forward(self, pil_images, layer_states, depth_mask): depth_preds = [] depth_embs = [] depth_loss = 0 depth_l1_loss = 0 depth_cont_loss = 0 if "depth" in self.mode and layer_states[0].shape[1] > self.NUM_SYS_TOKENS: 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.depth_layer_indices): depth_feats = self.forward_emb_predictor(layer_states, idx, i, "depth", self.image_depth_heads, self.depth_tokens) depth_embs.append(depth_feats) with torch.no_grad(): if self.use_intermediate_depth: depth_pred = self.da_v2_head(depth_feats) else: 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_mask is not None: depth_mask.zero_() if depth_targets is not None: layer_depth_loss, layer_l1_loss, layer_cont_loss = self._forward_depth(depth_feats, idx+1, depth_mask, depth_targets, depth_pred, depth_gts) depth_loss += layer_depth_loss * self.img_depth_loss_weight depth_l1_loss += layer_l1_loss * self.img_depth_loss_weight depth_cont_loss += layer_cont_loss * self.img_depth_loss_weight return depth_preds, depth_embs, depth_loss, depth_l1_loss, depth_cont_loss def seg_emb_forward(self, pil_images, hidden_states, layer_states, seg_mask): seg_embs = [] seg_loss = 0 seg_l1_loss = 0 seg_contrastive_loss = 0 if "seg" in self.mode and layer_states[0].shape[1] > self.NUM_SYS_TOKENS: 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.seg_layer_indices): seg_emb = self.forward_emb_predictor(layer_states, idx, i, "seg", self.image_seg_heads, self.seg_tokens) seg_embs.append(seg_emb) if seg_mask is not None: seg_mask.zero_() if seg_targets is not None: layer_seg_loss, seg_l1_loss, seg_contrastive_loss = self._forward_seg(seg_emb, idx+1, pil_images, seg_targets, seg_mask) seg_loss += layer_seg_loss * self.img_seg_loss_weight seg_l1_loss += seg_l1_loss * self.img_seg_loss_weight seg_contrastive_loss += seg_contrastive_loss * self.img_seg_loss_weight return seg_embs, seg_loss, seg_l1_loss, seg_contrastive_loss def gen_emb_forward(self, pil_images, hidden_states, layer_states, gen_mask): img_embs = [] gen_loss = 0 gen_con_loss = 0 gen_mse_loss = 0 if "gen" in self.mode and layer_states[0].shape[1] > self.NUM_SYS_TOKENS: 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.img_layer_indices): img_emb = self.forward_emb_predictor(layer_states, idx, i, "gen", self.image_gen_heads, self.gen_tokens) img_embs.append(img_emb) if gen_mask is not None: gen_mask.zero_() if gen_targets is not None: layer_gen_loss, layer_gen_con_loss, layer_gen_mse_loss = self._forward_gen(img_emb, idx+1, pil_images, gen_mask, gen_targets) gen_loss += layer_gen_loss * self.img_gen_loss_weight gen_con_loss += layer_gen_con_loss * self.img_gen_loss_weight gen_mse_loss += layer_gen_mse_loss * self.img_gen_loss_weight return img_embs, gen_loss, gen_mse_loss, gen_con_loss @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) depth_mask = kwargs.pop("seg_mask", None) gen_mask = kwargs.pop("seg_mask", None) seg_mask = kwargs.pop("seg_mask", 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 if depth_mask is not None: inputs['depth_mask'] = depth_mask if gen_mask is not None: inputs['gen_mask'] = gen_mask if seg_mask is not None: inputs['seg_mask'] = seg_mask return inputs