[Bugfix] fix: handle Qwen-Image-Layered layered RGBA output for jpeg edits (#2297)

Signed-off-by: David Chen <530634352@qq.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

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

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