[Perf][Frontend] Chunk in-memory image file responses (#7459)

Signed-off-by: QI JIA <qi.jia@shengshu.ai>
Co-authored-by: QI JIA <qi.jia@shengshu.ai>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/entrypoints/openai_api/test_image_protocol.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import base64
+import io
+import zipfile
+
+import pytest
+from fastapi.responses import StreamingResponse
+
+from vllm_omni.entrypoints.openai.protocol.images import (
+    _FILE_RESPONSE_CHUNK_SIZE,
+    ImageData,
+    ImageGenerationResponse,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_NEWLINE_HEAVY_PAYLOAD = (b"\x89PNG\n\x1a\n" * 64) + bytes(range(256)) * 8
+
+
+def _response(*payloads: bytes) -> ImageGenerationResponse:
+    return ImageGenerationResponse(
+        created=0,
+        data=[ImageData(b64_json=base64.b64encode(payload).decode()) for payload in payloads],
+        output_format="png",
+        size="1x1",
+    )
+
+
+async def _collect(response: StreamingResponse) -> list[bytes]:
+    return [bytes(chunk) async for chunk in response.body_iterator]
+
+
+@pytest.mark.asyncio
+async def test_single_file_response_chunks_ignore_newlines() -> None:
+    payload = _NEWLINE_HEAVY_PAYLOAD
+    assert b"\n" in payload
+
+    response = _response(payload).stream_response()
+    chunks = await _collect(response)
+
+    assert len(chunks) == 1
+    assert b"".join(chunks) == payload
+    assert response.headers["content-length"] == str(len(payload))
+    assert response.headers["content-type"] == "image/png"
+
+
+@pytest.mark.asyncio
+async def test_single_file_response_chunks_are_bounded() -> None:
+    payload = _NEWLINE_HEAVY_PAYLOAD * 200
+    expected_chunks = -(-len(payload) // _FILE_RESPONSE_CHUNK_SIZE)
+    assert expected_chunks > 1
+
+    chunks = await _collect(_response(payload).stream_response())
+
+    assert len(chunks) == expected_chunks
+    assert max(len(chunk) for chunk in chunks) <= _FILE_RESPONSE_CHUNK_SIZE
+    assert b"".join(chunks) == payload
+
+
+@pytest.mark.asyncio
+async def test_zip_file_response_streams_whole_archive() -> None:
+    payloads = (_NEWLINE_HEAVY_PAYLOAD, b"second image\nwith a newline")
+
+    response = _response(*payloads).stream_response()
+    chunks = await _collect(response)
+
+    zip_bytes = b"".join(chunks)
+    assert len(chunks) == -(-len(zip_bytes) // _FILE_RESPONSE_CHUNK_SIZE)
+    assert response.headers["content-length"] == str(len(zip_bytes))
+    assert response.headers["content-type"] == "application/zip"
+    with zipfile.ZipFile(io.BytesIO(zip_bytes)) as archive:
+        assert archive.read("image_0.png") == payloads[0]
+        assert archive.read("image_1.png") == payloads[1]