diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index cc2b1002e46f..fd4ac127407a 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -26,7 +26,8 @@ is_npu, is_xpu, load_audio, - load_image, + batch_decode_jpeg_gpu, + load_image_tensor, load_video, logger, ) @@ -529,22 +530,17 @@ def _load_single_item( Load a single multimodal data. If data is processor_output or precomputed embedding, return directly. + For images: + - JPEG: returns (img_tensor_bytes, 'jpeg') + - Non-JPEG: returns (img_tensor, None) - Class method that can be pickled for multiprocessing + Class method that can be pickled for multiprocessing. """ if cls._is_preprocessed_input(data): return data try: if modality == Modality.IMAGE: - img, _ = load_image(data, cls.gpu_image_decode) - if isinstance(img, torch.Tensor): - return img # JPEG already decoded on GPU by nvJPEG - # PIL decodes lazily; do it here in the io worker so the decode - # doesn't run later on the event-loop thread. - if discard_alpha_channel and img.mode != "RGB": - return img.convert("RGB") - img.load() - return img + return load_image_tensor(data, discard_alpha_channel) elif modality == Modality.VIDEO: return load_video(data, frame_count_limit) elif modality == Modality.AUDIO: @@ -658,7 +654,10 @@ def submit_data_loading_tasks( audio_sample_rate: Optional[int] = None, ) -> Tuple[List, List]: """ - load multimodal data parallelly using iterators. + Load multimodal data parallelly using iterators. + For images, will return either: + - JPEG: (img_tensor_bytes, 'jpeg') for batch decoding + - Non-JPEG: (img_tensor, None) already decoded """ futures = [] task_info = [] @@ -913,6 +912,11 @@ async def fast_load_mm_data( images: List[Any] = [None] * len(image_data) if image_data else [] videos: List[Any] = [None] * len(video_data) if video_data else [] audios: List[Any] = [None] * len(audio_data) if audio_data else [] + target_device = torch.device(f"cuda:{self.server_args.base_gpu_id}") + + # Track JPEG images for batch decoding + jpeg_indices: List[int] = [] # indices into images list + jpeg_bytes_list: List[Any] = [] # preprocessed JPEG data for modality, idx, future in futures: try: @@ -928,15 +932,36 @@ async def fast_load_mm_data( ) if modality == Modality.IMAGE: - images[idx] = result + # Check if this is a JPEG tuple result for batch decoding + if isinstance(result, tuple) and len(result) == 2: + img_data, format_type = result + if format_type == "jpeg": + # JPEG image - record position for batch decoding + jpeg_indices.append(idx) + jpeg_bytes_list.append(img_data) + images[idx] = None # Placeholder, will be replaced later + else: + # Non-JPEG image - already decoded, move to cuda directly + images[idx] = img_data.to(target_device) + else: + images[idx] = result elif modality == Modality.VIDEO: videos[idx] = result elif modality == Modality.AUDIO: audios[idx] = result + # Batch decode all JPEG images on GPU + if jpeg_bytes_list: + decoded_images = batch_decode_jpeg_gpu( + jpeg_bytes_list, device=target_device + ) + for img_idx, decoded_img in zip(jpeg_indices, decoded_images): + images[img_idx] = decoded_img.to(target_device) + logger.debug( - "[load_mm_data(simple)] loaded counts: images=%d, videos=%d, audios=%d", + "[load_mm_data(simple)] loaded counts: images=%d (jpeg_batch=%d), videos=%d, audios=%d", len(images), + len(jpeg_bytes_list), len(videos), len(audios), ) @@ -999,8 +1024,49 @@ async def legacy_load_mm_data( discard_alpha_channel=discard_alpha_channel, audio_sample_rate=audio_sample_rate, ) + final_results = [] + jpeg_indices = [] # Track which positions are JPEG images + jpeg_bytes_list = [] # Collect preprocessed JPEG data + target_device = torch.device(f"cuda:{self.server_args.base_gpu_id}") + + for idx, (future, (modality, raw_data, frame_limit)) in enumerate( + zip(futures, task_info) + ): + result = await asyncio.wrap_future(future) + + # Check if this is an image tuple result + if ( + modality == Modality.IMAGE + and isinstance(result, tuple) + and len(result) == 2 + ): + img_data, format_type = result + + if format_type == "jpeg": + # JPEG image - record position and preprocessed data + jpeg_indices.append(idx) + jpeg_bytes_list.append(img_data) + final_results.append(None) # Placeholder, will be replaced later + else: + # Non-JPEG image - already decoded, move to cuda directly + final_results.append(img_data.to(target_device)) + else: + # Non-image data or precomputed data + final_results.append(result) + + # Batch decode all JPEG images + if jpeg_bytes_list: + decoded_images = batch_decode_jpeg_gpu( + jpeg_bytes_list, device=target_device + ) + + # Put decoded images back to their original positions + for img_idx, decoded_img in zip(jpeg_indices, decoded_images): + final_results[img_idx] = decoded_img.to(target_device) + + # Create result iterators task_info_iter = iter(task_info) - futures_iter = iter(futures) + results_iter = iter(final_results) # Process results images, videos, audios = [], [], [] @@ -1010,7 +1076,9 @@ async def legacy_load_mm_data( try: if multimodal_tokens_pattern.match(text_part): modality, raw_data, frame_limit = next(task_info_iter) - result = await asyncio.wrap_future(next(futures_iter)) + result = next( + results_iter + ) # Get from results (already a complete tensor) is_precomputed, new_imgs, new_vids, new_auds = ( self._process_loaded_mm_data(modality, raw_data, result) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 82db90da3bfb..8b4328c235b3 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -91,7 +91,8 @@ from torch import nn from torch.library import Library from torch.utils._contextlib import _DecoratorContextManager -from torchvision.io import decode_jpeg +from torchvision.io import ImageReadMode, decode_image, decode_jpeg, read_image +from torchvision.transforms.v2 import functional as F from typing_extensions import Literal from sglang.srt.environ import envs @@ -1680,6 +1681,221 @@ def load_image( return image, image_size +def _is_jpeg(data): + """ + Detect if data is in JPEG format. + + Args: + data: bytes or torch.Tensor + + Returns: + bool: Whether the data is in JPEG format. + """ + if isinstance(data, torch.Tensor): + if len(data) < 3: + return False + return data[0] == 0xFF and data[1] == 0xD8 and data[2] == 0xFF + elif isinstance(data, (bytes, bytearray)): + if len(data) < 3: + return False + return data[0] == 0xFF and data[1] == 0xD8 and data[2] == 0xFF + return False + + +def decode_single_image_opencl(tensor_bytes): + import cv2 + + np_arr = tensor_bytes.numpy() + bgr_numpy = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) + if bgr_numpy is None: + raise ValueError("OpenCV imdecode returned None") + rgb_numpy = cv2.cvtColor(bgr_numpy, cv2.COLOR_BGR2RGB) + rgb_tensor = torch.from_numpy(rgb_numpy).permute(2, 0, 1).contiguous() + return rgb_tensor + + +def batch_decode_jpeg_gpu(img_tensor_bytes_list: list, device=None): + """ + Batch decode multiple JPEG images. + + Args: + img_tensor_bytes_list: List of encoded JPEG byte data. + device: Target GPU device. + + Returns: + List of decoded image tensors. + """ + if not img_tensor_bytes_list: + return [] + if device is None: + device = f"cuda:{torch.cuda.current_device()}" + + start = time.time() + try: + decoded_images = decode_jpeg( + img_tensor_bytes_list, mode=ImageReadMode.RGB, device=device + ) + total = (time.time() - start) * 1000 + logger.info( + "rocjpeg decode %d images in %.2f ms, avg: %.2f ms", + len(img_tensor_bytes_list), + total, + total / len(img_tensor_bytes_list), + ) + return decoded_images + except Exception: + import cv2 + from concurrent.futures import ThreadPoolExecutor + + cv2.ocl.setUseOpenCL(True) + if not cv2.ocl.haveOpenCL(): + logger.warning("OpenCL is not available, using opencv CPU for image decode") + else: + logger.info("Use_opencv_ocl %s", cv2.ocl.useOpenCL()) + results = [None] * len(img_tensor_bytes_list) + with ThreadPoolExecutor(max_workers=len(img_tensor_bytes_list)) as executor: + future_to_index = { + executor.submit(decode_single_image_opencl, img_bytes): idx + for idx, img_bytes in enumerate(img_tensor_bytes_list) + } + for future in future_to_index: + index = future_to_index[future] + try: + result = future.result() + results[index] = result + except Exception as e: + logger.error("Image at index %d failed to decode: %s", index, e) + results[index] = None + + total = (time.time() - start) * 1000 + logger.warning( + "opencv-ocl decode %d images in %.2f ms, avg: %.2f ms", + len(img_tensor_bytes_list), + total, + total / len(img_tensor_bytes_list), + ) + return results + + +def load_image_tensor( + image_file: Union[Image.Image, str, ImageData, bytes], + discard_alpha_channel: bool = True, +) -> tuple[Image.Image, tuple[int, int]]: + """ + Load image, return preprocessed result for JPEG format. + + For JPEG, return encoded bytes for batch decoding. For non-JPEG, decode + directly and return the image tensor. + """ + + if isinstance(image_file, ImageData): + image_file = image_file.url + if isinstance(image_file, str) and image_file.startswith("file://"): + image_file = unquote(urlparse(image_file).path) + + image = image_size = None + + if isinstance(image_file, Image.Image): + image = image_file + image_size = (image.width, image.height) + img_tensor = F.pil_to_tensor(image) + return img_tensor, None + + elif isinstance(image_file, bytes): + img_tensor_bytes = torch.frombuffer(bytearray(image_file), dtype=torch.uint8) + + if _is_jpeg(img_tensor_bytes): + return img_tensor_bytes, "jpeg" + else: + try: + img_tensor = decode_image(img_tensor_bytes, mode=ImageReadMode.RGB) + except Exception: + image = Image.open(BytesIO(image_file)) + if discard_alpha_channel and image.mode != "RGB": + image = image.convert("RGB") + img_tensor = F.pil_to_tensor(image) + return img_tensor, None + + elif image_file.startswith("http://") or image_file.startswith("https://"): + timeout = int(os.getenv("REQUEST_TIMEOUT", "3")) + response = requests.get(image_file, stream=True, timeout=timeout) + try: + response.raise_for_status() + img_bytes = response.content + img_tensor_bytes = torch.frombuffer(bytearray(img_bytes), dtype=torch.uint8) + + if _is_jpeg(img_tensor_bytes): + return img_tensor_bytes, "jpeg" + else: + try: + img_tensor = decode_image(img_tensor_bytes, mode=ImageReadMode.RGB) + except Exception: + image = Image.open(BytesIO(img_bytes)) + if discard_alpha_channel and image.mode != "RGB": + image = image.convert("RGB") + img_tensor = F.pil_to_tensor(image) + return img_tensor, None + finally: + response.close() + + elif image_file.lower().endswith(("png", "jpg", "jpeg", "webp", "gif")): + is_jpeg = image_file.lower().endswith(("jpg", "jpeg")) + + if is_jpeg: + from torchvision.io import read_file + + img_tensor_bytes = read_file(image_file) + return img_tensor_bytes, "jpeg" + else: + try: + img_tensor = read_image(image_file, mode=ImageReadMode.RGB) + except Exception: + image = Image.open(image_file) + if discard_alpha_channel and image.mode != "RGB": + image = image.convert("RGB") + img_tensor = F.pil_to_tensor(image) + return img_tensor, None + + elif image_file.startswith("data:"): + mime_type = image_file.split(";")[0].split(":")[1] + is_jpeg = "jpeg" in mime_type.lower() or "jpg" in mime_type.lower() + + base64_str = image_file.split(",")[1] + img_bytes = pybase64.b64decode(base64_str, validate=True) + img_tensor_bytes = torch.frombuffer(bytearray(img_bytes), dtype=torch.uint8) + + if is_jpeg: + return img_tensor_bytes, "jpeg" + else: + try: + img_tensor = decode_image(img_tensor_bytes, mode=ImageReadMode.RGB) + except Exception: + image = Image.open(BytesIO(img_bytes)) + if discard_alpha_channel and image.mode != "RGB": + image = image.convert("RGB") + img_tensor = F.pil_to_tensor(image) + + return img_tensor, None + + elif isinstance(image_file, str): + img_bytes = pybase64.b64decode(image_file, validate=True) + img_tensor_bytes = torch.frombuffer(bytearray(img_bytes), dtype=torch.uint8) + + if _is_jpeg(img_tensor_bytes): + return img_tensor_bytes, "jpeg" + else: + try: + img_tensor = decode_image(img_tensor_bytes, mode=ImageReadMode.RGB) + except Exception: + image = Image.open(BytesIO(img_bytes)) + if discard_alpha_channel and image.mode != "RGB": + image = image.convert("RGB") + img_tensor = F.pil_to_tensor(image) + return img_tensor, None + else: + raise ValueError(f"Invalid image: {image_file}") + + def get_image_bytes(image_file: Union[str, bytes]) -> bytes: """Normalize various image inputs into raw bytes.""" if isinstance(image_file, bytes): @@ -1693,7 +1909,9 @@ def get_image_bytes(image_file: Union[str, bytes]) -> bytes: finally: response.close() return result - if image_file.startswith(("file://", "/")): + if image_file.startswith("file://"): + image_file = unquote(urlparse(image_file).path) + if image_file.startswith("/"): with open(image_file, "rb") as f: return f.read() if isinstance(image_file, str) and image_file.startswith("data:"):