[Bugfix] fix diffusion quantization benchmarking for Omni outputs (#3653)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
This commit is contained in:
1 parent
990566aef1
commit
dbc589dbca
2 files changed
+49
-2
No files matched your search
@@ -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
|
||||
|
||||
|
||||
Reference in new issue
Block a user