import io
import logging
import json
import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import MSELoss
from transformers.modeling_outputs import (
CausalLMOutputWithPast,
)
from typing import List, Optional, Tuple, Union
from transformers import LlamaForCausalLM
from transformers.modeling_outputs import (
CausalLMOutputWithPast,
)
from torch.cuda.amp import autocast as autocast
import torch.nn.functional as F
import numpy as np
from .modeling_vit import build_vit, MLP, PostProcess
from .modeling_qformer import build_qformer
from .modeling_base import BaseMLLM
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
logger = logging.getLogger(__name__)
import pycocotools.mask as mask_util
from .modeling_base import VID_TOKEN, IMG_TOKEN
class MultiModalLLM_PT(BaseMLLM):
def __init__(
self,
config,
_tokenizer=None
):
super().__init__(config=config, _tokenizer=_tokenizer)
self.use_clip = False
self.num_frames = 16
self.num_clips = 1
self.token_merge_len = 4
self.per_clip_frames = self.num_frames // self.num_clips
print(self.config)
self.merge_proj = nn.Linear(
self.qformer.config.hidden_size*self.token_merge_len, self.config.hidden_size
)
if config.build_decoder:
self.track_embed = MLP(self.config.hidden_size, self.config.hidden_size, 3 * 256, 2, dropout=0)
self.track_embed_decode2 = MLP(4096, 4096, 4, 2, dropout=0)
self.temporal_embed = MLP(self.config.hidden_size, self.config.hidden_size, 2, 2, dropout=0.3)
self.action_embed = MLP(self.config.hidden_size, self.config.hidden_size, 1, 2, dropout=0.3)
self.postprocess = PostProcess()
self.track_token = nn.Parameter(torch.randn((1, 1, 4096)))
self.temporal_token = nn.Parameter(torch.randn((1, 1, 4096)))
self.box_token = nn.Parameter(torch.randn((1, 1, 4096)))
def forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
labels: Optional[torch.LongTensor] = None,
image: Optional[torch.Tensor] = None,
video: Optional[torch.Tensor] = None,
instruction = None,
video_idx = None,
image_idx = None,
output_boxes = None, # REC
input_boxes = None, # tracking inputs
text_input = None,
video_info = None,
temporal_labels = None,
gt_masks = None,
sam_images = None,
size_hw = None,
path = None,
mask_path = None,
tvg_inputs = None,
tvg_targets = None,
):
if text_input is not None:
time_instructions = self.get_clip_time_instruct(text_input)
else:
time_instructions = None
text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, return_visual=False,
video_idx=video_idx, image_idx=image_idx, instruction = instruction,
output_boxes = output_boxes, input_boxes=input_boxes, time_instructions = time_instructions)
outputs = self.lm(
inputs_embeds=text_embeds,
attention_mask=attention_mask,
labels=labels,
output_hidden_states=True,
return_dict=True,
)
loss = outputs.loss
logger.info(f'llm loss:{loss}')
if output_boxes is not None and self.use_bbox_loss:
last_hidden_states = outputs.hidden_states[-1]
pred_locs = []
for idx in range(last_hidden_states.shape[0]):
loc_positions = ( (input_ids[idx].flatten() == self.tokenizer.box_token) ).nonzero().flatten()
selected_hidden_states = last_hidden_states[idx][loc_positions]
pred_locs.append(self.loc_decoder(selected_hidden_states))
box_loss = self.box_loss(pred_locs, output_boxes)
logger.info(f'box loss:{box_loss}')
loss += box_loss
if (gt_masks is not None or input_boxes is not None) and self.use_mask_loss:
last_hidden_states = outputs.hidden_states[-1]
pred_masks = []
sam_losses = []
box_losses = []
for idx in range(last_hidden_states.shape[0]):
loc_positions = ( (input_ids[idx].flatten() == self.tokenizer.track_token) ).nonzero().flatten()
selected_hidden_states = last_hidden_states[idx][loc_positions]
embed_sam_boxes = self.track_embed(selected_hidden_states).reshape(1, 3, 256)
inference_state = self.sam.init_state_images(sam_images, size_hw[idx][0], size_hw[idx][1])
if input_boxes is not None:
gt_embeds = self.sam.get_prompt_embeding(inference_state, None, None, False, input_boxes[idx], device = text_embeds.device)
else:
input_boxes = self.find_boundaries_torch(gt_masks.squeeze(0)[:,:,:1].squeeze(2).cpu()).to(text_embeds.device)
gt_embeds = self.sam.get_prompt_embeding(inference_state, None, None, False, input_boxes, device = text_embeds.device)
pred_locs = [self.track_embed_decode2((selected_hidden_states))[0]]
target_boxes = [input_boxes[idx]]
src_boxes = pred_locs
loss_bbox = self.box_loss2(src_boxes, target_boxes)
loss_bbox = self.masked_loss(loss_bbox, 0)
box_losses.append(loss_bbox)
sam_losses.append( F.l1_loss(embed_sam_boxes, gt_embeds))
logger.info(f'refering sam loss:{sam_losses}')
sam_losses = torch.stack(sam_losses)
box_losses = torch.stack(box_losses)
loss += torch.mean(sam_losses)
loss += torch.mean(box_losses)
if tvg_inputs is not None and self.use_temporal_loss:
last_hidden_states = outputs.hidden_states[-1] # [bsz,1024, 4096]
last_hidden_states = last_hidden_states.view(-1, last_hidden_states.size(-1)) # [bsz*1024, 4096]
loc_positions = (input_ids.flatten()==self.tokenizer.temp_token).nonzero().flatten() # [bsz]
prompt_token = last_hidden_states[loc_positions]
prompt_token = prompt_token.view(input_ids.shape[0], -1 ,prompt_token.shape[-1]) # [bsz, 1, 4096]
cg_outputs = self.cg_model(**tvg_inputs, targets=tvg_targets, prompt_token=prompt_token)
loss_dict = self.cg_criterion(cg_outputs, tvg_targets)
weight_dict = self.cg_criterion.weight_dict
tvg_loss = 0.05*sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
logger.info(f'tvg_loss:{tvg_loss}')
loss += tvg_loss
logger.info(f'all loss:{loss}')
return CausalLMOutputWithPast(
loss=loss,
logits=outputs.logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
def pad_text_embeds(
self,
input_ids: torch.LongTensor = None,
image: Optional[torch.Tensor] = None,
video: Optional[torch.Tensor] = None,
image_idx = None,
video_idx = None,
return_visual: bool = False,
instruction = None,
output_boxes = None, # boxes for REC
input_boxes = None, # boxes for tracking
time_instructions = None,
):
text_embeds = self.lm.get_input_embeddings()(input_ids.long()).detach()
if input_boxes is not None:
input_boxes = input_boxes[0].to(dtype=text_embeds.dtype)
boxes_emb = self.loc_encoder(input_boxes)
boxes_emb = boxes_emb.view(-1, 4096)
text_embeds[input_ids == torch.full_like(input_ids, self.tokenizer.track_box_token)] = text_embeds[input_ids == torch.full_like(input_ids, self.tokenizer.track_box_token)] * 0 + boxes_emb.to(text_embeds.device)
logger.info(f'embedings:{text_embeds[input_ids == torch.full_like(input_ids, self.tokenizer.track_box_token)].shape}')
visual = None
visual_idx = None
if image is not None:
B, T, C, H, W = image.shape
image = image.permute(0, 2, 1, 3, 4)
instruction = None
prompt_image_embeds = self.encode_vision(image, instruction)
visual = prompt_image_embeds
prompt_image_embeds = self.project_up(prompt_image_embeds) # 768 -> 4096
prompt_image_embeds = prompt_image_embeds.view(-1, prompt_image_embeds.shape[-1])
visual_idx = image_idx
prompt_image_embeds = prompt_image_embeds.to(dtype=text_embeds.dtype)
text_embeds[image_idx == 1] = torch.zeros_like(text_embeds[image_idx == 1]) + prompt_image_embeds.to(text_embeds.device)
elif video is not None:
if len(video.shape) == 5:
B, T, C, H, W = video.shape
N = 1
if self.use_clip:
video = video.reshape(B*self.num_clips, T//self.num_clips, C, H, W) # [16, 8, 3, 224, 224]
else:
B, N, T, C, H, W = video.shape
video = video.permute(0,2,1,3,4) #
prompt_video_embeds = self.encode_vision(video, instruction=time_instructions) # [2, 96, 768]
if self.use_clip:
prompt_video_embeds = prompt_video_embeds.reshape(B,-1,prompt_video_embeds.shape[-1]) # [2,8*96,768]
batch_size, img_len, token_dim = prompt_video_embeds.shape
prompt_video_embeds = prompt_video_embeds.view(batch_size, img_len // self.token_merge_len, self.token_merge_len * token_dim) # [B, 768//4, 4*768] = [2, 192, 3072]
prompt_video_embeds = self.merge_proj(prompt_video_embeds) # [2, 192, 4096]
prompt_video_embeds = prompt_video_embeds.view(-1, prompt_video_embeds.shape[-1]) # [2*192, 4096]
else:
prompt_video_embeds = self.project_up(prompt_video_embeds) # [2, 96, 4096]
prompt_video_embeds = prompt_video_embeds.view(-1, prompt_video_embeds.shape[-1])
visual_idx = video_idx
text_embeds[video_idx == 1] = torch.zeros_like(text_embeds[video_idx == 1]) + prompt_video_embeds.to(text_embeds.device).to(text_embeds.dtype)
else:
logger.warn(f"don't get visual input, input_ids: {input_ids}")
for idx, text_embed in enumerate(text_embeds):
if text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.box_token].shape[0] != 0:
text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.box_token] = torch.zeros_like(text_embeds[idx][input_ids[idx] == self.tokenizer.box_token]) + torch.cat([self.box_token.squeeze(0)] * (text_embeds[idx][input_ids[idx] == self.tokenizer.box_token]).shape[0]).to(text_embeds.dtype)
if text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.temp_token].shape[0] != 0:
text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.temp_token] = torch.zeros_like(text_embeds[idx][input_ids[idx] == self.tokenizer.temp_token]) + self.temporal_token
if text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.track_token].shape[0] != 0:
text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.track_token] = torch.zeros_like(text_embeds[idx][input_ids[idx] == self.tokenizer.track_token]) + self.track_token
if return_visual:
return text_embeds, visual, visual_idx
return text_embeds
def temporal_decode(self, temporal_embedding):
pred_sted = self.temporal_embed(temporal_embedding)
pred_actioness = self.action_embed(temporal_embedding)
return pred_sted, pred_actioness
def box_loss2(self, src_boxes, target_boxes):
src_boxes = torch.cat(src_boxes)
target_boxes = torch.cat(target_boxes)
loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')
loss_bbox = self.masked_loss(loss_bbox, 0)
mask = (src_boxes[2:] >= src_boxes[:2]).all(-1)
src_boxes = src_boxes[mask]
target_boxes = target_boxes[mask]
return loss_bbox
def box_loss(self, src_boxes, target_boxes):
src_boxes = torch.cat(src_boxes)
target_boxes = torch.cat(target_boxes)
loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')
loss_bbox = self.masked_loss(loss_bbox, 0)
mask = (src_boxes[:, 2:] >= src_boxes[ :, :2]).all(-1)
src_boxes = src_boxes[mask]
target_boxes = target_boxes[mask]
if src_boxes.shape[0] > 0:
loss_giou = 1 - torch.diag(generalized_box_iou(
src_boxes,
target_boxes))
loss_giou = self.masked_loss(loss_giou, 0)
else:
loss_giou = torch.tensor(2, dtype=src_boxes.dtype)
iou, union = box_iou(src_boxes, target_boxes)
return loss_bbox * 2 + loss_giou / 5
def find_boundaries_torch(self, mask):
from skimage.segmentation import find_boundaries
mask_np = mask.to(torch.bool).numpy()
boundaries = find_boundaries(mask_np, mode='outer')
boundary_points = np.argwhere(boundaries)
if boundary_points.size == 0:
return torch.tensor([-1, -1, -1, -1], dtype = torch.bfloat16)
h0, w0 = boundary_points.min(axis=0)
h1, w1 = boundary_points.max(axis=0)
return torch.tensor([w0 / mask.shape[1], h0 / mask.shape[0], w1 / mask.shape[1], h1 / mask.shape[0]], dtype = torch.bfloat16)
def sam_loss(self, sam_outputs, gt_masks):
bound1 = self.find_boundaries_torch(gt_masks[:,:,:1].squeeze(2).cpu())
bound2 = self.find_boundaries_torch(sam_outputs[:,:,:1].squeeze(2).cpu())
lossl1 = F.l1_loss(bound1, bound2, reduction='none')
lossl1 = self.masked_loss(lossl1, 0)
loss_iou = self.iou_loss(sam_outputs, gt_masks)
loss_dice = self.dice_loss(sam_outputs, gt_masks)
# print(f'mask loss:{loss_iou, loss_dice}')
return loss_iou + loss_dice + lossl1
def masked_loss(self, loss, n):
mask = torch.ones_like(loss)
# mask[-n:] = 1e-10
loss = (loss*mask).sum()/(mask.sum())
return loss
def encode_vision(
self,
image,
instruction
):
device = image.device
B = image.shape[0]
T = image.shape[2]
use_image = True if T == 1 else False
image_embeds = self.vision_encoder(image, use_image=use_image)
C = image_embeds.shape[-1]
image_embeds = image_embeds.reshape(B, -1, C)
image_embeds = self.vision_layernorm(image_embeds).to(device) # [B, T*L, C]
image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
if self.extra_num_query_token > 0:
query_tokens = torch.cat([self.query_tokens, self.extra_query_tokens], dim=1)
query_tokens = query_tokens.expand(image_embeds.shape[0], -1, -1)
if instruction is not None:
text_Qformer = self.qformer_tokenizer(
instruction,
padding='longest',
truncation=True,
max_length=512,
return_tensors="pt",
).to(image_embeds.device)
query_atts = torch.ones(query_tokens.size()[:-1], dtype=torch.long).to(image_embeds.device)
Qformer_atts = torch.cat([query_atts, text_Qformer.attention_mask], dim=1)
query_output = self.qformer.bert(
text_Qformer.input_ids,
attention_mask=Qformer_atts,
query_embeds=query_tokens,
encoder_hidden_states=image_embeds,
encoder_attention_mask=image_atts,
return_dict=True,
)
else:
query_output = self.qformer.bert(
query_embeds=query_tokens,
encoder_hidden_states=image_embeds,
encoder_attention_mask=image_atts,
return_dict=True,
)
return query_output.last_hidden_state[:, :query_tokens.size(1), :]
def generate_caption(
self,
input_ids,
attention_mask,
image_idx = None,
video_idx = None,
image: Optional[torch.Tensor] = None,
video: Optional[torch.Tensor] = None,
num_beams=1,
max_new_tokens=200,
do_sample=True,
top_p=0.9,
top_k=None,
temperature=1.0,
length_penalty=1,
repetition_penalty=1.0,
):
text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, image_idx=image_idx, video_idx=video_idx)
outputs = self.lm.generate(
inputs_embeds=text_embeds,
attention_mask=attention_mask,
num_beams=num_beams,
max_new_tokens=max_new_tokens,
do_sample=do_sample,
min_length=1,
top_p=top_p,
top_k=top_k,
temperature=temperature,
length_penalty=length_penalty,
repetition_penalty=repetition_penalty,
)
return outputs
def generate_caption_bbox(
self,
input_ids,
attention_mask,
labels,
image_idx = None,
video_idx = None,
image: Optional[torch.Tensor] = None,
video: Optional[torch.Tensor] = None,
num_beams=1,
max_new_tokens=200,
do_sample=True,
top_p=0.9,
top_k=None,
temperature=0.9,
length_penalty=1,
repetition_penalty=1.0,
):
text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, image_idx=image_idx, video_idx=video_idx)
outputs = self.lm.generate(
inputs_embeds=text_embeds,
attention_mask=attention_mask,
num_beams=num_beams,
max_new_tokens=max_new_tokens,
do_sample=do_sample,
min_length=1,
top_p=top_p,
top_k=top_k,
temperature=temperature,
length_penalty=length_penalty,
repetition_penalty=repetition_penalty,
)
decoded_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
# torch.save({'text':decoded_text, 'output':{outputs}}, 'tmp.pth')
# print(decoded_text)
return outputs
def generate_temporal(self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
labels: Optional[torch.LongTensor] = None,
image: Optional[torch.Tensor] = None,
video: Optional[torch.Tensor] = None,
instruction = None,
video_idx = None,
image_idx = None,
boxes = None,
text_input = None,
video_info = None,
temporal_labels = None):
if text_input is not None:
time_instructions = self.get_clip_time_instruct(text_input)
else:
time_instructions = None
text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, return_visual=False,
video_idx=video_idx, image_idx=image_idx, instruction = instruction,
boxes = boxes, time_instructions = time_instructions)
# TODO
outputs = self.lm(
inputs_embeds=text_embeds,
attention_mask=attention_mask,
labels=labels,
output_hidden_states=True,
return_dict=True,
)
if temporal_labels is not None:
start_sec = temporal_labels["start_sec"]
end_sec = temporal_labels["end_sec"]
fps = video_info['fps']
frame_indices = video_info['frame_indices']
last_hidden_states = outputs.hidden_states[-1] # [2,1024, 4096]
last_hidden_states = last_hidden_states.view(-1, last_hidden_states.size(-1)) # [2048, 4096]
loc_positions = (input_ids.flatten()==self.tokenizer.temp_place_ids).nonzero().flatten() #
selected_hidden_states = last_hidden_states[loc_positions]
selected_hidden_states = selected_hidden_states.view(input_ids.shape[0], -1 ,selected_hidden_states.shape[-1]) # [2, 64, 4096]
# just for debug
# vis_embed = vis_embed[:,:64,:]
pred_sted, pred_actionness = self.temporal_decode(selected_hidden_states) # [2,64,2] [2,64,1]
pred_sted = self.postprocess(pred_sted, frame_indices)
pred_sec_s = pred_sted[0][0] / fps[0][0].item()
pred_sec_e = pred_sted[0][1] / fps[0][0].item()
output_file = "predictions2.jsonl"
prediction = {"pred_sec_s": round(pred_sec_s, 1), "pred_sec_e": round(pred_sec_e, 1), "start_sec":float(start_sec[0]), "end_sec": float(end_sec[0])}
with open(output_file, 'a') as f:
json.dump(prediction, f)
f.write('\n')
return outputs
def generate_seg(self, input_ids, attention_mask, labels, image, image_idx, video, video_idx, input_boxes, size_hw, sam_images):
device = input_ids.device
prompt = input_ids
l_prompt = len(input_ids)
temperature = 1e-5
max_new_tokens = 20
guide_w = 5
stop_str = '</s>'
bbox = []
output_ids = list(input_ids[0])
text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, image_idx=image_idx, video_idx=video_idx, return_visual=False,
instruction = None, output_boxes=None, input_boxes=input_boxes)
for i in range(max_new_tokens):
if i == 0:
outputs = self.lm(
inputs_embeds=text_embeds,
attention_mask=attention_mask,
output_hidden_states=True,
return_dict=True,
)
logits = outputs.logits
past_key_values = outputs.past_key_values
else:
attention_mask = torch.ones(1, past_key_values[0][0].shape[-2] + 1, device=device)
last_text_embeds = self.lm.get_input_embeddings()(torch.tensor(output_ids[-1], device=device).long()).detach().unsqueeze(0)
last_text_embeds = last_text_embeds.unsqueeze(0)
out = self.lm(
input_ids=None,
use_cache=True,
attention_mask=attention_mask,
output_hidden_states=True,
inputs_embeds=last_text_embeds,
past_key_values=past_key_values,
)
logits = out.logits
past_key_values = out.past_key_values
if logits is not None:
last_token_logits = logits[0][-1]
if temperature < 1e-4:
token = int(torch.argmax(last_token_logits))
else:
probs = torch.softmax(last_token_logits / temperature, dim=-1)
token = int(torch.multinomial(probs, num_samples=1))
output_ids.append(token)
ret = self.tokenizer.decode(token)
if ret == '<box_begin>':
attention_mask = torch.ones(1, past_key_values[0][0].shape[-2] + 1, device=device)
bbox_embeds = self.box_token.bfloat16()
out = self.lm(
inputs_embeds=bbox_embeds,
use_cache=True,
attention_mask=attention_mask,
output_hidden_states=True,
past_key_values=past_key_values
)
last_hidden_states = out.hidden_states[-1]
selected_hidden_states = last_hidden_states[0][0]
bbox.append(self.loc_decoder(selected_hidden_states))
last_token_logits = logits[0][-1]
if temperature < 1e-4:
token = int(torch.argmax(last_token_logits))
else:
probs = torch.softmax(last_token_logits / temperature, dim=-1)
token = int(torch.multinomial(probs, num_samples=1))
if ret == '<track_begin>':
attention_mask = torch.ones(1, past_key_values[0][0].shape[-2] + 1, device=device)
tracking_embeds = self.track_token
out = self.lm(
inputs_embeds=tracking_embeds,
use_cache=True,
attention_mask=attention_mask,
output_hidden_states=True,
past_key_values=past_key_values
)
last_hidden_states = out.hidden_states[-1]
selected_hidden_states = last_hidden_states[0][0].to(dtype = torch.bfloat16)
embed_sam_boxes = self.track_embed(selected_hidden_states).reshape(1, 3, 256)
inference_state = self.sam.init_state_images(sam_images, size_hw[0][0], size_hw[0][1])
gt_embeds = self.sam.get_prompt_embeding(inference_state, None, None, False, input_boxes[0].cuda(), device = text_embeds.device)
ann_frame_idx = 0
ann_obj_id = 0
box = np.array([0, 0, 0, 0], dtype=np.float32)
_, out_obj_ids, out_mask_logits = self.sam.add_new_box_embeding(
inference_state=inference_state,
frame_idx=ann_frame_idx,
obj_id=ann_obj_id,
box=box,
box_embeding=embed_sam_boxes,
)
video_segments = {} # video_segments contains the per-frame segmentation results
for out_frame_idx, out_obj_ids, out_mask_logits in self.sam.propagate_in_video(inference_state):
video_segments[out_frame_idx] = {
out_obj_id: (out_mask_logits[i] > 0.0)
for i, out_obj_id in enumerate(out_obj_ids)
}
video_segments = [video_segments[tt][0] for tt in video_segments]
# bbox = model.find_boundaries_torch(video_segments[0].squeeze(0).cpu())
# return ret, [], video_segments
if (ret == '</s>'):
break
ret = self.tokenizer.decode(output_ids)
del past_key_values
return ret, bbox, video_segments
def generate_answer(self, tokenizer, instruction, msg, user_prompt, media_type="video",video_tensor=None, image_tensor=None, answer_prompt=None, chat_history=[],return_history=False, debug=False, generation_config={}):
input_ids, attention_masks, labels = [], [], []
conversation = ""
if instruction:
conversation += instruction
conversation += (
"[INST]" + " "
)
if media_type == 'image':
conversation +=( "<Image>" + IMG_TOKEN + "</Image>")
else:
conversation += ("<Video>" + VID_TOKEN + "</Video>")
conversation += ( msg.rstrip() + "[/INST]")
for q,a in chat_history:
conversation += (" [INST] " + q + " [/INST]")
conversation += (a + "</s>")
conversation += (" [INST] " + user_prompt + " [/INST]")
conversation += ("")
if answer_prompt:
conversation += ("Best Option: (")
total_len = 0
indexs = []
if debug:
print(conversation)
tokenized = tokenizer.build_input_ids([conversation],
max_length=1024,
add_special_tokens=True,
truncation=False,
padding=False,
return_tensors='pt',
image=image_tensor,
video=video_tensor,
require_video=True)
if video_tensor is not None:
generation_output = self.generate_caption(
tokenized['input_ids'].unsqueeze(0).to(self.device),
tokenized['attention_mask'].unsqueeze(0).to(self.device),
video_idx = tokenized['video_index'].unsqueeze(0),
video = video_tensor.unsqueeze(0).to(self.device,dtype=torch.bfloat16),
do_sample=False
)
elif image_tensor is not None:
generation_output = self.generate_caption(
tokenized['input_ids'].unsqueeze(0).to(self.device),
tokenized['attention_mask'].unsqueeze(0).to(self.device),
image_idx = tokenized['image_index'].unsqueeze(0),
image = image_tensor.unsqueeze(0).to(self.device,dtype=torch.bfloat16),
do_sample=False
)
response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
if debug:
print(response)
return response, chat_history