913 lines
42 KiB
Python
913 lines
42 KiB
Python
import copy
|
|
import os
|
|
import re
|
|
import uuid
|
|
from typing import Dict, List, Optional, Union
|
|
|
|
import numpy as np
|
|
import PIL
|
|
from PIL import Image
|
|
import torch
|
|
from transformers.feature_extraction_utils import BatchFeature
|
|
from transformers.image_utils import ImageInput, load_image
|
|
from transformers.processing_utils import (
|
|
AllKwargsForChatTemplate,
|
|
ChatTemplateLoadKwargs,
|
|
ProcessingKwargs,
|
|
ProcessorMixin,
|
|
Unpack,
|
|
)
|
|
from transformers.tokenization_utils_base import AudioInput, TextInput
|
|
from transformers.utils import (
|
|
is_torch_device,
|
|
is_torch_dtype,
|
|
logging,
|
|
requires_backends,
|
|
)
|
|
from transformers.utils.chat_template_utils import render_jinja_template
|
|
from transformers.video_utils import VideoInput, VideoMetadata, load_video
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
class HCXBatchFeature(BatchFeature):
|
|
def to(self, *args, **kwargs) -> "BatchFeature":
|
|
"""
|
|
Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in
|
|
different `dtypes` and sending the `BatchFeature` to a different `device`.
|
|
|
|
Args:
|
|
args (`Tuple`):
|
|
Will be passed to the `to(...)` function of the tensors.
|
|
kwargs (`Dict`, *optional*):
|
|
Will be passed to the `to(...)` function of the tensors.
|
|
To enable asynchronous data transfer, set the `non_blocking` flag in `kwargs` (defaults to `False`).
|
|
|
|
Returns:
|
|
[`BatchFeature`]: The same instance after modification.
|
|
"""
|
|
requires_backends(self, ["torch"])
|
|
import torch # noqa
|
|
|
|
new_data = {}
|
|
device = kwargs.get("device")
|
|
non_blocking = kwargs.get("non_blocking", False)
|
|
# Check if the args are a device or a dtype
|
|
if device is None and len(args) > 0:
|
|
# device should be always the first argument
|
|
arg = args[0]
|
|
if is_torch_dtype(arg):
|
|
# The first argument is a dtype
|
|
pass
|
|
elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
|
|
device = arg
|
|
else:
|
|
# it's something else
|
|
raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
|
|
# We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
|
|
for k, v in self.items():
|
|
# check if v is a floating point
|
|
if isinstance(v, torch.Tensor) and torch.is_floating_point(v):
|
|
# cast and send to device
|
|
new_data[k] = v.to(*args, **kwargs)
|
|
elif isinstance(v, torch.Tensor) and device is not None:
|
|
new_data[k] = v.to(device=device, non_blocking=non_blocking)
|
|
elif "pixel_values" in k:
|
|
new_pixel_values_batch = []
|
|
for _v in v:
|
|
pixel_values = [pixel_value.to(device=device, non_blocking=non_blocking) for pixel_value in _v]
|
|
new_pixel_values_batch.append(pixel_values)
|
|
new_data[k] = new_pixel_values_batch
|
|
else:
|
|
new_data[k] = v
|
|
self.data = new_data
|
|
return self
|
|
|
|
|
|
class HCXProcessorKwargs(ProcessingKwargs, total=False):
|
|
_defaults = {
|
|
"text_kwargs": {
|
|
"return_tensors": "pt",
|
|
"calc_non_vision_query_lengths": False,
|
|
},
|
|
"images_kwargs": {},
|
|
"audio_kwargs": {},
|
|
"videos_kwargs": {
|
|
"max_image_cnt": 12,
|
|
"max_num_grids": 9,
|
|
},
|
|
}
|
|
|
|
|
|
class HCXProcessor(ProcessorMixin):
|
|
attributes = ["image_processor", "tokenizer"]
|
|
valid_kwargs = ["chat_template"]
|
|
|
|
image_processor_class = "AutoImageProcessor"
|
|
tokenizer_class = ("GPT2Tokenizer", "GPT2TokenizerFast")
|
|
|
|
def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
|
|
self.image_token = "<|dummy3|>"
|
|
self.video_token = "<|_unuse_missing_100270|>"
|
|
self.image_token_pattern = re.compile(r"<\|dummy3\|>")
|
|
self.video_token_pattern = re.compile(r"<\|_unuse_missing_100270\|>")
|
|
self.image_video_token_pattern = re.compile(r"<\|dummy3\|>|<\|_unuse_missing_100270\|>")
|
|
self.image_token_id = (
|
|
tokenizer.image_token_id
|
|
if getattr(tokenizer, "image_token_id", None)
|
|
else tokenizer.convert_tokens_to_ids(self.image_token)
|
|
)
|
|
self.video_token_id = (
|
|
tokenizer.video_token_id
|
|
if getattr(tokenizer, "video_token_id", None)
|
|
else tokenizer.convert_tokens_to_ids(self.video_token)
|
|
)
|
|
super().__init__(image_processor, tokenizer, chat_template=chat_template)
|
|
|
|
def apply_chat_template(
|
|
self,
|
|
conversation: Union[list[dict[str, str]], list[list[dict[str, str]]]],
|
|
chat_template: Optional[str] = None,
|
|
**kwargs: Unpack[AllKwargsForChatTemplate],
|
|
) -> str:
|
|
"""
|
|
Similar to the `apply_chat_template` method on tokenizers, this method applies a Jinja template to input
|
|
conversations to turn them into a single tokenizable string.
|
|
|
|
The input is expected to be in the following format, where each message content is a list consisting of text and
|
|
optionally image or video inputs. One can also provide an image, video, URL or local path which will be used to form
|
|
`pixel_values` when `return_dict=True`. If not provided, one will get only the formatted text, optionally tokenized text.
|
|
|
|
conversation = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image", "image": "https://www.ilankelman.org/stopsigns/australia.jpg"},
|
|
{"type": "text", "text": "Please describe this image in detail."},
|
|
],
|
|
},
|
|
]
|
|
|
|
Args:
|
|
conversation (`Union[List[Dict, [str, str]], List[List[Dict[str, str]]]]`):
|
|
The conversation to format.
|
|
chat_template (`Optional[str]`, *optional*):
|
|
The Jinja template to use for formatting the conversation. If not provided, the tokenizer's
|
|
chat template is used.
|
|
"""
|
|
|
|
if chat_template is None:
|
|
if isinstance(self.chat_template, dict) and "default" in self.chat_template:
|
|
chat_template = self.chat_template["default"]
|
|
elif isinstance(self.chat_template, dict):
|
|
raise ValueError(
|
|
'The processor has multiple chat templates but none of them are named "default". You need to specify'
|
|
" which one to use by passing the `chat_template` argument. Available templates are: "
|
|
f"{', '.join(self.chat_template.keys())}"
|
|
)
|
|
elif self.chat_template is not None:
|
|
chat_template = self.chat_template
|
|
else:
|
|
raise ValueError(
|
|
"Cannot use apply_chat_template because this processor does not have a chat template."
|
|
)
|
|
else:
|
|
if isinstance(self.chat_template, dict) and chat_template in self.chat_template:
|
|
# It's the name of a template, not a full template string
|
|
chat_template = self.chat_template[chat_template]
|
|
else:
|
|
# It's a template string, render it directly
|
|
chat_template = chat_template
|
|
|
|
if kwargs.get("continue_final_message", False):
|
|
if kwargs.get("add_generation_prompt", False):
|
|
raise ValueError(
|
|
"continue_final_message and add_generation_prompt are not compatible. Use continue_final_message when you want the model to continue the final message, and add_generation_prompt when you want to add a header that will prompt it to start a new assistant message instead."
|
|
)
|
|
if kwargs.get("return_assistant_tokens_mask", False):
|
|
raise ValueError("continue_final_message is not compatible with return_assistant_tokens_mask.")
|
|
|
|
# Fill sets of kwargs that should be used by different parts of template
|
|
processed_kwargs = {
|
|
"mm_load_kwargs": {},
|
|
"template_kwargs": {},
|
|
}
|
|
|
|
for kwarg_type in processed_kwargs:
|
|
for key in AllKwargsForChatTemplate.__annotations__[kwarg_type].__annotations__.keys():
|
|
kwarg_type_defaults = AllKwargsForChatTemplate.__annotations__[kwarg_type]
|
|
default_value = getattr(kwarg_type_defaults, key, None)
|
|
value = kwargs.pop(key, default_value)
|
|
if value is not None and not isinstance(value, dict):
|
|
processed_kwargs[kwarg_type][key] = value
|
|
|
|
# Pass unprocessed custom kwargs
|
|
processed_kwargs["template_kwargs"].update(kwargs)
|
|
|
|
if isinstance(conversation, (list, tuple)) and (
|
|
isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "content")
|
|
):
|
|
is_batched = True
|
|
conversations = conversation
|
|
else:
|
|
is_batched = False
|
|
conversations = [conversation]
|
|
|
|
tokenize = processed_kwargs["template_kwargs"].pop("tokenize", False)
|
|
return_dict = processed_kwargs["template_kwargs"].pop("return_dict", False)
|
|
mm_load_kwargs = processed_kwargs["mm_load_kwargs"]
|
|
|
|
if tokenize:
|
|
batch_images, batch_videos = [], []
|
|
batch_audios = []
|
|
batch_video_metadata = []
|
|
for conversation in conversations:
|
|
images, videos = [], []
|
|
video_metadata = []
|
|
for message in conversation:
|
|
visuals = [content for content in message["content"] if content["type"] in ["image", "video"]]
|
|
audio_fnames = [
|
|
content[key]
|
|
for content in message["content"]
|
|
for key in ["audio", "url", "path"]
|
|
if key in content and content["type"] == "audio"
|
|
]
|
|
image_fnames = [
|
|
vision_info[key]
|
|
for vision_info in visuals
|
|
for key in ["image", "url", "path", "base64"]
|
|
if key in vision_info and vision_info["type"] == "image"
|
|
]
|
|
video_fnames = [
|
|
vision_info[key]
|
|
for vision_info in visuals
|
|
for key in ["video", "url", "path"]
|
|
if key in vision_info and vision_info["type"] == "video"
|
|
]
|
|
|
|
for fname in image_fnames:
|
|
images.append(load_image(fname))
|
|
|
|
# Audio models do not accept nested list of audios (yet!) so we construct a flat input audio list
|
|
if not mm_load_kwargs["load_audio_from_video"]:
|
|
for fname in audio_fnames:
|
|
batch_audios.append(load_audio(fname, sampling_rate=mm_load_kwargs["sampling_rate"]))
|
|
else:
|
|
for fname in video_fnames:
|
|
batch_audios.append(load_audio(fname, sampling_rate=mm_load_kwargs["sampling_rate"]))
|
|
|
|
for fname in video_fnames:
|
|
if isinstance(fname, (list, tuple)) and isinstance(fname[0], str):
|
|
video = [np.array(load_image(image_fname)) for image_fname in fname]
|
|
# create a 4D video because `load_video` always returns a 4D array
|
|
video = np.stack(video)
|
|
metadata = None
|
|
logger.warning(
|
|
"When loading the video from list of images, we cannot infer metadata such as `fps` or `duration`. "
|
|
"If your model uses this metadata during processing, please load the whole video and let the model sample frames instead."
|
|
)
|
|
else:
|
|
# TODO: raushan, should be `self.video_processor.load_video_for_model` when API is added
|
|
video, metadata = self._load_video_for_model(
|
|
fname,
|
|
num_frames=mm_load_kwargs.get("num_frames", None),
|
|
fps=mm_load_kwargs.get("video_fps", None),
|
|
backend=mm_load_kwargs["video_load_backend"],
|
|
**kwargs,
|
|
)
|
|
videos.append(video)
|
|
video_metadata.append(metadata)
|
|
|
|
# Currently all processors can accept nested list of batches, but not flat list of visuals
|
|
# So we'll make a batched list of images and let the processor handle it
|
|
if images:
|
|
batch_images.append(images)
|
|
if videos:
|
|
batch_videos.append(videos)
|
|
batch_video_metadata.append(video_metadata)
|
|
|
|
# Process conversation with video/image information if needed. Then convert into a prompt using Jinja template
|
|
conversations = self._process_messages_for_chat_template(
|
|
conversations,
|
|
batch_images=batch_images,
|
|
batch_videos=batch_videos,
|
|
batch_video_metadata=batch_video_metadata,
|
|
**processed_kwargs["mm_load_kwargs"],
|
|
)
|
|
|
|
prompt, generation_indices = render_jinja_template(
|
|
conversations=conversations,
|
|
chat_template=chat_template,
|
|
**processed_kwargs["template_kwargs"], # different flags such as `return_assistant_mask`
|
|
**self.tokenizer.special_tokens_map, # tokenizer special tokens are used by some templates
|
|
)
|
|
|
|
if not is_batched:
|
|
prompt = prompt[0]
|
|
|
|
if tokenize:
|
|
# Tokenizer's `apply_chat_template` never adds special tokens when tokenizing
|
|
# But processor's `apply_chat_template` didn't have an option to tokenize, so users had to format the prompt
|
|
# and pass it to the processor. Users thus never worried about special tokens relying on processor handling
|
|
# everything internally. The below line is to keep BC for that and be able to work with model that have
|
|
# special tokens in the template (consistent with tokenizers). We dont want to raise warning, it will flood command line
|
|
# without actionable solution for users
|
|
single_prompt = prompt[0] if is_batched else prompt
|
|
if self.tokenizer.bos_token is not None and single_prompt.startswith(self.tokenizer.bos_token):
|
|
kwargs["add_special_tokens"] = False
|
|
|
|
out = self(
|
|
text=prompt,
|
|
images=batch_images if batch_images else None,
|
|
videos=batch_videos if batch_videos else None,
|
|
audio=batch_audios if batch_audios else None,
|
|
**kwargs,
|
|
)
|
|
if return_dict:
|
|
if processed_kwargs["template_kwargs"].get("return_assistant_tokens_mask", False):
|
|
assistant_masks = []
|
|
input_ids = out["input_ids"]
|
|
for i in range(len(input_ids)):
|
|
current_mask = [0] * len(input_ids[i])
|
|
for assistant_start_char, assistant_end_char in generation_indices[i]:
|
|
start_token = out.char_to_token(i, assistant_start_char)
|
|
end_token = out.char_to_token(i, assistant_end_char - 1)
|
|
if start_token is None:
|
|
# start_token is out of bounds maybe due to truncation.
|
|
break
|
|
for token_id in range(start_token, end_token + 1 if end_token else len(input_ids[i])):
|
|
current_mask[token_id] = 1
|
|
assistant_masks.append(current_mask)
|
|
out["assistant_masks"] = assistant_masks
|
|
out.convert_to_tensors(tensor_type=kwargs.get("return_tensors", None))
|
|
|
|
# vllm needs vision_query_lengths, but hf model doesn't need it
|
|
del out["vision_query_lengths_images"]
|
|
del out["vision_query_lengths_videos"]
|
|
return out
|
|
else:
|
|
return out["input_ids"]
|
|
|
|
def repeat_dummy_tokens(self, input_ids, target_token_id, vision_query_lengths):
|
|
input_ids = input_ids.clone().detach()
|
|
batch_indices, target_indices = torch.where(input_ids == target_token_id)
|
|
batch_size = input_ids.shape[0]
|
|
|
|
new_input_ids = [[] for _ in range(batch_size)]
|
|
start_indices = [0 for _ in range(batch_size)]
|
|
counter = [0 for _ in range(batch_size)]
|
|
for batch_idx, target_idx in zip(batch_indices, target_indices):
|
|
start_idx = start_indices[batch_idx]
|
|
new_input_ids[batch_idx].append(input_ids[batch_idx][start_idx:target_idx])
|
|
query_length = vision_query_lengths[batch_idx][counter[batch_idx]]
|
|
new_input_ids[batch_idx].append(input_ids[batch_idx][target_idx].repeat(query_length))
|
|
start_indices[batch_idx] = target_idx + 1
|
|
counter[batch_idx] += 1
|
|
|
|
for batch_idx in range(batch_size):
|
|
start_idx = start_indices[batch_idx]
|
|
new_input_ids[batch_idx].append(input_ids[batch_idx][start_idx:]) # append remaining tokens
|
|
new_input_ids[batch_idx] = torch.cat(new_input_ids[batch_idx], dim=0)
|
|
|
|
new_input_ids = torch.stack(new_input_ids)
|
|
return new_input_ids
|
|
|
|
def _load_video_for_model(
|
|
self,
|
|
video: str,
|
|
num_frames: Optional[int] = None,
|
|
fps: Optional[int] = None,
|
|
backend: str = "opencv",
|
|
**kwargs: Unpack[HCXProcessorKwargs],
|
|
) -> List[ImageInput]:
|
|
"""
|
|
Overrided function.
|
|
|
|
Loads `video` to a List[PIL.Image] (llava style)
|
|
|
|
Args:
|
|
video (`str`):
|
|
The video to convert to the numpy array format. Can be a link to video or local path.
|
|
num_frames (`int`, *optional*):
|
|
Number of frames to sample uniformly. If not passed, the whole video is loaded.
|
|
fps (`int`, *optional*):
|
|
Number of frames to sample per second. Should be passed only when `num_frames=None`.
|
|
If not specified and `num_frames==None`, all frames are sampled.
|
|
backend (`str`, *optional*, defaults to `"opencv"`):
|
|
The backend to use when loading the video. Can be any of ["decord", "pyav", "opencv", "torchvision"]. Defaults to "opencv".
|
|
|
|
Returns:
|
|
Tuple[`np.array`, Dict]: A tuple containing:
|
|
- List[PIL.Image] of frames in RGB.
|
|
- Metadata dictionary.
|
|
"""
|
|
output_kwargs = self._merge_kwargs(
|
|
HCXProcessorKwargs,
|
|
tokenizer_init_kwargs=self.tokenizer.init_kwargs,
|
|
**kwargs,
|
|
)
|
|
|
|
logger.warning_once(f"num_frames control via argument is not supported yet. Ignored num_frames: {num_frames}.")
|
|
logger.warning_once(f"fps control via argument is not supported yet. Ignored fps: {fps}.")
|
|
logger.warning_once(f"backend control via argument is not supported yet. Ignored backend: {backend}.")
|
|
|
|
# video_loaded, video_metadata = load_video(
|
|
# video, backend="decord", num_frames=32
|
|
# )
|
|
# frame_interval = int(video_metadata.total_num_frames / 32)
|
|
# time_interval = frame_interval / video_metadata.fps
|
|
# video_metadata.time_interval = time_interval
|
|
|
|
def _hcx_sample_indices_fn(metadata: VideoMetadata, num_frames=None, fps=None, **kwargs):
|
|
max_num_grids = output_kwargs["videos_kwargs"]["max_num_grids"]
|
|
max_image_cnt = output_kwargs["videos_kwargs"]["max_image_cnt"]
|
|
frame_indices, time_interval = extract_frame_indices(
|
|
metadata.duration,
|
|
metadata.total_num_frames,
|
|
metadata.fps,
|
|
max_num_grids,
|
|
max_image_cnt,
|
|
default_interval=0.4,
|
|
)
|
|
metadata.time_interval = time_interval
|
|
return np.array(frame_indices)
|
|
|
|
video_loaded, video_metadata = None, None
|
|
for backend in ["decord", "pyav", "opencv", "torchvision"]:
|
|
try:
|
|
video_loaded, video_metadata = load_video(
|
|
video, sample_indices_fn=_hcx_sample_indices_fn, backend=backend
|
|
)
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error loading video with {backend} backend: {e}")
|
|
continue
|
|
|
|
assert video_loaded is not None, "Failed to load video with any backend"
|
|
|
|
return video_loaded, video_metadata
|
|
|
|
def _process_messages_for_chat_template(
|
|
self,
|
|
conversation: List[List[Dict[str, str]]],
|
|
batch_images: List[List[ImageInput]],
|
|
batch_videos: List[List[VideoInput]],
|
|
batch_video_metadata: List[List[Dict[str, any]]],
|
|
**mm_load_kwargs: Unpack[ChatTemplateLoadKwargs],
|
|
):
|
|
"""
|
|
Overrided function.
|
|
Used within `apply_chat_template` when a model has a special way to process conversation history. For example,
|
|
video models might want to specify in the prompt the duration of video or which frame indices at which timestamps
|
|
were sampled. This information cannot be accessed before the video is loaded.
|
|
|
|
For most models it is a no-op, and must be overridden by model processors which require special processing.
|
|
|
|
Args:
|
|
conversation (`List[Dict, str, str]`):
|
|
The conversation to process. Always comes in batched format.
|
|
batch_images (`List[List[ImageInput]]`):
|
|
Batch of images that were loaded from url/path defined in the conversation. The images
|
|
are ordered in the same way as in the conversation. Comes in nested list format, one list of `PIL` images
|
|
per batch.
|
|
batch_videos (`List[List[ImageInput]]`):
|
|
Batch of videos that were loaded from url/path defined in the conversation. The videos
|
|
are ordered in the same way as in the conversation. Comes in nested list format, one list of `PIL.Image`
|
|
per batch.
|
|
batch_video_metadata (`List[List[Dict[[str, any]]]]`):
|
|
Batch of metadata returned from loading videos. That includes video fps, duration and total number of framer in original video.
|
|
Metadata are ordered in the same way as `batch_videos`. Comes in nested list format, one list of `Dict`
|
|
per batch.
|
|
"""
|
|
|
|
is_video_in_conversation = False
|
|
for batch_idx, messages in enumerate(conversation):
|
|
is_video_in_messages = False
|
|
is_image_in_messages = False
|
|
for message in messages:
|
|
for content in message["content"]:
|
|
if content["type"] == "video":
|
|
is_video_in_messages = True
|
|
elif content["type"] == "image":
|
|
is_image_in_messages = True
|
|
if not is_video_in_messages:
|
|
batch_videos.insert(batch_idx, [])
|
|
batch_video_metadata.insert(batch_idx, [])
|
|
if not is_image_in_messages:
|
|
batch_images.insert(batch_idx, [])
|
|
|
|
is_video_in_conversation = is_video_in_conversation or is_video_in_messages
|
|
|
|
if not is_video_in_conversation:
|
|
return conversation
|
|
|
|
# conversation processing
|
|
new_conversation = []
|
|
for batch_idx, messages in enumerate(conversation):
|
|
video_counter = 0
|
|
new_messages = []
|
|
|
|
for message in messages:
|
|
new_message = {
|
|
"role": message["role"],
|
|
"content": [],
|
|
}
|
|
for content in message["content"]:
|
|
if content["type"] == "video":
|
|
video = batch_videos[batch_idx][video_counter]
|
|
video_meta = batch_video_metadata[batch_idx][video_counter]
|
|
|
|
time_stamps = calc_timestamp_video_grids(video, video_meta.time_interval, max_grid_shape=(3, 3))
|
|
video_counter += 1
|
|
|
|
if "filename" in content:
|
|
filename = content["filename"]
|
|
else:
|
|
filename = content["video"].split("/")[-1]
|
|
if len(filename) > 50:
|
|
filename = f"{uuid.uuid4().hex}.mp4"
|
|
basename, ext = os.path.splitext(filename)
|
|
if ext == "":
|
|
ext = ".mp4"
|
|
|
|
for frame_idx, time_stamp in enumerate(time_stamps):
|
|
if frame_idx == len(video) - 1:
|
|
# final_grid
|
|
new_content = {
|
|
"filename": f"{basename}-{frame_idx}{ext}",
|
|
"video": content["video"],
|
|
"type": "video",
|
|
"video_time_stamp": time_stamp,
|
|
"lens_keywords": content["lens_keywords"],
|
|
"lens_local_keywords": content["lens_local_keywords"],
|
|
"speech_to_text": content["speech_to_text"],
|
|
"is_final_grid": True,
|
|
}
|
|
new_message["content"].append(new_content)
|
|
else:
|
|
new_content = {
|
|
"filename": f"{basename}-{frame_idx}{ext}",
|
|
"video": content["video"],
|
|
"type": "video",
|
|
"video_time_stamp": time_stamp,
|
|
}
|
|
new_message["content"].append(new_content)
|
|
else:
|
|
new_message["content"].append(copy.deepcopy(content))
|
|
new_messages.append(new_message)
|
|
new_conversation.append(new_messages)
|
|
|
|
return new_conversation
|
|
|
|
def __call__(
|
|
self,
|
|
text: TextInput = None,
|
|
images: List[List[ImageInput]] = None,
|
|
videos: List[List[VideoInput]] = None,
|
|
audio: AudioInput = None,
|
|
**kwargs: Unpack[HCXProcessorKwargs],
|
|
):
|
|
output_kwargs = self._merge_kwargs(
|
|
HCXProcessorKwargs,
|
|
tokenizer_init_kwargs=self.tokenizer.init_kwargs,
|
|
**kwargs,
|
|
)
|
|
|
|
# prepare model inputs
|
|
mm_inputs = {
|
|
"pixel_values_images": [],
|
|
"image_sizes_images": [],
|
|
"vision_query_lengths_images": [],
|
|
"pixel_values_videos": [],
|
|
# "image_sizes_videos": [],
|
|
"vision_query_lengths_videos": [],
|
|
}
|
|
calc_non_vision_query_lengths = output_kwargs["text_kwargs"].pop("calc_non_vision_query_lengths")
|
|
if calc_non_vision_query_lengths:
|
|
mm_inputs["non_vision_query_lengths"] = []
|
|
|
|
# video processing
|
|
if videos is not None:
|
|
vit_input_size = self.image_processor.crop_size["width"]
|
|
|
|
video_kwargs = copy.deepcopy(output_kwargs["videos_kwargs"])
|
|
|
|
for videos_in_single_conversation in videos:
|
|
pixel_values_videos = []
|
|
vision_query_lengths_videos = []
|
|
|
|
for video_frames in videos_in_single_conversation:
|
|
if len(video_frames) == 0:
|
|
mm_inputs["pixel_values_videos"].append([])
|
|
mm_inputs["vision_query_lengths_videos"].append([])
|
|
continue
|
|
video_frames_combined = combine_frames_into_images(
|
|
video_frames, max_grid_shape=(3, 3), vit_input_size=vit_input_size
|
|
)
|
|
video_kwargs["is_video"] = True
|
|
video_kwargs["return_tensors"] = None
|
|
|
|
frames_processed = self.image_processor(images=video_frames_combined, **video_kwargs)
|
|
sizes = [(size["width"], size["height"]) for size in frames_processed["image_sizes"]]
|
|
|
|
pixel_values_videos.extend(frames_processed["pixel_values"])
|
|
vision_query_lengths_videos.extend(frames_processed["vision_query_lengths"])
|
|
|
|
mm_inputs["pixel_values_videos"].append(pixel_values_videos)
|
|
mm_inputs["vision_query_lengths_videos"].append(vision_query_lengths_videos)
|
|
|
|
# image processing
|
|
if images is not None:
|
|
image_kwargs = copy.deepcopy(output_kwargs["images_kwargs"])
|
|
image_kwargs["is_video"] = False
|
|
image_kwargs["return_tensors"] = None
|
|
|
|
for images_in_single_conversation in images:
|
|
if isinstance(images_in_single_conversation, PIL.Image.Image): # single item to batch
|
|
images_in_single_conversation = [images_in_single_conversation, ]
|
|
if len(images_in_single_conversation) == 0:
|
|
mm_inputs["pixel_values_images"].append([])
|
|
mm_inputs["image_sizes_images"].append([])
|
|
mm_inputs["vision_query_lengths_images"].append([])
|
|
continue
|
|
images_processed = self.image_processor(images=images_in_single_conversation, **image_kwargs)
|
|
sizes = [(size["width"], size["height"]) for size in images_processed["image_sizes"]]
|
|
|
|
mm_inputs["pixel_values_images"].append(images_processed["pixel_values"])
|
|
mm_inputs["image_sizes_images"].append(sizes)
|
|
mm_inputs["vision_query_lengths_images"].append(images_processed["vision_query_lengths"])
|
|
|
|
# text processing
|
|
def _create_replacer(_target_token, _replacements):
|
|
_iterator = iter(_replacements)
|
|
|
|
def _replacer(match_obj):
|
|
# return self.image_token
|
|
num_query_tokens = next(_iterator)
|
|
return "".join([_target_token for _ in range(num_query_tokens)])
|
|
return _replacer
|
|
|
|
text_inputs = {}
|
|
if text is not None:
|
|
if not isinstance(text, list):
|
|
text = [text]
|
|
|
|
if images is not None:
|
|
new_texts = []
|
|
for batch_idx, text_in_single_conversation in enumerate(text):
|
|
new_text = self.image_token_pattern.sub(
|
|
_create_replacer(self.image_token, mm_inputs["vision_query_lengths_images"][batch_idx]),
|
|
text_in_single_conversation,
|
|
)
|
|
new_texts.append(new_text)
|
|
text = new_texts
|
|
|
|
if videos is not None:
|
|
new_texts = []
|
|
for batch_idx, text_in_single_conversation in enumerate(text):
|
|
new_text = self.video_token_pattern.sub(
|
|
_create_replacer(self.video_token, mm_inputs["vision_query_lengths_videos"][batch_idx]),
|
|
text_in_single_conversation,
|
|
)
|
|
new_texts.append(new_text)
|
|
text = new_texts
|
|
|
|
text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
|
|
|
|
# audio processing
|
|
if audio is not None:
|
|
raise NotImplementedError("Audio processing is not supported yet.")
|
|
|
|
return HCXBatchFeature(data={**text_inputs, **mm_inputs})
|
|
|
|
def decode(self, *args, **kwargs):
|
|
"""
|
|
This method forwards all its arguments to Siglip2Tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to
|
|
the docstring of this method for more information.
|
|
"""
|
|
return self.tokenizer.decode(*args, **kwargs)
|
|
|
|
def batch_decode(self, *args, **kwargs):
|
|
"""
|
|
This method forwards all its arguments to Siglip2Tokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please
|
|
refer to the docstring of this method for more information.
|
|
"""
|
|
return self.tokenizer.batch_decode(*args, **kwargs)
|
|
|
|
def post_process_image_text_to_text(
|
|
self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
|
|
):
|
|
"""
|
|
Post-process the output of the model to decode the text.
|
|
|
|
Args:
|
|
generated_outputs (`torch.Tensor` or `np.ndarray`):
|
|
The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
|
|
or `(sequence_length,)`.
|
|
skip_special_tokens (`bool`, *optional*, defaults to `True`):
|
|
Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
|
|
Clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
|
|
Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
|
|
**kwargs:
|
|
Additional arguments to be passed to the tokenizer's `batch_decode method`.
|
|
|
|
Returns:
|
|
`List[str]`: The decoded text.
|
|
"""
|
|
return self.tokenizer.batch_decode(
|
|
generated_outputs,
|
|
skip_special_tokens=skip_special_tokens,
|
|
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
|
**kwargs,
|
|
)
|
|
|
|
@property
|
|
def model_input_names(self):
|
|
tokenizer_input_names = self.tokenizer.model_input_names
|
|
image_processor_input_names = self.image_processor.model_input_names
|
|
names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
|
|
return names_from_processor + []
|
|
|
|
|
|
def extract_frame_indices(play_time, total_frames, fps, max_num_grids, max_image_cnt, default_interval=0.4):
|
|
"""
|
|
Extracts specific frame indices from a video based on duration, frame count, and sampling strategy.
|
|
|
|
The function determines which frames to extract given the video duration (`play_time`),
|
|
total frame count, and frame rate. It samples frames at regular intervals (default: 0.4s),
|
|
but if the number of frames exceeds the limit defined by `max_num_grids * max_image_cnt`,
|
|
it performs uniform sampling to stay within that limit.
|
|
|
|
Args:
|
|
play_time (float): Total play time of the video in seconds.
|
|
total_frames (int): Total number of frames in the video.
|
|
fps (float): Frames per second of the video.
|
|
max_num_grids (int): Maximum number of grids to display.
|
|
max_image_cnt (int): Maximum number of images per grid.
|
|
default_interval (float, optional): Interval in seconds between frame samples. Defaults to 0.4.
|
|
|
|
Returns:
|
|
Tuple:
|
|
frame_indices (List[int]): A list of selected frame indices.
|
|
time_interval (float): Time interval between selected frames (in seconds).
|
|
"""
|
|
|
|
# Calculate how many frames to extract with the default interval
|
|
default_frame_count = int(play_time / default_interval)
|
|
|
|
# Maximum frames allowed based on max_num_grids and max_image_cnt
|
|
max_frames_allowed = max_num_grids * max_image_cnt
|
|
|
|
# Determine whether we can use the default interval or need uniform sampling
|
|
if default_frame_count <= max_frames_allowed:
|
|
# Default interval is sufficient, extract frames every 0.4 seconds
|
|
frame_interval = int(total_frames / default_frame_count)
|
|
else:
|
|
# Use uniform sampling to fit within max_frames_allowed
|
|
frame_interval = int(total_frames / max_frames_allowed)
|
|
|
|
# Extract frame indices at the calculated interval
|
|
selected_indices = list(range(0, total_frames, frame_interval))
|
|
|
|
time_interval = frame_interval / fps
|
|
|
|
# Ensure the number of selected indices does not exceed max_frames_allowed
|
|
return selected_indices[:max_frames_allowed], time_interval
|
|
|
|
|
|
def calc_timestamp_video_grids(frames, time_interval, max_grid_shape=(3, 3)):
|
|
"""
|
|
Calculates the time range labels for each grid in a video.
|
|
|
|
Args:
|
|
frames (List[PIL.Image.Image]): A list of frames extracted from a video.
|
|
time_interval (float): Time interval (in seconds) between consecutive frames.
|
|
max_grid_shape (Tuple[int, int], optional): The maximum grid shape as (rows, cols). Defaults to (3, 3).
|
|
vit_input_size (int, optional): The target size (height and width) for the Vision Transformer input. Defaults to 378.
|
|
|
|
Returns:
|
|
Tuple:
|
|
image_time_stamps (List[str]): A list of time span labels for each combined image,
|
|
e.g., ["0.00s~1.50s", "1.50s~3.00s", ...].
|
|
"""
|
|
max_num_grids = max_grid_shape[0] * max_grid_shape[1]
|
|
# assert (
|
|
# max_grid_shape[1] == 1
|
|
# ), f"For video processing, decided to concatenate frames horizontally into a wide image."
|
|
|
|
# Calculate the number of canvases needed.
|
|
num_frames = len(frames)
|
|
num_canvases = num_frames // max_num_grids
|
|
leftover_frames = num_frames % max_num_grids
|
|
|
|
time_stamp = 0 # second
|
|
image_time_stamps = []
|
|
|
|
for canvas_idx in range(num_canvases):
|
|
# Determine the frames to fill in the current canvas.
|
|
start_idx = canvas_idx * max_num_grids
|
|
end_idx = min(start_idx + max_num_grids, num_frames)
|
|
|
|
# Append the current canvas to the result list.
|
|
frame_cnt = end_idx - start_idx
|
|
image_time_stamps.append(f"{time_stamp:.2f}s~{time_stamp + frame_cnt * time_interval:.2f}s")
|
|
time_stamp += frame_cnt * time_interval
|
|
|
|
if leftover_frames > 0:
|
|
# Add the current canvas to the list of combined images.
|
|
frame_cnt = leftover_frames
|
|
image_time_stamps.append(f"{time_stamp:.2f}s~{time_stamp + frame_cnt * time_interval:.2f}s")
|
|
time_stamp += frame_cnt * time_interval
|
|
|
|
return image_time_stamps
|
|
|
|
|
|
def combine_frames_into_images(frames, max_grid_shape=(3, 3), vit_input_size=378):
|
|
"""
|
|
Combines a sequence of video frames into grid-based images and generates corresponding time range labels.
|
|
|
|
Frames are grouped and arranged into a grid (e.g., 3x3) such that each combined image contains up to
|
|
`max_grid_shape[0] * max_grid_shape[1]` frames. Each combined image is resized to the given ViT input size.
|
|
|
|
Args:
|
|
frames (NDArray): (num_frames, H, W, C) shape. A list of frames extracted from a video.
|
|
time_interval (float): Time interval (in seconds) between consecutive frames.
|
|
max_grid_shape (Tuple[int, int], optional): The maximum grid shape as (rows, cols). Defaults to (3, 3).
|
|
vit_input_size (int, optional): The target size (height and width) for the Vision Transformer input. Defaults to 378.
|
|
|
|
Returns:
|
|
Tuple:
|
|
image_list (List[PIL.Image.Image]): A list of grid-combined images.
|
|
"""
|
|
max_num_grids = max_grid_shape[0] * max_grid_shape[1]
|
|
# assert (
|
|
# max_grid_shape[1] == 1
|
|
# ), f"For video processing, decided to concatenate frames horizontally into a wide image."
|
|
|
|
# List to store the resulting combined images.
|
|
image_list = []
|
|
|
|
# Calculate the number of canvases needed.
|
|
num_frames = len(frames)
|
|
num_canvases = num_frames // max_num_grids
|
|
leftover_frames = num_frames % max_num_grids
|
|
|
|
# change frames (4d numpy tensor) to List[PIL.Image.Image]
|
|
frames = [Image.fromarray(frame) for frame in frames]
|
|
|
|
for canvas_idx in range(num_canvases):
|
|
# Initialize the current canvas.
|
|
combined_image = Image.new(
|
|
"RGB", (vit_input_size * max_grid_shape[0], vit_input_size * max_grid_shape[1]), color=(0, 0, 0)
|
|
)
|
|
|
|
# Determine the frames to fill in the current canvas.
|
|
start_idx = canvas_idx * max_num_grids
|
|
end_idx = min(start_idx + max_num_grids, num_frames)
|
|
|
|
for idx in range(start_idx, end_idx):
|
|
img = frames[idx]
|
|
|
|
# Resize each frame to a square shape.
|
|
img_resized = img.resize((vit_input_size, vit_input_size))
|
|
|
|
# Calculate the (row, column) position to place the frame within the grid layout.
|
|
local_idx = idx - start_idx
|
|
x_offset = (local_idx % max_grid_shape[0]) * vit_input_size
|
|
y_offset = (local_idx // max_grid_shape[0]) * vit_input_size
|
|
|
|
# Calculate the position to place the frame in the grid.
|
|
combined_image.paste(img_resized, (x_offset, y_offset))
|
|
|
|
# Append the current canvas to the result list.
|
|
image_list.append(combined_image)
|
|
|
|
if leftover_frames > 0:
|
|
# canvas_idx might be undefined; default to 0 if not previously assigned to avoid "referenced before assignment" error.
|
|
canvas_idx = num_canvases
|
|
# Add the remaining frames to the final canvas.
|
|
# combined_image = Image.new("RGB", (vit_input_size * leftover_frames, vit_input_size * 1), color=(0, 0, 0)) # hsk
|
|
combined_image = Image.new(
|
|
"RGB", (vit_input_size * max_grid_shape[0], vit_input_size * max_grid_shape[1]), color=(0, 0, 0)
|
|
)
|
|
|
|
for idx in range(leftover_frames):
|
|
img = frames[num_canvases * max_num_grids + idx]
|
|
|
|
# Resize the frame to a square (equal width and height).
|
|
img_resized = img.resize((vit_input_size, vit_input_size))
|
|
|
|
# Calculate the (row, column) position to place the frame within the grid layout.
|
|
# x_offset = (idx % leftover_frames) * vit_input_size # hsk
|
|
# y_offset = (idx // leftover_frames) * vit_input_size # hsk
|
|
x_offset = (idx % max_grid_shape[0]) * vit_input_size
|
|
y_offset = (idx // max_grid_shape[0]) * vit_input_size
|
|
|
|
# Calculate the position to place the frame within the grid layout.
|
|
combined_image.paste(img_resized, (x_offset, y_offset))
|
|
|
|
# Add the current canvas to the list of combined images.
|
|
image_list.append(combined_image)
|
|
|
|
return image_list
|