Llava-hd Support (#92)

Co-authored-by: Haotian Liu <liuhaotian.cn@gmail.com>

python/sglang/srt/managers/io_struct.py

@@ -62,6 +62,7 @@ class TokenizedGenerateReqInput:
     input_ids: List[int]
     pixel_values: List[float]
     image_hash: int
+    image_size: List[int]
     sampling_params: SamplingParams
     return_logprob: bool
     logprob_start_len: int

python/sglang/srt/managers/router/infer_batch.py

@@ -26,6 +26,7 @@ class Req:
         self.input_ids = []
         self.output_ids = []
         self.pixel_values = None
+        self.image_size = None
         self.image_offset = 0
         self.sampling_params = None
         self.return_logprob = False
@@ -104,6 +105,7 @@ class Batch:
 
     # for multimodal
     pixel_values: List[torch.Tensor] = None
+    image_sizes: List[List[int]] = None
     image_offsets: List[int] = None
 
     # other arguments for control
@@ -195,6 +197,7 @@ class Batch:
             flatten_input_ids, dtype=torch.int32, device=device
         )
         self.pixel_values = [r.pixel_values for r in reqs]
+        self.image_sizes = [r.image_size for r in reqs]
         self.image_offsets = [
             r.image_offset - p_len for r, p_len in zip(reqs, prefix_lens)
         ]

python/sglang/srt/managers/router/model_rpc.py

@@ -203,6 +203,7 @@ class ModelRpcServer(rpyc.Service):
         req = Req(recv_req.rid)
         req.input_ids = recv_req.input_ids
         req.pixel_values = recv_req.pixel_values
+        req.image_size = recv_req.image_size
         if req.pixel_values is not None:
             pad_value = [
                 (recv_req.image_hash) % self.model_config.vocab_size,
@@ -211,7 +212,7 @@ class ModelRpcServer(rpyc.Service):
                 (recv_req.image_hash >> 64) % self.model_config.vocab_size,
             ]
             req.input_ids, req.image_offset = self.model_runner.model.pad_input_ids(
-                req.input_ids, pad_value
+                req.input_ids, pad_value, req.pixel_values.shape, req.image_size
             )
         req.sampling_params = recv_req.sampling_params
         req.return_logprob = recv_req.return_logprob

python/sglang/srt/managers/router/model_runner.py

@@ -409,6 +409,7 @@ class ModelRunner:
         self,
         input_ids,
         pixel_values,
+        image_sizes,
         image_offsets,
         req_pool_indices,
         seq_lens,
@@ -433,6 +434,7 @@ class ModelRunner:
             input_metadata.positions,
             input_metadata,
             pixel_values,
+            image_sizes,
             image_offsets,
         )
 
@@ -441,6 +443,7 @@ class ModelRunner:
             kwargs = {
                 "input_ids": batch.input_ids,
                 "pixel_values": batch.pixel_values,
+                "image_sizes": batch.image_sizes,
                 "image_offsets": batch.image_offsets,
                 "req_pool_indices": batch.req_pool_indices,
                 "seq_lens": batch.seq_lens,

python/sglang/srt/managers/tokenizer_manager.py

@@ -20,6 +20,7 @@ from sglang.srt.managers.io_struct import (
     GenerateReqInput,
     TokenizedGenerateReqInput,
 )
+from sglang.srt.mm_utils import expand2square, process_anyres_image
 from sglang.srt.sampling_params import SamplingParams
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import get_exception_traceback, is_multimodal_model, load_image
@@ -48,14 +49,25 @@ def init_global_processor(server_args: ServerArgs):
     )
 
 
-def get_pixel_values(image_data, processor=None):
+def get_pixel_values(image_data, model_cfg, processor=None):
+    image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
     try:
         processor = processor or global_processor
         image = load_image(image_data)
         image_hash = hash(image_data)
-        pixel_values = processor.image_processor(image)["pixel_values"][0]
+        if image_aspect_ratio == "pad":
+            image = expand2square(
+                image, tuple(int(x * 255) for x in processor.image_processor.image_mean)
+            )
+            pixel_values = processor.image_processor(image)["pixel_values"][0]
+        elif image_aspect_ratio == "anyres":
+            pixel_values = process_anyres_image(
+                image, processor.image_processor, model_cfg.image_grid_pinpoints
+            )
+        else:
+            pixel_values = processor.image_processor(image)["pixel_values"][0]
         pixel_values = pixel_values.astype(np.float16)
-        return pixel_values, image_hash
+        return pixel_values, image_hash, image.size
     except Exception:
         print("Exception in TokenizerManager:\n" + get_exception_traceback())
 
@@ -77,6 +89,7 @@ class TokenizerManager:
         self.hf_config = get_config(
             self.model_path, trust_remote_code=server_args.trust_remote_code
         )
+
         self.context_len = get_context_length(self.hf_config)
 
         if is_multimodal_model(self.model_path):
@@ -104,10 +117,10 @@ class TokenizerManager:
         if self.executor is not None:
             loop = asyncio.get_event_loop()
             return await loop.run_in_executor(
-                self.executor, get_pixel_values, image_data
+                self.executor, get_pixel_values, image_data, self.hf_config
             )
         else:
-            return get_pixel_values(image_data, self.processor)
+            return get_pixel_values(image_data, self.hf_config, self.processor)
 
     async def generate_request(self, obj: GenerateReqInput):
         if self.to_create_loop:
@@ -123,14 +136,17 @@ class TokenizerManager:
                 sampling_params.normalize(self.tokenizer)
                 sampling_params.verify()
             if obj.image_data is None:
-                pixel_values, image_hash = None, None
+                pixel_values, image_hash, image_size = None, None, None
             else:
-                pixel_values, image_hash = await self.get_pixel_values(obj.image_data)
+                pixel_values, image_hash, image_size = await self.get_pixel_values(
+                    obj.image_data
+                )
             tokenized_obj = TokenizedGenerateReqInput(
                 rid=rid,
                 input_ids=input_ids,
                 pixel_values=pixel_values,
                 image_hash=image_hash,
+                image_size=image_size,
                 sampling_params=sampling_params,
                 return_logprob=obj.return_logprob,
                 logprob_start_len=obj.logprob_start_len,
@@ -162,9 +178,9 @@ class TokenizerManager:
                     sampling_params.normalize(self.tokenizer)
                     sampling_params.verify()
                 if obj.image_data[i] is None:
-                    pixel_values, image_hash = None, None
+                    pixel_values, image_hash, image_size = None, None, None
                 else:
-                    pixel_values, image_hash = await self.get_pixel_values(
+                    pixel_values, image_hash, image_size = await self.get_pixel_values(
                         obj.image_data[i]
                     )
                 tokenized_obj = TokenizedGenerateReqInput(
@@ -172,6 +188,7 @@ class TokenizerManager:
                     input_ids=input_ids,
                     pixel_values=pixel_values,
                     image_hash=image_hash,
+                    image_size=image_size,
                     sampling_params=sampling_params,
                     return_logprob=obj.return_logprob[i],
                     logprob_start_len=obj.logprob_start_len[i],