diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md index fef28142e..a38512342 100644 --- a/docs/serving/image_edit_api.md +++ b/docs/serving/image_edit_api.md @@ -123,7 +123,14 @@ When `stream=false` or omitted, the endpoint returns the standard image edit res } ], "output_format": null, - "size": null + "size": null, + "metrics": { + "stage_durations": { + "queue_wait_ms": 1.0, + "stage_0_gen_ms": 2.0 + }, + "peak_memory_mb": 1024.0 + } } ``` diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 2136d84f8..b2eb2edfc 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -117,10 +117,23 @@ def test_encode_image_base64(): class MockGenerationResult: """Mock result object compatible with current diffusion output shape.""" - def __init__(self, images): + def __init__(self, images, stage_durations=None, peak_memory_mb=0.0): self.images = images - self.stage_durations = {} - self.peak_memory_mb = 0.0 + self.stage_durations = {} if stage_durations is None else stage_durations + self.peak_memory_mb = peak_memory_mb + + +def _stub_engine_generate_with_metrics(engine, *, stage_durations, peak_memory_mb): + """Stub engine.generate to return profiler metrics.""" + + async def generate(*args, **kwargs): + yield MockGenerationResult( + [Image.new("RGB", (16, 16), color="blue")], + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + engine.generate = generate class MockStageResult: @@ -972,6 +985,53 @@ def test_image_edits_streaming_rejects_single_stage_before_loading_url(test_clie assert "multi-stage" in response.json()["detail"] +def test_generate_images_returns_metrics_single_stage(test_client): + """Single-stage /v1/images/generations copies getattr metrics onto the response.""" + stage_durations = {"queue_wait_ms": 1.0, "stage_0_gen_ms": 2.0} + peak_memory_mb = 1024.0 + _stub_engine_generate_with_metrics( + test_client.app.state.engine_client, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + response = test_client.post( + "/v1/images/generations", + json={"prompt": "a cat", "n": 1, "size": "256x256"}, + ) + assert response.status_code == 200 + assert response.json()["metrics"] == { + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + } + + +def test_generate_images_returns_metrics_multistage(async_omni_test_client): + """Multi-stage /v1/images/generations unpacks generate_diffusion_images metrics.""" + stage_durations = { + "queue_wait_ms": 1.0, + "preprocess_ms": 2.0, + "stage_0_gen_ms": 3.0, + "stage_1_gen_ms": 4.0, + } + peak_memory_mb = 2048.0 + _stub_engine_generate_with_metrics( + async_omni_test_client.app.state.engine_client, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + response = async_omni_test_client.post( + "/v1/images/generations", + json={"prompt": "a cat", "n": 1, "size": "256x256"}, + ) + assert response.status_code == 200 + assert response.json()["metrics"] == { + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + } + + def test_generate_images_max_size_rejected(async_omni_test_client): """Test that a size exceeding max_generated_image_size returns 400.""" response = async_omni_test_client.post( @@ -1987,6 +2047,56 @@ def test_image_edit_with_seed_zero(async_omni_test_client): ) +def test_image_edits_returns_metrics_single_stage(test_client): + """Single-stage /v1/images/edits copies getattr metrics onto the response.""" + stage_durations = {"queue_wait_ms": 1.0, "stage_0_gen_ms": 2.0} + peak_memory_mb = 1024.0 + _stub_engine_generate_with_metrics( + test_client.app.state.engine_client, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + response = test_client.post( + "/v1/images/edits", + files=[("image", make_test_image_bytes((16, 16)))], + data={"prompt": "edit me"}, + ) + assert response.status_code == 200 + assert response.json()["metrics"] == { + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + } + + +def test_image_edits_returns_metrics_multistage(async_omni_test_client): + """Multi-stage /v1/images/edits unpacks generate_diffusion_images metrics.""" + stage_durations = { + "queue_wait_ms": 1.0, + "preprocess_ms": 2.0, + "ar2diffusion_ms": 3.0, + "stage_0_gen_ms": 4.0, + "stage_1_gen_ms": 5.0, + } + peak_memory_mb = 2048.0 + _stub_engine_generate_with_metrics( + async_omni_test_client.app.state.engine_client, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", make_test_image_bytes((16, 16)))], + data={"prompt": "edit me"}, + ) + assert response.status_code == 200 + assert response.json()["metrics"] == { + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + } + + def test_image_edit_with_seed_zero_single_stage(test_client): """Test that seed=0 is correctly handled in image editing (single stage). diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 71020bd84..57adaad08 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -2377,7 +2377,7 @@ async def edit_images( status_code=generation_result.error.code if generation_result.error else 400, detail=generation_result.message, ) - images, _, _, cot_output = generation_result + images, stage_durations, peak_memory_mb, cot_output = generation_result else: # Single-stage diffusion: use the direct path. result = await _generate_with_async_omni( @@ -2388,6 +2388,8 @@ async def edit_images( request_id=request_id, ) images = _extract_images_from_result(result) + stage_durations = getattr(result, "stage_durations", None) + peak_memory_mb = getattr(result, "peak_memory_mb", None) logger.debug(f"Successfully generated {len(images)} image(s)") @@ -2408,6 +2410,10 @@ async def edit_images( output_format=output_format, size=size_str, cot_output=cot_output, + metrics={ + "stage_durations": stage_durations or None, + "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, + }, ) except (EngineGenerateError, EngineDeadError) as exc: