[Refactor][OutputProcessor 2/8]: OmniRequestOutput should inherit RequestOutput and no nested wrap-up (#5146)

Signed-off-by: Boao Shi <aoibosh@connect.hku.hk>
Signed-off-by: bowieshi <aoibosh@connect.hku.hk>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: R2-Y <ruiruyang2@gmail.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: Shi Boao Bowie <aoibosh@connect.hku.hk>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: ruirui(rein) yang <73573651+R2-Y@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
authored and GitHub committed 2026-08-12 22:58:23 +08:00
1 parent 38a2500bdf
commit fcf1085395
89 files changed
+944 -639

No files matched your search

+10 -11
View File
@@ -158,7 +158,6 @@ 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)
@@ -176,7 +175,7 @@ def _generate_image(omni, args, prompt, seed):
elapsed = time.perf_counter() - start
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
req_out = OmniRequestOutput.unwrap_result(outputs)
req_out = outputs[0] if isinstance(outputs, list) else outputs
if not req_out.images:
raise ValueError("Could not extract image output from result.")
img = req_out.images[0]
@@ -206,15 +205,15 @@ def _generate_video(omni, args, prompt, seed):
elapsed = time.perf_counter() - start
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
first = outputs[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
frames = inner.images[0] if inner.images else None
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images[0]
if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput):
first = outputs[0]
elif isinstance(outputs, OmniRequestOutput):
first = outputs
else:
raise ValueError("Could not extract video frames from output.")
if hasattr(first, "images"):
frames = first.images[0] if first.images else None
else:
raise ValueError("Could not extract video frames from output.")
+1 -1
View File
@@ -300,7 +300,7 @@ def benchmark(args: argparse.Namespace) -> None:
output_idx = 0
for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
if stage_outputs.final_output_type == "image":
request_output = stage_outputs.request_output
request_output = stage_outputs
request_id = getattr(request_output, "request_id", "")
images = getattr(request_output, "images", [])
@@ -131,7 +131,7 @@ def build_sampling_params() -> list[SamplingParams]:
def save_stage_output(stage_outputs, output_dir: str) -> None:
output = stage_outputs.request_output
output = stage_outputs
request_id = output.request_id
if stage_outputs.final_output_type == "text":
text_output = output.outputs[0].text
@@ -104,7 +104,7 @@ def main(args):
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -228,7 +228,7 @@ async def main():
if not outputs:
raise ValueError("No output produced from omni.generate()")
first_out = outputs[0].request_output
first_out = outputs[0]
req_out: OmniRequestOutput = first_out
# Verify trajectory data (from custom pipeline)
@@ -741,7 +741,7 @@ def make_mmu_prompt(
def iter_mm_outputs(outputs: list[Any]):
for omni_out in outputs:
req_out = getattr(omni_out, "request_output", None)
req_out = omni_out
req_list = req_out if isinstance(req_out, list) else [req_out]
for item in req_list:
if item is None:
@@ -303,7 +303,7 @@ def main():
)
img_idx = 0
for req_output in omni_outputs:
ro = getattr(req_output, "request_output", None)
ro = req_output
stage_id = getattr(req_output, "stage_id", None)
# AR stage text — each CompletionOutput.text is already a delta when
@@ -621,7 +621,7 @@ def main():
images = getattr(output, "images", None)
if images:
break
req_out = getattr(output, "request_output", None)
req_out = output
images = getattr(req_out, "images", None) if req_out is not None else None
if images:
break
@@ -646,8 +646,8 @@ def main():
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
if frames.is_pipeline_output and frames.request_output is not None:
inner_output = frames.request_output
if frames.is_pipeline_output and frames is not None:
inner_output = frames
if isinstance(inner_output, OmniRequestOutput):
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
audio = inner_output.multimodal_output["audio"]
@@ -0,0 +1,117 @@
"""
Offline inference example: MammothModa2 image summarization (single AR stage).
Example:
uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
--model path/to/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
--image /path/to/input.jpg \
--question "Please summarize the content of this image."
"""
from __future__ import annotations
import argparse
import os
from PIL import Image
from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
parser.add_argument("--image", type=str, required=True, help="Path to input image.")
parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
parser.add_argument(
"--max-tokens",
type=int,
default=512,
help="Max new tokens to generate.",
)
parser.add_argument("--temperature", type=float, default=0.2)
parser.add_argument("--top-p", type=float, default=0.9)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--trust-remote-code", action="store_true")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
return parser.parse_args()
def build_prompt(system: str, question: str) -> str:
return (
f"<|im_start|>system\n{system}<|im_end|>\n"
"<|im_start|>user\n"
"<|vision_start|><|image_pad|><|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
def main() -> None:
args = parse_args()
if not os.path.exists(args.image):
raise FileNotFoundError(f"Image file not found: {args.image}")
pil_image = Image.open(args.image)
image_data = convert_image_mode(pil_image, "RGB")
prompt = build_prompt(args.system, args.question)
omni = Omni(
model=args.model,
deploy_config=args.deploy_config,
trust_remote_code=args.trust_remote_code,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
)
try:
sp = SamplingParams(
temperature=float(args.temperature),
top_p=float(args.top_p),
top_k=-1,
max_tokens=int(args.max_tokens),
seed=int(args.seed),
detokenize=True,
)
# NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
# Consume it inside the try block so the worker isn't closed early.
outputs = list(
omni.generate(
[
{
"prompt": prompt,
"multi_modal_data": {"image": image_data},
"additional_information": {"omni_task": ["chat"]},
}
],
[sp],
)
)
finally:
omni.close()
lines: list[str] = []
for stage_outputs in outputs:
ro = stage_outputs
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
lines.append("answer:\n")
lines.append(text.strip() + "\n")
lines.append("\n")
print("\n".join(lines))
if __name__ == "__main__":
main()
@@ -297,7 +297,7 @@ def main(args):
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -347,7 +347,7 @@ def main(args):
omni.start_profile(stages=args.profiler_stages)
for stage_outputs in omni.generate(prompts, sampling_params_list):
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -392,7 +392,7 @@ def main(args):
total_requests = len(prompts)
processed_count = 0
for stage_outputs in omni_generator:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -367,7 +367,7 @@ def main(args):
print(f"query type: {args.query_type}")
for stage_outputs in omni_generator:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -237,7 +237,7 @@ async def run_single_request(
sampling_params_list=sampling_params_list,
output_modalities=output_modalities,
):
output = omni_output.request_output
output = omni_output
if omni_output.final_output_type == "text":
if stage_0_first_output_ts is None:
stage_0_first_output_ts = time.perf_counter()
@@ -26,7 +26,6 @@ import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -352,8 +351,8 @@ def main():
else:
print("[Profiler] No valid profiling data returned.")
# Extract output from result
output = OmniRequestOutput.unwrap_result(result)
# omni.generate() returns a list for sync calls; unwrap single-result list.
output = result[0] if isinstance(result, list) else result
if not output.images:
raise ValueError("No video frames found in OmniRequestOutput.")
@@ -345,7 +345,7 @@ def main(args):
# Process outputs from each stage
for stage_idx, stage_outputs in enumerate(omni_outputs):
# request_output may be a single object or a list
raw_output = stage_outputs.request_output
raw_output = stage_outputs
outputs_list = raw_output if isinstance(raw_output, list) else [raw_output]
if stage_outputs.final_output_type == "text":
@@ -43,6 +43,7 @@ from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
@@ -388,9 +389,9 @@ def main():
raise ValueError("No output generated from omni.generate()")
output = outputs[0]
if not hasattr(output, "request_output") or not output.request_output:
if not isinstance(output, OmniRequestOutput) or not output:
raise ValueError("No request_output found in OmniRequestOutput")
request_output = output.request_output
request_output = output
if not hasattr(request_output, "multimodal_output"):
raise ValueError("No multimodal_output found in request_output")
@@ -115,7 +115,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
)
images_outputs = []
for output in outputs:
req_out = output.request_output
req_out = output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out.images
@@ -656,7 +656,7 @@ def main():
images = getattr(output, "images", None)
if images:
break
req_out = getattr(output, "request_output", None)
req_out = output
images = getattr(req_out, "images", None) if req_out is not None else None
if images:
break
@@ -168,7 +168,7 @@ def run_e2e():
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro = output.request_output
ro = output
if ro is None:
print("No request_output found.")
continue
@@ -179,7 +179,7 @@ def main(args):
t_start = time.perf_counter()
for stage_outputs in omni.generate(inputs):
request_output = stage_outputs.request_output
request_output = stage_outputs
if request_output is None or not request_output.outputs:
continue
_save_wav(
@@ -215,7 +215,7 @@ async def main_streaming(args):
sample_rate = None
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = stage_output.request_output.outputs[0].multimodal_output
mm = stage_output.outputs[0].multimodal_output
if not stage_output.finished:
t_now = time.perf_counter()
sample_rate, n = _collect_audio_chunk(mm, all_audio_chunks, sample_rate)
@@ -137,7 +137,7 @@ def main(args) -> None:
params_list = sampling_params
for stage_outputs in omni.generate(inputs, params_list):
for i, req_output in enumerate(stage_outputs.request_output):
for i, req_output in enumerate(stage_outputs):
for j, out in enumerate(req_output.outputs):
mm = out.multimodal_output
if mm is None:
@@ -135,7 +135,7 @@ def run_e2e():
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro = output.request_output
ro = output
if ro is None:
print("No request_output found.")
continue
@@ -391,7 +391,7 @@ def main(args):
for batch_start in range(0, len(inputs), batch_size):
batch = inputs[batch_start : batch_start + batch_size]
for stage_outputs in omni.generate(batch):
output = stage_outputs.request_output
output = stage_outputs
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
@@ -412,7 +412,7 @@ async def main_streaming(args):
t_prev = t_start
chunk_idx = 0
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = stage_output.request_output.outputs[0].multimodal_output
mm = stage_output.outputs[0].multimodal_output
if not stage_output.finished:
t_now = time.perf_counter()
audio = mm.get("audio")
@@ -0,0 +1,172 @@
"""Offline Qwen3-TTS example with word-level timestamps.
This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with
the same shared forced aligner used by the streaming server path, then writes
the WAV and JSON sidecar:
python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0``
to avoid FlashInfer sampler JIT during warmup.
"""
from __future__ import annotations
import asyncio
import json
import os
from pathlib import Path
from typing import Any
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
# The in-process offline aligner consumes token_classify outputs. Keeping V1
# multiprocessing off avoids msgspec serialization issues for token-wise logits.
os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
import numpy as np
import soundfile as sf
import torch
from end2end import _estimate_prompt_len
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
def _default_stage_config() -> str:
repo_root = Path(__file__).resolve().parents[4]
return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
def _build_custom_voice_input(args: Any) -> dict[str, Any]:
additional_information = {
"task_type": ["CustomVoice"],
"text": [args.text],
"language": [args.language],
"speaker": [args.speaker],
"instruct": [args.instructions],
"max_new_tokens": [args.max_new_tokens],
}
return {
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model),
"additional_information": additional_information,
}
def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]:
audio_data = mm["audio"]
sr_raw = mm["sr"]
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val)
audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
return audio_tensor.float().detach().cpu().flatten(), sample_rate
def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes:
samples = audio.numpy()
pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype("<i2")
return pcm.tobytes()
async def _run_alignment(args: Any, audio: torch.Tensor, sample_rate: int) -> list[dict[str, Any]] | None:
aligner_config = build_forced_aligner_config(args)
if aligner_config is None:
raise ValueError("--forced-aligner or --forced-aligner-config is required")
timestamps = await align(
audio=_float_audio_to_pcm16_bytes(audio),
text=args.text,
sample_rate=sample_rate,
config=aligner_config,
language=args.language,
)
if timestamps is None:
return None
return [
{
"word": item.word,
"start_ms": item.start_ms,
"end_ms": item.end_ms,
}
for item in timestamps
]
def main(args: Any) -> None:
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
omni_kwargs = vars(args).copy()
for key in (
"forced_aligner",
"text",
"language",
"speaker",
"instructions",
"max_new_tokens",
"output_dir",
):
omni_kwargs.pop(key, None)
omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
omni_kwargs["log_stats"] = args.log_stats
omni = Omni(**omni_kwargs)
prompt = _build_custom_voice_input(args)
final_output = None
for stage_outputs in omni.generate([prompt]):
final_output = stage_outputs
if final_output is None:
raise RuntimeError("Qwen3-TTS did not produce an output.")
mm = final_output.outputs[0].multimodal_output
audio, sample_rate = _audio_tensor_and_sample_rate(mm)
wav_path = output_dir / "qwen3_tts_word_timestamps.wav"
# Align the in-memory tensor, not the saved WAV. The WAV is only a demo
# artifact so users can listen to the same audio referenced by the sidecar.
timestamps = asyncio.run(_run_alignment(args, audio, sample_rate))
sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV")
sidecar = {
"text": args.text,
"sample_rate": sample_rate,
"audio_path": str(wav_path),
"timestamps": timestamps,
}
json_path = output_dir / "qwen3_tts_word_timestamps.json"
json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
print(f"Saved audio: {wav_path}")
print(f"Saved timestamps: {json_path}")
print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2))
def parse_args() -> Any:
parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example")
parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name")
parser.add_argument(
"--forced-aligner",
default=None,
help="Qwen3 forced aligner model path/name",
)
parser.add_argument(
"--forced-aligner-config",
default=None,
help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
)
parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
parser.add_argument("--instructions", default="", help="Optional speaking style instruction")
parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
nullify_stage_engine_defaults(parser)
return parser.parse_args()
if __name__ == "__main__":
main(parse_args())
@@ -38,7 +38,7 @@ def _stage_durations_from_output(outputs) -> dict[str, float]:
durations = getattr(omni_out, "stage_durations", None) or {}
if durations:
return dict(durations)
ro = getattr(omni_out, "request_output", None)
ro = omni_out
if ro is not None:
return dict(getattr(ro, "stage_durations", None) or {})
return {}
@@ -99,7 +99,7 @@ def extract_audio(outputs) -> tuple[np.ndarray, int]:
out = outputs[0]
mm = getattr(out, "multimodal_output", None)
if mm is None:
ro = getattr(out, "request_output", None)
ro = out
if ro is not None:
mm = getattr(ro, "multimodal_output", None)
if mm is None and ro.outputs:
@@ -390,7 +390,7 @@ def _extract_peak_memory_mb(result: Any) -> float:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = getattr(result, "request_output", None)
inner = result
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
@@ -575,8 +575,8 @@ def main():
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
if frames.is_pipeline_output and frames.request_output is not None:
inner_output = frames.request_output
if frames.is_pipeline_output and frames is not None:
inner_output = frames
if isinstance(inner_output, OmniRequestOutput):
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
audio = inner_output.multimodal_output["audio"]
@@ -138,7 +138,7 @@ def _extract_peak_memory_mb(result: Any) -> float:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = getattr(result, "request_output", None)
inner = result
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
@@ -125,7 +125,7 @@ def _extract_images(output: OmniRequestOutput) -> list:
return output.images
# When the output comes from the orchestrator pipeline, images may be
# nested inside request_output.
inner = getattr(output, "request_output", None)
inner = output
if inner is not None and hasattr(inner, "images") and inner.images:
return inner.images
return []
+2 -2
View File
@@ -56,10 +56,10 @@ def test_teacache(model_name: str):
# Extract images from request_output['images']
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
if not isinstance(first_output, OmniRequestOutput) or not first_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output
req_out = first_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
@@ -131,7 +131,7 @@ def _run_inference(
elapsed_ms = (time.time() - start) * 1000
return InferenceResult(
images=outputs[0].request_output.images,
images=outputs[0].images,
elapsed_ms=elapsed_ms,
)
finally:
@@ -49,10 +49,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
def _extract_single_image(outputs) -> Image.Image:
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
if not isinstance(first_output, OmniRequestOutput) or not first_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output
req_out = first_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
+2 -2
View File
@@ -29,10 +29,10 @@ def test_diffusion_model(tmp_path: Path, omni_runner):
raise ValueError("Empty outputs from Omni.generate()")
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
if not isinstance(first_output, OmniRequestOutput) or not first_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output
req_out = first_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
return req_out.images
@@ -122,8 +122,8 @@ def test_cpu_offload_diffusion_model(model_name: str):
print(f"No offload peak memory: {no_offload_peak_memory} MB")
if model_name == "stabilityai/stable-audio-open-1.0":
audio_offload = output_offload[0].request_output.multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio")
audio_offload = output_offload[0].multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].multimodal_output.get("audio")
check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2)
# Set platform-specific VRAM saving thresholds to account
@@ -20,18 +20,10 @@ IMAGE_VIDEO_MODELS = {
MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
LAYERWISE_MODEL_CASES = [
pytest.param(
"stabilityai/stable-audio-open-1.0",
marks=[pytest.mark.full_model, pytest.mark.diffusion],
id="stable_audio",
),
pytest.param(
"riverclouds/qwen_image_random",
marks=[pytest.mark.core_model, pytest.mark.diffusion],
id="qwen_image_random",
),
]
MODEL_MARKS = {
"riverclouds/qwen_image_random": pytest.mark.core_model,
"stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
}
AUDIO_MODEL_PARAMS = {
"runner_params": {},
@@ -101,8 +93,9 @@ def run_inference(
return peak, output
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
@pytest.mark.parametrize("model_name", LAYERWISE_MODEL_CASES)
@pytest.mark.parametrize("model_name", list(MODELS.keys()))
def test_layerwise_offload_diffusion_model(model_name: str):
"""Test that layerwise offloading reduces GPU memory usage.
@@ -134,8 +127,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):
print(f"No offload peak memory: {no_offload_peak_memory} MB")
if model_name == "stabilityai/stable-audio-open-1.0":
audio_offload = output_offload[0].request_output.multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio")
audio_offload = output_offload[0].multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].multimodal_output.get("audio")
# Match the sibling cpu-offload test's tolerance: layerwise offload moves
# blocks across the PCIe bus on a side stream, which can perturb cuBLAS
# algorithm selection and produce ~ULP-level drift larger than 1e-3.
@@ -94,8 +94,8 @@ def _generate_single_stage_image(
if hasattr(first_output, "images") and first_output.images:
images = first_output.images
else:
assert hasattr(first_output, "request_output") and first_output.request_output
request_output = first_output.request_output
assert isinstance(first_output, OmniRequestOutput) and first_output
request_output = first_output
if isinstance(request_output, list):
req_out = request_output[0]
else:
@@ -88,8 +88,8 @@ def _generate_single_stage_image(
if hasattr(first_output, "images") and first_output.images:
images = first_output.images
else:
assert hasattr(first_output, "request_output") and first_output.request_output
request_output = first_output.request_output
assert isinstance(first_output, OmniRequestOutput) and first_output
request_output = first_output
if isinstance(request_output, list):
req_out = request_output[0]
else:
@@ -147,10 +147,10 @@ def _generate_single_stage_video(
first = outputs[0]
# Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput.request_output).
# Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput).
frames: Any = None
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(first, OmniRequestOutput) and isinstance(first, list):
inner = first[0]
if isinstance(inner, OmniRequestOutput) and inner.images:
frames = inner.images[0]
if frames is None and hasattr(first, "images") and first.images:
@@ -234,8 +234,8 @@ def _generate_bagel_image(
if images := getattr(req_output, "images", None):
generated_image = images[0]
break
if hasattr(req_output, "request_output") and req_output.request_output:
stage_outputs = req_output.request_output
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_outputs = req_output
if not isinstance(stage_outputs, list):
stage_outputs = [stage_outputs]
for stage_out in stage_outputs:
@@ -250,8 +250,8 @@ def _generate_bagel_image(
# Check LLM stage output — should have finish_reason=stop (not length)
for req_output in omni_outputs:
if hasattr(req_output, "request_output") and req_output.request_output:
stage_outputs = req_output.request_output
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_outputs = req_output
if not isinstance(stage_outputs, list):
stage_outputs = [stage_outputs]
for stage_out in stage_outputs:
@@ -310,7 +310,7 @@ def _generate_image(omni, config: QualityTestConfig):
first = outputs[0]
if hasattr(first, "images") and first.images:
return first.images[0], peak_mem
inner = first.request_output
inner = first
if inner is not None and hasattr(inner, "images") and inner.images:
return inner.images[0], peak_mem
raise ValueError("Could not extract image from output.")
@@ -341,15 +341,15 @@ def _generate_video(omni, config: QualityTestConfig):
)
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
first = outputs[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
frames = inner.images[0] if inner.images else None
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images[0]
if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput):
first = outputs[0]
elif isinstance(outputs, OmniRequestOutput):
first = outputs
else:
raise ValueError("Could not extract video frames from output.")
if hasattr(first, "images"):
frames = first.images[0] if first.images else None
else:
raise ValueError("Could not extract video frames from output.")
@@ -444,11 +444,11 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch
image = Image.new("RGB", (2, 2))
inner = OmniRequestOutput.from_diffusion(request_id="req", images=[image])
outer = OmniRequestOutput(
outer = OmniRequestOutput.from_stage_output(
inner,
request_id="req",
stage_id=0,
final_output_type="image",
request_output=inner,
finished=True,
)
@@ -73,11 +73,9 @@ def test_run_summary_reports_worker_peak_memory():
assert summary["per_run_peak_memory_mb"] == [100.0, 150.0]
def test_request_peak_memory_prefers_inner_diffusion_output():
inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})()
outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})()
assert _request_peak_memory_mb(outer) == 321.0
def test_request_peak_memory_reads_peak_memory_mb_directly():
result = type("Output", (), {"peak_memory_mb": 123.0})()
assert _request_peak_memory_mb(result) == 123.0
def test_metric_guidance_describes_thresholds():
@@ -273,7 +273,7 @@ class TestPipelineStreamingOutputToEntrypoint:
await _wait_for(lambda: len(pipeline.requests) == 1)
assert pipeline.requests[0].request_id.startswith("req-omni-")
assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True]
assert [output.finished for output in outputs] == [False, True]
finally:
await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
@@ -25,6 +25,7 @@ from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")]
@@ -95,8 +96,8 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
for req_output in omni_outputs:
if images := getattr(req_output, "images", None):
return images[0]
if hasattr(req_output, "request_output") and req_output.request_output:
stage_out = req_output.request_output
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_out = req_output
if hasattr(stage_out, "images") and stage_out.images:
return stage_out.images[0]
return None
@@ -24,6 +24,7 @@ from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")]
@@ -122,8 +123,8 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
for req_output in omni_outputs:
if images := getattr(req_output, "images", None):
return images[0]
if hasattr(req_output, "request_output") and req_output.request_output:
stage_out = req_output.request_output
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_out = req_output
if hasattr(stage_out, "images") and stage_out.images:
return stage_out.images[0]
return None
+2 -2
View File
@@ -328,7 +328,7 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag
images = None
cot_text = ""
for out in outputs:
ro = getattr(out, "request_output", None)
ro = out
if ro and getattr(ro, "outputs", None):
cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
if not cot_text:
@@ -467,7 +467,7 @@ def _extract_image(outputs) -> Image.Image:
assert outputs, "Pipeline produced no outputs"
for output in outputs:
images = getattr(output, "images", None)
request_output = getattr(output, "request_output", None)
request_output = output
if not images and request_output is not None:
images = getattr(request_output, "images", None)
if images:
@@ -141,8 +141,8 @@ def _extract_frames(result) -> list[Image.Image]:
if isinstance(result, list):
result = result[0] if result else None
if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None:
result = result.request_output
if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result is not None:
result = result
images = getattr(result, "images", None)
if not images:
raise AssertionError("Pipeline output had no images")
@@ -171,11 +171,11 @@ def _build_text_output(text: str = "This is a test response.") -> OmniRequestOut
metrics=None,
lora_request=None,
)
return OmniRequestOutput(
return OmniRequestOutput.from_stage_output(
request_output,
request_id="test_req_text",
finished=True,
final_output_type="text",
request_output=request_output,
)
@@ -200,11 +200,11 @@ def _build_audio_chat_output(num_samples: int = 24000) -> OmniRequestOutput:
metrics=None,
lora_request=None,
)
return OmniRequestOutput(
return OmniRequestOutput.from_stage_output(
request_output,
request_id="test_req_audio_chat",
finished=True,
final_output_type="audio",
request_output=request_output,
)
@@ -240,11 +240,11 @@ def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
images=[_build_image_output(color="blue")],
finished=True,
)
return OmniRequestOutput(
return OmniRequestOutput.from_stage_output(
request_output,
request_id="test_req_img_chat",
finished=True,
images=[_build_image_output(color="blue")],
final_output_type="image",
request_output=request_output,
)
@@ -515,8 +515,19 @@ def mock_async_omni(
@pytest.fixture
def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni, tmp_path):
"""Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
# Override the STORAGE_MANAGER path to a writable temp directory before
# forking the server subprocess. The default /tmp/storage may be owned
# by a different user (e.g. from a previous CI run) and cause
# PermissionError when the server tries to save generated video files.
from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
old_storage_path = STORAGE_MANAGER.storage_path
new_storage_path = tmp_path / "omni_storage"
new_storage_path.mkdir(exist_ok=True)
STORAGE_MANAGER.storage_path = str(new_storage_path)
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="command")
cmd = OmniServeCommand()
@@ -555,6 +566,9 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
yield f"http://127.0.0.1:{port}/v1"
# Restore the original storage path after the test.
STORAGE_MANAGER.storage_path = old_storage_path
if server_process.is_alive():
server_process.terminate()
server_process.join(timeout=10)
@@ -294,7 +294,7 @@ class vLLMOmniHttpServerLocal:
"global_steps": self.global_steps,
}
req_output = final_res.request_output
req_output = final_res
if req_output is not None and hasattr(req_output, "finish_reason"):
finish_reason = req_output.finish_reason or "stop"
else:
@@ -53,9 +53,9 @@ def test_audiox_model(omni_runner: OmniRunner) -> None:
assert outputs is not None
first_output = outputs[0]
assert first_output.final_output_type == "audio"
assert hasattr(first_output, "request_output") and first_output.request_output
assert isinstance(first_output, OmniRequestOutput) and first_output
req_out = first_output.request_output
req_out = first_output
assert isinstance(req_out, OmniRequestOutput)
assert req_out.final_output_type == "audio"
assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output
+1 -1
View File
@@ -55,7 +55,7 @@ pytestmark = [
def _extract_text(omni_outputs: list) -> str:
"""Extract generated text from OmniRequestOutput list."""
for req_output in omni_outputs:
ro = getattr(req_output, "request_output", None)
ro = req_output
if ro and getattr(ro, "outputs", None):
return "".join(getattr(o, "text", "") or "" for o in ro.outputs)
return ""
@@ -28,7 +28,6 @@ from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
@@ -176,12 +175,12 @@ def test_cosyvoice3_offline_reference_zero_shot(omni_runner: OmniRunner) -> None
assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
# Code2wav is ``final_output`` (stage 1); ``Omni.generate`` yields its ``request_output`` / completions.
pipeline_out = OmniRequestOutput.unwrap_result(outputs[0])
pipeline_out = outputs[0]
assert pipeline_out.stage_id == 1, f"expected final stage 1, got {pipeline_out.stage_id}"
assert pipeline_out.final_output_type == "audio", (
f"expected audio final_output_type, got {pipeline_out.final_output_type!r}"
)
ro = pipeline_out.request_output
ro = pipeline_out
assert ro is not None, "stage 1 should include request_output"
completions = getattr(ro, "outputs", None) or []
assert completions, "stage 1 request_output should include at least one completion"
@@ -175,7 +175,7 @@ def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None:
# Prefer the latest finished chunk to avoid picking an intermediate stream output.
for stage_output in reversed(matched):
if _is_finished_request_output(getattr(stage_output, "request_output", None)):
if _is_finished_request_output(stage_output):
return stage_output
return matched[-1]
@@ -205,7 +205,7 @@ def _to_token_list(value: Any) -> list[int]:
def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
request_output = getattr(stage_output, "request_output", None)
request_output = stage_output
if request_output is None:
return ""
req_list = request_output if isinstance(request_output, list) else [request_output]
@@ -240,7 +240,7 @@ def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
def _extract_audio(stage_output: Any) -> Any | None:
request_output = getattr(stage_output, "request_output", None)
request_output = stage_output
if request_output is None:
return None
req_list = request_output if isinstance(request_output, list) else [request_output]
@@ -256,7 +256,7 @@ def _extract_audio(stage_output: Any) -> Any | None:
def _extract_image(stage_output: Any) -> Any | None:
request_output = getattr(stage_output, "request_output", None)
request_output = stage_output
if request_output is None:
return None
req_list = request_output if isinstance(request_output, list) else [request_output]
@@ -54,7 +54,7 @@ def _sampling_params() -> OmniDiffusionSamplingParams:
def _first_request_images(outputs) -> list:
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output.request_output
req_out = first_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
return req_out.images
@@ -161,7 +161,7 @@ def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunn
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output.request_output
req_out = first_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
images = req_out.images
@@ -224,7 +224,7 @@ def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunne
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output.request_output
req_out = first_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
images = req_out.images
@@ -288,7 +288,7 @@ def _extract_generated_image(outputs: list[object]) -> Image.Image:
if images := getattr(first_output, "images", None):
return images[0]
request_output = getattr(first_output, "request_output", None)
request_output = first_output
if request_output is not None and (images := getattr(request_output, "images", None)):
return images[0]
@@ -74,7 +74,7 @@ def test_i2t_generates_text(omni: Omni) -> None:
outputs = omni.generate(prompts=[prompt_dict])
assert outputs, "No outputs returned from Omni.generate()"
request_output = outputs[0].request_output
request_output = outputs[0]
assert request_output.outputs, "No completion outputs"
completion = request_output.outputs[0]
@@ -172,7 +172,7 @@ def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner):
found_image = False
for out in outputs:
ro_list = getattr(out, "request_output", out)
ro_list = out
if not isinstance(ro_list, list):
ro_list = [ro_list]
for ro in ro_list:
@@ -137,14 +137,11 @@ def _extract_final_audio_outputs(outputs):
if getattr(item, "final_output_type", None) == "audio":
final_outputs.append(item)
continue
request_output = getattr(item, "request_output", None)
if request_output is None:
continue
multimodal_output = getattr(request_output, "multimodal_output", None)
multimodal_output = getattr(item, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
final_outputs.append(item)
continue
completions = getattr(request_output, "outputs", None) or []
completions = getattr(item, "outputs", None) or []
if any(isinstance(getattr(completion, "multimodal_output", None), Mapping) for completion in completions):
final_outputs.append(item)
return final_outputs
@@ -155,16 +152,11 @@ def _extract_multimodal_output(output) -> Mapping[str, Any]:
if isinstance(multimodal_output, Mapping):
return multimodal_output
request_output = getattr(output, "request_output", None)
if request_output is not None:
multimodal_output = getattr(request_output, "multimodal_output", None)
completions = getattr(output, "outputs", None) or []
for completion in completions:
multimodal_output = getattr(completion, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
return multimodal_output
completions = getattr(request_output, "outputs", None) or []
for completion in completions:
multimodal_output = getattr(completion, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
return multimodal_output
raise AssertionError("No multimodal audio output found in Ming generate results")
@@ -108,7 +108,7 @@ def _build_request(
def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
"""Run a single request and return (waveform, sample_rate)."""
for stage_outputs in omni.generate(request, DEFAULT_SAMPLING):
req_output = stage_outputs.request_output
req_output = stage_outputs
if req_output is not None:
mm = req_output.outputs[0].multimodal_output
assert mm is not None, "Expected multimodal_output to be non-None"
@@ -167,7 +167,7 @@ def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
results = []
# Single-stage model (num_stages=1): one sampling param for all requests.
for stage_outputs in omni_runner.omni.generate(requests, [DEFAULT_SAMPLING]):
req_output = stage_outputs.request_output
req_output = stage_outputs
if req_output is not None:
mm = req_output.outputs[0].multimodal_output
assert mm is not None
@@ -58,7 +58,7 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None:
# Check final output has audio
final_output = outputs[-1]
ro = final_output.request_output
ro = final_output
assert ro is not None, "No request_output"
mm = getattr(ro, "multimodal_output", None)
@@ -42,7 +42,7 @@ def _sampling_params() -> OmniDiffusionSamplingParams:
def _first_request_images(outputs) -> list:
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output.request_output
req_out = first_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
return req_out.images
@@ -62,9 +62,9 @@ def generate_stable_audio_short_clip(
# Audio-output diffusion pipelines (those with ``support_audio_output = True``) now have
# ``final_output_type="audio"`` set on the outer stage metadata as well as the inner request.
assert first_output.final_output_type == "audio"
assert hasattr(first_output, "request_output") and first_output.request_output
assert isinstance(first_output, OmniRequestOutput) and first_output
req_out = first_output.request_output
req_out = first_output
assert isinstance(req_out, OmniRequestOutput)
assert req_out.final_output_type == "audio"
assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output
@@ -63,8 +63,8 @@ def _assert_non_empty_text_output(outputs: list[OmniRequestOutput]) -> None:
text_output = next((o for o in outputs if o.final_output_type == "text"), None)
assert text_output is not None
assert text_output.request_output is not None
text_content = text_output.request_output.outputs[0].text
assert text_output is not None
text_content = text_output.outputs[0].text
assert text_content is not None
assert len(text_content.strip()) > 0
@@ -81,8 +81,8 @@ def test_audio_to_text_and_audio(omni_runner: OmniRunner) -> None:
audio_output = next((o for o in outputs if o.final_output_type == "audio"), None)
if audio_output is not None:
assert audio_output.request_output is not None
audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio")
assert audio_output is not None
audio_tensor = audio_output.outputs[0].multimodal_output.get("audio")
if audio_tensor is not None:
assert audio_tensor.numel() > 0
@@ -65,7 +65,7 @@ def test_realtime_duplex_multi_session_script_is_directly_executable():
cwd=MULTI_DEMO_PATH.parents[3],
capture_output=True,
text=True,
timeout=20,
timeout=60,
check=False,
)
@@ -78,7 +78,7 @@ def test_realtime_duplex_demo_pair_script_is_directly_executable():
cwd=PAIR_DEMO_PATH.parents[3],
capture_output=True,
text=True,
timeout=20,
timeout=60,
check=False,
)
@@ -91,7 +91,7 @@ def test_realtime_duplex_soft_interrupt_script_is_directly_executable():
cwd=SOFT_INTERRUPT_DEMO_PATH.parents[3],
capture_output=True,
text=True,
timeout=20,
timeout=60,
check=False,
)
+1 -1
View File
@@ -50,7 +50,7 @@ async def generate(
output_modalities=["text"],
):
stage_id = omni_output.stage_id
out = omni_output.request_output
out = omni_output
if stage_id == 0:
num_tokens = sum(len(output.token_ids) for output in out.outputs)
count = num_tokens
+246 -148
View File
@@ -4,171 +4,269 @@
import pytest
from PIL import Image
from pytest_mock import MockerFixture
from vllm.outputs import CompletionOutput, RequestOutput
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestOmniRequestOutput:
"""Tests for OmniRequestOutput class."""
# ── helpers ────────────────────────────────────────────────────────────────
def test_from_diffusion(self):
"""Test creating output from diffusion model."""
images = [Image.new("RGB", (64, 64), color="red")]
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=images,
prompt="a cat",
def _make_text_request_output(
request_id: str = "req-1",
text: str = "hello world",
token_ids: list[int] | None = None,
finished: bool = True,
) -> RequestOutput:
"""Build a minimal vLLM RequestOutput for testing."""
return RequestOutput(
request_id=request_id,
prompt="test prompt",
prompt_token_ids=[1, 2, 3],
prompt_logprobs=None,
outputs=[
CompletionOutput(
index=0,
text=text,
token_ids=token_ids or [10, 11, 12],
cumulative_logprob=-1.0,
logprobs=None,
finish_reason="stop" if finished else None,
stop_reason=None,
)
],
finished=finished,
metrics=None,
lora_request=None,
)
# ── from_stage_output ──────────────────────────────────────────────────────
class TestFromStageOutput:
"""Tests for OmniRequestOutput.from_stage_output()."""
def test_copies_request_output_content(self):
"""Content from a vLLM RequestOutput is flattened onto the new object."""
source = _make_text_request_output(text="copied text")
out = OmniRequestOutput.from_stage_output(source, request_id="req-1", stage_id=0, final_output_type="text")
assert out.prompt == "test prompt"
assert out.prompt_token_ids == [1, 2, 3]
assert len(out.outputs) == 1
assert out.outputs[0].text == "copied text"
assert out.finished is True
assert out.stage_id == 0
assert out.final_output_type == "text"
def test_copies_omni_content(self):
"""Diffusion fields are copied when source is another OmniRequestOutput."""
image = Image.new("RGB", (64, 64), color="blue")
inner = OmniRequestOutput.from_diffusion(request_id="req-diff", images=[image], prompt="a cat")
outer = OmniRequestOutput.from_stage_output(inner, request_id="req-diff", stage_id=0, final_output_type="image")
assert len(outer.images) == 1
assert outer.images[0] is image
assert outer.prompt == "a cat"
assert outer.is_diffusion_output
def test_with_duck_typed_source(self):
"""A duck-typed source (SimpleNamespace-style) copies correctly."""
from types import SimpleNamespace
source = SimpleNamespace(
outputs=[],
prompt="duck prompt",
prompt_token_ids=[7, 8, 9],
finished=False,
)
out = OmniRequestOutput.from_stage_output(source, request_id="duck-1", stage_id=2)
assert out.prompt == "duck prompt"
assert out.prompt_token_ids == [7, 8, 9]
assert out.finished is False
assert out.stage_id == 2
def test_request_id_always_copied_from_source(self):
"""request_id is always copied from source, even when passed in kwargs."""
source = _make_text_request_output(request_id="source-id")
out = OmniRequestOutput.from_stage_output(source, request_id="kwargs-id", stage_id=1)
# Source wins — request_id is content, not a control field.
assert out.request_id == "source-id"
def test_passes_control_fields(self):
"""Control fields (metrics, stage_durations, peak_memory_mb) are set."""
source = _make_text_request_output()
out = OmniRequestOutput.from_stage_output(
source,
request_id="r1",
stage_id=3,
metrics={"ttft": 0.5},
stage_durations={"prefill": 0.12},
peak_memory_mb=1024.0,
)
assert out.metrics == {"ttft": 0.5}
assert out.stage_durations == {"prefill": 0.12}
assert out.peak_memory_mb == 1024.0
def test_finished_field_copied(self):
"""The *finished* flag is copied from the source when present."""
source = _make_text_request_output(finished=False)
out = OmniRequestOutput.from_stage_output(source, request_id="r1")
assert out.finished is False
source2 = _make_text_request_output(finished=True)
out2 = OmniRequestOutput.from_stage_output(source2, request_id="r2")
assert out2.finished is True
def test_outputs_field_copied(self):
"""The *outputs* list is copied from the source."""
source = _make_text_request_output(text="hello", token_ids=[1, 2, 3])
out = OmniRequestOutput.from_stage_output(source, request_id="r1")
assert len(out.outputs) == 1
assert out.outputs[0].text == "hello"
assert out.outputs[0].token_ids == [1, 2, 3]
# ── dataclass safety ────────────────────────────────────────────────────────
class TestDataclassSafety:
def test_from_stage_output_is_classmethod(self):
"""from_stage_output must be a classmethod."""
assert callable(OmniRequestOutput.from_stage_output)
# ── msgpack round-trip ─────────────────────────────────────────────────────
class TestMsgpackRoundTrip:
"""Tests that OmniRequestOutput survives msgpack encode → decode without
infinite recursion or data loss."""
def test_round_trip_pipeline_output(self):
"""A pipeline output round-trips without data loss."""
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
OmniMsgpackEncoder,
)
source = _make_text_request_output(text="round trip")
out = OmniRequestOutput.from_stage_output(
source,
request_id="rt-1",
stage_id=0,
final_output_type="text",
)
encoder = OmniMsgpackEncoder()
decoder = OmniMsgpackDecoder()
encoded = encoder.encode(out)
decoded = decoder.decode(encoded)
assert isinstance(decoded, OmniRequestOutput)
assert decoded.request_id == "req-1"
assert decoded.stage_id == 0
assert decoded.final_output_type == "text"
assert decoded.finished is True
assert decoded.outputs[0].text == "round trip"
assert decoded.prompt == "test prompt"
assert decoded.prompt_token_ids == [1, 2, 3]
def test_round_trip_diffusion_output(self):
"""A diffusion output round-trips without data loss."""
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
OmniMsgpackEncoder,
)
image = Image.new("RGB", (64, 64), color="red")
out = OmniRequestOutput.from_diffusion(
request_id="diff-1",
images=[image],
prompt="a red image",
metrics={"steps": 50},
)
assert output.request_id == "test-123"
assert output.images == images
assert output.prompt == "a cat"
assert output.metrics == {"steps": 50}
assert output.is_diffusion_output
assert output.num_images == 1
def test_from_pipeline(self, mocker: MockerFixture):
"""Test creating output from pipeline stage."""
mock_request_output = mocker.Mock()
mock_request_output.request_id = "pipeline-123"
mock_request_output.prompt_token_ids = [1, 2, 3]
mock_request_output.outputs = [mocker.Mock()]
mock_request_output.encoder_prompt_token_ids = None
mock_request_output.prompt_logprobs = None
mock_request_output.num_cached_tokens = 10
mock_request_output.kv_transfer_params = None
mock_request_output.multimodal_output = {"image": mocker.Mock()}
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
multimodal_output={"audio": None},
)
assert output.request_id == "pipeline-123"
assert output.stage_id == 0
assert output.final_output_type == "text"
assert output.is_pipeline_output
encoder = OmniMsgpackEncoder()
decoder = OmniMsgpackDecoder()
def test_prompt_token_ids_property(self, mocker: MockerFixture):
"""Test prompt_token_ids property for streaming compatibility."""
mock_request_output = mocker.Mock()
mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5]
encoded = encoder.encode(out)
decoded = decoder.decode(encoded)
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
assert isinstance(decoded, OmniRequestOutput)
assert decoded.request_id == "diff-1"
assert decoded.prompt == "a red image"
assert decoded.num_images == 1
assert decoded.images[0].size == (64, 64)
assert decoded.metrics == {"steps": 50}
def test_round_trip_does_not_recurse(self):
"""Encoding a pipeline output must not trigger infinite recursion."""
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackEncoder,
)
assert output.prompt_token_ids == [1, 2, 3, 4, 5]
source = _make_text_request_output(text="no recursion")
out = OmniRequestOutput.from_stage_output(source, request_id="rec-1", stage_id=0)
def test_prompt_token_ids_none_when_no_request_output(self):
"""Test prompt_token_ids returns None when no request_output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[],
prompt="a cat",
)
assert output.prompt_token_ids is None
encoder = OmniMsgpackEncoder()
encoded = encoder.encode(out)
assert isinstance(encoded, bytes)
assert len(encoded) > 0
def test_outputs_property(self, mocker: MockerFixture):
"""Test outputs property for chat completion compatibility."""
mock_output = mocker.Mock()
mock_request_output = mocker.Mock()
mock_request_output.outputs = [mock_output]
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
def test_legacy_wire_format_deserialized(self):
"""When the wire format contains a nested ``request_output`` key
(pre-inheritance legacy), it is merged via _copy_content_from."""
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
)
assert output.outputs == [mock_output]
# Build a legacy-style dict: OmniRequestOutput fields + a nested
# request_output dict containing RequestOutput-like content.
legacy_dict = {
"request_id": "legacy-1",
"finished": True,
"stage_id": 1,
"final_output_type": "text",
"images": [],
"_multimodal_output": {},
"_custom_output": {},
"stage_durations": {},
"peak_memory_mb": 0.0,
# --- nested request_output key (legacy) ---
"request_output": {
"request_id": "legacy-1",
"prompt": "legacy prompt",
"prompt_token_ids": [1, 2],
"prompt_logprobs": None,
"outputs": [],
"finished": True,
"metrics": {},
"lora_request": None,
"encoder_prompt": None,
"encoder_prompt_token_ids": None,
"num_cached_tokens": None,
"kv_transfer_params": None,
},
}
def test_outputs_empty_when_no_request_output(self):
"""Test outputs returns empty list when no request_output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[],
prompt="a cat",
)
assert output.outputs == []
# Manually run the decoder's post_process to test legacy path.
# The decoder's _decode_omni_request_output pops "request_output"
# and merges it via _copy_content_from.
decoder = OmniMsgpackDecoder()
result = decoder._post_process(legacy_dict)
def test_encoder_prompt_token_ids_property(self, mocker: MockerFixture):
"""Test encoder_prompt_token_ids property."""
mock_request_output = mocker.Mock()
mock_request_output.encoder_prompt_token_ids = [10, 20, 30]
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.encoder_prompt_token_ids == [10, 20, 30]
def test_num_cached_tokens_property(self, mocker: MockerFixture):
"""Test num_cached_tokens property."""
mock_request_output = mocker.Mock()
mock_request_output.num_cached_tokens = 42
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.num_cached_tokens == 42
def test_multimodal_output_property(self, mocker: MockerFixture):
"""Test multimodal_output property."""
mock_request_output = mocker.Mock()
mock_audio = mocker.Mock()
expected_output = {"audio": mock_audio}
mock_request_output.outputs = []
mock_request_output.multimodal_output = expected_output
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="audio",
request_output=mock_request_output,
)
assert output.multimodal_output is expected_output
def test_to_dict_diffusion(self):
"""Test to_dict for diffusion output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[Image.new("RGB", (64, 64), color="red")],
prompt="a cat",
metrics={"steps": 50},
)
result = output.to_dict()
assert result["request_id"] == "test-123"
assert result["finished"] is True
assert result["final_output_type"] == "image"
assert result["num_images"] == 1
assert result["prompt"] == "a cat"
def test_to_dict_pipeline(self, mocker: MockerFixture):
"""Test to_dict for pipeline output."""
mock_request_output = mocker.Mock()
mock_request_output.request_id = "pipeline-123"
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
result = output.to_dict()
assert result["request_id"] == "pipeline-123"
assert result["finished"] is True
assert result["final_output_type"] == "text"
assert result["stage_id"] == 0
assert isinstance(result, OmniRequestOutput)
assert result.request_id == "legacy-1"
assert result.finished is True
# Content merged from legacy nested key
assert result.prompt == "legacy prompt"
assert result.prompt_token_ids == [1, 2]
@@ -2156,12 +2156,12 @@ def test_direct_listen_decision_survives_inner_completion_metadata():
)
assert decision is not None
output = attach_duplex_output_decision(
OmniRequestOutput(
OmniRequestOutput.from_stage_output(
inner_output,
request_id="duplex-direct-listen",
finished=True,
stage_id=0,
final_output_type=decision.final_output_type,
request_output=inner_output,
metrics={
"stage_metrics": {
"0": {
@@ -3049,12 +3049,12 @@ async def test_minicpmo_auto_response_listen_without_response_does_not_defer_com
)
assert decision is not None
listen_output = attach_duplex_output_decision(
OmniRequestOutput(
OmniRequestOutput.from_stage_output(
inner_output,
request_id=request_id,
finished=True,
stage_id=0,
final_output_type=decision.final_output_type,
request_output=inner_output,
),
decision,
)
@@ -119,7 +119,6 @@ class MockGenerationResult:
def __init__(self, images):
self.images = images
self.request_output = SimpleNamespace(images=images)
self.stage_durations = {}
self.peak_memory_mb = 0.0
@@ -137,10 +136,7 @@ class MockStageResult:
outputs = [SimpleNamespace(text=text, index=0)]
else:
outputs = []
self.request_output = SimpleNamespace(
outputs=outputs,
images=self.images,
)
self.outputs = outputs
self.stage_durations = {}
self.peak_memory_mb = 0.0
@@ -2080,22 +2076,20 @@ def test_extract_images_from_result():
assert all(isinstance(img, Image.Image) for img in images)
assert all(img.size == (64, 64) for img in images)
# Test dict path: result.request_output["images"]
# Test result with "images" attribute set in __init__
class DictRequestOutput:
def __init__(self):
self.request_output = {"images": [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]}
self.images = [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]
result = DictRequestOutput()
images = _extract_images_from_result(result)
assert len(images) == 1
assert isinstance(images[0], Image.Image)
# Test attribute path: result.request_output.images
# Test result with "images" attribute from an inner object
class AttrRequestOutput:
def __init__(self):
self.request_output = type(
"obj", (), {"images": [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]}
)()
self.images = [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]
result = AttrRequestOutput()
images = _extract_images_from_result(result)
@@ -180,10 +180,10 @@ def create_mock_audio_output_for_test(
audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor)
return OmniRequestOutput(
return OmniRequestOutput.from_stage_output(
mock_request_output,
stage_id=0,
final_output_type="audio",
request_output=mock_request_output,
)
@@ -2292,10 +2292,10 @@ class TestStreamingResponse:
self.prompt_logprobs = None
self.kv_transfer_params = None
return OmniRequestOutput(
return OmniRequestOutput.from_stage_output(
MockRequestOutput(audio_tensor=chunk),
stage_id=0,
final_output_type="audio",
request_output=MockRequestOutput(audio_tensor=chunk),
finished=finished,
)
@@ -3441,10 +3441,10 @@ class TestWAVStreaming:
self.prompt_logprobs = None
self.kv_transfer_params = None
return OmniRequestOutput(
return OmniRequestOutput.from_stage_output(
MockRequestOutput(audio_tensor=chunk),
stage_id=0,
final_output_type="audio",
request_output=MockRequestOutput(audio_tensor=chunk),
finished=finished,
)
@@ -47,7 +47,7 @@ def _text_result(text: str) -> OmniRequestOutput:
output.text = text
request_output = RequestOutput()
request_output.outputs = [output]
return OmniRequestOutput(final_output_type="text", request_output=request_output)
return OmniRequestOutput.from_stage_output(request_output, final_output_type="text")
def _audio_result(audio_data: Any) -> OmniRequestOutput:
@@ -61,7 +61,7 @@ def _audio_result(audio_data: Any) -> OmniRequestOutput:
output.multimodal_output = {"audio": audio_data}
request_output = RequestOutput()
request_output.outputs = [output]
return OmniRequestOutput(final_output_type="audio", request_output=request_output)
return OmniRequestOutput.from_stage_output(request_output, final_output_type="audio")
class MockWebSocket:
+4 -6
View File
@@ -480,11 +480,9 @@ def test_async_omni_duplex_direct_output_prefers_outer_control_metadata():
OmniRequestOutput(
request_id="duplex-sid-e0-stage0",
stage_id=0,
request_output=SimpleNamespace(
outputs=[],
multimodal_output=SimpleNamespace(kind="processed-payload"),
_custom_output={"inner": "completion"},
),
outputs=[],
_multimodal_output={"kind": "processed-payload"},
_custom_output={"inner": "completion"},
),
decision,
)
@@ -557,7 +555,7 @@ async def test_async_omni_duplex_collect_wraps_raw_response_stage_output():
assert outputs[0].finished is False
assert outputs[0].stage_id == 1
assert outputs[0].final_output_type == "audio"
assert outputs[0].request_output is raw_stage1_output
assert outputs[0].peak_memory_mb == raw_stage1_output.peak_memory_mb
def test_async_omni_duplex_request_info_includes_response_stage():
+9 -10
View File
@@ -72,8 +72,8 @@ def make_output_msg(
final_output_type=final_output_type,
images=images or [],
stage_durations={},
outputs=[SimpleNamespace(text=payload, index=0)],
)
engine_output.payload = payload
return OutputMessage(
request_id=request_id,
stage_id=stage_id,
@@ -474,7 +474,7 @@ async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.Mo
app.shutdown()
assert [output.stage_id for output in outputs] == [2]
assert [output.request_output.payload for output in outputs] == ["final"]
assert [output.outputs[0].text for output in outputs] == ["final"]
assert "req-1" not in app.request_states
@@ -492,7 +492,7 @@ async def test_async_omni_accepts_multiple_final_stage_streams(monkeypatch: pyte
app.shutdown()
assert [output.stage_id for output in outputs] == [0, 0, 0, 2, 2, 2]
assert [output.request_output.payload for output in outputs] == [
assert [output.outputs[0].text for output in outputs] == [
"req-1-stage0-0",
"req-1-stage0-1",
"req-1-stage0-2",
@@ -517,7 +517,7 @@ async def test_async_omni_stops_on_final_stage_finished(monkeypatch: pytest.Monk
finally:
app.shutdown()
assert [output.request_output.payload for output in outputs] == [
assert [output.outputs[0].text for output in outputs] == [
"req-1-stage0",
"req-1-stage2-final",
]
@@ -544,7 +544,7 @@ async def test_async_omni_diffusion_only_yields_single_image_output(monkeypatch:
assert outputs[0].stage_id == 0
assert outputs[0].final_output_type == "image"
assert outputs[0].images == ["req-1-image"]
assert outputs[0].request_output.payload == "req-1-diffusion-final"
assert outputs[0].outputs[0].text == "req-1-diffusion-final"
@pytest.mark.asyncio
@@ -565,7 +565,7 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc
assert [output.stage_id for output in outputs] == [0, 0, 0, 1]
assert [output.final_output_type for output in outputs] == ["text", "text", "text", "image"]
assert [output.request_output.payload for output in outputs] == [
assert [output.outputs[0].text for output in outputs] == [
"req-1-text-0",
"req-1-text-1",
"req-1-text-2",
@@ -695,7 +695,7 @@ def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkey
assert len(outputs) == 4
assert [output.stage_id for output in outputs] == [0, 2, 0, 2]
assert [output.request_output.payload for output in outputs] == [
assert [output.outputs[0].text for output in outputs] == [
f"{engine.submitted[0]['request_id']}-stage0-0",
f"{engine.submitted[0]['request_id']}-stage2-final",
f"{engine.submitted[1]['request_id']}-stage0-0",
@@ -739,7 +739,7 @@ def test_omni_generate_diffusion_only_yields_single_image_per_request(monkeypatc
assert len(outputs) == 2
assert [output.stage_id for output in outputs] == [0, 0]
assert [output.final_output_type for output in outputs] == ["image", "image"]
assert [output.request_output.payload for output in outputs] == [
assert [output.outputs[0].text for output in outputs] == [
f"{engine.submitted[0]['request_id']}-diffusion-final",
f"{engine.submitted[1]['request_id']}-diffusion-final",
]
@@ -767,7 +767,7 @@ def test_omni_generate_llm_diffusion_yields_final_text_then_image_per_request(
assert len(outputs) == 4
assert [output.stage_id for output in outputs] == [0, 1, 0, 1]
assert [output.final_output_type for output in outputs] == ["text", "image", "text", "image"]
assert [output.request_output.payload for output in outputs] == [
assert [output.outputs[0].text for output in outputs] == [
f"{engine.submitted[0]['request_id']}-text-0",
f"{engine.submitted[0]['request_id']}-image-final",
f"{engine.submitted[1]['request_id']}-text-0",
@@ -1046,7 +1046,6 @@ def _enqueue_stage_error(
if kill_engine:
engine._alive = False
engine_output = OmniRequestOutput.from_error(msg["request_id"], error_text)
engine_output.payload = ""
engine.output_q.put_nowait(
OutputMessage(
request_id=msg["request_id"],
@@ -69,10 +69,10 @@ def _make_text_omni_output(
],
finished=finish_reason is not None,
)
return OmniRequestOutput(
return OmniRequestOutput.from_stage_output(
res,
request_id=request_id,
final_output_type="text",
request_output=res,
finished=finish_reason is not None,
)
@@ -108,12 +108,12 @@ def _make_audio_omni_output(
outputs=[completion],
finished=True,
)
return OmniRequestOutput(
return OmniRequestOutput.from_stage_output(
res,
request_id=request_id,
stage_id=stage_id,
replica_id=replica_id,
final_output_type="audio",
request_output=res,
finished=True,
)
@@ -151,7 +151,7 @@ def _build_serving_chat():
)
instance._create_audio_choice = MagicMock(
side_effect=lambda omni_res, role, request, stream=False: _mock_audio_choices(
index=omni_res.request_output.outputs[0].index,
index=omni_res.outputs[0].index,
role=role,
)
)
@@ -489,7 +489,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
request = _make_request(modalities=["text", "audio"])
empty_audio = _make_audio_omni_output()
empty_audio.request_output.outputs[0].multimodal_output = {"audio": []}
empty_audio.outputs[0].multimodal_output = {"audio": []}
def create_audio_choice(omni_res, role, request, stream=False):
return OmniOpenAIServingChat._create_audio_choice(serving_chat, omni_res, role, request, stream=stream)
+19 -1
View File
@@ -26,10 +26,28 @@ from PIL import Image
logger = logging.getLogger(__name__)
_synthetic_media_fallback_dir: Path | None = None
def _resolve_synthetic_media_cache_dir(cache_dir: Path | str | None) -> Path:
if cache_dir is not None:
return Path(cache_dir).expanduser().resolve()
return Path(tempfile.gettempdir()) / "vllm_omni_test_synthetic_media"
default = Path(tempfile.gettempdir()) / "vllm_omni_test_synthetic_media"
try:
default.mkdir(parents=True, exist_ok=True)
# Verify write access: the directory may exist but belong to
# another user (e.g. a previous CI job), causing PermissionError
# later when individual files are saved.
canary = default / ".write_test"
canary.touch()
canary.unlink()
return default
except (PermissionError, OSError):
global _synthetic_media_fallback_dir
if _synthetic_media_fallback_dir is None:
_synthetic_media_fallback_dir = Path(tempfile.mkdtemp(prefix="vllm_omni_test_synthetic_media_"))
return _synthetic_media_fallback_dir
def _np_array_from_mp4_bytes(video_bytes: bytes) -> np.ndarray:
+3 -3
View File
@@ -2938,9 +2938,9 @@ class OmniRunnerHandler:
audio_content = None
for stage_output in outputs:
if getattr(stage_output, "final_output_type", None) == "text":
text_content = stage_output.request_output.outputs[0].text
text_content = stage_output.outputs[0].text
if getattr(stage_output, "final_output_type", None) == "audio":
audio_content = stage_output.request_output.outputs[0].multimodal_output["audio"]
audio_content = stage_output.outputs[0].multimodal_output["audio"]
result.audio_content = audio_content
result.text_content = text_content
result.success = True
@@ -3080,7 +3080,7 @@ class OmniRunnerHandler:
mm_out: dict[str, Any] | None = None
for stage_out in outputs:
if getattr(stage_out, "final_output_type", None) == "audio":
mm_out = stage_out.request_output.outputs[0].multimodal_output
mm_out = stage_out.outputs[0].multimodal_output
break
if mm_out is None:
raise AssertionError("No audio output from pipeline")
+7 -2
View File
@@ -51,8 +51,13 @@ def _iter_multimodal_image_payloads(output: OmniRequestOutput | RequestOutput) -
if isinstance(output, OmniRequestOutput):
if output.multimodal_output is not None:
yield from _image_values_from_mapping_like(output.multimodal_output)
if output.request_output is not None:
yield from _iter_request_output_payloads(output.request_output)
# OmniRequestOutput IS the RequestOutput now; iterate completion
# outputs directly for multimodal payloads that the property
# (which returns only the first match) may have missed.
for completion in getattr(output, "outputs", []):
mm = getattr(completion, "multimodal_output", None)
if mm is not None:
yield from _image_values_from_mapping_like(mm)
else:
yield from _iter_request_output_payloads(output)
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Mapping
from dataclasses import asdict, is_dataclass
from dataclasses import asdict, fields, is_dataclass
from typing import Any
import msgspec
@@ -59,7 +59,10 @@ class OmniMsgpackEncoder:
if isinstance(obj, Image.Image):
return self._encode_pil_image(obj)
# RequestOutput (not a dataclass, needs special handling)
# RequestOutput (not a dataclass, needs special handling).
# Note: OmniRequestOutput subclasses RequestOutput but is a dataclass,
# so msgspec encodes it natively (field map) and it never reaches
# this hook.
if isinstance(obj, RequestOutput):
return self._encode_request_output(obj)
@@ -148,7 +151,6 @@ class OmniMsgpackEncoder:
"encoder_prompt": obj.encoder_prompt,
"encoder_prompt_token_ids": obj.encoder_prompt_token_ids,
"num_cached_tokens": obj.num_cached_tokens,
"multi_modal_placeholders": getattr(obj, "multi_modal_placeholders", None),
"kv_transfer_params": obj.kv_transfer_params,
}
# Handle multimodal_output attribute (MultimodalPayload or dict)
@@ -252,23 +254,27 @@ class OmniMsgpackDecoder:
def _decode_omni_request_output(self, obj: dict[str, Any]) -> Any:
"""Decode dict to OmniRequestOutput.
OmniRequestOutput is a dataclass, so we can use msgspec.convert
or construct it directly.
By the time this runs, ``_post_process`` has already reconstructed
nested values bottom-up (CompletionOutput, tensors, images, ...), so
the object is built directly from the known dataclass fields.
"""
from vllm_omni.outputs import OmniRequestOutput
data = dict(obj)
# Legacy wire format (pre-inheritance): the stage output was nested
# under "request_output". Merge its content into the new object.
legacy_inner = data.pop("request_output", None)
field_names = {f.name for f in fields(OmniRequestOutput)}
known = {k: v for k, v in data.items() if k in field_names}
try:
# Use msgspec.convert for dataclass reconstruction
return msgspec.convert(obj, OmniRequestOutput)
out = OmniRequestOutput(**known)
except Exception:
try:
# Fallback: construct directly if msgspec.convert fails
# (e.g., if some fields are missing or have wrong types)
return OmniRequestOutput(**obj)
except Exception:
# If both attempts fail, return dict as-is (defensive fallback)
# This should rarely happen if _is_omni_request_output is correct
return obj
# Defensive fallback; should rarely happen if
# _is_omni_request_output is correct.
return obj
if legacy_inner is not None and not isinstance(legacy_inner, dict):
out._copy_content_from(legacy_inner)
return out
def _decode_tensor(self, obj: dict[str, Any]) -> torch.Tensor:
"""Decode dict to torch.Tensor."""
@@ -314,19 +320,16 @@ class OmniMsgpackDecoder:
We construct it manually using only the known __init__ parameters to
avoid triggering the "Ignoring extra arguments" warning in vllm.
Fields that are not part of RequestOutput.__init__ (e.g.
multi_modal_placeholders, multimodal_output) are extracted first and
multimodal_output) are extracted first and
then restored as dynamic attributes after construction.
"""
# Extract dynamically-added / non-init fields before constructing so
# they are not passed as unknown **kwargs to RequestOutput.__init__.
mm_output = obj.pop("multimodal_output", None)
multi_modal_placeholders = obj.pop("multi_modal_placeholders", None)
ro = RequestOutput(**obj)
# Restore dynamic attributes that are not part of __init__.
if multi_modal_placeholders is not None:
setattr(ro, "multi_modal_placeholders", multi_modal_placeholders)
if mm_output is not None:
setattr(ro, "multimodal_output", mm_output)
return ro
+2 -2
View File
@@ -1633,12 +1633,12 @@ class Orchestrator:
from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision
engine_output = attach_duplex_output_decision(
OmniRequestOutput(
OmniRequestOutput.from_stage_output(
output,
request_id=req_id,
finished=True,
stage_id=stage_id,
final_output_type=decision.final_output_type,
request_output=output,
),
decision,
)
+5 -9
View File
@@ -628,7 +628,6 @@ class OmniBase(PDDisaggregationMixin):
self.prom_metrics.set_running(running)
self.prom_metrics.set_waiting(max(0, total - running))
images = getattr(engine_outputs, "images", []) if output_type == "image" else []
response_metrics: dict[str, Any] = {}
stage_metrics: dict[str, dict[str, Any]] = {}
rid_key = str(req_id)
@@ -653,19 +652,16 @@ class OmniBase(PDDisaggregationMixin):
response_metrics["final_output_type"] = current_stage_metrics["final_output_type"]
response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"]
response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"]
return OmniRequestOutput(
# Generation content (outputs, prompt, images, trajectory_*, ...) is
# copied from engine_outputs onto the returned object by
# OmniRequestOutput.from_stage_output().
return OmniRequestOutput.from_stage_output(
engine_outputs,
request_id=req_id or "",
finished=finished,
stage_id=stage_id,
replica_id=result.replica_id,
final_output_type=output_type,
request_output=engine_outputs,
images=images,
trajectory_latents=getattr(engine_outputs, "trajectory_latents", None),
trajectory_timesteps=getattr(engine_outputs, "trajectory_timesteps", None),
trajectory_log_probs=getattr(engine_outputs, "trajectory_log_probs", None),
trajectory_decoded=getattr(engine_outputs, "trajectory_decoded", None),
_custom_output=getattr(engine_outputs, "_custom_output", {}),
metrics=response_metrics,
stage_durations=stage_durations,
peak_memory_mb=peak_memory_mb,
@@ -2649,12 +2649,6 @@ def _extract_images_from_result(result: Any) -> list[Any]:
images = []
if hasattr(result, "images") and result.images:
images = result.images
elif hasattr(result, "request_output"):
request_output = result.request_output
if isinstance(request_output, dict) and request_output.get("images"):
images = request_output["images"]
elif hasattr(request_output, "images") and request_output.images:
images = request_output.images
# Handle when generate more than one image
if images and isinstance(images[0], np.ndarray) and images[0].shape[0] > 1 and images[0].ndim == 5:
# Unwrap batch: (N, T, H, W, C) -> [img1, img2, ...]
@@ -119,9 +119,6 @@ class OmniOpenAIServingAudioGenerate(OpenAIServing, AudioMixin):
audio_output = None
if hasattr(final_output, "multimodal_output") and final_output.multimodal_output:
audio_output = final_output.multimodal_output
if not audio_output and hasattr(final_output, "request_output"):
if final_output.request_output and hasattr(final_output.request_output, "multimodal_output"):
audio_output = final_output.request_output.multimodal_output
# Check for audio data using either "audio" or "model_outputs" key
audio_key = None
+18 -20
View File
@@ -1468,7 +1468,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
try:
async for omni_res in result_generator:
final_output_type = omni_res.final_output_type
res = omni_res.request_output
res = omni_res
if final_output_type not in first_iteration_dict:
logger.warning(f"final output type: {final_output_type} is not needed by the request")
continue
@@ -2338,7 +2338,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
for omni_outputs in final_outputs:
choices_data = []
if omni_outputs.request_output is not None and not getattr(omni_outputs.request_output, "finished", False):
if not omni_outputs.finished:
continue
# Filter outputs based on requested modalites
@@ -2347,7 +2347,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
continue
if omni_outputs.final_output_type == "text":
if omni_outputs.request_output is not None:
if omni_outputs.stage_id is not None or omni_outputs.outputs:
(
choices_data,
usage,
@@ -2362,7 +2362,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
role,
reasoning_parser,
)
final_res = omni_outputs.request_output
final_res = omni_outputs
else:
# Diffusion pipeline text output (e.g. single-stage
# img2text / text2text) — no AR request_output, so build
@@ -2463,7 +2463,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
role: str,
reasoning_parser: ReasoningParser | None = None,
):
final_res = omni_outputs.request_output
final_res = omni_outputs
if self.tool_call_id_type == "kimi_k2":
history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
else:
@@ -2711,7 +2711,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
self, omni_outputs: OmniRequestOutput, role: str, request: ChatCompletionRequest, stream: bool = False
) -> list[ChatCompletionResponseChoice] | list[ChatCompletionResponseStreamChoice] | ErrorResponse:
choices: list[ChatCompletionResponseChoice] = []
final_res = omni_outputs.request_output
final_res = omni_outputs
# OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
# Reference: examples/offline_inference/qwen3_omni/end2end.py line 421
# The attribute is attached dynamically when stage audio arrives; fall
@@ -2834,7 +2834,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
from PIL import Image
choices: list[ChatCompletionResponseChoice] = []
final_res = omni_outputs.request_output
# Handle profiling data
stage_durations = omni_outputs.stage_durations
@@ -2846,10 +2845,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
# First check omni_outputs.images directly (for diffusion mode via from_diffusion)
if omni_outputs.images:
images = omni_outputs.images
# Fall back to request_output for pipeline mode
# OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
elif final_res is not None and final_res.outputs:
completion_output = final_res.outputs[0]
# Fall back to completion outputs for pipeline mode (multimodal_output
# is attached to CompletionOutput by AR stages).
elif omni_outputs.outputs:
completion_output = omni_outputs.outputs[0]
final_res = omni_outputs
if hasattr(completion_output, "multimodal_output") and completion_output.multimodal_output:
image_data = completion_output.multimodal_output.get("image")
if image_data is not None:
@@ -3311,12 +3311,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
request_id=request_id,
)
images = getattr(result.request_output, "images", [])
images = getattr(result, "images", [])
stage_durations = result.stage_durations
peak_memory_mb = result.peak_memory_mb
cot_output = None
req_out = getattr(result, "request_output", None)
req_out = result
if req_out:
prompt_obj = getattr(req_out, "prompt", None)
if isinstance(prompt_obj, dict):
@@ -3326,7 +3326,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
if isinstance(ar_text, str) and ar_text.strip():
cot_output = ar_text
req_out = getattr(result, "request_output", None)
req_out = result
if req_out:
prompt_obj = getattr(req_out, "prompt", None)
if isinstance(prompt_obj, dict):
@@ -3336,7 +3336,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
if isinstance(ar_text, str) and ar_text.strip():
cot_output = ar_text
req_out = getattr(result, "request_output", None)
req_out = result
if req_out:
prompt_obj = getattr(req_out, "prompt", None)
if isinstance(prompt_obj, dict):
@@ -3368,8 +3368,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
stage_id = getattr(output, "stage_id", None)
metrics = getattr(output, "metrics", None) if return_stage_metrics else None
if final_output_type == "text" and stage_id == 0:
request_output = output.request_output
for completion in request_output.outputs:
for completion in output.outputs:
text = completion.text or ""
if not text:
continue
@@ -3382,7 +3381,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
)
yield f"data: {chunk.model_dump_json()}\n\n"
elif final_output_type == "image":
images = self._flatten_diffusion_images(getattr(output.request_output, "images", []))
images = self._flatten_diffusion_images(getattr(output, "images", []))
if not images:
raise RuntimeError("Streaming image edit produced an empty final image output.")
image_data = [
@@ -3686,8 +3685,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
# Image output path (text2img / img2img)
final_output_type = getattr(result, "final_output_type", "image")
# Handle nested OmniRequestOutput structure where images might be in request_output
images = getattr(result.request_output, "images", [])
images = getattr(result, "images", [])
multimodal_output = getattr(result, "multimodal_output", {}) or {}
stage_durations = result.stage_durations
peak_memory_mb = result.peak_memory_mb
@@ -2429,7 +2429,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
mm = getattr(res, "multimodal_output", None)
ro = None
if not mm:
ro = getattr(res, "request_output", None)
ro = res
mm = getattr(ro, "multimodal_output", None) if ro else None
if not mm:
# MultimodalOutputProcessor attaches mm_accumulated on per-completion outputs.
+3 -23
View File
@@ -505,16 +505,6 @@ class OmniOpenAIServingVideo:
videos = None
if hasattr(result, "images") and result.images:
videos = result.images
elif hasattr(result, "request_output"):
request_output = result.request_output
if isinstance(request_output, dict) and request_output.get("images"):
videos = request_output["images"]
elif hasattr(request_output, "images") and request_output.images:
videos = request_output.images
else:
request_multimodal_output = getattr(request_output, "multimodal_output", None)
if isinstance(request_multimodal_output, Mapping):
videos = request_multimodal_output.get("video")
if videos is None:
multimodal_output = getattr(result, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
@@ -534,16 +524,6 @@ class OmniOpenAIServingVideo:
multimodal_output = getattr(result, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
audio = multimodal_output.get("audio")
elif hasattr(result, "request_output"):
request_output = result.request_output
if isinstance(request_output, dict) and request_output.get("multimodal_output"):
mm_output = request_output.get("multimodal_output") or {}
if isinstance(mm_output, Mapping):
audio = mm_output.get("audio")
else:
request_multimodal_output = getattr(request_output, "multimodal_output", None)
if isinstance(request_multimodal_output, Mapping):
audio = request_multimodal_output.get("audio")
if audio is None:
return [None] * expected_count
@@ -599,7 +579,7 @@ class OmniOpenAIServingVideo:
if isinstance(multimodal_output, Mapping):
return dict(multimodal_output)
request_output = getattr(result, "request_output", None)
request_output = result
if isinstance(request_output, dict):
multimodal_output = request_output.get("multimodal_output")
if multimodal_output is None:
@@ -716,7 +696,7 @@ class OmniOpenAIServingVideo:
except (TypeError, ValueError):
pass
request_output = getattr(result, "request_output", None)
request_output = result
if isinstance(request_output, dict):
mm = request_output.get("multimodal_output") or {}
if isinstance(mm, Mapping):
@@ -755,7 +735,7 @@ class OmniOpenAIServingVideo:
if sample_rate is not None:
return sample_rate
request_output = getattr(result, "request_output", None)
request_output = result
if isinstance(request_output, dict):
multimodal_output = request_output.get("multimodal_output") or {}
if isinstance(multimodal_output, Mapping):
@@ -671,17 +671,8 @@ class OmniStreamingVideoOutputHandler:
if result.images:
# Video frames may be list[list[Image.Image]] despite images: list[Image.Image].
videos = result.images # type: ignore[assignment]
elif result.request_output is not None:
request_output = result.request_output
if isinstance(request_output, dict) and request_output.get("images"):
videos = request_output["images"]
elif isinstance(request_output, OmniRequestOutput):
if request_output.images:
# Same nested-frames mismatch as result.images above.
videos = request_output.images # type: ignore[assignment]
else:
request_multimodal_output = request_output.multimodal_output
videos = request_multimodal_output.get("video")
# OmniRequestOutput IS the RequestOutput now — images and
# multimodal_output are checked directly on result.
if videos is None:
videos = result.multimodal_output.get("video")
@@ -870,7 +870,7 @@ class OmniStreamingVideoHandler:
@staticmethod
def _get_audio_data(result: OmniRequestOutput):
"""Navigate OmniRequestOutput → multimodal_output['audio']. None on miss."""
request_output = getattr(result, "request_output", None)
request_output = result
if request_output is None:
return None
outputs = getattr(request_output, "outputs", None)
@@ -1000,7 +1000,7 @@ class OmniStreamingVideoHandler:
if result.final_output_type != "text":
return "", previous_text
request_output = getattr(result, "request_output", None)
request_output = result
if request_output is None:
return "", previous_text
+128 -177
View File
@@ -3,7 +3,7 @@ from typing import Any
import torch
from PIL import Image
from vllm.outputs import RequestOutput
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm_omni.inputs.data import OmniPromptType
@@ -71,12 +71,55 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
)
# RequestOutput attributes that hold generation content. When wrapping a
# stage output, these are copied onto the OmniRequestOutput itself so that
# the object IS the RequestOutput instead of holding one.
_REQUEST_OUTPUT_CONTENT_ATTRS = (
"request_id",
"prompt",
"prompt_token_ids",
"prompt_logprobs",
"outputs",
"finished",
"lora_request",
"encoder_prompt",
"encoder_prompt_token_ids",
"num_cached_tokens",
"kv_transfer_params",
)
# Omni-specific content copied when the wrapped stage output is itself an
# OmniRequestOutput (e.g. a diffusion stage inside a multi-stage pipeline).
_OMNI_CONTENT_ATTRS = (
"images",
"latents",
"trajectory_latents",
"trajectory_timesteps",
"trajectory_log_probs",
"trajectory_decoded",
"_multimodal_output",
"_custom_output",
)
@dataclass
class OmniRequestOutput:
class OmniRequestOutput(RequestOutput):
"""Unified request output for both pipeline stages and diffusion models.
Extends vLLM's ``RequestOutput`` so that omni outputs can flow directly
through vLLM serving codepaths (which expect ``prompt_token_ids``,
``outputs``, etc. as real attributes). The inherited fields store the
LLM generation content; omni-specific fields store pipeline/diffusion
extras.
Note: ``RequestOutput`` is a plain class (not a dataclass), so all of its
attributes are redeclared below as dataclass fields with defaults — the
dataclass-generated ``__init__`` replaces ``RequestOutput.__init__`` and
must set them itself.
This class handles outputs from:
1. Multi-stage LLM pipelines (with stage_id, final_output_type, request_output)
1. Multi-stage LLM pipelines (with stage_id, final_output_type, and the
inherited RequestOutput fields carrying the stage's generation content)
2. Diffusion models (with images, prompt, metrics)
Attributes:
@@ -85,31 +128,39 @@ class OmniRequestOutput:
stage_id: Identifier of the stage that produced this output (pipeline mode)
replica_id: Identifier of the stage replica that produced this output
final_output_type: Type of output ("text", "image", "audio", "latents")
request_output: The underlying RequestOutput from the stage (pipeline mode)
images: List of generated PIL images (diffusion mode)
prompt: The prompt used for generation (diffusion mode)
prompt: The prompt used for generation
latents: Optional tensor of latent representations (diffusion mode)
metrics: Optional dictionary of generation metrics
metrics: Generation metrics. A plain dict for omni outputs; may carry
vLLM's request stats object when copied from a raw RequestOutput.
"""
# --- Inherited RequestOutput attributes (redeclared as fields) ---
request_id: str = ""
prompt: OmniPromptType | None = None
prompt_token_ids: list[int] | None = None
prompt_logprobs: Any = None
outputs: list[CompletionOutput] = field(default_factory=list)
finished: bool = True
metrics: Any = field(default_factory=dict)
lora_request: Any = None
encoder_prompt: str | None = None
encoder_prompt_token_ids: list[int] | None = None
num_cached_tokens: int | None = None
kv_transfer_params: dict[str, Any] | None = None
# Pipeline stage fields
# --- Pipeline stage fields ---
stage_id: int | None = None
replica_id: int | None = None
final_output_type: str = "text"
request_output: RequestOutput | None = None
# Diffusion model fields
# --- Diffusion model fields ---
images: list[Image.Image] = field(default_factory=list)
prompt: OmniPromptType | None = None
latents: torch.Tensor | None = None
trajectory_latents: torch.Tensor | None = None
trajectory_timesteps: torch.Tensor | None = None
trajectory_log_probs: torch.Tensor | None = None
trajectory_decoded: list | None = None
metrics: dict[str, Any] = field(default_factory=dict)
_multimodal_output: dict[str, Any] = field(default_factory=dict)
_custom_output: dict[str, Any] = field(default_factory=dict)
@@ -124,6 +175,64 @@ class OmniRequestOutput:
error_status_code: int | None = None
error_type: str | None = None
def _copy_content_from(self, source: RequestOutput) -> None:
"""Copy generation content from a stage output into this object.
*source* should be a vLLM ``RequestOutput`` or a subclass such as
``OmniRequestOutput``. Control fields (``stage_id``,
``final_output_type``, ``stage_durations``, ``peak_memory_mb``,
``error``, ...) are intentionally not copied.
"""
# RequestOutput attributes — guaranteed on any RequestOutput.
# Only copy attributes that actually exist on the source, so that
# duck-typed mocks can omit optional fields without overwriting the
# dataclass defaults (e.g. ``outputs=[]``, ``finished=True``).
for name in _REQUEST_OUTPUT_CONTENT_ATTRS:
if hasattr(source, name):
setattr(self, name, getattr(source, name))
# Omni-specific fields — only present when the source is itself an
# OmniRequestOutput (e.g. a diffusion stage inside a pipeline).
if isinstance(source, OmniRequestOutput):
for name in _OMNI_CONTENT_ATTRS:
if hasattr(source, name):
setattr(self, name, getattr(source, name))
# Propagate multimodal_output from duck-typed sources (e.g. test mocks
# that set multimodal_output directly on a non-OmniRequestOutput).
if not isinstance(source, OmniRequestOutput):
src_mm = getattr(source, "multimodal_output", None)
if src_mm:
self._multimodal_output = src_mm
@classmethod
def from_stage_output(cls, source: RequestOutput, **kwargs: Any) -> "OmniRequestOutput":
"""Create an OmniRequestOutput from a stage's raw output.
Copies generation content (``outputs``, ``prompt``, ``prompt_token_ids``,
``finished``, ``images``, ``latents``, etc.) from *source* onto the
returned object. *source* may be a vLLM ``RequestOutput``, another
``OmniRequestOutput`` (which inherits from ``RequestOutput``).
This is the **preferred** way to construct an ``OmniRequestOutput``
that wraps a stage result.
Args:
source: The stage output whose content is copied onto the new object.
**kwargs: Passed through to the dataclass constructor (``request_id``,
``stage_id``, ``final_output_type``, ``metrics``,
``stage_durations``, ``peak_memory_mb``, ``finished``, etc.).
Typed as ``Any`` because the exact set of valid keys is the
dataclass field list, which is validated by ``cls(**kwargs)``
at call time.
Returns:
A new ``OmniRequestOutput`` with the stage's content flattened onto it.
"""
obj = cls(**kwargs)
obj._copy_content_from(source)
return obj
@classmethod
def from_error(
cls,
@@ -150,33 +259,6 @@ class OmniRequestOutput:
error_type=error_type,
)
@classmethod
def from_pipeline(
cls,
stage_id: int,
final_output_type: str,
request_output: RequestOutput,
replica_id: int | None = None,
) -> "OmniRequestOutput":
"""Create output from pipeline stage.
Args:
stage_id: Stage identifier
final_output_type: Type of output
request_output: The stage's output
Returns:
OmniRequestOutput configured for pipeline mode
"""
return cls(
request_id=getattr(request_output, "request_id", ""),
stage_id=stage_id,
replica_id=replica_id,
final_output_type=final_output_type,
request_output=request_output,
finished=True,
)
@classmethod
def from_diffusion(
cls,
@@ -236,35 +318,21 @@ class OmniRequestOutput:
@property
def multimodal_output(self) -> Any:
"""Return multimodal output from the underlying request output or local field.
"""Return the multimodal output payload.
For pipeline outputs, this checks completion outputs first, then request_output.
For diffusion outputs, this returns the local _multimodal_output field.
Checks completion outputs first (where multimodal_output is attached
by AR stages), then the local _multimodal_output field.
Returns either a MultimodalPayload (Phase 3+) or a plain dict (legacy).
"""
if self.request_output is None:
return self._multimodal_output
# Check completion outputs first (where multimodal_output is attached)
for output in getattr(self.request_output, "outputs", []):
for output in self.outputs:
if mm := getattr(output, "multimodal_output", None):
return mm
if mm := getattr(self.request_output, "multimodal_output", None):
return mm
return self._multimodal_output
@property
def custom_output(self) -> dict[str, Any]:
"""Return custom output data from diffusion pipelines.
For diffusion outputs, returns the local _custom_output field.
For pipeline outputs with an inner OmniRequestOutput, forwards
the custom_output from the inner request output.
"""
if self.request_output is not None:
if isinstance(self.request_output, OmniRequestOutput):
return self.request_output._custom_output
"""Return custom output data from diffusion pipelines."""
return self._custom_output
@custom_output.setter
@@ -276,58 +344,6 @@ class OmniRequestOutput:
"""Return the number of generated images."""
return len(self.images)
# Pass-through properties keep vLLM serving codepaths compatible with
# OmniRequestOutput for pipeline outputs (Issue #345).
@property
def prompt_token_ids(self) -> list[int] | None:
"""Return prompt token IDs from the underlying request output.
This property is required for compatibility with vLLM's streaming
chat completion generator which checks res.prompt_token_ids.
"""
if self.request_output is not None:
return getattr(self.request_output, "prompt_token_ids", None)
return None
@property
def outputs(self) -> list[Any]:
"""Return outputs from the underlying request output.
This property is required for compatibility with vLLM's streaming
and non-streaming chat completion generators.
"""
if self.request_output is not None:
return getattr(self.request_output, "outputs", [])
return []
@property
def encoder_prompt_token_ids(self) -> list[int] | None:
"""Return encoder prompt token IDs from the underlying request output."""
if self.request_output is not None:
return getattr(self.request_output, "encoder_prompt_token_ids", None)
return None
@property
def prompt_logprobs(self) -> Any:
"""Return prompt logprobs from the underlying request output."""
if self.request_output is not None:
return getattr(self.request_output, "prompt_logprobs", None)
return None
@property
def num_cached_tokens(self) -> int | None:
"""Return number of cached tokens from the underlying request output."""
if self.request_output is not None:
return getattr(self.request_output, "num_cached_tokens", None)
return None
@property
def kv_transfer_params(self) -> Any:
"""Return KV transfer params from the underlying request output."""
if self.request_output is not None:
return getattr(self.request_output, "kv_transfer_params", None)
return None
@property
def is_diffusion_output(self) -> bool:
"""Check if this is a diffusion model output."""
@@ -336,73 +352,7 @@ class OmniRequestOutput:
@property
def is_pipeline_output(self) -> bool:
"""Check if this is a pipeline stage output."""
return self.stage_id is not None and self.request_output is not None
def unwrap(self) -> "OmniRequestOutput":
"""Unwrap nested OmniRequestOutput to get the innermost result.
This helper handles the common pattern where pipeline outputs may wrap
other OmniRequestOutput instances. It recursively unwraps until it reaches
the final output with actual content (images, text, etc.).
Returns:
The innermost OmniRequestOutput containing the actual generation results.
Example:
```python
result = omni.generate(...)
output = OmniRequestOutput.unwrap_result(result)
if output.images:
# Access images directly
video_frames = output.images
```
"""
current = self
# Unwrap nested pipeline outputs
while current.is_pipeline_output and current.request_output is not None:
if isinstance(current.request_output, OmniRequestOutput):
current = current.request_output
else:
break
return current
@staticmethod
def unwrap_result(result: Any) -> "OmniRequestOutput":
"""Unwrap result from omni.generate() to get the final OmniRequestOutput.
This static helper handles the full unwrapping pattern including:
1. Extracting from list if needed
2. Type validation
3. Recursive unwrapping of nested pipeline outputs
Args:
result: The result from omni.generate() - may be a list or OmniRequestOutput
Returns:
The innermost OmniRequestOutput with actual content
Raises:
ValueError: If result is not an OmniRequestOutput or list containing one
Example:
```python
result = omni.generate(...)
output = OmniRequestOutput.unwrap_result(result)
# output is guaranteed to be the final OmniRequestOutput
```
"""
# Handle list wrapper
if isinstance(result, list):
if not result:
raise ValueError("Result list is empty")
result = result[0]
# Validate type
if not isinstance(result, OmniRequestOutput):
raise ValueError(f"Expected OmniRequestOutput, got {type(result)}")
# Unwrap nested outputs
return result.unwrap()
return self.stage_id is not None
def to_dict(self) -> dict[str, Any]:
"""Convert to dictionary for JSON serialization."""
@@ -441,7 +391,8 @@ class OmniRequestOutput:
f"finished={self.finished}",
f"stage_id={self.stage_id}",
f"final_output_type={self.final_output_type!r}",
f"request_output={self.request_output}",
f"prompt_token_ids={self.prompt_token_ids}",
f"outputs={self.outputs}",
f"images={images_repr}",
f"prompt={self.prompt!r}",
f"latents={self.latents}",
@@ -259,17 +259,14 @@ def _extract_inner_output(outputs: Any) -> Any:
if len(outputs) != 1:
raise ValueError(f"Expected one request output, got {len(outputs)}.")
outputs = outputs[0]
inner = getattr(outputs, "request_output", None)
inner = outputs
if inner is not None:
return inner
return outputs
def _request_peak_memory_mb(result: Any) -> float | None:
diffusion_output = getattr(result, "request_output", None)
peak_memory_mb = getattr(diffusion_output, "peak_memory_mb", None)
if peak_memory_mb is None:
peak_memory_mb = getattr(result, "peak_memory_mb", None)
peak_memory_mb = getattr(result, "peak_memory_mb", None)
if peak_memory_mb is None:
return None
return float(peak_memory_mb)
+2 -2
View File
@@ -79,7 +79,7 @@ def audio_chunk_pcm_bytes(omni_res: OmniRequestOutput) -> int:
sample rather than recording a wrong byte count.
"""
try:
final_res = omni_res.request_output
final_res = omni_res
mm_output = final_res.outputs[0].multimodal_output
audio_data = mm_output.get("audio")
if isinstance(audio_data, list):
@@ -101,7 +101,7 @@ def audio_chunk_pcm_bytes(omni_res: OmniRequestOutput) -> int:
def audio_chunk_sample_rate(omni_res: OmniRequestOutput) -> int:
"""Resolve audio sample rate for the request's audio stream."""
try:
mm_output = omni_res.request_output.outputs[0].multimodal_output
mm_output = omni_res.outputs[0].multimodal_output
except Exception:
return _metric_defs.DEFAULT_AUDIO_SAMPLE_RATE
return _metric_defs.resolve_audio_sample_rate(mm_output)