[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:
89 files changed
+944
-639
No files matched your search
@@ -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.")
|
||||
|
||||
|
||||
@@ -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
@@ -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")
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
Reference in new issue
Block a user