[Bugfix] fix diffusion quantization benchmarking for Omni outputs (#3653)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
This commit is contained in:
Lancer authored and GitHub committed 2026-05-18 17:23:40 +08:00
1 parent 990566aef1
commit dbc589dbca
2 files changed
+49 -2

No files matched your search

+4 -2
View File
@@ -148,6 +148,7 @@ def _build_omni_kwargs(args, quantization=None):
def _generate_image(omni, args, prompt, seed):
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
@@ -165,8 +166,9 @@ def _generate_image(omni, args, prompt, seed):
elapsed = time.perf_counter() - start
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
first = outputs[0]
req_out = first.request_output[0] if hasattr(first, "request_output") else first
req_out = OmniRequestOutput.unwrap_result(outputs)
if not req_out.images:
raise ValueError("Could not extract image output from result.")
img = req_out.images[0]
return img, elapsed, peak_mem