model: Support Janus-pro (#3203)

python/sglang/srt/managers/image_processors/base_image_processor.py

@@ -13,6 +13,7 @@ from PIL import Image
 
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import load_image
+from sglang.utils import logger
 
 global global_processor
 
@@ -22,6 +23,13 @@ def get_global_processor():
     return global_processor
 
 
+def init_global_processor(sglang_image_processor, server_args: ServerArgs):
+    """Init the global processor for multi-modal models."""
+    global global_processor
+    transformers.logging.set_verbosity_error()
+    global_processor = sglang_image_processor._build_processor(server_args=server_args)
+
+
 @dataclasses.dataclass
 class BaseImageProcessorOutput:
     image_hashes: list[int]
@@ -119,6 +127,11 @@ class BaseImageProcessor(ABC):
     ) -> BaseImageProcessorOutput:
         """
         Each frame of video/image will be replaced by a single image token
+
+        Args:
+
+            discard_alpha_channel: if True, discards the alpha channel in the returned images
+
         """
         image_hashes, image_sizes = [], []
         all_frames = []
@@ -133,7 +146,7 @@ class BaseImageProcessor(ABC):
         if return_text:
             text_parts = input_text.split(image_token)
 
-        # roughly calculate the max number of frames under the max_req_input_len limit
+        # TODO(mick): load from server_args, env, or sampling_params
         MAX_NUM_FRAMES = 30
         estimated_frames_list = self.get_estimated_frames_list(image_data=image_data)
         total_frame_count = sum(estimated_frames_list)

python/sglang/srt/managers/image_processors/janus_pro.py

@@ -0,0 +1,79 @@
+import asyncio
+from typing import List, Union
+
+from sglang.srt.managers.image_processors.base_image_processor import (
+    BaseImageProcessor as SGLangBaseImageProcessor,
+)
+from sglang.srt.managers.image_processors.base_image_processor import (
+    get_global_processor,
+)
+from sglang.srt.models.deepseek_janus_pro import MultiModalityCausalLM
+
+
+class JanusProProcessor(SGLangBaseImageProcessor):
+    def __init__(self, hf_config, server_args, _processor):
+        super().__init__(hf_config, server_args, _processor)
+
+    @staticmethod
+    def _process_images_task(images, input_text):
+        processor = get_global_processor()
+        result = processor.__call__(
+            prompt=input_text, images=images, return_tensors="pt"
+        )
+        return {
+            "input_ids": result["input_ids"],
+            "pixel_values": result["pixel_values"],
+            "images_emb_mask": result["images_emb_mask"],
+            "im_start_id": processor.image_start_id,
+            "im_end_id": processor.image_end_id,
+            "im_token_id": processor.image_id,
+        }
+
+    async def _process_images(self, images, input_text):
+        if self.executor is not None:
+            loop = asyncio.get_event_loop()
+            image_inputs = await loop.run_in_executor(
+                self.executor,
+                JanusProProcessor._process_images_task,
+                images,
+                input_text,
+            )
+        else:
+            image_inputs = self._processor(
+                images=images, text=input_text, return_tensors="pt"
+            )
+
+        return image_inputs
+
+    async def process_images_async(
+        self,
+        image_data: List[Union[str, bytes]],
+        input_ids,
+        request_obj,
+        max_req_input_len,
+        **kwargs,
+    ):
+        if not image_data:
+            return None
+
+        if not isinstance(image_data, list):
+            image_data = [image_data]
+
+        base_out = self.load_images(
+            input_ids, image_data, "<image_placeholder>", max_req_input_len
+        )
+        images = base_out.all_frames
+        res = await self._process_images(images=images, input_text=base_out.input_text)
+
+        return {
+            "input_ids": res["input_ids"].flatten().tolist(),
+            "pixel_values": res["pixel_values"],
+            "images_emb_mask": res["images_emb_mask"],
+            "image_hashes": base_out.image_hashes,
+            "im_start_id": res["im_start_id"],
+            "im_end_id": res["im_end_id"],
+            "im_token_id": res["im_token_id"],
+        }
+
+
+ImageProcessorMapping = {MultiModalityCausalLM: JanusProProcessor}