Merge d11fda69877939c590c42b26c5522bc00dc80b6b into 553741557a57df44a46fb3ab72b6aafeaf945337

deepinfer/ve/processor/image_processor.py

@@ -152,6 +152,16 @@ def safe_resize(
     return layout
 
 
+def exceeds_max_wh_ratio(width: int, height: int, max_wh_ratio: int | None) -> bool:
+    """该图是否超出宽高比上限（即其宽度会被裁掉）。
+
+    compute_image_layout 会把这类图的宽度夹到 height*max_wh_ratio，resize 时要把
+    裁掉的宽度压回来（stretch）而不是等比内接（contain）。判定必须严格一致：恰好
+    等于上限的图没有被裁剪，应当保持长宽比，否则它会被无端拉伸变形。
+    """
+    return max_wh_ratio is not None and width > max_wh_ratio * height
+
+
 def compute_image_layout(
     width: int,
     height: int,
@@ -163,7 +173,7 @@ def compute_image_layout(
 ) -> ImageLayout:
     """把原图尺寸规划成 patch 对齐的 best_width/best_height 及 LLM token 数。"""
     # 超宽图先把宽度夹到 height*max_wh_ratio；该图的 resize 相应走 stretch 而非等比内接
-    if max_wh_ratio is not None and width > height * max_wh_ratio:
+    if exceeds_max_wh_ratio(width, height, max_wh_ratio):
         width = height * max_wh_ratio
 
     current_pixels = width * height
@@ -264,8 +274,13 @@ class ImageProcessorAnyRes(BaseProcessor):
         self.max_wh_ratio = max_wh_ratio
 
     def _is_stretched(self, width: int, height: int) -> bool:
-        """该图是否触发了 compute_image_layout 的超宽裁剪，裁掉的宽度靠 stretch resize 压回来。"""
-        return self.max_wh_ratio is not None and width >= self.max_wh_ratio * height
+        """该图是否触发了 compute_image_layout 的超宽裁剪，裁掉的宽度靠 stretch resize 压回来。
+
+        判定与 compute_image_layout 的裁剪条件共用 exceeds_max_wh_ratio：恰好等于上限
+        的图没有被裁剪，必须走等比内接，否则会被拉伸变形（patch 对齐带来的长宽比误差
+        全都压到图上，而只差一个像素的邻居却不会）。
+        """
+        return exceeds_max_wh_ratio(width, height, self.max_wh_ratio)
 
     def _prepare_image(self, image: Image.Image, target: tuple[int, int], *, stretch: bool = False) -> ResizedImage:
         """CPU resize 到 target 尺寸，保留 uint8 像素（pad / normalize 由 GPU 变换完成）。"""

tests/ve/test_image_stretch_boundary.py

@@ -0,0 +1,112 @@
+"""CPU-only regression test for the ultra-wide resize boundary.
+
+``compute_image_layout`` clamps images that are wider than ``height * max_wh_ratio``
+and the CPU resize then has to squeeze that removed width back (stretch) instead of
+keeping the aspect ratio (contain).  The predicate that picks between the two has to
+agree with the clamp condition exactly: an image whose ratio sits *at* the limit is
+not clamped, so it must keep its aspect ratio.  Otherwise it is distorted by the
+patch-alignment rounding, while a one-pixel narrower neighbour is not.
+
+Only CPU geometry is exercised: no checkpoint, no accelerator.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from PIL import Image
+
+from deepinfer.ve.processor.image_processor import (
+    BasicImageTransform,
+    ImageLayout,
+    ImageProcessorAnyRes,
+    compute_image_layout,
+    exceeds_max_wh_ratio,
+)
+
+PATCH_SIZE = 14
+DOWNSAMPLE_RATIO = 3
+MAX_N_TOKEN = 4096
+MAX_WH_RATIOS = [4, 8, 16]
+HEIGHTS = [100, 250]
+
+
+def _processor(max_wh_ratio: int | None) -> ImageProcessorAnyRes:
+    return ImageProcessorAnyRes(
+        image_transform=BasicImageTransform(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+        patch_size=PATCH_SIZE,
+        downsample_ratio=DOWNSAMPLE_RATIO,
+        dtype=torch.bfloat16,
+        max_n_token=MAX_N_TOKEN,
+        max_wh_ratio=max_wh_ratio,
+    )
+
+
+def _layout(width: int, height: int, max_wh_ratio: int | None = None) -> ImageLayout:
+    return compute_image_layout(
+        width,
+        height,
+        patch_size=PATCH_SIZE,
+        downsample_ratio=DOWNSAMPLE_RATIO,
+        max_n_token=MAX_N_TOKEN,
+        max_wh_ratio=max_wh_ratio,
+    )
+
+
+@pytest.mark.parametrize("max_wh_ratio", MAX_WH_RATIOS)
+@pytest.mark.parametrize("height", HEIGHTS)
+def test_stretch_predicate_mirrors_the_layout_clamp(max_wh_ratio: int, height: int) -> None:
+    """The resize must stretch exactly the images whose width the layout clamped."""
+    processor = _processor(max_wh_ratio)
+    limit = height * max_wh_ratio
+
+    # The layout clamps the width to `height * max_wh_ratio`, so a wider image ends up
+    # with the very same geometry as the clamped one.
+    assert _layout(limit + 4 * PATCH_SIZE, height, max_wh_ratio) == _layout(limit, height)
+
+    # Only images strictly beyond the limit are clamped, so only they may be stretched.
+    assert processor._is_stretched(limit - 1, height) is False
+    assert processor._is_stretched(limit, height) is False
+    assert processor._is_stretched(limit + 1, height) is True
+
+
+@pytest.mark.parametrize("max_wh_ratio", MAX_WH_RATIOS)
+@pytest.mark.parametrize("height", HEIGHTS)
+def test_image_at_the_ratio_limit_keeps_its_aspect_ratio(max_wh_ratio: int, height: int) -> None:
+    """An image exactly at the limit is not clamped, so it must not be stretched."""
+    processor = _processor(max_wh_ratio)
+    limit = height * max_wh_ratio
+
+    layout = _layout(limit, height, max_wh_ratio)
+    resized = processor._prepare_image(
+        Image.new("RGB", (limit, height)),
+        (layout.best_width, layout.best_height),
+        stretch=processor._is_stretched(limit, height),
+    )
+    width, out_height = resized.pixels.shape[1], resized.pixels.shape[0]
+    # Contain keeps the aspect ratio (the remainder is padded by the GPU transform).
+    assert abs(width / out_height - limit / height) / (limit / height) < 0.02
+
+
+@pytest.mark.parametrize("max_wh_ratio", MAX_WH_RATIOS)
+@pytest.mark.parametrize("height", HEIGHTS)
+def test_image_beyond_the_ratio_limit_is_squeezed_into_the_canvas(max_wh_ratio: int, height: int) -> None:
+    """A clamped image keeps the layout canvas and is squeezed into it on purpose."""
+    processor = _processor(max_wh_ratio)
+    limit = height * max_wh_ratio
+    width = limit + 4 * PATCH_SIZE
+
+    layout = _layout(width, height, max_wh_ratio)
+    resized = processor._prepare_image(
+        Image.new("RGB", (width, height)),
+        (layout.best_width, layout.best_height),
+        stretch=processor._is_stretched(width, height),
+    )
+    assert (resized.pixels.shape[1], resized.pixels.shape[0]) == (layout.best_width, layout.best_height)
+
+
+def test_no_ratio_limit_never_stretches() -> None:
+    """Without a limit there is no clamp, so the resize always keeps the aspect ratio."""
+    processor = _processor(None)
+    assert processor._is_stretched(100_000, 1) is False
+    assert exceeds_max_wh_ratio(100_000, 1, None) is False