anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Frontend] Add omni benchmark support for image and video endpoints (#4728)

Browse Source
This commit is contained in:
ZacheryAU authored and GitHub committed 2026-09-08 09:28:21 +08:00
1 parent e1c59edec6
commit b1ec136bd7
21 files changed
+1626 -72
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/entrypoints/openai/api_server.py
Executable → Regular
+84 -25
View File
@@ -15,7 +15,7 @@ import tempfile
import time
import uuid
from argparse import Namespace
from collections.abc import AsyncIterator, Mapping
from collections.abc import AsyncIterator, Mapping, Sequence
from contextlib import asynccontextmanager, suppress
from http import HTTPStatus
from numbers import Integral
@@ -131,6 +131,7 @@ from vllm_omni.entrypoints.openai.protocol.images import (
from vllm_omni.entrypoints.openai.protocol.videos import (
SecondStr,
SizeStr,
VideoAction,
VideoDeleteResponse,
VideoError,
VideoGenerationRequest,
@@ -1946,12 +1947,26 @@ async def show_available_models(raw_request: Request) -> JSONResponse:
# Image generation API endpoints
def _build_image_response_metrics(
*,
response_metrics: Any,
stage_durations: Any,
peak_memory_mb: Any,
) -> dict[str, Any]:
"""Merge detailed stage metrics with the legacy image timing fields."""
metrics = dict(response_metrics) if isinstance(response_metrics, dict) else {}
metrics["stage_durations"] = stage_durations or None
metrics["peak_memory_mb"] = float(peak_memory_mb) if peak_memory_mb else None
return metrics
def _build_image_generation_response(
*,
images: list[Image.Image],
request: ImageGenerationRequest,
stage_durations: Any,
peak_memory_mb: Any,
response_metrics: Any = None,
) -> ImageGenerationResponse | StreamingResponse:
"""Encode generated images and apply the requested response format."""
output_format = _choose_output_format(request.output_format or "png", None)
@@ -1966,10 +1981,11 @@ def _build_image_generation_response(
"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,
},
"metrics": _build_image_response_metrics(
response_metrics=response_metrics,
stage_durations=stage_durations,
peak_memory_mb=peak_memory_mb,
),
}
if request.size is not None:
response_kwargs["size"] = request.size
@@ -2066,6 +2082,8 @@ async def generate_images(
extra_body["use_system_prompt"] = request.use_system_prompt
if request.system_prompt is not None:
extra_body["system_prompt"] = request.system_prompt
if request.return_stage_metrics is not None:
extra_body["return_stage_metrics"] = request.return_stage_metrics
generation_result = await chat_handler.generate_diffusion_images(
prompt=request.prompt,
@@ -2079,12 +2097,13 @@ async def generate_images(
status_code=generation_result.error.code if generation_result.error else 400,
content=generation_result.model_dump(),
)
flat_images, stage_durations, peak_memory_mb, _ = generation_result
flat_images, stage_durations, peak_memory_mb, _, response_metrics = generation_result
return _build_image_generation_response(
images=flat_images,
request=request,
stage_durations=stage_durations,
peak_memory_mb=peak_memory_mb,
response_metrics=response_metrics,
)
# Build params - pass through user values directly
@@ -2174,11 +2193,13 @@ async def generate_images(
stage_durations = getattr(result, "stage_durations", None)
peak_memory_mb = getattr(result, "peak_memory_mb", None)
response_metrics = getattr(result, "metrics", None) if request.return_stage_metrics else None
return _build_image_generation_response(
images=images,
request=request,
stage_durations=stage_durations,
peak_memory_mb=peak_memory_mb,
response_metrics=response_metrics,
)
except (EngineGenerateError, EngineDeadError) as exc:
@@ -2510,7 +2531,7 @@ async def edit_images(
status_code=generation_result.error.code if generation_result.error else 400,
detail=generation_result.message,
)
images, stage_durations, peak_memory_mb, cot_output = generation_result
images, stage_durations, peak_memory_mb, cot_output, response_metrics = generation_result
else:
# Single-stage diffusion: use the direct path.
result = await _generate_with_async_omni(
@@ -2523,6 +2544,7 @@ async def edit_images(
images = _extract_images_from_result(result)
stage_durations = getattr(result, "stage_durations", None)
peak_memory_mb = getattr(result, "peak_memory_mb", None)
response_metrics = getattr(result, "metrics", None) if return_stage_metrics else None
logger.debug(f"Successfully generated {len(images)} image(s)")
@@ -2543,10 +2565,11 @@ 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,
},
metrics=_build_image_response_metrics(
response_metrics=response_metrics,
stage_durations=stage_durations,
peak_memory_mb=peak_memory_mb,
),
)
except (EngineGenerateError, EngineDeadError) as exc:
@@ -3010,12 +3033,22 @@ def _reference_video_decode_spec(
def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse:
video_params = req.resolve_video_params()
duration_s = None
if req.seconds is not None:
duration_s = float(req.seconds)
elif video_params.num_frames is not None and video_params.fps is not None:
duration_s = video_params.num_frames / video_params.fps
resp = VideoResponse(
model=model_name,
status=VideoGenerationStatus.QUEUED,
size=req.size,
prompt=req.prompt,
quality=req.quality or "default",
fps=video_params.fps,
num_frames=video_params.num_frames,
duration_s=duration_s,
)
resp.seconds = str(req.seconds or resp.seconds)
return resp
@@ -3091,6 +3124,25 @@ def _cleanup_video_references(
os.unlink(control_path)
def _unpack_video_generation_result(
result: Sequence[object],
) -> tuple[bytes, dict[str, float], float, VideoAction | None, dict[str, object]]:
video_metadata: dict[str, object] = {}
if len(result) == 5:
video_bytes, stage_durations, peak_memory_mb, action, raw_metadata = result
if isinstance(raw_metadata, dict):
video_metadata = {str(key): value for key, value in raw_metadata.items()}
else:
video_bytes, stage_durations, peak_memory_mb, action = result
return (
cast(bytes, video_bytes),
cast(dict[str, float], stage_durations),
float(cast(float, peak_memory_mb)),
cast(VideoAction | None, action),
video_metadata,
)
async def _run_video_generation_job(
handler: OmniOpenAIServingVideo,
request: VideoGenerationRequest,
@@ -3110,12 +3162,14 @@ async def _run_video_generation_job(
await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS})
started_at = time.perf_counter()
try:
video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes(
request,
video_id,
reference_image=reference_image,
reference_video=reference_video,
reference_audio=reference_audio,
video_bytes, stage_durations, peak_memory_mb, action, video_metadata = _unpack_video_generation_result(
await handler.generate_video_bytes(
request,
video_id,
reference_image=reference_image,
reference_video=reference_video,
reference_audio=reference_audio,
)
)
save_context = await STORAGE_MANAGER.save(video_bytes, video_id)
@@ -3131,6 +3185,7 @@ async def _run_video_generation_job(
"peak_memory_mb": peak_memory_mb,
"action": action,
}
updated_fields.update({key: value for key, value in video_metadata.items() if value is not None})
if save_context.expires_at is not None:
updated_fields["expires_at"] = save_context.expires_at
@@ -3494,6 +3549,7 @@ async def _parse_video_form(
frame_interpolation_model_path: str | None = Form(default=None),
lora: str | None = Form(default=None),
extra_params: str | None = Form(default=None),
return_stage_metrics: bool | None = Form(default=None),
) -> tuple[
VideoGenerationRequest,
"OmniOpenAIServingVideo",
@@ -3564,6 +3620,7 @@ async def _parse_video_form(
"frame_interpolation_model_path": frame_interpolation_model_path,
"lora": _parse_form_json(lora, expected_type=dict),
"extra_params": _parse_form_json(extra_params, expected_type=dict),
"return_stage_metrics": return_stage_metrics,
}
request_data = {k: v for k, v in request_data.items() if v is not None}
request = VideoGenerationRequest(**request_data)
@@ -3844,15 +3901,17 @@ async def create_video_sync(
raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
started_at = time.perf_counter()
try:
video_bytes, stage_durations, peak_memory_mb, _action = await asyncio.wait_for(
handler.generate_video_bytes(
request,
request_id,
reference_image=reference_image,
reference_video=reference_video,
reference_audio=reference_audio,
video_bytes, stage_durations, peak_memory_mb, _action, _video_metadata = _unpack_video_generation_result(
await asyncio.wait_for(
handler.generate_video_bytes(
request,
request_id,
reference_image=reference_image,
reference_video=reference_video,
reference_audio=reference_audio,
),
timeout=VIDEO_SYNC_TIMEOUT_S,
),
timeout=VIDEO_SYNC_TIMEOUT_S,
)
except asyncio.TimeoutError:
raise HTTPException(