From 8d1d80e2b5e1ed4a3900b73d0e0b5a08952a50e6 Mon Sep 17 00:00:00 2001 From: Roger Wang Date: Sun, 8 Dec 2024 01:55:51 -0800 Subject: [PATCH] iterate Signed-off-by: Roger Wang --- vllm/model_executor/models/internvl.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index 064a7e0bdcd3b..42c769f79e202 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -743,11 +743,12 @@ def forward( "intermediate_tensors": intermediate_tensors, "inputs_embeds": inputs_embeds, } - if self.img_context_token_id is not None: - visual_token_mask = self._get_visual_token_mask(input_ids) - if self.is_mono: - forward_kwargs.update({"visual_token_mask": visual_token_mask}) + # Only required if the model is mono-architecture + if self.visual_token_mask is not None: + forward_kwargs.update( + {"visual_token_mask": self.visual_token_mask}) + self.visual_token_mask = None hidden_states = self.language_model.model(**forward_kwargs) return hidden_states