diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 036ca003b..61ede2283 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -115,7 +115,7 @@ class MockGenerationResult: class FakeAsyncOmni: """Fake AsyncOmni that yields a single diffusion output.""" - def __init__(self): + def __init__(self, images=None): self.stage_configs = [ SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion"), @@ -123,11 +123,12 @@ class FakeAsyncOmni: self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()] self.captured_sampling_params_list = None self.captured_prompt = None + self._images = images or [Image.new("RGB", (64, 64), color="green")] async def generate(self, prompt, request_id, sampling_params_list): self.captured_sampling_params_list = sampling_params_list self.captured_prompt = prompt - images = [Image.new("RGB", (64, 64), color="green")] + images = [img.copy() for img in self._images] yield MockGenerationResult(images) @@ -205,6 +206,28 @@ def async_omni_test_client(): return TestClient(app) +@pytest.fixture +def async_omni_rgba_test_client(): + """Create test client with mocked AsyncOmni engine returning RGBA output.""" + from fastapi import FastAPI + + from vllm_omni.entrypoints.openai.api_server import router + + app = FastAPI() + app.include_router(router) + + app.state.engine_client = FakeAsyncOmni(images=[Image.new("RGBA", (64, 64), color=(0, 255, 0, 128))]) + app.state.stage_configs = [ + SimpleNamespace(stage_type="llm"), + SimpleNamespace(stage_type="diffusion"), + ] + app.state.args = Namespace( + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}', + max_generated_image_size=1048576, + ) + return TestClient(app) + + @pytest.fixture def async_omni_stage_configs_only_client(): """Create test client with refactored AsyncOmni compatibility surface only.""" @@ -983,6 +1006,27 @@ def test_image_edit_compression_jpeg(test_client): assert len(img_bytes_50) < len(img_bytes_100) +def test_image_edit_rgba_output_converts_to_jpeg(async_omni_rgba_test_client): + img_bytes_1 = make_test_image_bytes((16, 16)) + + response = async_omni_rgba_test_client.post( + "/v1/images/edits", + files=[("image", img_bytes_1)], + data={ + "prompt": "hello world.", + "output_format": "jpeg", + }, + ) + assert response.status_code == 200 + + data = response.json() + img_bytes = base64.b64decode(data["data"][0]["b64_json"]) + img = Image.open(io.BytesIO(img_bytes)) + assert img.format.lower() == "jpeg" + assert img.mode == "RGB" + assert data["output_format"] == "jpeg" + + def test_image_edit_compression_png(async_omni_test_client): img_bytes_1 = make_test_image_bytes((16, 16)) # uploadfile with image key diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index da7361acc..9f5d465b3 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1753,10 +1753,27 @@ def _choose_output_format(output_format: str | None, background: str | None) -> return "jpeg" +def _prepare_image_for_output_format(image: Image.Image, format: str) -> Image.Image: + fmt = format.lower() + if fmt not in {"jpg", "jpeg"}: + return image + + if image.mode == "RGB": + return image + + if image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info): + alpha_image = image.convert("RGBA") + flattened = Image.new("RGB", alpha_image.size, (255, 255, 255)) + flattened.paste(alpha_image, mask=alpha_image.getchannel("A")) + return flattened + + return image.convert("RGB") + + def _encode_image_base64_with_compression( image: Image.Image, format: str = "png", output_compression: int = 100 ) -> str: - """Encode PIL Image to base64 PNG string. + """Encode PIL Image to a base64 image string. Args: image: PIL Image object @@ -1766,6 +1783,7 @@ def _encode_image_base64_with_compression( Base64-encoded image as string """ buffer = io.BytesIO() + image = _prepare_image_for_output_format(image, format) save_kwargs = {} if format in ("jpg", "jpeg", "webp"): save_kwargs["quality"] = output_compression