[FIX][1/2] fix step3-vl/deepseek-ocr image processor error (#24701)

Co-authored-by: wanghanpei <wanghanpei@bytedance.com>

python/sglang/srt/multimodal/processors/step3_vl.py

@@ -8,6 +8,7 @@ import torch
 from PIL import Image
 from torchvision import transforms
 from torchvision.transforms import InterpolationMode
+from torchvision.transforms import functional as F
 from transformers import BatchFeature, ProcessorMixin, TensorType
 
 from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput
@@ -20,14 +21,37 @@ from sglang.srt.multimodal.processors.base_processor import (
     MultimodalSpecialTokens,
 )
 
-ImageWithPatches = tuple[Image.Image, list[Image.Image], list[int] | None]
+Step3Image = Union[Image.Image, torch.Tensor]
+ImageWithPatches = tuple[Step3Image, list[Step3Image], list[int] | None]
 
 
 class GPUToTensor(torch.nn.Module):
 
-    def forward(self, raw_image: Union[np.ndarray, Image.Image]) -> torch.Tensor:
+    def forward(
+        self, raw_image: Union[np.ndarray, Image.Image, torch.Tensor]
+    ) -> torch.Tensor:
+        if isinstance(raw_image, torch.Tensor):
+            image_tensor = raw_image
+            if image_tensor.ndim != 3:
+                raise TypeError(
+                    f"Expected CHW image tensor, got shape {tuple(image_tensor.shape)}"
+                )
+            if image_tensor.shape[0] == 1:
+                image_tensor = image_tensor.repeat(3, 1, 1)
+            elif image_tensor.shape[0] != 3:
+                raise TypeError(
+                    f"Expected CHW image tensor with 1 or 3 channels, got shape {tuple(image_tensor.shape)}"
+                )
+            if image_tensor.dtype == torch.uint8:
+                image_tensor = image_tensor.to(torch.float32).div(255)
+            elif not image_tensor.is_floating_point():
+                image_tensor = image_tensor.to(torch.float32)
+            return image_tensor.contiguous()
         if isinstance(raw_image, Image.Image):
-            return transforms.ToTensor()(raw_image)
+            image_tensor = transforms.ToTensor()(raw_image)
+            if torch.cuda.is_available():
+                image_tensor = image_tensor.to(torch.device("cuda"))
+            return image_tensor
         if raw_image.ndim == 2:
             raw_image = raw_image[:, :, None].repeat(3, -1)
         if torch.cuda.is_available():
@@ -91,6 +115,16 @@ class Step3VisionProcessor:
 
 
 class ImagePatcher:
+    def get_image_size(self, img: Step3Image) -> tuple[int, int]:
+        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 determine_window_size(self, long: int, short: int) -> int:
         if long <= 728:
@@ -132,14 +166,16 @@ class ImagePatcher:
             for box in windows
         ], (x_num, y_num)
 
-    def square_pad(self, img: Image.Image) -> Image.Image:
-        w, h = img.size
+    def square_pad(self, img: Step3Image) -> Step3Image:
+        w, h = self.get_image_size(img)
         if w == h:
             return img
         size = max(w, h)
-        padded = Image.new(img.mode, (size, size), 0)
-        padded.paste(img, (0, 0))
-        return padded
+        if isinstance(img, Image.Image):
+            padded = Image.new(img.mode, (size, size), 0)
+            padded.paste(img, (0, 0))
+            return padded
+        return torch.nn.functional.pad(img, (0, size - w, 0, size - h), value=0)
 
     def get_image_size_for_padding(
         self, img_width: int, img_height: int
@@ -182,9 +218,22 @@ class ImagePatcher:
             height_new = window_size * h_ratio
         return int(width_new), int(height_new)
 
-    def patch_crop(self, img: Image.Image, i: int, j: int, th: int, tw: int):
-        target = img.crop((j, i, j + tw, i + th))
-        return target
+    def resize(self, img: Step3Image, size: tuple[int, int]) -> Step3Image:
+        if isinstance(img, Image.Image):
+            return img.resize(size, Image.Resampling.BILINEAR)
+        return F.resize(
+            img,
+            [size[1], size[0]],
+            interpolation=InterpolationMode.BILINEAR,
+            antialias=True,
+        ).contiguous()
+
+    def patch_crop(
+        self, img: Step3Image, i: int, j: int, th: int, tw: int
+    ) -> Step3Image:
+        if isinstance(img, Image.Image):
+            return img.crop((j, i, j + tw, i + th))
+        return img[:, i : i + th, j : j + tw].contiguous()
 
     def get_num_patches(self, img_width: int, img_height: int) -> tuple[int, int]:
         img_width, img_height = self.get_image_size_for_padding(img_width, img_height)
@@ -212,20 +261,20 @@ class ImagePatcher:
             return len(center_list), full_rows
 
     def __call__(
-        self, img: Image.Image
-    ) -> tuple[Image.Image, list[Image.Image], list[bool] | None]:
-        img_width, img_height = img.size
+        self, img: Step3Image
+    ) -> tuple[Step3Image, list[Step3Image], list[bool] | None]:
+        img_width, img_height = self.get_image_size(img)
         new_img_width, new_img_height = self.get_image_size_for_padding(
             img_width, img_height
         )
         if new_img_width != img_width or new_img_height != img_height:
             img = self.square_pad(img)
-            img_width, img_height = img.size
+            img_width, img_height = self.get_image_size(img)
 
         new_img_width, new_img_height = self.get_image_size_for_preprocess(
             img_width, img_height
         )
-        img = img.resize((new_img_width, new_img_height), Image.Resampling.BILINEAR)
+        img = self.resize(img, (new_img_width, new_img_height))
         window_size = self.determine_window_size(
             max(new_img_height, new_img_width), min(new_img_height, new_img_width)
         )
@@ -236,9 +285,7 @@ class ImagePatcher:
                 new_img_width, new_img_height, window_size
             )
             if (new_img_width, new_img_height) != (img_width, img_height):
-                img_for_crop = img.resize(
-                    (new_img_width, new_img_height), Image.Resampling.BILINEAR
-                )
+                img_for_crop = self.resize(img, (new_img_width, new_img_height))
             else:
                 img_for_crop = img
 
@@ -320,7 +367,7 @@ class Step3VLProcessor:
 
     def _convert_images_to_pixel_values(
         self,
-        images: list[Image.Image],
+        images: list[Step3Image],
         is_patch: bool = False,
     ) -> list[torch.Tensor]:
         return [