[MM] Keep more metadata tensors on CPU (#52827)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
@@ -268,7 +268,7 @@ class Cohere2VisionMultiModalProcessor(
|
||||
num_patches = hf_inputs.get("num_patches", torch.empty(0))
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
|
||||
num_patches=MultiModalFieldConfig.batched("image"),
|
||||
num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
image_embeds=MultiModalFieldConfig.batched("image"),
|
||||
)
|
||||
|
||||
|
||||
@@ -320,7 +320,9 @@ class DeepseekOCRMultiModalProcessor(
|
||||
patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.batched("image"),
|
||||
images_spatial_crop=MultiModalFieldConfig.batched("image"),
|
||||
images_spatial_crop=MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
),
|
||||
images_crop=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", patches_per_image
|
||||
),
|
||||
|
||||
@@ -191,7 +191,9 @@ class DeepseekOCR2MultiModalProcessor(
|
||||
patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.batched("image"),
|
||||
images_spatial_crop=MultiModalFieldConfig.batched("image"),
|
||||
images_spatial_crop=MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
),
|
||||
images_crop=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", patches_per_image
|
||||
),
|
||||
|
||||
@@ -270,7 +270,9 @@ class DeepseekVL2MultiModalProcessor(
|
||||
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
|
||||
images_spatial_crop=MultiModalFieldConfig.batched("image"),
|
||||
images_spatial_crop=MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
),
|
||||
image_embeds=MultiModalFieldConfig.batched("image"),
|
||||
)
|
||||
|
||||
|
||||
@@ -1246,11 +1246,11 @@ class Ernie4_5VLMultiModalProcessor(BaseMultiModalProcessor[Ernie4_5_VLProcessin
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_grid_sizes
|
||||
),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image"),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_grid_sizes
|
||||
),
|
||||
video_grid_thw=MultiModalFieldConfig.batched("video"),
|
||||
video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -260,7 +260,7 @@ class FireRedASR2MultiModalProcessor(
|
||||
return dict(
|
||||
input_features=MultiModalFieldConfig.batched("audio"),
|
||||
speech_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
fake_token_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -482,8 +482,8 @@ class FireRedLIDMultiModalProcessor(
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
return dict(
|
||||
input_features=MultiModalFieldConfig.batched("audio"),
|
||||
speech_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
fake_token_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
speech_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -779,7 +779,7 @@ class FunASRMultiModalProcessor(BaseMultiModalProcessor[FunASRProcessingInfo]):
|
||||
return dict(
|
||||
input_features=MultiModalFieldConfig.batched("audio"),
|
||||
speech_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
fake_token_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -695,7 +695,9 @@ class FunAudioChatMultiModalProcessor(
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
return {
|
||||
"speech_ids": MultiModalFieldConfig.batched("audio"),
|
||||
"speech_attention_mask": MultiModalFieldConfig.batched("audio"),
|
||||
"speech_attention_mask": MultiModalFieldConfig.batched(
|
||||
"audio", keep_on_cpu=True
|
||||
),
|
||||
"input_features": MultiModalFieldConfig.batched("audio"),
|
||||
"feature_attention_mask": MultiModalFieldConfig.batched("audio"),
|
||||
"feature_exist_mask": MultiModalFieldConfig.batched("audio"),
|
||||
|
||||
@@ -625,13 +625,13 @@ def _glmasr_field_config(
|
||||
feature_attention_mask=MultiModalFieldConfig.flat_from_sizes(
|
||||
"audio", chunk_counts, dim=0
|
||||
),
|
||||
chunk_counts=MultiModalFieldConfig.batched("audio"),
|
||||
chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
return dict(
|
||||
audio_embeds=MultiModalFieldConfig.batched("audio"),
|
||||
input_features=MultiModalFieldConfig.batched("audio"),
|
||||
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
|
||||
chunk_counts=MultiModalFieldConfig.batched("audio"),
|
||||
chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -398,7 +398,7 @@ class Granite4VisionMultiModalProcessor(
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.batched("image"),
|
||||
image_sizes=MultiModalFieldConfig.batched("image"),
|
||||
image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -324,10 +324,14 @@ class HCXVisionMultiModalProcessor(BaseMultiModalProcessor[HCXVisionProcessingIn
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
fields = dict(
|
||||
pixel_values_images=MultiModalFieldConfig.batched("image"),
|
||||
image_sizes_images=MultiModalFieldConfig.batched("image"),
|
||||
vision_query_lengths_images=MultiModalFieldConfig.batched("image"),
|
||||
image_sizes_images=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
vision_query_lengths_images=MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
),
|
||||
pixel_values_videos=MultiModalFieldConfig.batched("video"),
|
||||
vision_query_lengths_videos=MultiModalFieldConfig.batched("video"),
|
||||
vision_query_lengths_videos=MultiModalFieldConfig.batched(
|
||||
"video", keep_on_cpu=True
|
||||
),
|
||||
)
|
||||
|
||||
return fields
|
||||
|
||||
@@ -440,7 +440,7 @@ class IsaacMultiModalProcessor(BaseMultiModalProcessor):
|
||||
"pixel_values": MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_grid_sizes
|
||||
),
|
||||
"image_grid_thw": MultiModalFieldConfig.batched("image"),
|
||||
"image_grid_thw": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
}
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -561,8 +561,8 @@ class KananaVMultiModalProcessor(BaseMultiModalProcessor[KananaVProcessingInfo])
|
||||
|
||||
mm_fields_config = dict(
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", pixel_sizes),
|
||||
vision_grid_thw=MultiModalFieldConfig.batched("image"),
|
||||
image_token_thw=MultiModalFieldConfig.batched("image"),
|
||||
vision_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
image_token_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
)
|
||||
return mm_fields_config
|
||||
|
||||
|
||||
@@ -309,13 +309,15 @@ def _keye_field_config(
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
|
||||
image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image"),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_num_patches
|
||||
),
|
||||
video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_num_patches),
|
||||
video_grid_thw=MultiModalFieldConfig.flat_from_sizes("video", video_num_grids),
|
||||
num_frames=MultiModalFieldConfig.batched("video"),
|
||||
video_grid_thw=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_num_grids, keep_on_cpu=True
|
||||
),
|
||||
num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -201,7 +201,7 @@ class KimiAudioDummyInputsBuilder(BaseDummyInputsBuilder[KimiAudioProcessingInfo
|
||||
# Field config for Kimi-Audio multimodal data
|
||||
_KIMIAUDIO_FIELD_CONFIG = {
|
||||
"whisper_input_features": MultiModalFieldConfig.batched("audio"),
|
||||
"feature_attention_mask": MultiModalFieldConfig.batched("audio"),
|
||||
"feature_attention_mask": MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -251,7 +251,7 @@ class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]):
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_grid_sizes
|
||||
),
|
||||
image_grid_hws=MultiModalFieldConfig.batched("image"),
|
||||
image_grid_hws=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -431,20 +431,20 @@ def _create_field_factory(
|
||||
image_embeds=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_embed_grid_sizes
|
||||
),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image"),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
# OV2 first-class MM kwarg: per-patch (t,h,w)
|
||||
# positions required by the 3-D vision RoPE.
|
||||
patch_positions=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_pixel_grid_sizes
|
||||
"image", image_pixel_grid_sizes, keep_on_cpu=True
|
||||
),
|
||||
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_patch_sizes
|
||||
),
|
||||
video_grid_thw=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_num_frames
|
||||
"video", video_num_frames, keep_on_cpu=True
|
||||
),
|
||||
patch_positions_videos=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_patch_sizes
|
||||
"video", video_patch_sizes, keep_on_cpu=True
|
||||
),
|
||||
video_num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
frame_timestamps=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
|
||||
@@ -893,17 +893,21 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing
|
||||
merge_size = self.info.get_hf_config().vision_config.spatial_merge_size
|
||||
fields: dict[str, MultiModalFieldConfig] = dict(
|
||||
**_create_qwen2vl_field_factory(merge_size)(hf_inputs),
|
||||
second_per_grid_ts=MultiModalFieldConfig.batched("video"),
|
||||
video_start_times=MultiModalFieldConfig.batched("video"),
|
||||
second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
video_start_times=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
audio_features=MultiModalFieldConfig.batched("audio"),
|
||||
audio_token_lens=MultiModalFieldConfig.batched("audio"),
|
||||
audio_token_lens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
# video_audio fields: only present when video_audio content was processed
|
||||
if "video_audio_n_segs" in hf_inputs:
|
||||
fields["video_audio_n_segs"] = MultiModalFieldConfig.batched("video")
|
||||
fields["video_audio_n_segs"] = MultiModalFieldConfig.batched(
|
||||
"video", keep_on_cpu=True
|
||||
)
|
||||
# video_audio_seg_lens: list of per-video 1D tensors, batched("video")
|
||||
if "video_audio_seg_lens" in hf_inputs:
|
||||
fields["video_audio_seg_lens"] = MultiModalFieldConfig.batched("video")
|
||||
fields["video_audio_seg_lens"] = MultiModalFieldConfig.batched(
|
||||
"video", keep_on_cpu=True
|
||||
)
|
||||
if "va_audio_features" in hf_inputs:
|
||||
fields["va_audio_features"] = MultiModalFieldConfig.batched("va_audio")
|
||||
return fields
|
||||
|
||||
@@ -472,12 +472,12 @@ def get_version_by_config(config: PretrainedConfig) -> tuple[int, ...]:
|
||||
def _minicpmv_field_config(hf_inputs: Mapping[str, torch.Tensor]):
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.batched("image"),
|
||||
image_sizes=MultiModalFieldConfig.batched("image"),
|
||||
tgt_sizes=MultiModalFieldConfig.batched("image"),
|
||||
image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
image_embeds=MultiModalFieldConfig.batched("image"),
|
||||
video_pixel_values=MultiModalFieldConfig.batched("video"),
|
||||
video_image_sizes=MultiModalFieldConfig.batched("video"),
|
||||
video_tgt_sizes=MultiModalFieldConfig.batched("video"),
|
||||
video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
video_embeds=MultiModalFieldConfig.batched("video"),
|
||||
)
|
||||
|
||||
|
||||
@@ -75,15 +75,17 @@ from .vision import is_vit_use_data_parallel
|
||||
def _minicpmv4_6_field_config(hf_inputs: Mapping[str, torch.Tensor]):
|
||||
fields = dict(
|
||||
pixel_values=MultiModalFieldConfig.batched("image"),
|
||||
tgt_sizes=MultiModalFieldConfig.batched("image"),
|
||||
tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
image_embeds=MultiModalFieldConfig.batched("image"),
|
||||
video_pixel_values=MultiModalFieldConfig.batched("video"),
|
||||
video_image_sizes=MultiModalFieldConfig.batched("video"),
|
||||
video_tgt_sizes=MultiModalFieldConfig.batched("video"),
|
||||
video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
video_embeds=MultiModalFieldConfig.batched("video"),
|
||||
)
|
||||
if "use_vit_merger" in hf_inputs:
|
||||
fields["use_vit_merger"] = MultiModalFieldConfig.batched("image")
|
||||
fields["use_vit_merger"] = MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
)
|
||||
return fields
|
||||
|
||||
|
||||
|
||||
@@ -653,8 +653,8 @@ class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo])
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", patches_per_image
|
||||
),
|
||||
patches_per_image=MultiModalFieldConfig.batched("image"),
|
||||
aspect_ratios=MultiModalFieldConfig.batched("image"),
|
||||
patches_per_image=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
aspect_ratios=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -1208,8 +1208,10 @@ class MolmoMultiModalProcessor(BaseMultiModalProcessor[MolmoProcessingInfo]):
|
||||
images=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
|
||||
image_masks=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
|
||||
image_input_idx=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
|
||||
num_crops=MultiModalFieldConfig.batched("image"),
|
||||
img_patch_id=MultiModalFieldConfig.shared("image", num_images),
|
||||
num_crops=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
img_patch_id=MultiModalFieldConfig.shared(
|
||||
"image", num_images, keep_on_cpu=True
|
||||
),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -2127,26 +2127,30 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):
|
||||
image_token_pooling=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_num_pooled_patches
|
||||
),
|
||||
image_num_crops=MultiModalFieldConfig.batched("image"),
|
||||
image_num_pooled_patches=MultiModalFieldConfig.batched("image"),
|
||||
image_num_patches=MultiModalFieldConfig.batched("image"),
|
||||
image_num_crops=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
image_num_pooled_patches=MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
),
|
||||
image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
image_tokens=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", num_image_tokens
|
||||
),
|
||||
num_image_tokens=MultiModalFieldConfig.batched("image"),
|
||||
num_image_tokens=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_num_crops
|
||||
),
|
||||
video_token_pooling=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_num_pooled_patches
|
||||
),
|
||||
video_num_crops=MultiModalFieldConfig.batched("video"),
|
||||
video_num_pooled_patches=MultiModalFieldConfig.batched("video"),
|
||||
video_num_patches=MultiModalFieldConfig.batched("video"),
|
||||
video_num_crops=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
video_num_pooled_patches=MultiModalFieldConfig.batched(
|
||||
"video", keep_on_cpu=True
|
||||
),
|
||||
video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
video_tokens=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", num_video_tokens
|
||||
),
|
||||
num_video_tokens=MultiModalFieldConfig.batched("video"),
|
||||
num_video_tokens=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -352,9 +352,12 @@ def _mtd_field_config(
|
||||
audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes(
|
||||
"audio",
|
||||
audio_chunk_counts,
|
||||
keep_on_cpu=True,
|
||||
),
|
||||
audio_chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
audio_token_lengths=MultiModalFieldConfig.batched(
|
||||
"audio", keep_on_cpu=True
|
||||
),
|
||||
audio_chunk_counts=MultiModalFieldConfig.batched("audio"),
|
||||
audio_token_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
)
|
||||
return fields
|
||||
|
||||
|
||||
@@ -378,12 +378,16 @@ class MuseGlimmerMultiModalProcessor(
|
||||
if "image_pixel_values" in hf_inputs:
|
||||
fields.update(
|
||||
image_pixel_values=MultiModalFieldConfig.batched("image"),
|
||||
image_feature_sizes=MultiModalFieldConfig.batched("image"),
|
||||
image_feature_sizes=MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
),
|
||||
)
|
||||
if "video_pixel_values" in hf_inputs:
|
||||
fields.update(
|
||||
video_pixel_values=MultiModalFieldConfig.batched("video"),
|
||||
video_feature_sizes=MultiModalFieldConfig.batched("video"),
|
||||
video_feature_sizes=MultiModalFieldConfig.batched(
|
||||
"video", keep_on_cpu=True
|
||||
),
|
||||
)
|
||||
return fields
|
||||
|
||||
|
||||
@@ -386,10 +386,12 @@ class NanoNemotronVLMultiModalProcessor(
|
||||
|
||||
return dict(
|
||||
pixel_values_flat=pixel_values_flat,
|
||||
image_num_patches=MultiModalFieldConfig.batched("image"),
|
||||
image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
image_embeds=MultiModalFieldConfig.batched("image"),
|
||||
num_tokens_per_image=MultiModalFieldConfig.batched("image"),
|
||||
imgs_sizes=MultiModalFieldConfig.batched("image"),
|
||||
num_tokens_per_image=MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
),
|
||||
imgs_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_video_fields_config(self, hf_inputs: BatchFeature):
|
||||
@@ -399,9 +401,9 @@ class NanoNemotronVLMultiModalProcessor(
|
||||
pixel_values_flat_video=MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", video_num_patches
|
||||
),
|
||||
video_num_patches=MultiModalFieldConfig.batched("video"),
|
||||
frames_indices=MultiModalFieldConfig.batched("video"),
|
||||
frame_duration_ms=MultiModalFieldConfig.batched("video"),
|
||||
video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
frames_indices=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
frame_duration_ms=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_audio_fields_config(self, hf_inputs: BatchFeature):
|
||||
|
||||
@@ -386,7 +386,7 @@ class OvisMultiModalProcessor(BaseMultiModalProcessor[OvisProcessingInfo]):
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.batched("image"),
|
||||
grids=MultiModalFieldConfig.batched("image"),
|
||||
grids=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
indicator_tokens=MultiModalFieldConfig.batched("image"),
|
||||
)
|
||||
|
||||
|
||||
@@ -279,7 +279,7 @@ class PaddleOCRVLMultiModalProcessor(
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
return dict(
|
||||
pixel_values=MultiModalFieldConfig.batched("image"),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image"),
|
||||
image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
@@ -918,7 +918,7 @@ class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]):
|
||||
input_image_embeds=MultiModalFieldConfig.batched("image"),
|
||||
image_attention_mask=MultiModalFieldConfig.batched("image"),
|
||||
image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
num_img_tokens=MultiModalFieldConfig.batched("image"),
|
||||
num_img_tokens=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
input_audio_embeds=MultiModalFieldConfig.batched("audio"),
|
||||
)
|
||||
|
||||
|
||||
@@ -301,8 +301,12 @@ def create_qwen2_5_omni_thinker_field_factory(
|
||||
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
|
||||
"audio", audio_feature_lengths, dim=1
|
||||
),
|
||||
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
|
||||
audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
feature_attention_mask=MultiModalFieldConfig.batched(
|
||||
"audio", keep_on_cpu=True
|
||||
),
|
||||
audio_feature_lengths=MultiModalFieldConfig.batched(
|
||||
"audio", keep_on_cpu=True
|
||||
),
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", image_pixel_grid_sizes
|
||||
),
|
||||
@@ -318,7 +322,9 @@ def create_qwen2_5_omni_thinker_field_factory(
|
||||
),
|
||||
video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
|
||||
use_audio_in_video=MultiModalFieldConfig.shared("video", num_videos),
|
||||
use_audio_in_video=MultiModalFieldConfig.shared(
|
||||
"video", num_videos, keep_on_cpu=True
|
||||
),
|
||||
)
|
||||
|
||||
return _qwen2_5_omni_thinker_field_config
|
||||
@@ -1031,7 +1037,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
|
||||
audio_input: Qwen2_5OmniAudioFeatureInputs,
|
||||
) -> torch.Tensor:
|
||||
input_features = audio_input["input_features"]
|
||||
audio_feature_lengths = audio_input["audio_feature_lengths"]
|
||||
# audio_feature_lengths is keep_on_cpu; the audio tower derives
|
||||
# device placement from feature_lens, so move it explicitly.
|
||||
audio_feature_lengths = audio_input["audio_feature_lengths"].to(
|
||||
input_features.device, non_blocking=True
|
||||
)
|
||||
|
||||
audio_feat_lengths, audio_output_lengths = (
|
||||
self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths)
|
||||
|
||||
@@ -130,7 +130,7 @@ def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]):
|
||||
return dict(
|
||||
audio_embeds=MultiModalFieldConfig.batched("audio"),
|
||||
input_features=MultiModalFieldConfig.batched("audio"),
|
||||
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
|
||||
feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -254,8 +254,8 @@ def _qwen3asr_field_config(hf_inputs: Mapping[str, torch.Tensor]):
|
||||
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
|
||||
"audio", audio_feature_lengths, dim=1
|
||||
),
|
||||
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
|
||||
audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
|
||||
feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
audio_feature_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
|
||||
@@ -459,7 +459,11 @@ class Qwen3ASRForConditionalGeneration(
|
||||
audio_input: Qwen2_5OmniAudioFeatureInputs,
|
||||
) -> torch.Tensor:
|
||||
input_features = audio_input["input_features"]
|
||||
audio_feature_lengths = audio_input["audio_feature_lengths"]
|
||||
# audio_feature_lengths is keep_on_cpu; the audio tower derives
|
||||
# device placement from feature_lens, so move it explicitly.
|
||||
audio_feature_lengths = audio_input["audio_feature_lengths"].to(
|
||||
input_features.device, non_blocking=True
|
||||
)
|
||||
|
||||
audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)
|
||||
|
||||
|
||||
@@ -1553,7 +1553,11 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
|
||||
audio_input: Qwen2_5OmniAudioFeatureInputs,
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
input_features = audio_input["input_features"]
|
||||
audio_feature_lengths = audio_input["audio_feature_lengths"]
|
||||
# audio_feature_lengths is keep_on_cpu; the audio tower derives
|
||||
# device placement from feature_lens, so move it explicitly.
|
||||
audio_feature_lengths = audio_input["audio_feature_lengths"].to(
|
||||
input_features.device, non_blocking=True
|
||||
)
|
||||
|
||||
audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)
|
||||
|
||||
|
||||
@@ -202,7 +202,7 @@ class Step3VLMultiModalProcessor(BaseMultiModalProcessor[Step3VLProcessingInfo])
|
||||
),
|
||||
num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
patch_newline_mask=MultiModalFieldConfig.flat_from_sizes(
|
||||
"image", num_patches
|
||||
"image", num_patches, keep_on_cpu=True
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -300,11 +300,17 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
|
||||
|
||||
# Keep these as batched, as they always have batch size as first dim
|
||||
if "audio" in modalities:
|
||||
mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched("audio")
|
||||
mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched(
|
||||
"audio", keep_on_cpu=True
|
||||
)
|
||||
if "image" in modalities:
|
||||
mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image")
|
||||
mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
)
|
||||
# TODO: route to "video" once the video modality is supported
|
||||
mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image")
|
||||
mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
)
|
||||
mm_fields["num_image_patches"] = MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
)
|
||||
@@ -755,6 +761,12 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
|
||||
|
||||
num_image_patches = kwargs.pop("num_image_patches")
|
||||
|
||||
# grid_thw fields are registered keep_on_cpu; restore the on-device
|
||||
# placement that HF get_image_features implementations expect.
|
||||
for key, value in kwargs.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
kwargs[key] = value.to(pixel_values.device, non_blocking=True)
|
||||
|
||||
# The underlying HuggingFace `get_image_features` implementations
|
||||
# contain model-internal syncs (e.g. Idefics3 filters all-zero
|
||||
# padding images via boolean-mask indexing, LlavaOnevision
|
||||
|
||||
@@ -251,7 +251,7 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo
|
||||
"audio", num_chunks, keep_on_cpu=True
|
||||
),
|
||||
# num_chunks can convert audio_chunked to audio batch dimension
|
||||
audio_num_chunks=MultiModalFieldConfig.batched("audio"),
|
||||
audio_num_chunks=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
audio_embeds=MultiModalFieldConfig.batched("audio"),
|
||||
)
|
||||
|
||||
|
||||
@@ -307,12 +307,12 @@ class InklingMultiModalProcessor(BaseMultiModalProcessor[InklingProcessingInfo])
|
||||
return dict(
|
||||
# Ragged per-image patches, grouped by num_patches.
|
||||
pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
|
||||
num_patches=MultiModalFieldConfig.batched("image"),
|
||||
num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
|
||||
# Ragged per-audio frames, grouped by num_audio_tokens.
|
||||
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
|
||||
"audio", num_audio_tokens
|
||||
),
|
||||
num_audio_tokens=MultiModalFieldConfig.batched("audio"),
|
||||
num_audio_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
|
||||
)
|
||||
|
||||
def _get_prompt_updates(
|
||||
|
||||
Reference in New Issue
Block a user