[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>

docs/user_guide/examples/online_serving/bagel.md

@@ -247,12 +247,18 @@ curl http://localhost:8091/v1/chat/completions \
   -d @payload.json
 ```
 
-BAGEL img2img derives its generated image size from the input image. It
-preserves the input aspect ratio, aligns dimensions to the latent stride, and
-applies the checkpoint size limit; explicit `height` and `width` values do not
-override this model-specific resize policy. In step mode, the effective size is
-resolved before scheduler admission so requests with different resulting
-shapes are not combined.
+When no output size is requested (`size=auto` on `/v1/images/edits`, or no
+`height`/`width` in a chat request), BAGEL img2img derives the canvas from the
+input image: it preserves the input aspect ratio, aligns dimensions to the
+latent stride (16 px), and applies the checkpoint size limit. An explicit
+`size`, `height`, or `width` is honored instead, in both the single-stage and
+the two-stage deployment; the input image is still resized to a stride-aligned
+size for the VAE/ViT prefill. The latent grid needs multiples of 16, so a
+requested side that is not one is floored to the next multiple (1000x700
+generates 992x688) and the server logs a warning; the response reports the
+generated size. A side above the checkpoint limit is rejected with HTTP 400.
+In step mode, the effective size is resolved before scheduler admission so
+requests with different resulting shapes are not combined.
 
 ### Image to Text (img2text)
 
@@ -318,8 +324,8 @@ curl http://localhost:8091/v1/chat/completions \
 | `--server` / `-s` | `http://localhost:8091` | Server URL |
 | `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) |
 | `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
-| `--height` | `512` | Text-to-image output height; BAGEL img2img derives its size from the input image |
-| `--width` | `512` | Text-to-image output width; BAGEL img2img derives its size from the input image |
+| `--height` | `512` | Output height, a multiple of 16; omit both `--height` and `--width` to let img2img derive the size from the input image |
+| `--width` | `512` | Output width, a multiple of 16; omit both `--height` and `--width` to let img2img derive the size from the input image |
 | `--steps` | `25` | Number of inference steps |
 | `--seed` | `42` | Random seed |
 | `--negative` | `None` | Negative prompt for CFG |

recipes/Bagel/BAGEL-7B-MoT.md

@@ -120,12 +120,19 @@ python examples/offline_inference/image_to_image/image_edit.py \
 The `--extra-args` JSON forwards BAGEL-specific parameters (e.g. `cfg_text_scale`,
 `cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`) into
 `OmniDiffusionSamplingParams.extra_args` via the model-extras registry.
-For BAGEL img2img, the pipeline derives the generated image size from the input
-image: it preserves the input aspect ratio, aligns dimensions to the latent
-stride, and applies the checkpoint size limit. Explicit `--height` and
-`--width` values do not override that BAGEL-specific resize policy. Step-mode
-admission resolves this effective size before batching so differently shaped
-requests are scheduled separately.
+For BAGEL img2img, an explicit `--height`/`--width` (or `size` on
+`/v1/images/edits`) sets the output canvas, in both the single-stage and the
+two-stage deployment; the input image is still resized to a stride-aligned size
+for the VAE/ViT prefill. The latent grid needs multiples of 16, so a requested
+side that is not one is floored to the next multiple (1000x700 generates
+992x688) with a server-side warning, and the response reports the generated
+size. A side above the checkpoint limit is rejected with HTTP 400. When no size
+is requested
+(`size=auto`, or no `height`/`width`), the pipeline derives the canvas from the
+input image: it preserves the input aspect ratio, aligns dimensions to the
+latent stride, and applies the checkpoint size limit. Step-mode admission
+resolves this effective size before batching so differently shaped requests are
+scheduled separately.
 
 Run text-to-text with the shared understanding example. BAGEL's default
 `bagel.yaml` deploy config is discovered from the checkpoint, so no model-specific

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

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -35,6 +35,7 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.errors import OmniClientError
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
 from .autoencoder import AutoEncoder, AutoEncoderParams, DistributedAutoEncoder
@@ -184,6 +185,45 @@ def _bagel_effective_image_size(
     return min(resized_height, max_image_size), min(resized_width, max_image_size)
 
 
+def _bagel_canvas_requested(sampling: OmniDiffusionSamplingParams) -> bool:
+    return (sampling.height is not None and not sampling.height_not_provided) or (
+        sampling.width is not None and not sampling.width_not_provided
+    )
+
+
+def _bagel_align_requested_canvas(
+    sampling: OmniDiffusionSamplingParams,
+    *,
+    latent_downsample: int,
+    max_latent_size: int,
+) -> None:
+    """Floor an explicitly requested canvas to the latent stride; reject sides over the checkpoint limit."""
+
+    max_image_size = int(max_latent_size * latent_downsample)
+    resized: list[str] = []
+    for name in ("width", "height"):
+        value = getattr(sampling, name)
+        if value is None or getattr(sampling, f"{name}_not_provided"):
+            continue
+        value = int(value)
+        if value > max_image_size:
+            raise OmniClientError(
+                f"Requested {name}={value} exceeds the BAGEL checkpoint limit of {max_image_size} "
+                f"(max_latent_size={max_latent_size}, latent_downsample={latent_downsample})."
+            )
+        aligned = max(latent_downsample, (value // latent_downsample) * latent_downsample)
+        if aligned != value:
+            setattr(sampling, name, aligned)
+            resized.append(f"{name} {value}->{aligned}")
+    if resized:
+        logger.warning(
+            "BAGEL: requested canvas is not a multiple of the latent stride %d and cannot be "
+            "generated as-is; resizing %s.",
+            latent_downsample,
+            ", ".join(resized),
+        )
+
+
 def get_bagel_pre_process_func(od_config: OmniDiffusionConfig):
     """Resolve BAGEL execution mode and step-batch compatibility."""
 
@@ -194,12 +234,23 @@ def get_bagel_pre_process_func(od_config: OmniDiffusionConfig):
         nonlocal image_geometry
 
         sampling = request.sampling_params
+        modalities = (request.prompt.get("modalities") or []) if isinstance(request.prompt, dict) else []
         kv_metadata = getattr(sampling, "kv_metadata", None) or {}
         image_shape = kv_metadata.get("image_shape")
-        if image_shape is not None:
+        if _bagel_canvas_requested(sampling) and "text" not in modalities:
+            if image_geometry is None:
+                image_geometry = _resolve_bagel_image_geometry(od_config)
+            latent_downsample, max_latent_size = image_geometry
+            _bagel_align_requested_canvas(
+                sampling,
+                latent_downsample=latent_downsample,
+                max_latent_size=max_latent_size,
+            )
+        elif image_shape is not None:
             sampling.height, sampling.width = (int(value) for value in image_shape)
+            sampling.height_not_provided = False
+            sampling.width_not_provided = False
         elif isinstance(request.prompt, dict):
-            modalities = request.prompt.get("modalities") or []
             multi_modal_data = request.prompt.get("multi_modal_data") or {}
             image_input = multi_modal_data.get("img2img")
             if image_input is None and "text" not in modalities:
@@ -220,6 +271,8 @@ def get_bagel_pre_process_func(od_config: OmniDiffusionConfig):
                     latent_downsample=latent_downsample,
                     max_latent_size=max_latent_size,
                 )
+                sampling.height_not_provided = True
+                sampling.width_not_provided = True
 
         extra_args = request.sampling_params.extra_args or {}
         cfg_interval = tuple(extra_args.get("cfg_interval", (0.4, 1.0)))
@@ -528,6 +581,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                 f"latent_downsample={self.bagel.latent_downsample})."
             )
         image_shape = (height, width)
+        image_shape_requested = _bagel_canvas_requested(sampling)
 
         extra_args = getattr(sampling, "extra_args", {}) or {}
         cfg_text_scale = extra_args.get("cfg_text_scale", 4.0)
@@ -568,7 +622,15 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                 gen_context["ropes"] = [seq_len]
 
             if sampling.kv_metadata and "image_shape" in sampling.kv_metadata:
+                if not image_shape_requested:
                     image_shape = tuple(sampling.kv_metadata["image_shape"])
+                elif tuple(sampling.kv_metadata["image_shape"]) != image_shape:
+                    logger.info(
+                        "Generating at requested %dx%d instead of KV image shape %s",
+                        image_shape[1],
+                        image_shape[0],
+                        sampling.kv_metadata["image_shape"],
+                    )
 
             branch_kvs = getattr(sampling, "cfg_branch_past_key_values", None) or {}
             branch_metadata = getattr(sampling, "cfg_branch_kv_metadata", None) or {}
@@ -673,8 +735,15 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                     image_input = [_resize_to_stride(img) for img in image_input]
 
                     resized_w, resized_h = image_input[0].size
+                    if not image_shape_requested:
                         image_shape = (resized_h, resized_w)
-                    logger.info(f"img2img: resized image to {resized_w}x{resized_h}")
+                    logger.info(
+                        "img2img: resized source to %dx%d, generating at %dx%d",
+                        resized_w,
+                        resized_h,
+                        image_shape[1],
+                        image_shape[0],
+                    )
 
                     def vae_transforms(img):
                         if img.mode != "RGB":

vllm_omni/entrypoints/openai/api_server.py

@@ -128,6 +128,7 @@ from vllm_omni.entrypoints.openai.images.helpers import (
     _check_max_generated_image_size,
     _choose_output_format,
     _extract_images_from_result,
+    _generated_size_str,
     _get_max_edit_input_images,
     _load_input_images,
     _update_if_not_none,
@@ -2002,8 +2003,9 @@ def _build_image_generation_response(
             peak_memory_mb=peak_memory_mb,
         ),
     }
-    if request.size is not None:
-        response_kwargs["size"] = request.size
+    size = _generated_size_str(images, request.size)
+    if size is not None:
+        response_kwargs["size"] = size
     response = ImageGenerationResponse(**response_kwargs)
     if request.response_format == ResponseFormat.FILE:
         return response.stream_response()
@@ -2428,6 +2430,8 @@ async def edit_images(
 
         _update_if_not_none(gen_params, "width", width)
         _update_if_not_none(gen_params, "height", height)
+        gen_params.width_not_provided = size_was_auto
+        gen_params.height_not_provided = size_was_auto
 
         # 3.4 Add optional parameters ONLY if provided
         _update_if_not_none(gen_params, "num_inference_steps", num_inference_steps)
@@ -2580,7 +2584,7 @@ async def edit_images(
             created=int(time.time()),
             data=image_data,
             output_format=output_format,
-            size=size_str,
+            size=_generated_size_str(images, size_str),
             cot_output=cot_output,
             metrics=_build_image_response_metrics(
                 response_metrics=response_metrics,

vllm_omni/entrypoints/openai/images/helpers.py

@@ -169,6 +169,14 @@ def _extract_images_from_result(result: Any) -> list[Any]:
     return [_normalize_image(img) for img in flattened]
 
 
+def _generated_size_str(images: list[Any], fallback: str | None) -> str | None:
+    """``"WxH"`` of the first generated image, or ``fallback`` when it is unknown."""
+    size = getattr(images[0], "size", None) if images else None
+    if isinstance(size, tuple) and len(size) == 2 and all(isinstance(value, Integral) for value in size):
+        return f"{size[0]}x{size[1]}"
+    return fallback
+
+
 async def _load_input_images(
     inputs: list[str],
     *,