[Bugfix] Honor the requested output canvas in BAGEL img2img (#7287)

Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/diffusion/models/bagel/test_step_execution.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Focused CPU tests for BAGEL step execution and packed batching."""
 
 from __future__ import annotations
@@ -18,12 +18,14 @@ from vllm_omni.diffusion.models.bagel.bagel_transformer import Bagel, NaiveCache
 from vllm_omni.diffusion.models.bagel.pipeline_bagel import (
     BagelGenParams,
     BagelPipeline,
+    _bagel_canvas_requested,
     get_bagel_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import StepScheduler
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.utils import StepRequestState
+from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -118,7 +120,7 @@ def test_bagel_step_preprocessor_separates_incompatible_cfg_settings(first_extra
     assert first.batch_compatibility_key != second.batch_compatibility_key
 
 
-def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
+def _pre_process_with_stub_checkpoint(tmp_path):
     (tmp_path / "config.json").write_text(
         '{"vae_config":{"downsample":8},"latent_patch_size":2,"max_latent_size":32}',
         encoding="utf-8",
@@ -131,9 +133,22 @@ def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
         '{"weight_map":{"latent_pos_embed.pos_embed":"ema.safetensors"}}',
         encoding="utf-8",
     )
-    pre_process = get_bagel_pre_process_func(
-        types.SimpleNamespace(model=str(tmp_path), revision=None, step_execution=True)
+    return get_bagel_pre_process_func(types.SimpleNamespace(model=str(tmp_path), revision=None, step_execution=True))
+
+
+def _edit_request(source_size, sampling_params, request_id="req"):
+    return OmniDiffusionRequest(
+        prompt={
+            "prompt": "edit this image",
+            "multi_modal_data": {"image": [Image.new("RGB", source_size)]},
+        },
+        sampling_params=sampling_params,
+        request_id=request_id,
     )
+
+
+def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
     requests = [
         OmniDiffusionRequest(
             prompt={
@@ -163,6 +178,172 @@ def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
     assert schedule.num_waiting_reqs == 1
 
 
+def test_bagel_pre_process_keeps_explicitly_requested_canvas(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request((512, 1024), OmniDiffusionSamplingParams(num_inference_steps=2, height=1024, width=1024))
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
+    assert _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_aligns_source_derived_canvas_for_size_auto(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (500, 700),
+        OmniDiffusionSamplingParams(
+            num_inference_steps=2,
+            height=700,
+            width=500,
+            height_not_provided=True,
+            width_not_provided=True,
+        ),
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (704, 496)
+    assert not _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_derives_canvas_when_caller_asked_for_none(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request((800, 400), OmniDiffusionSamplingParams(num_inference_steps=2))
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (400, 800)
+    assert (request.sampling_params.height_not_provided, request.sampling_params.width_not_provided) == (True, True)
+    assert not _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_marks_kv_metadata_canvas_as_requested(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (500, 700),
+        OmniDiffusionSamplingParams(
+            num_inference_steps=2,
+            height=700,
+            width=500,
+            height_not_provided=True,
+            width_not_provided=True,
+            kv_metadata={"image_shape": (1024, 1024)},
+        ),
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
+    assert (request.sampling_params.height_not_provided, request.sampling_params.width_not_provided) == (False, False)
+    assert _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_prefers_explicit_canvas_over_kv_metadata(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (1280, 720),
+        OmniDiffusionSamplingParams(
+            num_inference_steps=2,
+            height=1024,
+            width=1024,
+            kv_metadata={"image_shape": (576, 1024)},
+        ),
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (1024, 1024)
+    assert _bagel_canvas_requested(request.sampling_params)
+
+
+def test_bagel_pre_process_floors_unaligned_explicit_canvas_with_warning(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (1280, 720),
+        OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
+    )
+
+    with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.logger.warning") as warning:
+        pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (688, 992)
+    assert _bagel_canvas_requested(request.sampling_params)
+    warning.assert_called_once()
+    assert "width 1000->992, height 700->688" in warning.call_args.args[-1]
+
+
+def test_bagel_pre_process_keeps_aligned_explicit_canvas_silent(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (1280, 720),
+        OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=1024),
+    )
+
+    with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.logger.warning") as warning:
+        pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (512, 1024)
+    warning.assert_not_called()
+
+
+def test_bagel_pre_process_rejects_explicit_canvas_over_checkpoint_limit(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = _edit_request(
+        (1280, 720),
+        OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=2048),
+    )
+
+    with pytest.raises(OmniClientError, match="width=2048 exceeds the BAGEL checkpoint limit of 1024") as excinfo:
+        pre_process(request)
+
+    assert excinfo.value.status_code == 400
+
+
+def test_bagel_pre_process_aligns_text2img_canvas_too(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = OmniDiffusionRequest(
+        prompt="a cat",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
+        request_id="req",
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (688, 992)
+
+
+def test_bagel_pre_process_skips_canvas_alignment_for_text_output(tmp_path):
+    pre_process = _pre_process_with_stub_checkpoint(tmp_path)
+    request = OmniDiffusionRequest(
+        prompt={
+            "prompt": "describe this image",
+            "modalities": ["text"],
+            "multi_modal_data": {"image": [Image.new("RGB", (1280, 720))]},
+        },
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=700, width=1000),
+        request_id="req",
+    )
+
+    pre_process(request)
+
+    assert (request.sampling_params.height, request.sampling_params.width) == (700, 1000)
+    assert request.use_step_execution is False
+
+
+@pytest.mark.parametrize(
+    ("sampling_kwargs", "requested"),
+    [
+        ({"height": 1024, "width": 1024}, True),
+        ({"height": 1024, "width": 512, "height_not_provided": True, "width_not_provided": True}, False),
+        ({}, False),
+    ],
+    ids=["explicit_size", "size_auto", "unset"],
+)
+def test_bagel_canvas_requested_separates_asked_for_from_derived(sampling_kwargs, requested):
+    assert _bagel_canvas_requested(OmniDiffusionSamplingParams(**sampling_kwargs)) is requested
+
+
 def _cache(length: int, value: float, num_layers: int = 1) -> NaiveCache:
     cache = NaiveCache(num_layers)
     if length:

tests/e2e/online_serving/test_bagel.py

@@ -25,9 +25,11 @@ import os
 from io import BytesIO
 
 import pytest
+from PIL import Image
 from vllm.assets.image import ImageAsset
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_image
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
@@ -120,3 +122,55 @@ def test_bagel_img2img_online(omni_server, online_client) -> None:
     }
 
     online_client.send_diffusion_request(request_config)
+
+
+EDIT_SOURCE_SIZE = (1280, 720)
+EDIT_REQUESTED_SIZE = "1024x1024"
+
+
+def _edit_source_jpeg() -> tuple[str, bytes, str]:
+    synthetic = generate_synthetic_image(*EDIT_SOURCE_SIZE, seed=7287)
+    return ("source.jpg", base64.b64decode(synthetic["base64"]), "image/jpeg")
+
+
+def _build_edit_request(size: str) -> dict:
+    return {
+        "files": [("image", _edit_source_jpeg())],
+        "data": {
+            "prompt": IMG2IMG_PROMPT,
+            "size": size,
+            "num_inference_steps": 2,
+            "seed": 42,
+            "response_format": "b64_json",
+        },
+    }
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_bagel_image_edit_honors_explicit_size_two_stage(omni_server, online_client) -> None:
+    """An explicit ``size`` wins over the AR stage's source-derived KV image shape (#7283)."""
+    (response,) = online_client.send_images_edits_http_request(_build_edit_request(EDIT_REQUESTED_SIZE))
+
+    payload = response.json_body
+    assert isinstance(payload, dict)
+    assert payload["size"] == EDIT_REQUESTED_SIZE
+    image = Image.open(BytesIO(base64.b64decode(payload["data"][0]["b64_json"])))
+    assert image.size == (1024, 1024)
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_bagel_image_edit_floors_unaligned_size_two_stage(omni_server, online_client) -> None:
+    """A side that is not a multiple of the latent stride is floored to one (1000x700 -> 992x688)."""
+    (response,) = online_client.send_images_edits_http_request(_build_edit_request("1000x700"))
+
+    payload = response.json_body
+    assert isinstance(payload, dict)
+    assert payload["size"] == "992x688"
+    image = Image.open(BytesIO(base64.b64decode(payload["data"][0]["b64_json"])))
+    assert image.size == (992, 688)

tests/entrypoints/openai_api/test_image_server.py

@@ -1719,7 +1719,7 @@ def test_image_edit_parameter_pass(async_omni_test_client):
         img = Image.open(io.BytesIO(img_bytes))
         assert img.format.lower() == "jpeg"
         assert data["output_format"] == "jpeg"
-        assert data["size"] == "16x24"
+        assert data["size"] == "64x64"
 
 
 def test_image_edit_layers_and_resolution(async_omni_test_client):
@@ -1896,6 +1896,7 @@ def test_image_edit_parameter_default_single_stage(test_client):
 
     assert captured_sampling_params.width == 24
     assert captured_sampling_params.height == 16
+    assert (captured_sampling_params.height_not_provided, captured_sampling_params.width_not_provided) == (True, True)
     assert captured_sampling_params.num_outputs_per_prompt == 1
     assert captured_sampling_params.num_inference_steps == 4
     assert captured_sampling_params.guidance_scale == 7.5
@@ -1913,6 +1914,40 @@ def test_image_edit_parameter_default_single_stage(test_client):
     assert response.status_code == 400
 
 
+def test_image_edit_explicit_size_marks_canvas_provided_single_stage(test_client):
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", make_test_image_bytes((24, 16)))],
+        data={"prompt": "hello world.", "size": "16x24"},
+    )
+
+    assert response.status_code == 200
+    sampling = test_client.app.state.engine_client.captured_sampling_params_list[0]
+    assert (sampling.height, sampling.width) == (24, 16)
+    assert (sampling.height_not_provided, sampling.width_not_provided) == (False, False)
+
+
+def test_image_edit_response_size_reports_generated_image_single_stage(test_client):
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", make_test_image_bytes((24, 16)))],
+        data={"prompt": "hello world.", "size": "16x24"},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["size"] == "64x64"
+
+
+def test_generate_images_response_size_reports_generated_image(test_client):
+    response = test_client.post(
+        "/v1/images/generations",
+        json={"prompt": "a cat", "n": 1, "size": "1024x1024"},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["size"] == "64x64"
+
+
 def test_image_edit_compression_jpeg(test_client):
     img_bytes_1 = make_test_image_bytes((16, 16))
     # uploadfile with image key