[Bugfix] Return diffusion metrics from /v1/images/generations (#5278)

Signed-off-by: Zhiyu0603 <Zhiyu0603@users.noreply.github.com>
Co-authored-by: Zhiyu0603 <Zhiyu0603@users.noreply.github.com>

vllm_omni/entrypoints/openai/api_server.py

@@ -1830,10 +1830,17 @@ async def generate_images(
                     status_code=generation_result.error.code if generation_result.error else 400,
                     content=generation_result.model_dump(),
                 )
-            flat_images, _, _, _ = generation_result
+            flat_images, stage_durations, peak_memory_mb, _ = generation_result
             image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in flat_images]
 
-            return ImageGenerationResponse(created=int(time.time()), data=image_data)
+            return ImageGenerationResponse(
+                created=int(time.time()),
+                data=image_data,
+                metrics={
+                    "stage_durations": stage_durations or None,
+                    "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None,
+                },
+            )
 
         # Build params - pass through user values directly
         prompt: OmniTextPrompt = {"prompt": request.prompt, "modalities": ["image"]}
@@ -1899,7 +1906,7 @@ async def generate_images(
 
         logger.debug(f"Generating {request.n} image(s) {size_str}")
 
-        # Generate images using AsyncOmni (multi-stage mode)
+        # Generate images using AsyncOmni.
         result = await _generate_with_async_omni(
             engine_client=engine_client,
             gen_params=gen_params,
@@ -1929,10 +1936,16 @@ async def generate_images(
             for img in images
         ]
 
+        stage_durations = getattr(result, "stage_durations", None)
+        peak_memory_mb = getattr(result, "peak_memory_mb", None)
         response_kwargs = {
             "created": int(time.time()),
             "data": image_data,
             "output_format": output_format,
+            "metrics": {
+                "stage_durations": stage_durations or None,
+                "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None,
+            },
         }
         if request.size:
             response_kwargs["size"] = size_str

vllm_omni/entrypoints/openai/protocol/images.py

@@ -190,6 +190,10 @@ class ImageGenerationResponse(BaseModel):
         description="Chain-of-thought text output from the AR stage. "
         "Only present for image editing (IT2I) with CoT-enabled models.",
     )
+    metrics: dict[str, Any] | None = Field(
+        default=None,
+        description="Per-request generation metrics.",
+    )
 
     def stream_response(self) -> StreamingResponse:
         if not self.data or not self.data[0].b64_json: