From 2de74035a5d2f674ff3cb55e9c019984f7209ca5 Mon Sep 17 00:00:00 2001 From: kousakawang Date: Sun, 24 May 2026 01:36:39 +0800 Subject: [PATCH] [FIX][2/2] fix step3-vl/deepseek-ocr image processor error (#25403) Co-authored-by: kousakawang Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- python/sglang/srt/configs/deepseek_ocr.py | 112 +++++++++++++++++++--- 1 file changed, 100 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/configs/deepseek_ocr.py b/python/sglang/srt/configs/deepseek_ocr.py index b742ff036..9a4041773 100644 --- a/python/sglang/srt/configs/deepseek_ocr.py +++ b/python/sglang/srt/configs/deepseek_ocr.py @@ -1,9 +1,11 @@ import math from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union import torch from PIL import Image, ImageOps +from torchvision.transforms import InterpolationMode +from torchvision.transforms import functional as TF from transformers import ( AutoProcessor, LlamaTokenizerFast, @@ -18,6 +20,8 @@ from sglang.srt.sampling.custom_logit_processor import ( DeepseekOCRNoRepeatNGramLogitProcessor, ) +DeepseekOCRImage = Union[Image.Image, torch.Tensor] + BASE_SIZE = 1024 IMAGE_SIZE = 640 CROP_MODE = True @@ -50,6 +54,77 @@ def get_default_ngram_custom_params() -> Dict[str, Any]: PROMPT = "\n<|grounding|>Convert the document to markdown." +def get_image_size(img: DeepseekOCRImage) -> Tuple[int, int]: + """Return (width, height) for both PIL.Image and torch.Tensor (CHW).""" + if isinstance(img, Image.Image): + return img.size + if isinstance(img, torch.Tensor): + if img.ndim != 3: + raise TypeError(f"Expected CHW image tensor, got shape {tuple(img.shape)}") + return int(img.shape[-1]), int(img.shape[-2]) + raise TypeError(f"Unsupported image type: {type(img)}") + + +def resize_image(img: DeepseekOCRImage, size: Tuple[int, int]) -> DeepseekOCRImage: + """Resize image to (width, height) for both PIL and tensor.""" + if isinstance(img, Image.Image): + return img.resize(size, Image.BICUBIC) + return TF.resize( + img, + [size[1], size[0]], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).contiguous() + + +def crop_image( + img: DeepseekOCRImage, box: Tuple[int, int, int, int] +) -> DeepseekOCRImage: + """Crop image with box=(left, upper, right, lower) for both PIL and tensor.""" + if isinstance(img, Image.Image): + return img.crop(box) + left, upper, right, lower = box + return img[:, upper:lower, left:right].contiguous() + + +def pad_image( + img: DeepseekOCRImage, + target_size: Tuple[int, int], + fill_color: Tuple[int, int, int], +) -> DeepseekOCRImage: + """Fit-and-center-pad image to target_size=(width, height). + + Replaces ImageOps.pad for tensor inputs. + """ + if isinstance(img, Image.Image): + return ImageOps.pad(img, target_size, color=fill_color) + # tensor path: CHW format + _, h, w = img.shape + target_w, target_h = target_size + scale = min(target_w / w, target_h / h) + new_w = int(w * scale) + new_h = int(h * scale) + resized = TF.resize( + img, + [new_h, new_w], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ) + pad_left = (target_w - new_w) // 2 + pad_top = (target_h - new_h) // 2 + if img.dtype == torch.uint8: + fill_tensor = torch.tensor( + list(fill_color), device=img.device, dtype=torch.uint8 + ).view(3, 1, 1) + else: + fill_tensor = torch.tensor( + [c / 255.0 for c in fill_color], device=img.device, dtype=img.dtype + ).view(3, 1, 1) + result = fill_tensor.expand(3, target_h, target_w).clone() + result[:, pad_top : pad_top + new_h, pad_left : pad_left + new_w] = resized + return result.contiguous() + + class DictOutput(object): def items(self): return self.__dict__.items() @@ -110,8 +185,20 @@ class ImageTransform(object): self.transform = T.Compose(transform_pipelines) - def __call__(self, pil_img: Image.Image): - x = self.transform(pil_img) + def __call__(self, img): + if isinstance(img, torch.Tensor): + x = img + if x.dtype == torch.uint8: + x = x.to(torch.float32).div(255) + elif not x.is_floating_point(): + x = x.to(torch.float32) + if self.normalize: + + import torchvision.transforms as T + + x = T.Normalize(self.mean, self.std)(x) + return x + x = self.transform(img) return x @@ -134,7 +221,7 @@ def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_ def dynamic_preprocess( image, min_num=MIN_CROPS, max_num=MAX_CROPS, image_size=640, use_thumbnail=False ): - orig_width, orig_height = image.size + orig_width, orig_height = get_image_size(image) aspect_ratio = orig_width / orig_height # calculate the existing image aspect ratio @@ -158,7 +245,7 @@ def dynamic_preprocess( blocks = target_aspect_ratio[0] * target_aspect_ratio[1] # resize the image - resized_img = image.resize((target_width, target_height)) + resized_img = resize_image(image, (target_width, target_height)) processed_images = [] for i in range(blocks): box = ( @@ -168,11 +255,11 @@ def dynamic_preprocess( ((i // (target_width // image_size)) + 1) * image_size, ) # split the image - split_img = resized_img.crop(box) + split_img = crop_image(resized_img, box) processed_images.append(split_img) assert len(processed_images) == blocks if use_thumbnail and len(processed_images) != 1: - thumbnail_img = image.resize((image_size, image_size)) + thumbnail_img = resize_image(image, (image_size, image_size)) processed_images.append(thumbnail_img) return processed_images, target_aspect_ratio @@ -454,9 +541,10 @@ class DeepseekOCRProcessor(ProcessorMixin): tokenized_str += tokenized_sep images_seq_mask += [False] * len(tokenized_sep) - image_shapes.append(image.size) + img_w, img_h = get_image_size(image) + image_shapes.append((img_w, img_h)) - if image.size[0] <= 640 and image.size[1] <= 640: + if img_w <= 640 and img_h <= 640: crop_ratio = [1, 1] else: if cropping: @@ -468,12 +556,12 @@ class DeepseekOCRProcessor(ProcessorMixin): """process the global view""" if self.image_size <= 640 and not cropping: - image = image.resize((self.image_size, self.image_size)) + image = resize_image(image, (self.image_size, self.image_size)) - global_view = ImageOps.pad( + global_view = pad_image( image, (self.base_size, self.base_size), - color=tuple(int(x * 255) for x in self.image_transform.mean), + tuple(int(x * 255) for x in self.image_transform.mean), ) images_list.append(self.image_transform(global_view))