diff --git a/docs/user_guide/examples/online_serving/bagel.md b/docs/user_guide/examples/online_serving/bagel.md index 6c10c4349..6b1e20391 100644 --- a/docs/user_guide/examples/online_serving/bagel.md +++ b/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 | diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md index 19ace3d33..17e10fb69 100644 --- a/recipes/Bagel/BAGEL-7B-MoT.md +++ b/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 diff --git a/tests/diffusion/models/bagel/test_step_execution.py b/tests/diffusion/models/bagel/test_step_execution.py index 40d4154e3..e5f8e5d5a 100644 --- a/tests/diffusion/models/bagel/test_step_execution.py +++ b/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: diff --git a/tests/e2e/online_serving/test_bagel.py b/tests/e2e/online_serving/test_bagel.py index 84c2302cc..1a11dc037 100644 --- a/tests/e2e/online_serving/test_bagel.py +++ b/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) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 60850248e..1c124a19b 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index d805fd3b8..25f31679c 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/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: - image_shape = tuple(sampling.kv_metadata["image_shape"]) + 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 - image_shape = (resized_h, resized_w) - logger.info(f"img2img: resized image to {resized_w}x{resized_h}") + if not image_shape_requested: + image_shape = (resized_h, resized_w) + 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": diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 6541c11df..3bf3f3ff5 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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, diff --git a/vllm_omni/entrypoints/openai/images/helpers.py b/vllm_omni/entrypoints/openai/images/helpers.py index dd790ca3d..4001409fe 100644 --- a/vllm_omni/entrypoints/openai/images/helpers.py +++ b/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], *,