[Refactor] simplify multimodal data processing (#8107)

Signed-off-by: Xinyuan Tong <justinning0323@outlook.com>
This commit is contained in:
Xinyuan Tong
2025-07-20 21:43:09 -07:00
committed by GitHub
parent c9e8613c97
commit 8430bfe3e9
30 changed files with 297 additions and 421 deletions

View File

@@ -5,7 +5,7 @@ import multiprocessing as mp
import os
import re
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple, Union
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
import numpy as np
import torch
@@ -155,17 +155,15 @@ class BaseMultimodalProcessor(ABC):
self.ATTR_NAME_TO_MODALITY = {
# Image-related attributes
"pixel_values": Modality.IMAGE,
"pixel_values_videos": Modality.VIDEO,
"image_sizes": Modality.IMAGE,
"image_grid_thw": Modality.IMAGE,
"image_attention_mask": Modality.IMAGE,
"image_emb_mask": Modality.IMAGE,
"image_spatial_crop": Modality.IMAGE,
"images_spatial_crop": Modality.IMAGE,
"tgt_size": Modality.IMAGE,
"image_grid_hws": Modality.IMAGE,
"aspect_ratio_id": Modality.IMAGE,
"aspect_ratio_ids": Modality.IMAGE,
"aspect_ratio_mask": Modality.IMAGE,
"second_per_grid_ts": Modality.IMAGE,
# Audio-related attributes
"audio_features": Modality.AUDIO,
"audio_feature_lens": Modality.AUDIO,
@@ -173,9 +171,11 @@ class BaseMultimodalProcessor(ABC):
"input_features_mask": Modality.AUDIO,
"audio_attention_mask": Modality.AUDIO,
# Video-related attributes
"pixel_values_videos": Modality.VIDEO,
"second_per_grid_ts": Modality.VIDEO,
"video_grid_thw": Modality.VIDEO,
# Generic attributes that could apply to multiple modalities
# "precomputed_features" - handled specially as it can be any modality
# "precomputed_embeddings" - handled specially as it can be any modality
}
# name of the feature filed
@@ -222,7 +222,6 @@ class BaseMultimodalProcessor(ABC):
audio_data,
input_text,
request_obj,
max_req_input_len,
**kwargs,
) -> Optional[Dict[str, Any]]:
pass
@@ -283,7 +282,7 @@ class BaseMultimodalProcessor(ABC):
self,
text_parts: List[str],
multimodal_tokens: MultimodalSpecialTokens,
data_iterators: dict,
data_iterators: dict[Modality, Iterator[Any]],
discard_alpha_channel: bool = True,
image_estimated_frames_iter: Optional[iter] = None,
image_scaling_factor: float = 1.0,
@@ -354,7 +353,6 @@ class BaseMultimodalProcessor(ABC):
self,
prompt: str,
multimodal_tokens: MultimodalSpecialTokens,
max_req_input_len: int,
image_data: Optional[list] = None,
video_data: Optional[list] = None,
audio_data: Optional[list] = None,
@@ -489,50 +487,11 @@ class BaseMultimodalProcessor(ABC):
return list(zip(indices_start.tolist(), indices_end.tolist()))
@staticmethod
def _extract_processor_features(
items: List[dict], attr_name: str
) -> Optional[torch.Tensor]:
"""
Helper function to concat extracted attributes from processor output.
"""
values = [value for item in items if (value := item.get(attr_name)) is not None]
return torch.cat(values) if values else None
# When we assume that all the items have the same attributes
def _extract_processor_features_from_all_attributes(
self, items: List[dict]
) -> dict:
values = {}
# Verify all items have the same keys
first_keys = set(items[0].keys())
for item in items[1:]:
if set(item.keys()) != first_keys:
raise ValueError(
f"All items must have the same attributes. "
f"First item has {first_keys}, but found {set(item.keys())}"
)
# Process each attribute
for k, v in items[0].items():
if isinstance(v, list):
values[k] = self._extract_processor_features(items, k)
else:
# Verify all items have the same value for non-list attributes
for item in items[1:]:
if item[k] != v:
raise ValueError(
f"All items must have the same value for attribute {k}. "
f"First item has {v}, but found {item[k]}"
)
values[k] = v
return values
def collect_mm_items_from_processor_output(
self, data_dict: dict
) -> List[MultimodalDataItem]:
"""Create mm_items directly from processor output."""
items = {} # modality -> MultimodalDataItem
items: dict[Modality, MultimodalDataItem] = {}
for attr_name, value in data_dict.items():
if attr_name == "input_ids":
@@ -541,16 +500,15 @@ class BaseMultimodalProcessor(ABC):
# Get modality for this attribute
modality = self.ATTR_NAME_TO_MODALITY.get(attr_name)
if not modality and attr_name == "precomputed_features":
if attr_name == "precomputed_embeddings":
modality_str = data_dict.get("modality")
try:
modality = (
Modality.from_str(modality_str)
if modality_str
else Modality.IMAGE
)
except ValueError:
modality = Modality.IMAGE
modality = Modality.IMAGE
if modality_str:
try:
modality = Modality.from_str(modality_str)
except ValueError:
pass
if modality:
# Create item if needed
if modality not in items:
@@ -559,8 +517,7 @@ class BaseMultimodalProcessor(ABC):
if attr_name in self.FEATURE_NAMES:
attr_name = "feature"
# Set attribute
setattr(items[modality], attr_name, value)
items[modality].set(attr_name, value)
return list(items.values())
@@ -586,6 +543,7 @@ class BaseMultimodalProcessor(ABC):
self,
base_output: BaseMultiModalProcessorOutput,
mm_tokens: MultimodalSpecialTokens,
**kwargs,
) -> Tuple[List[MultimodalDataItem], torch.Tensor, dict]:
"""
Process multimodal data and return the combined multimodal items and input_ids.
@@ -618,7 +576,7 @@ class BaseMultimodalProcessor(ABC):
else:
raise ValueError(f"Unknown multimodal item type: {type(item)}")
# Process items and get input_ids
all_collected_items = []
all_collected_items: list[MultimodalDataItem] = []
input_ids = None
# Handle dict items (already processed)
@@ -634,6 +592,7 @@ class BaseMultimodalProcessor(ABC):
images=raw_images,
audios=raw_audios,
videos=raw_videos,
**kwargs,
)
all_collected_items.extend(collected_items)
else: