chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Organize image inputs (#1531)

Browse Source
This commit is contained in:
Liangsheng Yin
2024-09-29 06:28:55 +00:00
committed by GitHub
parent e165a9fc1b
commit fd9ad817ec
8 changed files with 121 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/io_struct.py
+2 -8
View File
@@ -172,12 +172,8 @@ class TokenizedGenerateReqInput:
input_text: str
# The input token ids
input_ids: List[int]
# The pixel values for input images
pixel_values: List[float]
# The hash values of input images
image_hashes: List[int]
# The image sizes
image_sizes: List[List[int]]
# The image input
image_inputs: dict
# The sampling parameters
sampling_params: SamplingParams
# Whether to return the logprobs
@@ -188,8 +184,6 @@ class TokenizedGenerateReqInput:
top_logprobs_num: int
# Whether to stream output
stream: bool
# Modalities of the input images
modalites: Optional[List[str]] = None
# LoRA related
lora_path: Optional[str] = None # None means just use the base model
python/sglang/srt/managers/schedule_batch.py
+37 -14
View File
@@ -102,6 +102,39 @@ class FINISH_ABORT(BaseFinishReason):
}
@dataclass
class ImageInputs:
pixel_values: torch.Tensor
image_hash: int
image_sizes: Optional[list] = None
image_offsets: Optional[list] = None
pad_values: Optional[list] = None
modalities: Optional[list] = None
image_embeds: Optional[List[torch.Tensor]] = None
aspect_ratio_ids: Optional[List[torch.Tensor]] = None
aspect_ratio_mask: Optional[List[torch.Tensor]] = None
@staticmethod
def from_dict(obj, vocab_size):
# Use image hash as fake token_ids, which is then used for prefix matching
ret = ImageInputs(
pixel_values=obj["pixel_values"],
image_hash=hash(tuple(obj["image_hashes"])),
)
image_hash = ret.image_hash
ret.pad_values = [
(image_hash) % vocab_size,
(image_hash >> 16) % vocab_size,
(image_hash >> 32) % vocab_size,
(image_hash >> 64) % vocab_size,
]
ret.image_sizes = obj["image_sizes"]
# Only when pixel values is not None we have modalities
ret.modalities = obj["modalities"]
return ret
class Req:
"""Store all inforamtion of a request."""
@@ -147,11 +180,7 @@ class Req:
self.completion_tokens_wo_jump_forward = 0
# For vision inputs
self.pixel_values = None
self.image_sizes = None
self.image_offsets = None
self.pad_value = None
self.modalities = None
self.image_inputs: Optional[ImageInputs] = None
# Prefix info
self.prefix_indices = []
@@ -654,15 +683,9 @@ class ScheduleBatch:
self.tree_cache.cache_finished_req(req, cur_all_ids)
# re-applying image padding
if req.pixel_values is not None:
(
req.origin_input_ids,
req.image_offsets,
) = model_runner.model.pad_input_ids(
req.origin_input_ids_unpadded,
req.pad_value,
req.pixel_values,
req.image_sizes,
if req.image_inputs is not None:
req.origin_input_ids = model_runner.model.pad_input_ids(
req.origin_input_ids_unpadded, req.image_inputs
)
jump_forward_reqs.append(req)
python/sglang/srt/managers/tokenizer_manager.py
+16 -21
View File
@@ -194,10 +194,9 @@ class TokenizerManager:
)
if self.is_generation:
pixel_values, image_hashes, image_sizes = await self._get_pixel_values(
obj.image_data if not_use_index else obj.image_data[index]
image_inputs = await self._get_image_inputs(
obj, obj.image_data if not_use_index else obj.image_data[index]
)
modalities = obj.modalities
return_logprob = (
obj.return_logprob if not_use_index else obj.return_logprob[index]
)
@@ -248,10 +247,7 @@ class TokenizerManager:
sampling_params = SamplingParams(**obj.sampling_params[0])
sampling_params.max_new_tokens = 0
pixel_values, image_hashes, image_sizes = await self._get_pixel_values(
obj.image_data[0]
)
modalities = obj.modalities
image_inputs = await self._get_image_inputs(obj, obj.image_data[0])
return_logprob = obj.return_logprob[0]
logprob_start_len = obj.logprob_start_len[0]
top_logprobs_num = obj.top_logprobs_num[0]
@@ -262,15 +258,12 @@ class TokenizerManager:
rid,
input_text,
input_ids,
pixel_values,
image_hashes,
image_sizes,
image_inputs,
sampling_params,
return_logprob,
logprob_start_len,
top_logprobs_num,
obj.stream,
modalities,
(
obj.lora_path[index]
if isinstance(obj.lora_path, list)
@@ -369,24 +362,20 @@ class TokenizerManager:
sampling_params = self._get_sampling_params(obj.sampling_params[index])
if self.is_generation:
pixel_values, image_hashes, image_sizes = (
await self._get_pixel_values(obj.image_data[index])
image_inputs = await self._get_image_inputs(
obj, obj.image_data[index]
)
modalities = obj.modalities
tokenized_obj = TokenizedGenerateReqInput(
rid,
input_text,
input_ids,
pixel_values,
image_hashes,
image_sizes,
image_inputs,
sampling_params,
obj.return_logprob[index],
obj.logprob_start_len[index],
obj.top_logprobs_num[index],
obj.stream,
modalities,
(
obj.lora_path[index]
if isinstance(obj.lora_path, list)
@@ -697,10 +686,11 @@ class TokenizerManager:
)
return top_logprobs
async def _get_pixel_values(self, image_data: List[Union[str, bytes]]):
async def _get_image_inputs(self, obj, image_data: List[Union[str, bytes]]):
if not image_data:
return None, None, None
return None
# TODO: move this into a processor for each vision architecture
aspect_ratio = getattr(self.hf_config, "image_aspect_ratio", None)
grid_pinpoints = (
self.hf_config.image_grid_pinpoints
@@ -741,7 +731,12 @@ class TokenizerManager:
else:
raise ValueError(f"Invalid image data: {image_data}")
return pixel_values, image_hashes, image_sizes
return {
"pixel_values": pixel_values,
"image_hashes": image_hashes,
"image_sizes": image_sizes,
"modalities": obj.modalities,
}
async def _process_single_image(
self, image_data: Union[bytes, str], aspect_ratio: str, grid_pinpoints: str
python/sglang/srt/managers/tp_worker.py
+10 -22
View File
@@ -49,6 +49,7 @@ from sglang.srt.managers.policy_scheduler import PolicyScheduler, PrefillAdder
from sglang.srt.managers.schedule_batch import (
FINISH_ABORT,
BaseFinishReason,
ImageInputs,
Req,
ScheduleBatch,
)
@@ -340,29 +341,16 @@ class ModelTpServer:
req = Req(recv_req.rid, recv_req.input_text, recv_req.input_ids)
req.tokenizer = self.tokenizer
req.sampling_params = recv_req.sampling_params
req.pixel_values = recv_req.pixel_values
if req.pixel_values is not None:
# Use image hash as fake token_ids, which is then used
# for prefix matching
image_hash = hash(tuple(recv_req.image_hashes))
req.pad_value = [
(image_hash) % self.model_config.vocab_size,
(image_hash >> 16) % self.model_config.vocab_size,
(image_hash >> 32) % self.model_config.vocab_size,
(image_hash >> 64) % self.model_config.vocab_size,
]
req.image_sizes = recv_req.image_sizes
(
req.origin_input_ids,
req.image_offsets,
) = self.model_runner.model.pad_input_ids(
req.origin_input_ids_unpadded,
req.pad_value,
req.pixel_values,
req.image_sizes,
# Image inputs
if recv_req.image_inputs is not None:
req.image_inputs = ImageInputs.from_dict(
recv_req.image_inputs, self.model_config.vocab_size
)
# Only when pixel values is not None we have modalities
req.modalities = recv_req.modalites
req.origin_input_ids = self.model_runner.model.pad_input_ids(
req.origin_input_ids_unpadded, req.image_inputs
)
req.return_logprob = recv_req.return_logprob
req.top_logprobs_num = recv_req.top_logprobs_num
req.stream = recv_req.stream