A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
[Bugfix] Multi-video inference on LLaVA-Onevision (#15082)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Isotr0py <2037008807@qq.com>
Cyrus Leung committed
Mar 20, 2025 at 22:10 UTC
27261e40a6bf1cfd4ece68e55540a0388606ab47
1 file changed
+38
-45
vllm/model_executor/models/llava_onevision.py
+38
-45
@@ -25,7 +25,6 @@ from vllm.multimodal.parse import (ImageSize, MultiModalDataItems,
25
from vllm.multimodal.processing import PromptReplacement, PromptUpdate
26
from vllm.multimodal.profiling import ProcessorInputs
27
from vllm.sequence import IntermediateTensors
28
-from vllm.utils import is_list_of
28
29
from .clip import CLIPVisionModel
30
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
@@ -44,7 +43,7 @@ class LlavaOnevisionVideoPixelInputs(TypedDict):
43
type: Literal["pixel_values_videos"]
44
pixel_values_videos: Union[torch.Tensor, list[torch.Tensor]]
45
"""
47
- Shape: `(batch_size, num_videos, num_frames, num_channels, height, width)`
46
+ Shape: `(batch_size * num_videos, num_frames, num_channels, height, width)`
47
48
Note that `num_videos` may be different for each batch, and 'num_frames'
49
may be different for each video, in which case the data is passed as a
@@ -580,7 +579,7 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,
579
580
return LlavaOnevisionVideoPixelInputs(
581
type="pixel_values_videos",
583
- pixel_values_videos=pixel_values_videos,
582
+ pixel_values_videos=flatten_bn(pixel_values_videos),
583
)
584
585
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
@@ -768,22 +767,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,
767
for i, patch_features_batch in enumerate(patch_embeddings)
768
]
769
771
- def _add_image_newline(
772
- self,
773
- video_features: torch.Tensor,
774
- videos: int = 1,
775
- frames: int = 1,
776
- strategy: str = "one_token",
777
- ) -> torch.Tensor:
778
- if strategy == "one_token":
779
- video_features = video_features.reshape(
780
- videos, frames * video_features.shape[1], -1)
781
- image_newline = self.image_newline[None, None, :].repeat(
782
- videos, 1, 1).to(video_features.device)
783
- video_features = torch.cat((video_features, image_newline), dim=1)
784
- return video_features
785
- raise ValueError(f"Unexpected video newline strategy: {strategy}")
786
-
770
def _video_pixels_to_features(
771
self,
772
vision_tower: Union[CLIPVisionModel, SiglipVisionModel],
@@ -807,33 +790,43 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,
790
video_pixels = inputs["pixel_values_videos"]
791
792
if isinstance(video_pixels, torch.Tensor):
810
- b, num_videos, frames, c, h, w = video_pixels.shape
811
- pixel_values = video_pixels.view(b * num_videos * frames, c, h, w)
812
- stacked_embeddings = self._video_pixels_to_features(
813
- self.vision_tower, pixel_values)
814
- stacked_embeddings = self._add_image_newline(stacked_embeddings,
815
- videos=b * num_videos,
816
- frames=frames,
817
- strategy="one_token")
818
- return stacked_embeddings
819
- elif is_list_of(video_pixels, torch.Tensor):
820
- stacked_embeddings = []
821
- for video_pixel in video_pixels:
822
- num_videos, frames, c, h, w = video_pixel.shape
823
- pixel_values = video_pixel.view(num_videos * frames, c, h, w)
824
- embeddings = self._video_pixels_to_features(
825
- self.vision_tower, pixel_values)
826
- embeddings = self._add_image_newline(embeddings,
827
- videos=num_videos,
828
- frames=frames,
829
- strategy="one_token")
830
- stacked_embeddings.append(embeddings)
831
- return stacked_embeddings
832
- else:
833
- raise ValueError(
834
- f"Unsupported type of video input {type(video_pixels)}")
793
+ total_videos, frames, c, h, w = video_pixels.shape
794
+ video_pixels_flat = video_pixels.view(total_videos * frames, c, h,
795
+ w)
796
+
797
+ embeddings_flat = self._video_pixels_to_features(
798
+ self.vision_tower, video_pixels_flat)
799
+
800
+ embeddings_flat = embeddings_flat.reshape(
801
+ total_videos, frames * embeddings_flat.shape[1], -1)
802
+
803
+ image_newline = self.image_newline[None, None, :].expand(
804
+ total_videos, -1, -1)
805
+ return torch.cat((embeddings_flat, image_newline), dim=1)
806
+
807
+ frames_per_video = [len(video) for video in video_pixels]
808
+ video_pixels_flat = torch.cat(video_pixels)
809
+
810
+ embeddings_flat = self._video_pixels_to_features(
811
+ self.vision_tower, video_pixels_flat)
812
+
813
+ image_newline = self.image_newline[None, None, :]
814
+
815
+ return [
816
+ torch.cat(
817
+ (
818
+ embeds.reshape(1, num_frame * embeddings_flat.shape[1],
819
+ -1),
820
+ image_newline,
821
+ ),
822
+ dim=1,
823
+ ) for num_frame, embeds in zip(
824
+ frames_per_video,
825
+ torch.split(embeddings_flat, frames_per_video),
826
+ )
827
+ ]
828
836
- def apply_pooling(self, image_features, stride=2):
829
+ def apply_pooling(self, image_features: torch.Tensor, stride: int = 2):
830
vision_config = self.config.vision_config
831
height = width = vision_config.image_size // vision_config.patch_size
832
batch_frames, _, dim = image_features.shape