praeclarumjj3's picture
:zap: add code
9fa3d89
raw
history blame
26.7 kB
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