diff --git a/docs/user_guide/diffusion/frame_interpolation.md b/docs/user_guide/diffusion/frame_interpolation.md index ab03363cf..817ae67f7 100644 --- a/docs/user_guide/diffusion/frame_interpolation.md +++ b/docs/user_guide/diffusion/frame_interpolation.md @@ -47,7 +47,8 @@ For supported Wan2.2 pipelines, the execution order is: 1. Diffusion worker finishes denoising and decodes the raw video tensor. 2. Worker-side model-specific post-processing runs. 3. If frame interpolation is enabled, RIFE interpolates the decoded video - tensor on the worker side and records a FPS multiplier in `custom_output`. + tensor on the worker side and records a FPS multiplier in + `metadata.video.video_fps_multiplier`. 4. The API server receives the already-interpolated video and only performs MP4 export. diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py index 504dceec6..d26910161 100644 --- a/examples/offline_inference/lance/end2end.py +++ b/examples/offline_inference/lance/end2end.py @@ -293,16 +293,19 @@ def main(): if args.modality in ("img2text", "video2text"): for i, req_output in enumerate(outputs): - text = getattr(req_output, "output", None) or getattr(req_output, "text", None) - if text is None and hasattr(req_output, "custom_output"): - text = (req_output.custom_output or {}).get("text_output") + multimodal_output = getattr(req_output, "multimodal_output", {}) or {} + metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {} + text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {} + text = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or ( + text_metadata.get("text_output") if isinstance(text_metadata, dict) else None + ) + text = text or getattr(req_output, "output", None) or getattr(req_output, "text", None) print(f"[Output {i}] {text}") return if args.modality in ("text2video", "video2video", "image2video"): for i, req_output in enumerate(outputs): - custom = getattr(req_output, "custom_output", None) or {} - frames = custom.get("video_frames") + frames = getattr(req_output, "images", None) if not frames: print(f"[Output {i}] no video frames returned") continue diff --git a/examples/offline_inference/lance/gradio_demo.py b/examples/offline_inference/lance/gradio_demo.py index 7dfeb46c9..756b006cb 100644 --- a/examples/offline_inference/lance/gradio_demo.py +++ b/examples/offline_inference/lance/gradio_demo.py @@ -835,12 +835,15 @@ def run_task( print(f"[unified] task={task} {width}x{height} frames={num_frames} cfg={cfg} steps={steps} seed={seed}", flush=True) outputs = list(omni.generate(prompts=formatted, sampling_params_list=params)) - custom = getattr(outputs[0], "custom_output", None) or {} - # Lance/Bagel return text at outputs[0].output (string) or - # custom_output["text_output"]. + multimodal_output = getattr(outputs[0], "multimodal_output", {}) or {} + metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {} + text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {} text_out = None if task in {"x2t_image", "x2t_video"}: - text_out = custom.get("text_output") or getattr(outputs[0], "output", None) + text_out = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or ( + text_metadata.get("text_output") if isinstance(text_metadata, dict) else None + ) + text_out = text_out or getattr(outputs[0], "output", None) if not isinstance(text_out, str): text_out = None @@ -856,13 +859,11 @@ def run_task( gr.update(value=str(text_out or "(empty)"), visible=True), ) if task in {"t2v", "i2v", "video_edit"}: - frames = custom.get("video_frames") + frames = getattr(outputs[0], "images", None) if frames is None: raise gr.Error("No video frames in output.") return (gr.update(visible=False), gr.update(value=_save_video(frames), visible=True), gr.update(visible=False)) - # image tasks — orchestrator strips ``outputs[0].output`` across the - # IPC boundary but preserves ``custom_output``, so prefer the latter. - img_out = custom.get("image") or getattr(outputs[0], "output", None) + img_out = (getattr(outputs[0], "images", None) or [None])[0] or getattr(outputs[0], "output", None) if img_out is None: raise gr.Error("No image in output.") return (gr.update(value=_save_image(img_out), visible=True), gr.update(visible=False), gr.update(visible=False)) diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py index 7c727449e..9547c40b9 100644 --- a/examples/offline_inference/sensenova_u1/end2end.py +++ b/examples/offline_inference/sensenova_u1/end2end.py @@ -278,12 +278,17 @@ def main(): ) for req_output in outputs: - custom = getattr(req_output, "_custom_output", {}) or {} - if args.print_think and custom.get("think_text"): - print(f"[Think]\n{custom['think_text']}\n") + multimodal_output = getattr(req_output, "multimodal_output", {}) or {} + metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {} + text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {} + think_text = text_metadata.get("think_text") if isinstance(text_metadata, dict) else None + if args.print_think and think_text: + print(f"[Think]\n{think_text}\n") if is_text_output: - text = custom.get("text_output", "") + text = multimodal_output.get("text", "") if isinstance(multimodal_output, dict) else "" + if not text: + text = text_metadata.get("text_output", "") if isinstance(text_metadata, dict) else "" if not text: text = getattr(req_output, "text", "") or "" print(f"[Response]\n{text}") diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index f0159f82e..ba482e793 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -294,7 +294,6 @@ def test_pipeline_registered_and_exported() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.registry import ( - _DIFFUSION_ACTION_POST_PROCESS_FUNCS, _DIFFUSION_IR_OP_PRIORITY_FUNCS, _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS, @@ -311,9 +310,6 @@ def test_pipeline_registered_and_exported() -> None: ) assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func" assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func" - assert ( - _DIFFUSION_ACTION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_action_post_process_func" - ) assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_ir_op_priority_func" assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__ @@ -718,11 +714,11 @@ def test_postprocess_handles_image_video_audio_and_validation() -> None: def test_action_postprocess_handles_robolab_policy_outputs() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( RoboLabPolicyInputs, - get_cosmos3_action_post_process_func, + get_cosmos3_post_process_func, make_robolab_action_postprocess_inputs, ) - func = get_cosmos3_action_post_process_func(SimpleNamespace()) + func = get_cosmos3_post_process_func(SimpleNamespace()) inputs = RoboLabPolicyInputs( prompt="Pick the cube.", video_tensor=torch.zeros(1, 3, 3, 16, 16), @@ -746,13 +742,41 @@ def test_action_postprocess_handles_robolab_policy_outputs() -> None: ) action = torch.tensor([[[0.0, 0.25], [1.0, 0.75]]]) - custom_output = {"robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs)} - processed = func(action, custom_output=custom_output) + processed = func( + { + "payload": { + "actions": action, + }, + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + "common": { + "action_only_output": True, + }, + "internal": { + "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs), + }, + }, + } + ) - assert processed.shape == (1, 2) - assert processed.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype - torch.testing.assert_close(torch.from_numpy(processed), torch.tensor([[1.0, 0.25]])) - assert "robolab_action_postprocess" not in custom_output + processed_action = processed["payload"]["actions"] + assert processed_action.shape == (1, 2) + assert processed_action.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype + torch.testing.assert_close(torch.from_numpy(processed_action), torch.tensor([[1.0, 0.25]])) + assert processed["metadata"] == { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + "common": { + "action_only_output": True, + }, + } def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.MonkeyPatch) -> None: @@ -1208,8 +1232,8 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint assert captured["flow_shifts"] == [10.0] # Transfer uses the V2V flow-sigma schedule (karras off) so flow_shift applies. assert captured["scheduler_use_karras_sigmas"] == [False] - assert output.custom_output["fps"] == expected_fps - assert output.output["video"].device.type == "meta" + assert output.output["metadata"]["video"]["fps"] == expected_fps + assert output.output["payload"]["video"].device.type == "meta" def test_forward_transfer_runs_multichunk_overlap_path( @@ -1278,12 +1302,12 @@ def test_forward_transfer_runs_multichunk_overlap_path( assert captured["conditional_frames"] == [2, 1] assert len(captured["decode_calls"]) == 2 - assert output.output["video"].shape == (1, 3, 8, 16, 16) + assert output.output["payload"]["video"].shape == (1, 3, 8, 16, 16) torch.testing.assert_close( - output.output["video"][0, 0, :, 0, 0], + output.output["payload"]["video"][0, 0, :, 0, 0], torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2, 0.2, 0.3, 0.4]), ) - assert output.custom_output["transfer_controls"]["edge"].shape == (1, 3, 8, 16, 16) + assert output.output["metadata"]["transfer"]["controls"]["edge"].shape == (1, 3, 8, 16, 16) torch.testing.assert_close(captured["targets"][0][:, :, 0], torch.full((1, 3, 16, 16), -1.0)) torch.testing.assert_close(captured["targets"][0][:, :, 1], torch.full((1, 3, 16, 16), 1.0)) torch.testing.assert_close(captured["targets"][0][:, :, 2:], torch.full((1, 3, 3, 16, 16), 1.0)) @@ -1507,8 +1531,13 @@ class TestForwardRouting: ) ) assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7] - assert output.custom_output["action"].shape == (1, 2, 2) - assert "action_only_output" not in output.custom_output + assert output.output["payload"]["actions"].shape == (1, 2, 2) + assert output.output["metadata"]["actions"] == { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + } + assert "common" not in output.output["metadata"] def test_forward_dispatches_robolab_policy_flow( self, @@ -1592,12 +1621,15 @@ class TestForwardRouting: assert captured["prepare_action_video"]["kwargs"] == {"image_size": None} assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7] assert captured["diffuse_calls"][-1]["timesteps"].tolist() == [4, 3, 2, 1] - assert output.output == {} - assert output.custom_output["action_only_output"] is True - assert output.custom_output["action"].shape == (1, 2, 2) - assert "actions" not in output.custom_output - assert "robolab_action_postprocess" in output.custom_output - assert "robolab_policy_inputs" not in output.custom_output + assert output.output["payload"]["actions"].shape == (1, 2, 2) + assert output.output["metadata"]["actions"] == { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + } + assert output.output["metadata"]["common"]["action_only_output"] is True + assert "robolab_action_postprocess" in output.output["metadata"]["internal"] + assert "robolab_policy_inputs" not in output.output["metadata"] @pytest.mark.parametrize( ("prompt", "sampling_params", "message"), diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py new file mode 100644 index 000000000..62d17f705 --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image as qwen_image + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_qwen_image_postprocess_preserves_output_envelope( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + model_path = tmp_path / "model" + vae_path = model_path / "vae" + vae_path.mkdir(parents=True) + (vae_path / "config.json").write_text(json.dumps({}), encoding="utf-8") + + processed = object() + captured: dict[str, object] = {} + + class FakeImageProcessor: + def __init__(self, *, vae_scale_factor: int) -> None: + captured["vae_scale_factor"] = vae_scale_factor + + def postprocess(self, images: torch.Tensor) -> object: + captured["images"] = images + return processed + + monkeypatch.setattr(qwen_image, "VaeImageProcessor", FakeImageProcessor) + postprocess = qwen_image.get_qwen_image_post_process_func(SimpleNamespace(model=str(model_path))) + + image = torch.zeros(1, 3, 4, 4) + output = postprocess( + { + "payload": {"image": image}, + "metadata": {"prompt_embeddings": {"present": True}}, + } + ) + + assert captured["vae_scale_factor"] == 16 + assert captured["images"] is image + assert output == { + "payload": {"image": processed}, + "metadata": {"prompt_embeddings": {"present": True}}, + } diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py index 185d51b54..22b305eab 100644 --- a/tests/diffusion/test_diffusion_output_formatter.py +++ b/tests/diffusion/test_diffusion_output_formatter.py @@ -51,21 +51,23 @@ def _timings() -> DiffusionStepTimings: ) -def test_formatter_preserves_single_video_audio_actions_and_custom_output( +def test_formatter_preserves_single_video_audio_actions_and_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) - custom_output = {"from_output": "kept"} postprocess_output = normalize_diffusion_postprocess_output( { - "video": ["frame-0"], - "audio": "audio-0", - "actions": "action-0", - "custom_output": {"from_postprocess": "merged"}, - "audio_sample_rate": 48000, - "fps": 24.0, - }, - custom_output, + "payload": { + "video": ["frame-0"], + "audio": "audio-0", + "actions": "action-0", + }, + "metadata": { + "audio": {"sample_rate": 48000}, + "video": {"fps": 24.0}, + "actions": {"action_mode": "policy"}, + }, + } ) results = format_diffusion_outputs( @@ -73,7 +75,6 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output( od_config=_config(), diffusion_output=DiffusionOutput( output=None, - custom_output=custom_output, stage_durations={"execute": 1.25}, peak_memory_mb=321.0, ), @@ -87,12 +88,12 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output( assert result.images == ["frame-0"] assert result.prompt == "prompt-0" assert result.final_output_type == "image" - assert result.custom_output == { - "from_output": "kept", - "from_postprocess": "merged", - } - assert custom_output == {"from_output": "kept"} assert result.multimodal_output == { + "metadata": { + "audio": {"sample_rate": 48000}, + "video": {"fps": 24.0}, + "actions": {"action_mode": "policy"}, + }, "audio": "audio-0", "audio_sample_rate": 48000, "fps": 24.0, @@ -110,18 +111,123 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output( } -def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch) -> None: +def test_formatter_normalizes_payload_metadata_envelope( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) postprocess_output = normalize_diffusion_postprocess_output( - "caption", - {"text_output": "caption"}, + { + "payload": { + "video": ["frame-0"], + "actions": "action-0", + }, + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + "common": { + "action_only_output": True, + }, + "internal": { + "robolab_action_postprocess": object(), + }, + }, + } + ) + + assert postprocess_output.outputs == { + "video": ["frame-0"], + "actions": "action-0", + } + assert postprocess_output.primary_key == "video" + assert postprocess_output.metadata == { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + "common": { + "action_only_output": True, + }, + } + assert "internal" not in postprocess_output.metadata + + [result] = format_diffusion_outputs( + request=_request("prompt-0"), + od_config=_config(), + diffusion_output=DiffusionOutput(output=None), + output_data={"raw": "output"}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.multimodal_output == { + "metadata": postprocess_output.metadata, + "actions": "action-0", + } + + +def test_formatter_maps_trajectory_payload_to_request_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) + latents = torch.zeros(2, 1, 4) + timesteps = torch.tensor([1.0, 0.5]) + log_probs = torch.tensor([-0.1, -0.2]) + postprocess_output = normalize_diffusion_postprocess_output( + { + "payload": { + "image": "image-0", + "trajectory": { + "latents": latents, + "timesteps": timesteps, + "log_probs": log_probs, + }, + }, + "metadata": {"trajectory": {"type": "denoising"}}, + } + ) + + [result] = format_diffusion_outputs( + request=_request("prompt-0"), + od_config=_config(), + diffusion_output=DiffusionOutput(output=None), + output_data={"raw": "output"}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.images == ["image-0"] + assert result.latents is latents + assert result.trajectory_latents is latents + assert result.trajectory_timesteps is timesteps + assert result.trajectory_log_probs is log_probs + assert result.multimodal_output == { + "metadata": {"trajectory": {"type": "denoising"}}, + "trajectory": { + "latents": latents, + "timesteps": timesteps, + "log_probs": log_probs, + }, + } + + +def test_formatter_preserves_text_envelope_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) + postprocess_output = normalize_diffusion_postprocess_output( + { + "payload": {"text": "caption"}, + "metadata": {"text": {"text_output": "caption"}}, + } ) [result] = format_diffusion_outputs( request=_request("describe this"), od_config=_config(), - diffusion_output=DiffusionOutput(output="caption"), - output_data="caption", + diffusion_output=DiffusionOutput(output=None), + output_data={"payload": {"text": "caption"}}, postprocess_output=postprocess_output, timings=_timings(), ) @@ -129,7 +235,10 @@ def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch) assert result.images == [] assert result.prompt == "describe this" assert result.final_output_type == "text" - assert result.custom_output == {"text_output": "caption"} + assert result.multimodal_output == { + "metadata": {"text": {"text_output": "caption"}}, + "text": "caption", + } def test_formatter_preserves_audio_output_with_model_sample_rate_fallback( @@ -144,7 +253,7 @@ def test_formatter_preserves_audio_output_with_model_sample_rate_fallback( "_try_load_model_cls", lambda _: AudioModel, ) - postprocess_output = normalize_diffusion_postprocess_output(["waveform"], {}) + postprocess_output = normalize_diffusion_postprocess_output(["waveform"]) [result] = format_diffusion_outputs( request=_request("speak"), @@ -175,8 +284,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions( "actions": "action-0", "audio_sample_rate": 16000, "fps": 30.0, - }, - {}, + } ) [result] = format_diffusion_outputs( @@ -192,6 +300,10 @@ def test_formatter_preserves_audio_model_video_audio_and_actions( assert result.prompt == "watch and listen" assert result.final_output_type == "image" assert result.multimodal_output == { + "metadata": { + "audio": {"sample_rate": 16000}, + "video": {"fps": 30.0}, + }, "audio": "audio-0", "audio_sample_rate": 16000, "fps": 30.0, @@ -207,8 +319,7 @@ def test_formatter_preserves_audio_only_postprocess_dict( { "audio": "waveform", "audio_sample_rate": 24000, - }, - {}, + } ) [result] = format_diffusion_outputs( @@ -224,6 +335,9 @@ def test_formatter_preserves_audio_only_postprocess_dict( assert result.prompt == "speak" assert result.final_output_type == "audio" assert result.multimodal_output == { + "metadata": { + "audio": {"sample_rate": 24000}, + }, "audio": "waveform", "audio_sample_rate": 24000, } @@ -238,7 +352,7 @@ def test_formatter_preserves_single_prompt_multiple_audio_outputs( "_try_load_model_cls", lambda _: None, ) - postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"], {}) + postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"]) [result] = format_diffusion_outputs( request=_request("speak", num_outputs_per_prompt=2), @@ -264,10 +378,8 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads( "video": ["frame-0", "frame-1"], "audio": ["audio-0", "audio-1"], "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), - "custom_output": {"shared": True}, "fps": 12.5, - }, - {}, + } ) results = format_diffusion_outputs( @@ -282,7 +394,7 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads( assert len(results) == 1 assert results[0].images == ["frame-0", "frame-1"] assert results[0].prompt == "prompt-0" - assert results[0].custom_output == {"shared": True} + assert results[0].multimodal_output["metadata"] == {"video": {"fps": 12.5}} assert results[0].multimodal_output["audio"] == ["audio-0", "audio-1"] assert results[0].multimodal_output["fps"] == 12.5 torch.testing.assert_close( diff --git a/tests/diffusion/test_diffusion_output_metadata.py b/tests/diffusion/test_diffusion_output_metadata.py new file mode 100644 index 000000000..ee3559f47 --- /dev/null +++ b/tests/diffusion/test_diffusion_output_metadata.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.outputs.output_metadata import ( + strip_internal_metadata, + validate_diffusion_metadata, + validate_public_diffusion_metadata, +) + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu] + + +def test_validate_accepts_canonical_metadata_fields() -> None: + validate_diffusion_metadata( + { + "video": { + "fps": 24.0, + "video_fps_multiplier": 2, + "shape": [1, 3, 16, 64, 64], + }, + "audio": { + "sample_rate": 48000, + "audiox_task": "tts", + }, + "actions": { + "raw_action_dim": 7, + "action_mode": "policy", + "domain_id": 3, + }, + "transfer": { + "controls": {"edge": torch.zeros(1, 3, 4, 8, 8)}, + "hints": ["edge"], + }, + "text": { + "text_output": "caption", + "think_text": "reasoning", + }, + "common": { + "action_only_output": True, + }, + } + ) + + +@pytest.mark.parametrize( + ("metadata", "error_type"), + [ + ({"video": {"fps": 0}}, ValueError), + ({"video": {"fps": "24"}}, TypeError), + ({"video": {"video_fps_multiplier": 0}}, ValueError), + ({"video": {"shape": [1, 0, 64]}}, ValueError), + ({"audio": {"sample_rate": 0}}, ValueError), + ({"audio": {"audiox_task": ""}}, ValueError), + ({"actions": {"raw_action_dim": 0}}, ValueError), + ({"actions": {"action_mode": ""}}, ValueError), + ({"actions": {"domain_id": "7"}}, TypeError), + ({"transfer": {"controls": {"edge": "not-a-tensor"}}}, TypeError), + ({"transfer": {"hints": ["edge", 1]}}, TypeError), + ({"text": {"text_output": 1}}, TypeError), + ({"text": {"think_text": 1}}, TypeError), + ({"common": {"action_only_output": "true"}}, TypeError), + ], +) +def test_validate_rejects_invalid_canonical_metadata(metadata: dict[str, object], error_type: type[Exception]) -> None: + with pytest.raises(error_type): + validate_diffusion_metadata(metadata) + + +def test_validate_allows_unknown_model_metadata() -> None: + validate_diffusion_metadata( + { + "model_specific": {"opaque": object()}, + "video": {"model_key": object()}, + } + ) + + +def test_public_metadata_rejects_internal_group() -> None: + with pytest.raises(ValueError, match="internal"): + validate_public_diffusion_metadata({"internal": {"secret": object()}}) + + +def test_strip_internal_metadata_removes_private_group() -> None: + metadata = { + "actions": {"action_mode": "policy"}, + "internal": {"robolab_action_postprocess": object()}, + } + + public_metadata = strip_internal_metadata(metadata) + + assert public_metadata == {"actions": {"action_mode": "policy"}} + assert "internal" not in public_metadata diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py index 6804d550d..a437a572a 100644 --- a/tests/diffusion/test_diffusion_plugin_hooks.py +++ b/tests/diffusion/test_diffusion_plugin_hooks.py @@ -16,7 +16,6 @@ from unittest.mock import Mock, patch import pytest from vllm_omni.diffusion.registry import ( - _DIFFUSION_ACTION_POST_PROCESS_FUNCS, _DIFFUSION_IR_OP_PRIORITY_FUNCS, _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS, @@ -62,7 +61,6 @@ class TestRegisterDiffusionModel: original_models = _DIFFUSION_MODELS.copy() original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy() original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy() - original_action_post = _DIFFUSION_ACTION_POST_PROCESS_FUNCS.copy() original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy() yield _DIFFUSION_MODELS.clear() @@ -71,8 +69,6 @@ class TestRegisterDiffusionModel: _DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre) _DIFFUSION_POST_PROCESS_FUNCS.clear() _DIFFUSION_POST_PROCESS_FUNCS.update(original_post) - _DIFFUSION_ACTION_POST_PROCESS_FUNCS.clear() - _DIFFUSION_ACTION_POST_PROCESS_FUNCS.update(original_action_post) _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear() _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority) @@ -84,7 +80,6 @@ class TestRegisterDiffusionModel: class_name="TestPipeline", pre_process_func_name="test_pre_process", post_process_func_name="test_post_process", - action_post_process_func_name="test_action_post_process", ir_op_priority_func_name="test_ir_op_priority", ) assert "TestPipeline" in _DIFFUSION_MODELS @@ -95,9 +90,21 @@ class TestRegisterDiffusionModel: ) assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process" assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process" - assert _DIFFUSION_ACTION_POST_PROCESS_FUNCS["TestPipeline"] == "test_action_post_process" assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority" + def test_register_model_accepts_deprecated_action_postprocess_keyword(self): + """Deprecated action hook keyword is accepted but not registered.""" + register_diffusion_model( + model_arch="LegacyActionPipeline", + module_name="test_plugin.diffusion.pipeline", + class_name="LegacyActionPipeline", + post_process_func_name="test_post_process", + action_post_process_func_name="test_action_post_process", + ) + + assert "LegacyActionPipeline" in _DIFFUSION_MODELS + assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process" + class TestWorkerUsesHook: """Test that DiffusionWorker resolves model runner via platform hook.""" diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py index 50ed772c1..ae5f9eee8 100644 --- a/tests/diffusion/test_diffusion_streaming_output.py +++ b/tests/diffusion/test_diffusion_streaming_output.py @@ -85,19 +85,20 @@ class _PipelineBackedEngine: outputs = self.pipeline.step_outputs(request) except Exception as exc: outputs = [ - DiffusionOutput(custom_output={"chunk": 0}, finished=False), + _streaming_diffusion_output(chunk=0, finished=False), DiffusionOutput(error=str(exc), finished=True), ] for output in outputs: if output.error is not None: yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)] continue - custom_output = output.custom_output or {} + payload, metadata = _streaming_payload_and_metadata(output) + stream_metadata = metadata.get("stream", {}) yield [ OmniRequestOutput.from_diffusion( request_id=request.request_id, - images=list(custom_output.get("images") or []), - custom_output=custom_output, + images=list(payload.get("image") or []), + multimodal_output={"metadata": {"stream": stream_metadata}}, finished=output.finished, ) ] @@ -106,6 +107,30 @@ class _PipelineBackedEngine: del request_id +def _streaming_diffusion_output( + *, + chunk: int, + finished: bool, + images: list | None = None, +) -> DiffusionOutput: + payload = {"image": images or []} + metadata = {"stream": {"chunk": chunk}} + return DiffusionOutput( + output={ + "payload": payload, + "metadata": metadata, + }, + finished=finished, + ) + + +def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]: + envelope = output.output if isinstance(output.output, dict) else {} + payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {} + metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {} + return payload, metadata + + class TestPipelineStreamingOutputToStageDiffusionClient: """Streaming pipeline output over ZMQ into ``StageDiffusionClient``.""" @@ -114,8 +139,8 @@ class TestPipelineStreamingOutputToStageDiffusionClient: """Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags.""" pipeline = _StepStreamingPipeline( [ - DiffusionOutput(custom_output={"chunk": 0}, finished=False), - DiffusionOutput(custom_output={"chunk": 1}, finished=True), + _streaming_diffusion_output(chunk=0, finished=False), + _streaming_diffusion_output(chunk=1, finished=True), ] ) @@ -123,7 +148,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient: assert [output.request_id for output in outputs] == ["req-stream", "req-stream"] assert [output.finished for output in outputs] == [False, True] - assert [output.custom_output["chunk"] for output in outputs] == [0, 1] + assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1] assert pipeline.requests[0].request_id == "req-stream" @pytest.mark.asyncio @@ -136,7 +161,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient: assert len(outputs) == 2 assert outputs[0].request_id == "req-error" assert outputs[0].finished is False - assert outputs[0].custom_output["chunk"] == 0 + assert outputs[0].multimodal_output["metadata"]["stream"]["chunk"] == 0 assert outputs[1].request_id == "req-error" assert outputs[1].finished is True assert outputs[1].error is not None @@ -227,8 +252,8 @@ class TestPipelineStreamingOutputToEntrypoint: """Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator.""" pipeline = _StepStreamingPipeline( [ - DiffusionOutput(custom_output={"chunk": 0}, finished=False), - DiffusionOutput(custom_output={"chunk": 1}, finished=True), + _streaming_diffusion_output(chunk=0, finished=False), + _streaming_diffusion_output(chunk=1, finished=True), ] ) inline_client = self._make_inline_pipeline_client(pipeline) @@ -247,7 +272,7 @@ class TestPipelineStreamingOutputToEntrypoint: await _wait_for(lambda: len(pipeline.requests) == 1) assert pipeline.requests[0].request_id.startswith("req-omni-") - assert [output.custom_output["chunk"] for output in outputs] == [0, 1] + 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] finally: await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client) @@ -258,8 +283,8 @@ class TestPipelineStreamingOutputToEntrypoint: frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)] pipeline = _StepStreamingPipeline( [ - DiffusionOutput(custom_output={"chunk": 0, "images": [frames]}, finished=False), - DiffusionOutput(custom_output={"chunk": 1, "images": [frames]}, finished=True), + _streaming_diffusion_output(chunk=0, images=[frames], finished=False), + _streaming_diffusion_output(chunk=1, images=[frames], finished=True), ] ) inline_client = self._make_inline_pipeline_client(pipeline) diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index 753050c41..00e77fc8f 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -331,7 +331,10 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag if ro and getattr(ro, "outputs", None): cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs) if not cot_text: - ar_text = getattr(out, "custom_output", {}).get("ar_generated_text") + multimodal_output = getattr(out, "multimodal_output", {}) or {} + metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {} + text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {} + ar_text = text_metadata.get("ar_generated_text") if isinstance(text_metadata, dict) else None if isinstance(ar_text, list): cot_text = "\n".join(text for text in ar_text if text) else: diff --git a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py b/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py index 97efab0fe..193d1b1e2 100644 --- a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py +++ b/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py @@ -288,7 +288,7 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline): else: # Both prompt_ids and prompt_embeds are None (e.g. during warmup/dummy run). # Return a minimal dummy output to avoid crashing. - return DiffusionOutput(output=None, custom_output={}) + return DiffusionOutput(output=None) if isinstance(negative_prompt_ids, list): negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device) @@ -399,14 +399,20 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline): image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] return DiffusionOutput( - output=_maybe_to_cpu(image), + output={ + "payload": { + "image": _maybe_to_cpu(image), + }, + "metadata": { + "prompt_embeddings": { + "prompt_embeds": _maybe_to_cpu(prompt_embeds), + "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask), + "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds), + "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask), + }, + }, + }, trajectory_latents=_maybe_to_cpu(all_latents), trajectory_log_probs=_maybe_to_cpu(all_log_probs), trajectory_timesteps=_maybe_to_cpu(all_timesteps), - custom_output={ - "prompt_embeds": _maybe_to_cpu(prompt_embeds), - "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask), - "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds), - "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask), - }, ) diff --git a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py index 9bcaf0e40..69952435e 100644 --- a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py +++ b/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py @@ -390,20 +390,23 @@ class vLLMOmniHttpServerLocal: diffusion_output = self._to_tensor(final_res.images[0]).float() / 255.0 - mm_output = final_res.custom_output or {} + metadata = ( + final_res.multimodal_output.get("metadata", {}) if isinstance(final_res.multimodal_output, dict) else {} + ) + prompt_embeddings = metadata.get("prompt_embeddings", {}) if isinstance(metadata, dict) else {} if sampling_params.get("logprobs", False): - all_log_probs = mm_output.get("all_log_probs") + all_log_probs = final_res.trajectory_log_probs log_probs = all_log_probs[0] if all_log_probs is not None else None else: log_probs = None - all_latents = mm_output.get("all_latents") - all_timesteps = mm_output.get("all_timesteps") - prompt_embeds = mm_output.get("prompt_embeds") - prompt_embeds_mask = mm_output.get("prompt_embeds_mask") - negative_prompt_embeds = mm_output.get("negative_prompt_embeds") - negative_prompt_embeds_mask = mm_output.get("negative_prompt_embeds_mask") + all_latents = final_res.trajectory_latents + all_timesteps = final_res.trajectory_timesteps + prompt_embeds = prompt_embeddings.get("prompt_embeds") + prompt_embeds_mask = prompt_embeddings.get("prompt_embeds_mask") + negative_prompt_embeds = prompt_embeddings.get("negative_prompt_embeds") + negative_prompt_embeds_mask = prompt_embeddings.get("negative_prompt_embeds_mask") extra_fields = { "all_latents": all_latents[0] if all_latents is not None else None, diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 3f1c5f8e5..ff69a2bf2 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -44,23 +44,18 @@ class MockVideoResult: videos, audios=None, sample_rate=None, - custom_output=None, + multimodal_output=None, stage_durations=None, peak_memory_mb=0.0, ): - self.multimodal_output = {"video": videos} + self.multimodal_output = dict(multimodal_output or {"video": videos}) if audios is not None: self.multimodal_output["audio"] = audios if sample_rate is not None: self.multimodal_output["audio_sample_rate"] = sample_rate - self._custom_output = custom_output or {} self.stage_durations = stage_durations or {} self.peak_memory_mb = peak_memory_mb - @property - def custom_output(self): - return self._custom_output - class FakeAsyncOmni: def __init__(self): @@ -835,7 +830,13 @@ def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker: engine.captured_sampling_params_list = sampling_params_list import numpy as np - yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], custom_output={"video_fps_multiplier": 2}) + yield MockVideoResult( + [np.zeros((1, 64, 64, 3), dtype=np.uint8)], + multimodal_output={ + "video": [np.zeros((1, 64, 64, 3), dtype=np.uint8)], + "metadata": {"video": {"video_fps_multiplier": 2}}, + }, + ) engine.generate = _generate @@ -939,11 +940,16 @@ def test_video_generation_response_exposes_action_payload(mocker: MockerFixture) yield MockVideoResult( [object()], - custom_output={ - "action": np.array([[[1.5, 2.5], [3.5, 4.5]]], dtype=np.float32), - "raw_action_dim": 2, - "action_mode": "policy", - "domain_id": 7, + multimodal_output={ + "video": [object()], + "actions": np.array([[[1.5, 2.5], [3.5, 4.5]]], dtype=np.float32), + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + }, }, ) @@ -981,11 +987,16 @@ def test_video_job_persists_action_metadata(test_client, mocker: MockerFixture): engine.captured_sampling_params_list = sampling_params_list yield MockVideoResult( [object()], - custom_output={ - "action": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32), - "raw_action_dim": 2, - "action_mode": "policy", - "domain_id": 7, + multimodal_output={ + "video": [object()], + "actions": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32), + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + }, }, ) @@ -1019,19 +1030,51 @@ def test_action_extraction_accepts_unbatched_action(): result = MockVideoResult( [object()], - custom_output={ - "action": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), - "raw_action_dim": 2, - "action_mode": "policy", - "domain_id": 7, + multimodal_output={ + "video": [object()], + "actions": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + }, }, ) actions = OmniOpenAIServingVideo._extract_action_outputs(result, expected_count=1) + assert actions and actions[0] is not None + assert actions[0].data == [[1.0, 2.0], [3.0, 4.0]] + assert actions[0].shape == [2, 2] + + +def test_action_extraction_accepts_multimodal_actions_payload(): + import numpy as np + + result = MockVideoResult([object()]) + result.multimodal_output.update( + { + "actions": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32), + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + }, + } + ) + + actions = OmniOpenAIServingVideo._extract_action_outputs(result, expected_count=1) + assert actions[0] is not None assert actions[0].data == [[1.0, 2.0], [3.0, 4.0]] assert actions[0].shape == [2, 2] + assert actions[0].raw_action_dim == 2 + assert actions[0].action_mode == "policy" + assert actions[0].domain_id == 7 def test_missing_handler_returns_503(): @@ -1840,7 +1883,13 @@ def test_worker_fps_multiplier_is_applied_to_sync_encoding(test_client, mocker: async def _generate(prompt, request_id, sampling_params_list): engine.captured_prompt = prompt engine.captured_sampling_params_list = sampling_params_list - yield MockVideoResult([object()], custom_output={"video_fps_multiplier": 2}) + yield MockVideoResult( + [object()], + multimodal_output={ + "video": [object()], + "metadata": {"video": {"video_fps_multiplier": 2}}, + }, + ) engine.generate = _generate diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 3d7e1bcfc..0feab5b28 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1198,8 +1198,10 @@ class DiffusionOutput: Final output (after pipeline completion) """ - # Fields may be replaced with SHM handle dicts by ipc.pack_diffusion_output_shm output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None + + # Legacy compatibility fields. New pipeline-specific payloads should be + # carried by output["payload"] instead. trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None trajectory_latents: torch.Tensor | dict[str, Any] | None = None trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None @@ -1212,10 +1214,6 @@ class DiffusionOutput: post_process_func: Callable[..., Any] | None = None - # Extra custom output data (e.g. latent trajectories, prompt embeds) - # passed through to OmniRequestOutput.custom_output - custom_output: dict[str, Any] = field(default_factory=dict) - # logged timings info, directly from Req.timings # timings: Optional["RequestTimings"] = None @@ -1230,10 +1228,9 @@ class DiffusionOutput: # memory usage info peak_memory_mb: float = 0.0 - # When True, move all tensor fields (including tensors inside - # ``custom_output``) to CPU at construction time. Useful when the output - # is shipped across process boundaries (e.g. step-execution mode) and the - # receiving side must not initialise a stray CUDA context. + # When True, move tensor fields to CPU at construction time. Useful when + # the output is shipped across process boundaries (e.g. step-execution + # mode) and the receiving side must not initialise a stray CUDA context. to_cpu: bool = False def __post_init__(self) -> None: @@ -1243,14 +1240,18 @@ class DiffusionOutput: def _maybe_to_cpu(value: Any) -> Any: if isinstance(value, torch.Tensor): return value.detach().cpu() + if isinstance(value, dict): + return {key: _maybe_to_cpu(item) for key, item in value.items()} + if isinstance(value, list): + return [_maybe_to_cpu(item) for item in value] + if isinstance(value, tuple): + return tuple(_maybe_to_cpu(item) for item in value) return value self.output = _maybe_to_cpu(self.output) self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps) self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents) self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs) - if self.custom_output: - self.custom_output = {k: _maybe_to_cpu(v) for k, v in self.custom_output.items()} @classmethod def from_exception(cls, exc: BaseException) -> "DiffusionOutput": diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index f552ae293..d6e7b35de 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -10,7 +10,7 @@ import queue import threading import time from collections.abc import AsyncGenerator, Iterable -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any import numpy as np @@ -40,7 +40,6 @@ from vllm_omni.diffusion.output_formatter import ( ) from vllm_omni.diffusion.registry import ( DiffusionModelRegistry, - get_diffusion_action_post_process_func, get_diffusion_post_process_func, get_diffusion_pre_process_func, ) @@ -155,17 +154,10 @@ class DiffusionEngine: self.od_config = od_config self.post_process_func = get_diffusion_post_process_func(od_config) - self.action_post_process_func = get_diffusion_action_post_process_func(od_config) self.pre_process_func = get_diffusion_pre_process_func(od_config) # Cache whether the model-specific postprocess accepts request-level # sampling params so step() can support both legacy and extended hooks. self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params") - self._action_post_process_accepts_sampling_params = _func_accepts_parameter( - self.action_post_process_func, "sampling_params" - ) - self._action_post_process_accepts_custom_output = _func_accepts_parameter( - self.action_post_process_func, "custom_output" - ) self.step_execution = bool(getattr(od_config, "step_execution", False)) if self.od_config.streaming_output and not self.step_execution: @@ -307,46 +299,18 @@ class DiffusionEngine: if self.od_config.enable_cpu_offload: output_data = _move_tensor_tree_to_cpu(output_data) - custom_output = output.custom_output or {} - action_payload = None - action_only_output = bool(custom_output.get("action_only_output")) - postprocess_start_time = time.perf_counter() - if action_only_output: - outputs = [] - elif self.post_process_func is not None: + if self.post_process_func is not None: # Some video pipelines need request-level controls during # postprocess (for example worker-side frame interpolation). + postprocess_kwargs: dict[str, object] = {} if self._post_process_accepts_sampling_params: - outputs = self.post_process_func(output_data, sampling_params=request.sampling_params) - else: - outputs = self.post_process_func(output_data) + postprocess_kwargs["sampling_params"] = request.sampling_params + outputs = self.post_process_func(output_data, **postprocess_kwargs) else: outputs = output_data - postprocess_output = normalize_diffusion_postprocess_output(outputs, custom_output) - custom_output = postprocess_output.custom_output - action_payload = postprocess_output.action_payload - if action_payload is None: - action_payload = custom_output.get("actions") - if action_payload is not None: - postprocess_output = replace(postprocess_output, action_payload=action_payload) - action_post_process_func = getattr(self, "action_post_process_func", None) - if action_payload is None and action_post_process_func is not None: - raw_action_payload = custom_output.get("action") - if raw_action_payload is not None: - action_kwargs: dict[str, Any] = {} - if getattr(self, "_action_post_process_accepts_custom_output", False): - action_kwargs["custom_output"] = custom_output - if getattr(self, "_action_post_process_accepts_sampling_params", False): - action_kwargs["sampling_params"] = request.sampling_params - action_payload = action_post_process_func(raw_action_payload, **action_kwargs) - custom_output = {**custom_output, "actions": action_payload} - postprocess_output = replace( - postprocess_output, - custom_output=custom_output, - action_payload=action_payload, - ) + postprocess_output = normalize_diffusion_postprocess_output(outputs) postprocess_time = time.perf_counter() - postprocess_start_time logger.debug("Post-processing completed in %.4f seconds", postprocess_time) diff --git a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py b/vllm_omni/diffusion/models/audiox/pipeline_audiox.py index 02639f1af..395e06dc8 100644 --- a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py +++ b/vllm_omni/diffusion/models/audiox/pipeline_audiox.py @@ -60,6 +60,15 @@ def get_audiox_post_process_func(od_config: OmniDiffusionConfig): """Convert the pipeline's float audio tensor to a CPU numpy array for serving.""" def post_process_func(audio: torch.Tensor) -> Any: + if isinstance(audio, dict) and isinstance(audio.get("payload"), dict): + payload = dict(audio["payload"]) + audio_payload = payload.get("audio") + if isinstance(audio_payload, torch.Tensor): + payload["audio"] = audio_payload.detach().cpu().float().numpy() + return { + "payload": payload, + "metadata": audio.get("metadata") or {}, + } if isinstance(audio, torch.Tensor): return audio.detach().cpu().float().numpy() return audio @@ -915,9 +924,10 @@ class AudioXPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMix audio = audio[..., :target_samples] stage_durations = self.stage_durations if getattr(self, "enable_diffusion_pipeline_profiler", False) else None - custom_output = {"audiox_task": task_norm} return DiffusionOutput( - output=audio, - custom_output=custom_output, + output={ + "payload": {"audio": audio}, + "metadata": {"audio": {"audiox_task": task_norm}}, + }, stage_durations=stage_durations, ) diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 453622765..9786e2b47 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -739,8 +739,10 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf text_output = text_output.split("<|im_start|>")[-1] return DiffusionOutput( - output=text_output, - custom_output={"text_output": text_output}, + output={ + "payload": {"text": text_output}, + "metadata": {"text": {"text_output": text_output}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -847,22 +849,28 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf if trajectory_log_probs: trajectory_log_probs_stacked = torch.stack(trajectory_log_probs) - custom = {} + payload = {"image": img} + metadata = {} if think_text is not None: - custom["think_text"] = think_text - # Mirror the PIL image into ``custom_output`` so callers reading via - # the orchestrator IPC boundary (which strips the bare ``output`` - # field) can still recover the result. ``video_frames`` already - # uses this pattern. - custom["image"] = img + metadata["text"] = {"think_text": think_text} + trajectory_payload = {} + if trajectory_latents_stacked is not None: + trajectory_payload["latents"] = trajectory_latents_stacked + if trajectory_timesteps_stacked is not None: + trajectory_payload["timesteps"] = trajectory_timesteps_stacked + if trajectory_log_probs_stacked is not None: + trajectory_payload["log_probs"] = trajectory_log_probs_stacked + if trajectory_decoded is not None: + trajectory_payload["decoded"] = trajectory_decoded + if trajectory_payload: + payload["trajectory"] = trajectory_payload + metadata["trajectory"] = {"type": "denoising"} return DiffusionOutput( - output=img, - trajectory_latents=trajectory_latents_stacked, - trajectory_timesteps=trajectory_timesteps_stacked, - trajectory_log_probs=trajectory_log_probs_stacked, - trajectory_decoded=trajectory_decoded, - custom_output=custom, + output={ + "payload": payload, + "metadata": metadata, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) diff --git a/vllm_omni/diffusion/models/cosmos3/__init__.py b/vllm_omni/diffusion/models/cosmos3/__init__.py index 43be0c37a..6df062b5c 100644 --- a/vllm_omni/diffusion/models/cosmos3/__init__.py +++ b/vllm_omni/diffusion/models/cosmos3/__init__.py @@ -3,7 +3,6 @@ from .pipeline_cosmos3 import ( Cosmos3OmniDiffusersPipeline, - get_cosmos3_action_post_process_func, get_cosmos3_post_process_func, get_cosmos3_pre_process_func, ) @@ -11,7 +10,6 @@ from .transformer_cosmos3 import Cosmos3VFMTransformer __all__ = [ "Cosmos3OmniDiffusersPipeline", - "get_cosmos3_action_post_process_func", "get_cosmos3_post_process_func", "get_cosmos3_pre_process_func", "Cosmos3VFMTransformer", diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py index 27f6d36c2..c5fe41a91 100644 --- a/vllm_omni/diffusion/models/cosmos3/action.py +++ b/vllm_omni/diffusion/models/cosmos3/action.py @@ -4,8 +4,8 @@ These helpers cover action modes that use action tokens as an auxiliary output stream: ``policy``, ``forward_dynamics``, and ``inverse_dynamics``. The pipeline -returns predicted actions through ``custom_output`` rather than as a normal -``modalities`` entry. +returns predicted actions through the diffusion output payload/metadata +envelope rather than as a video frame stream. """ from __future__ import annotations diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 4153eaa67..cd8c12777 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -20,7 +20,8 @@ Video generation is further specialized by inputs and extra args: transfer video generation. * ``action_mode``: action-capable video generation. RoboLab/OpenPI observation payloads in ``extra_args["robot_obs"]`` or ``extra_args["observation"]`` - bypass normal video output and return action-only custom output. + bypass normal video output and return an action-only payload/metadata + envelope. Generated sound is video-only, cannot be combined with action or transfer, and is produced from sound latents rather than from ``multi_modal_data["audio"]``. @@ -133,6 +134,7 @@ from .utils import ( logger = init_logger(__name__) + COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = ( max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1 ) @@ -530,6 +532,37 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig): if output_type == "latent": return output + def _postprocess_action(action: Any, metadata: dict[str, Any]) -> Any: + internal_metadata = metadata.get("internal") + inputs = ( + internal_metadata.get("robolab_action_postprocess") if isinstance(internal_metadata, dict) else None + ) + if isinstance(inputs, RoboLabActionPostprocessInputs): + return postprocess_robolab_action(action, inputs) + return action + + pending_action = None + pending_action_metadata: dict[str, Any] = {} + envelope_public_metadata: dict[str, Any] = {} + if isinstance(output, dict) and isinstance(output.get("payload"), dict): + envelope_payload = dict(output.get("payload") or {}) + metadata = output.get("metadata") or {} + envelope_metadata = metadata if isinstance(metadata, dict) else {} + envelope_public_metadata = {key: value for key, value in envelope_metadata.items() if key != "internal"} + action = envelope_payload.pop("actions", None) + if action is not None: + pending_action = _postprocess_action(action, envelope_metadata) + pending_action_metadata = envelope_public_metadata + if not envelope_payload: + return { + "payload": { + "video": [], + "actions": pending_action, + }, + "metadata": pending_action_metadata, + } + output = envelope_payload + audio = None audio_sample_rate = None if isinstance(output, dict): @@ -568,11 +601,31 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig): # check_video_safety expects a 5D tensor; re-add T axis. checked = check_video_safety(image.unsqueeze(2)) image = checked.squeeze(2) - return video_processor.postprocess(image, output_type="pil") - if is_guardrails_enabled(od_config, sampling_params): + processed_image = video_processor.postprocess(image, output_type="pil") + if envelope_public_metadata: + return { + "payload": {"image": processed_image}, + "metadata": envelope_public_metadata, + } + return processed_image + guardrails_enabled = is_guardrails_enabled(od_config, sampling_params) + if guardrails_enabled: video = check_video_safety(video) processed_video = video_processor.postprocess_video(video, output_type=output_type) if audio is None: + if pending_action is not None: + return { + "payload": { + "video": processed_video, + "actions": pending_action, + }, + "metadata": pending_action_metadata, + } + if envelope_public_metadata: + return { + "payload": {"video": processed_video}, + "metadata": envelope_public_metadata, + } return processed_video if isinstance(audio, torch.Tensor): audio = audio.detach().cpu() @@ -583,33 +636,38 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig): } if audio_sample_rate is not None: result["audio_sample_rate"] = int(audio_sample_rate) + if pending_action is not None: + return { + "payload": { + "video": result["video"], + "audio": result["audio"], + "actions": pending_action, + }, + "metadata": { + **pending_action_metadata, + "video": {"fps": result["fps"]}, + "audio": {"sample_rate": result.get("audio_sample_rate")}, + }, + } + if envelope_public_metadata: + return { + "payload": { + "video": result["video"], + "audio": result["audio"], + }, + "metadata": { + **envelope_public_metadata, + "video": {"fps": result["fps"], **envelope_public_metadata.get("video", {})} + if isinstance(envelope_public_metadata.get("video"), dict) + else {"fps": result["fps"]}, + "audio": {"sample_rate": result.get("audio_sample_rate")}, + }, + } return result return post_process_func -def get_cosmos3_action_post_process_func(od_config: OmniDiffusionConfig): - """Build the custom-output postprocessor for Cosmos3 action predictions. - - Action modes return predicted action tensors in ``custom_output`` alongside - normal video output. RoboLab/OpenPI policy serving marks action-only output - and carries observation metadata used here to map model-space actions back - to the requested robot action representation. - """ - del od_config - - def action_post_process_func(action: Any, custom_output: dict[str, Any] | None = None, sampling_params=None): - del sampling_params - inputs = custom_output.get("robolab_action_postprocess") if isinstance(custom_output, dict) else None - if isinstance(inputs, RoboLabActionPostprocessInputs): - processed_action = postprocess_robolab_action(action, inputs) - custom_output.pop("robolab_action_postprocess", None) - return processed_action - return action - - return action_post_process_func - - def get_cosmos3_ir_op_priority_func(od_config: OmniDiffusionConfig): del od_config @@ -669,9 +727,10 @@ class Cosmos3OmniDiffusersPipeline( T2I, transfer, and action+sound are rejected. * **Action generation** when ``action_mode`` is provided. ``policy`` and ``forward_dynamics`` require an image or video input; ``inverse_dynamics`` - requires video input. Action predictions are returned in ``custom_output``. + requires video input. Action predictions are returned in the diffusion + output payload/metadata envelope. RoboLab/OpenPI observations in ``extra_args['robot_obs']`` or - ``extra_args['observation']`` return action-only custom output. + ``extra_args['observation']`` return an action-only envelope. * **T2V** otherwise (default video generation). """ @@ -1409,15 +1468,25 @@ class Cosmos3OmniDiffusersPipeline( logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start) action = action_latents[:, :, :raw_action_dim].detach().cpu() - custom_action_output: dict[str, Any] = { - "action": action, - "raw_action_dim": raw_action_dim, - "action_mode": action_mode, - "domain_id": domain_id, - "action_only_output": True, - "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs), + action_output: dict[str, Any] = { + "payload": { + "actions": action, + }, + "metadata": { + "actions": { + "raw_action_dim": raw_action_dim, + "action_mode": action_mode, + "domain_id": domain_id, + }, + "common": { + "action_only_output": True, + }, + "internal": { + "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs), + }, + }, } - return DiffusionOutput(output={}, custom_output=custom_action_output) + return DiffusionOutput(output=action_output) @staticmethod def _truthy(value) -> bool: @@ -2881,11 +2950,15 @@ class Cosmos3OmniDiffusersPipeline( full_output = torch.cat([normalized_input.to(full_output), full_output], dim=-1) return DiffusionOutput( - output={"video": full_output}, - custom_output={ - "transfer_controls": full_controls, - "transfer_hints": list(per_hint_frames), - "fps": frame_rate, + output={ + "payload": {"video": full_output}, + "metadata": { + "transfer": { + "controls": full_controls, + "hints": list(per_hint_frames), + }, + "video": {"fps": frame_rate}, + }, }, ) @@ -3287,12 +3360,18 @@ class Cosmos3OmniDiffusersPipeline( raise ValueError("Cosmos3 action generation finished without action latents.") action = action_latents[:, :, :raw_action_dim].detach().cpu() return DiffusionOutput( - output={"video": video}, - custom_output={ - "action": action, - "raw_action_dim": raw_action_dim, - "action_mode": action_mode, - "domain_id": domain_id, + output={ + "payload": { + "video": video, + "actions": action, + }, + "metadata": { + "actions": { + "raw_action_dim": raw_action_dim, + "action_mode": action_mode, + "domain_id": domain_id, + }, + }, }, ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 73e9ff666..c0476b879 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -2209,7 +2209,6 @@ class HunyuanImage3Pipeline( if output_type == "latent": return DiffusionOutput( output=latents, - custom_output={}, stage_durations=getattr(self, "stage_durations", None), ) @@ -2232,12 +2231,14 @@ class HunyuanImage3Pipeline( ) cot_text_list = state.extra.get(_STEP_COT_TEXT_LIST) or [] - custom_output = {} + metadata = {} if any(text is not None for text in cot_text_list): - custom_output["ar_generated_text"] = cot_text_list[0] + metadata["text"] = {"ar_generated_text": cot_text_list[0]} return DiffusionOutput( - output=image[0], - custom_output=custom_output, + output={ + "payload": {"image": image[0]}, + "metadata": metadata, + }, stage_durations=getattr(self, "stage_durations", None), ) @@ -2302,12 +2303,14 @@ class HunyuanImage3Pipeline( outputs = self._generate(**model_inputs, **kwargs) image = outputs[0] - custom_output = {} + metadata = {} if any(t is not None for t in cot_text_list): - custom_output["ar_generated_text"] = cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list + metadata["text"] = {"ar_generated_text": cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list} stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None return DiffusionOutput( - output=image, - custom_output=custom_output, + output={ + "payload": {"image": image}, + "metadata": metadata, + }, stage_durations=stage_durations, ) diff --git a/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py b/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py index 6c28e760c..10498620c 100644 --- a/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py +++ b/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py @@ -247,11 +247,13 @@ class InternVLAA1Pipeline(nn.Module, DiffusionPipelineProfilerMixin): noise=extra_args.get("noise"), decode_image=bool(extra_args.get("decode_image", False)), ) - custom_output: dict[str, Any] = {} + metadata: dict[str, Any] = {} if decoded is not None: - custom_output["decoded"] = decoded + metadata["actions"] = {"decoded": decoded} return DiffusionOutput( - output=output, - custom_output=custom_output, + output={ + "payload": {"actions": output}, + "metadata": metadata, + }, post_process_func=get_internvla_a1_post_process_func(self.od_config), ) diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py index 6f7e9b95e..dc885ac44 100644 --- a/vllm_omni/diffusion/models/lance/pipeline_lance.py +++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py @@ -836,8 +836,10 @@ class LancePipeline(BagelPipeline): frames = [Image.fromarray(f) for f in frames_np] logger.info("Lance t2v: decoded %d frames at %dx%d", len(frames), frames[0].width, frames[0].height) return DiffusionOutput( - output=frames[0], - custom_output={"video_frames": frames, "video_shape": video_shape}, + output={ + "payload": {"video": frames}, + "metadata": {"video": {"shape": video_shape}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1069,13 +1071,11 @@ class LancePipeline(BagelPipeline): ) img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape) - # Put the PIL image into ``custom_output`` too — the orchestrator - # serializes ``custom_output`` across the IPC boundary but strips - # the bare ``output`` field, so consumers that only see - # ``outputs[0].output`` (e.g. gradio_demo) get ``None`` otherwise. return DiffusionOutput( - output=img, - custom_output={"image": img, "image_shape": image_shape}, + output={ + "payload": {"image": img}, + "metadata": {"image": {"shape": image_shape}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1361,8 +1361,10 @@ class LancePipeline(BagelPipeline): frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], out_shape) frames = [Image.fromarray(f) for f in frames_np] return DiffusionOutput( - output=frames[0], - custom_output={"video_frames": frames, "video_shape": out_shape}, + output={ + "payload": {"video": frames}, + "metadata": {"video": {"shape": out_shape}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1622,8 +1624,10 @@ class LancePipeline(BagelPipeline): frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape) frames = [Image.fromarray(f) for f in frames_np] return DiffusionOutput( - output=frames[0], - custom_output={"video_frames": frames, "video_shape": video_shape}, + output={ + "payload": {"video": frames}, + "metadata": {"video": {"shape": video_shape}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1705,8 +1709,10 @@ class LancePipeline(BagelPipeline): text_output = text_output.lstrip("\n") logger.info("Lance x2t_video: generated %d tokens", token_ids.shape[0]) return DiffusionOutput( - output=text_output, - custom_output={"text_output": text_output}, + output={ + "payload": {"text": text_output}, + "metadata": {"text": {"text_output": text_output}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1801,8 +1807,10 @@ class LancePipeline(BagelPipeline): text_output = text_output.lstrip("\n") logger.info("Lance x2t_image: generated %d tokens", token_ids.shape[0]) return DiffusionOutput( - output=text_output, - custom_output={"text_output": text_output}, + output={ + "payload": {"text": text_output}, + "metadata": {"text": {"text_output": text_output}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 638316039..986c2ef86 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -73,8 +73,19 @@ def get_qwen_image_post_process_func( image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) def post_process_func( - images: torch.Tensor, + images: torch.Tensor | dict[str, Any], ): + if isinstance(images, dict) and isinstance(images.get("payload"), dict): + payload = dict(images["payload"]) + image_payload = payload.get("image") + if image_payload is None: + raise ValueError("Qwen-Image postprocess expected payload['image'] in output envelope.") + payload["image"] = image_processor.postprocess(image_payload) + metadata = images.get("metadata") or {} + return { + "payload": payload, + "metadata": metadata if isinstance(metadata, dict) else {}, + } return image_processor.postprocess(images) return post_process_func diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index 9c4a052ae..67e141e6a 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -938,8 +938,10 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli ) logger.info("Text generation: %d chars", len(text_output)) return DiffusionOutput( - output=text_output, - custom_output={"text_output": text_output}, + output={ + "payload": {"text": text_output}, + "metadata": {"text": {"text_output": text_output}}, + }, ) # ----------------------------------------------------------------------- @@ -1457,10 +1459,15 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli images = _to_pil(image_prediction) img = images[0] if images else None - custom = {} + metadata = {} if think_text: - custom["think_text"] = think_text - return DiffusionOutput(output=img, custom_output=custom) + metadata["text"] = {"think_text": think_text} + return DiffusionOutput( + output={ + "payload": {"image": img}, + "metadata": metadata, + } + ) # ----------------------------------------------------------------------- # Weight loading diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py index aed72a995..cf5f84e7a 100644 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py @@ -164,11 +164,11 @@ class SoulXPreprocessPipeline(nn.Module, SupportAudioInput, SupportsComponentDis sample_rate=self.target_sr, ) - # Non-None placeholder so DiffusionEngine forwards custom_output to stage-1 - # (output=None triggers an early return that drops handoff). return DiffusionOutput( - output=torch.zeros(0), - custom_output={SOULX_PREPROCESSED_KEY: payload}, + output={ + "payload": {}, + "metadata": {SOULX_PREPROCESSED_KEY: payload}, + }, to_cpu=True, ) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py index 64fa213ab..a48b5236d 100644 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py @@ -4,6 +4,7 @@ from collections.abc import Callable from contextlib import contextmanager from typing import Any, ClassVar +import numpy as np import torch import torch.nn as nn from vllm.logger import init_logger @@ -30,6 +31,21 @@ _SOULX_PROFILER_TARGETS = [ ] +def convert_soulx_audio_output_to_numpy(audio: torch.Tensor | dict[str, Any]) -> np.ndarray | dict[str, Any]: + if isinstance(audio, dict) and isinstance(audio.get("payload"), dict): + payload = dict(audio["payload"]) + audio_payload = payload.get("audio") + if isinstance(audio_payload, torch.Tensor): + payload["audio"] = audio_payload.detach().cpu().float().numpy() + return { + "payload": payload, + "metadata": audio.get("metadata") or {}, + } + if isinstance(audio, torch.Tensor): + return audio.detach().cpu().float().numpy() + return audio + + class FlowMatchingAudioPipeline( nn.Module, CFGParallelMixin, @@ -326,7 +342,7 @@ class FlowMatchingAudioPipeline( req: OmniDiffusionRequest, *, kind: str, - custom_output_key: str, + metadata_key: str, infer_batch_fn: Callable[..., tuple[torch.Tensor, int]], prepare_extra_args: Callable[[dict[str, Any], Any], dict[str, Any]] | None = None, ) -> DiffusionOutput: @@ -353,7 +369,9 @@ class FlowMatchingAudioPipeline( ) return DiffusionOutput( - output=audio, - custom_output={custom_output_key: pitch_shift}, + output={ + "payload": {"audio": audio}, + "metadata": {"audio": {metadata_key: pitch_shift}}, + }, stage_durations=self._profiler_stage_durations() or {}, ) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py index 829256134..85fed9ee7 100644 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py @@ -19,6 +19,7 @@ from vllm_omni.diffusion.models.soulx_singer.modules import ( ) from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( FlowMatchingAudioPipeline, + convert_soulx_audio_output_to_numpy, ) from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( SOULX_PREPROCESSED_KEY, @@ -96,7 +97,7 @@ def get_soulxsinger_svc_pre_process_func(od_config: OmniDiffusionConfig): def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig): def post_process_func(audio: torch.Tensor): - return audio.detach().cpu().float().numpy() + return convert_soulx_audio_output_to_numpy(audio) return post_process_func @@ -400,7 +401,7 @@ class PipelineSoulXSingerSVC(FlowMatchingAudioPipeline): return self._forward_batch_from_request( req, kind="svc", - custom_output_key="pitch_shift", + metadata_key="pitch_shift", infer_batch_fn=self.infer_svc_batch, ) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py index 9be70e651..fb2c086f2 100644 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py @@ -15,6 +15,7 @@ from vllm_omni.diffusion.models.soulx_singer.modules import ( ) from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( FlowMatchingAudioPipeline, + convert_soulx_audio_output_to_numpy, ) from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( SOULX_PREPROCESSED_KEY, @@ -97,7 +98,7 @@ def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig): """Convert pipeline audio tensor output for offline consumers.""" def post_process_func(audio: torch.Tensor): - return audio.detach().cpu().float().numpy() + return convert_soulx_audio_output_to_numpy(audio) return post_process_func @@ -343,7 +344,7 @@ class PipelineSoulXSingerSVS(FlowMatchingAudioPipeline): return self._forward_batch_from_request( req, kind="svs", - custom_output_key="f0_shift", + metadata_key="f0_shift", infer_batch_fn=self.infer_svs_batch, prepare_extra_args=lambda extra_args, _sampling: normalize_svs_control_extra_args(extra_args), ) diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py index 83f81e481..ff4e4b7ca 100644 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py @@ -13,7 +13,7 @@ SOULX_PRECOMPUTED_KEYS_BY_KIND: dict[str, tuple[str, ...]] = { # Key under which the SoulX preprocess payload is stored in -# prompt["additional_information"] or DiffusionOutput.custom_output. +# prompt["additional_information"] or diffusion output metadata. SOULX_PREPROCESSED_KEY = "soulx_preprocessed" diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 5d7c8a281..d3726e087 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -187,7 +187,7 @@ def get_wan22_post_process_func( ): if output_type == "latent": return video - custom_output = {} + video_metadata = {} if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): video, multiplier = interpolate_video_tensor( video, @@ -195,10 +195,10 @@ def get_wan22_post_process_func( scale=sampling_params.frame_interpolation_scale, model_path=sampling_params.frame_interpolation_model_path, ) - custom_output["video_fps_multiplier"] = multiplier + video_metadata["video_fps_multiplier"] = multiplier return { - "video": video_processor.postprocess_video(video, output_type=output_type), - "custom_output": custom_output, + "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)}, + "metadata": {"video": video_metadata} if video_metadata else {}, } return post_process_func diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 65f74b7f9..8b042b595 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -65,7 +65,7 @@ def get_wan22_i2v_post_process_func( ): if output_type == "latent": return video - custom_output = {} + video_metadata = {} if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): video, multiplier = interpolate_video_tensor( video, @@ -73,10 +73,10 @@ def get_wan22_i2v_post_process_func( scale=sampling_params.frame_interpolation_scale, model_path=sampling_params.frame_interpolation_model_path, ) - custom_output["video_fps_multiplier"] = multiplier + video_metadata["video_fps_multiplier"] = multiplier return { - "video": video_processor.postprocess_video(video, output_type=output_type), - "custom_output": custom_output, + "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)}, + "metadata": {"video": video_metadata} if video_metadata else {}, } return post_process_func diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 0611c6b04..a36305fce 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -3,8 +3,9 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import Any +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TypeGuard from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.io_support import supports_audio_output @@ -12,6 +13,18 @@ from vllm_omni.diffusion.registry import DiffusionModelRegistry from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniPromptType from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.output_metadata import ( + DiffusionMetadata, + DiffusionMultimodalOutput, + DiffusionOutputEnvelope, + DiffusionPayload, + DiffusionPayloadValue, + DiffusionPostprocessRawOutput, + DiffusionTrajectoryPayload, + strip_internal_metadata, + validate_diffusion_metadata, + validate_public_diffusion_metadata, +) @dataclass(frozen=True) @@ -24,53 +37,79 @@ class DiffusionStepTimings: @dataclass(frozen=True) class DiffusionPostprocessOutput: - outputs: Any - custom_output: dict[str, Any] - audio_payload: Any | None = None - action_payload: Any | None = None - audio_sample_rate: int | None = None - fps: float | None = None - has_video_payload: bool = False + outputs: DiffusionPayload + metadata: DiffusionMetadata = field(default_factory=dict) + primary_key: str | None = None + + +def _is_output_envelope(outputs: DiffusionPostprocessRawOutput) -> TypeGuard[DiffusionOutputEnvelope]: + return isinstance(outputs, dict) and isinstance(outputs.get("payload"), dict) def normalize_diffusion_postprocess_output( - outputs: Any, - custom_output: dict[str, Any], + outputs: DiffusionPostprocessRawOutput, ) -> DiffusionPostprocessOutput: - """Normalize the legacy postprocess dict shape used by diffusion models. + """Normalize diffusion postprocess output into payload and metadata.""" - The returned envelope owns a shallow merged copy of ``custom_output`` so - postprocess values keep legacy override precedence without mutating the - caller-owned dict. - """ + if _is_output_envelope(outputs): + payload: DiffusionPayload = outputs.get("payload") or {} + metadata = outputs.get("metadata") or {} + if not isinstance(metadata, dict): + metadata = {} + validate_diffusion_metadata(metadata) - audio_payload = None - action_payload = None - audio_sample_rate = None - fps = None - has_video_payload = False - merged_custom_output = dict(custom_output) + public_metadata = strip_internal_metadata(metadata) + if "text" in payload and "text" not in public_metadata: + public_metadata = {**public_metadata, "text": {}} + validate_public_diffusion_metadata(public_metadata) + + return DiffusionPostprocessOutput( + outputs=payload, + metadata=public_metadata, + primary_key=_infer_primary_payload_key(payload), + ) if isinstance(outputs, dict): - has_video_payload = "video" in outputs - audio_payload = outputs.get("audio") - action_payload = outputs.get("actions") - merged_custom_output.update(outputs.get("custom_output") or {}) - audio_sample_rate = outputs.get("audio_sample_rate") - fps = outputs.get("fps") - outputs = outputs.get("video", outputs) + payload = {key: value for key, value in outputs.items() if key not in {"audio_sample_rate", "fps"}} + metadata = _metadata_from_legacy_payload(outputs) + if "text" in payload and "text" not in metadata: + metadata["text"] = {} + if metadata: + validate_public_diffusion_metadata(metadata) + return DiffusionPostprocessOutput( + outputs=payload, + metadata=metadata, + primary_key=_infer_primary_payload_key(payload), + ) return DiffusionPostprocessOutput( - outputs=outputs, - custom_output=merged_custom_output, - audio_payload=audio_payload, - action_payload=action_payload, - audio_sample_rate=audio_sample_rate, - fps=fps, - has_video_payload=has_video_payload, + outputs={"output": outputs}, + primary_key="output", ) +def _metadata_from_legacy_payload(payload: dict[str, DiffusionPayloadValue]) -> DiffusionMetadata: + metadata: DiffusionMetadata = {} + audio_sample_rate = payload.get("audio_sample_rate") + if audio_sample_rate is not None: + metadata["audio"] = {"sample_rate": audio_sample_rate} + fps = payload.get("fps") + if fps is not None: + metadata["video"] = {"fps": fps} + return metadata + + +def _infer_primary_payload_key(payload: DiffusionPayload) -> str | None: + for key in ("video", "image", "text", "audio", "output"): + if key in payload: + return key + if set(payload).issubset({"actions", "trajectory"}): + return None + if "actions" in payload: + return None + return next(iter(payload), None) + + def format_empty_diffusion_outputs( request: OmniDiffusionRequest, *, @@ -93,13 +132,14 @@ def format_diffusion_outputs( request: OmniDiffusionRequest, od_config: OmniDiffusionConfig, diffusion_output: DiffusionOutput, - output_data: Any, + output_data: DiffusionPostprocessRawOutput, postprocess_output: DiffusionPostprocessOutput, timings: DiffusionStepTimings, ) -> list[OmniRequestOutput]: """Convert a finished diffusion model output into API-facing outputs.""" - outputs = _ensure_list(postprocess_output.outputs) + primary_payload = _primary_payload(postprocess_output) + outputs = _ensure_list(primary_payload) metrics = { "preprocess_time_ms": timings.preprocess_time_s * 1000, "diffusion_engine_exec_time_ms": timings.exec_time_s * 1000, @@ -112,10 +152,10 @@ def format_diffusion_outputs( # Detect text output: when the pipeline returns a string (e.g., # SenseNova-U1 / BAGEL single-stage img2text / text2text), wrap it # as a text-type response instead of an image. - is_text_output = isinstance(output_data, str) and postprocess_output.custom_output.get("text_output") is not None + is_text_output = postprocess_output.primary_key == "text" or "text" in postprocess_output.metadata is_audio_output = supports_audio_output(od_config.model_class_name) - audio_sample_rate = postprocess_output.audio_sample_rate + audio_sample_rate = _metadata_audio_sample_rate(postprocess_output.metadata) if is_audio_output and audio_sample_rate is None: model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name) audio_sample_rate = getattr(model_cls, "audio_sample_rate", None) @@ -134,14 +174,42 @@ def format_diffusion_outputs( ) -def _ensure_list(outputs: Any) -> list[Any]: +def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]: if isinstance(outputs, list): return outputs return [outputs] if outputs is not None else [] -def _format_audio_multimodal_output(payload: Any, audio_sample_rate: int | None) -> dict[str, Any]: - mm_output: dict[str, Any] = {"audio": payload} +def _primary_payload(postprocess_output: DiffusionPostprocessOutput) -> DiffusionPayloadValue | None: + if postprocess_output.primary_key is None: + return [] + return postprocess_output.outputs.get(postprocess_output.primary_key) + + +def _metadata_audio_sample_rate(metadata: DiffusionMetadata) -> int | None: + audio_metadata = metadata.get("audio") + if not isinstance(audio_metadata, dict): + return None + sample_rate = audio_metadata.get("sample_rate") + return sample_rate if isinstance(sample_rate, int) else None + + +def _metadata_video_fps(metadata: DiffusionMetadata) -> float | None: + video_metadata = metadata.get("video") + if not isinstance(video_metadata, dict): + return None + fps = video_metadata.get("fps") + return fps if isinstance(fps, (int, float)) and not isinstance(fps, bool) else None + + +def _format_audio_multimodal_output( + payload: DiffusionPayloadValue, + audio_sample_rate: int | None, + metadata: DiffusionMetadata | None = None, +) -> DiffusionMultimodalOutput: + mm_output: DiffusionMultimodalOutput = {"audio": payload} + if metadata: + mm_output["metadata"] = metadata if audio_sample_rate is not None: mm_output["audio_sample_rate"] = audio_sample_rate return mm_output @@ -149,25 +217,29 @@ def _format_audio_multimodal_output(payload: Any, audio_sample_rate: int | None) def _has_non_audio_postprocess_payload(postprocess_output: DiffusionPostprocessOutput) -> bool: return ( - postprocess_output.has_video_payload - or postprocess_output.action_payload is not None - or postprocess_output.fps is not None + "video" in postprocess_output.outputs + or "image" in postprocess_output.outputs + or "text" in postprocess_output.metadata + or "actions" in postprocess_output.outputs + or "trajectory" in postprocess_output.outputs + or _metadata_video_fps(postprocess_output.metadata) is not None ) def _build_multimodal_output( postprocess_output: DiffusionPostprocessOutput, audio_sample_rate: int | None, -) -> dict[str, Any]: - mm_output: dict[str, Any] = {} - if postprocess_output.audio_payload is not None: - mm_output["audio"] = postprocess_output.audio_payload +) -> DiffusionMultimodalOutput: + mm_output: DiffusionMultimodalOutput = {} + if postprocess_output.metadata: + mm_output["metadata"] = postprocess_output.metadata + for key, value in postprocess_output.outputs.items(): + if key in {"audio", "actions", "trajectory"}: + mm_output[key] = value if audio_sample_rate is not None: mm_output["audio_sample_rate"] = audio_sample_rate - if postprocess_output.fps is not None: - mm_output["fps"] = postprocess_output.fps - if postprocess_output.action_payload is not None: - mm_output["actions"] = postprocess_output.action_payload + if (fps := _metadata_video_fps(postprocess_output.metadata)) is not None: + mm_output["fps"] = fps return mm_output @@ -176,8 +248,8 @@ def _format_single_prompt_output( request: OmniDiffusionRequest, prompt: OmniPromptType, diffusion_output: DiffusionOutput, - outputs: list[Any], - metrics: dict[str, Any], + outputs: list[DiffusionPayloadValue], + metrics: dict[str, object], postprocess_output: DiffusionPostprocessOutput, is_text_output: bool, is_audio_output: bool, @@ -186,6 +258,13 @@ def _format_single_prompt_output( ) -> list[OmniRequestOutput]: request_id = request.request_id mm_output = _build_multimodal_output(postprocess_output, audio_sample_rate) + if is_text_output: + mm_output["text"] = outputs[0] if len(outputs) == 1 else outputs + trajectory_payload = _trajectory_payload(postprocess_output, diffusion_output) + trajectory_latents = trajectory_payload.get("latents") + trajectory_timesteps = trajectory_payload.get("timesteps") + trajectory_log_probs = trajectory_payload.get("log_probs") + trajectory_decoded = trajectory_payload.get("decoded") if is_text_output: return [ @@ -194,7 +273,6 @@ def _format_single_prompt_output( images=[], prompt=prompt, metrics=metrics, - custom_output=postprocess_output.custom_output, multimodal_output=mm_output, final_output_type="text", stage_durations=diffusion_output.stage_durations, @@ -204,7 +282,7 @@ def _format_single_prompt_output( ] if is_audio_output and not _has_non_audio_postprocess_payload(postprocess_output): - request_audio_payload = postprocess_output.audio_payload + request_audio_payload = postprocess_output.outputs.get("audio") if request_audio_payload is None: request_audio_payload = outputs[0] if len(outputs) == 1 else outputs return [ @@ -213,14 +291,15 @@ def _format_single_prompt_output( images=[], prompt=prompt, metrics=metrics, - latents=diffusion_output.trajectory_latents, - trajectory_latents=diffusion_output.trajectory_latents, - trajectory_timesteps=diffusion_output.trajectory_timesteps, - trajectory_log_probs=diffusion_output.trajectory_log_probs, - trajectory_decoded=diffusion_output.trajectory_decoded, + latents=trajectory_latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, multimodal_output=_format_audio_multimodal_output( request_audio_payload, audio_sample_rate, + postprocess_output.metadata, ), final_output_type="audio", stage_durations=diffusion_output.stage_durations, @@ -235,15 +314,43 @@ def _format_single_prompt_output( images=outputs, prompt=prompt, metrics=metrics, - latents=diffusion_output.trajectory_latents, - trajectory_latents=diffusion_output.trajectory_latents, - trajectory_timesteps=diffusion_output.trajectory_timesteps, - trajectory_log_probs=diffusion_output.trajectory_log_probs, - trajectory_decoded=diffusion_output.trajectory_decoded, - custom_output=postprocess_output.custom_output, + latents=trajectory_latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, multimodal_output=mm_output, stage_durations=diffusion_output.stage_durations, peak_memory_mb=diffusion_output.peak_memory_mb, finished=finished, ), ] + + +def _trajectory_payload( + postprocess_output: DiffusionPostprocessOutput, + diffusion_output: DiffusionOutput, +) -> DiffusionTrajectoryPayload: + trajectory: DiffusionTrajectoryPayload = {} + payload = postprocess_output.outputs.get("trajectory") + if isinstance(payload, Mapping): + for source_key, target_key in ( + ("latents", "latents"), + ("timesteps", "timesteps"), + ("log_probs", "log_probs"), + ("decoded", "decoded"), + ): + value = payload.get(source_key) + if value is not None: + trajectory[target_key] = value + + fallback_fields = ( + ("latents", diffusion_output.trajectory_latents), + ("timesteps", diffusion_output.trajectory_timesteps), + ("log_probs", diffusion_output.trajectory_log_probs), + ("decoded", diffusion_output.trajectory_decoded), + ) + for key, value in fallback_fields: + if key not in trajectory and value is not None: + trajectory[key] = value + return trajectory diff --git a/vllm_omni/diffusion/output_metadata.py b/vllm_omni/diffusion/output_metadata.py new file mode 100644 index 000000000..abea8759e --- /dev/null +++ b/vllm_omni/diffusion/output_metadata.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compatibility imports for diffusion output metadata. + +The shared output payload and metadata schema lives in +``vllm_omni.outputs.output_metadata``. Keep this module as a stable import +path for existing diffusion callers. +""" + +from vllm_omni.outputs.output_metadata import * # noqa: F403 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index c8a5ee974..1454a37b1 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -541,13 +541,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "Krea2Pipeline": "get_krea2_post_process_func", } -_DIFFUSION_ACTION_POST_PROCESS_FUNCS = { - # arch: action_post_process_func - # `action_post_process_func` function must be placed in {mod_folder}/{mod_relname}.py, - # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key. - "Cosmos3OmniDiffusersPipeline": "get_cosmos3_action_post_process_func", -} - _DIFFUSION_IR_OP_PRIORITY_FUNCS = { # arch: ir_op_priority_func # `ir_op_priority_func` function must be placed in {mod_folder}/{mod_relname}.py, @@ -588,8 +581,8 @@ def register_diffusion_model( class_name: str, pre_process_func_name: str | None = None, post_process_func_name: str | None = None, - action_post_process_func_name: str | None = None, ir_op_priority_func_name: str | None = None, + action_post_process_func_name: str | None = None, ) -> None: """Register a diffusion model pipeline from an out-of-tree plugin. @@ -608,13 +601,23 @@ def register_diffusion_model( post_process_func_name: Optional name of the post-process function located in *module_name*. Pass ``None`` to keep the existing entry when replacing a built-in model. - action_post_process_func_name: Optional name of the action post-process - function located in *module_name*. Pass ``None`` to keep the - existing entry when replacing a built-in model. ir_op_priority_func_name: Optional name of the IR op priority merge function located in *module_name*. Pass ``None`` to keep the existing entry when replacing a built-in model. + action_post_process_func_name: Deprecated compatibility-only keyword + for out-of-tree plugins. Action postprocess hooks are no longer + registered separately; move action handling into + ``post_process_func_name`` and return a payload/metadata envelope. """ + if action_post_process_func_name is not None: + logger.warning( + "Ignoring deprecated action_post_process_func_name=%r for diffusion " + "model %s. Move action postprocess logic into post_process_func_name " + "and return payload/metadata output.", + action_post_process_func_name, + model_arch, + ) + # Register model class in DiffusionModelRegistry DiffusionModelRegistry.register_model( model_arch, @@ -631,8 +634,6 @@ def register_diffusion_model( _DIFFUSION_PRE_PROCESS_FUNCS[model_arch] = pre_process_func_name if post_process_func_name is not None: _DIFFUSION_POST_PROCESS_FUNCS[model_arch] = post_process_func_name - if action_post_process_func_name is not None: - _DIFFUSION_ACTION_POST_PROCESS_FUNCS[model_arch] = action_post_process_func_name if ir_op_priority_func_name is not None: _DIFFUSION_IR_OP_PRIORITY_FUNCS[model_arch] = ir_op_priority_func_name @@ -665,13 +666,6 @@ def get_diffusion_post_process_func(od_config: OmniDiffusionConfig): return _load_process_func(od_config, func_name) -def get_diffusion_action_post_process_func(od_config: OmniDiffusionConfig): - if od_config.model_class_name not in _DIFFUSION_ACTION_POST_PROCESS_FUNCS: - return None - func_name = _DIFFUSION_ACTION_POST_PROCESS_FUNCS[od_config.model_class_name] - return _load_process_func(od_config, func_name) - - def get_diffusion_ir_op_priority_func(od_config: OmniDiffusionConfig): if od_config.model_class_name not in _DIFFUSION_IR_OP_PRIORITY_FUNCS: return None diff --git a/vllm_omni/engine/output_modality.py b/vllm_omni/engine/output_modality.py index ac72910dd..2536a8ade 100644 --- a/vllm_omni/engine/output_modality.py +++ b/vllm_omni/engine/output_modality.py @@ -9,9 +9,6 @@ from __future__ import annotations import re from enum import Enum, Flag, StrEnum, auto -from typing import Literal, TypeAlias - -FinalOutputModalityType: TypeAlias = Literal["text", "image", "audio", "video"] _MODALITY_ALIASES: dict[str, str] = { "speech": "audio", diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 276775f66..6db090955 100644 --- a/vllm_omni/engine/stage_client.py +++ b/vllm_omni/engine/stage_client.py @@ -11,8 +11,8 @@ if TYPE_CHECKING: from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniTokensPrompt from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.engine.output_modality import FinalOutputModalityType from vllm_omni.inputs.data import OmniSamplingParams +from vllm_omni.outputs.output_metadata import FinalOutputModalityType class StageClient(Protocol): diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 48afad0e5..8897043ca 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -32,7 +32,6 @@ from vllm_omni.engine.messages import ( EngineQueueMessage, RegisterRemoteReplicaMessage, ) -from vllm_omni.engine.output_modality import FinalOutputModalityType from vllm_omni.engine.stage_client import StageClient, StagePoolClient from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase from vllm_omni.engine.stage_engine_startup import ( @@ -63,6 +62,7 @@ from vllm_omni.engine.stage_init_utils import ( from vllm_omni.engine.stage_pool import StagePool from vllm_omni.entrypoints.stage_utils import resolve_stage_physical_devices from vllm_omni.entrypoints.utils import inject_omni_kv_config +from vllm_omni.outputs.output_metadata import FinalOutputModalityType from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 6896e3641..a1a51c303 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -120,6 +120,7 @@ from vllm_omni.entrypoints.openai.utils import ( from vllm_omni.errors import OmniClientError from vllm_omni.lora.request import LoRARequest from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.output_metadata import DiffusionMetadataMapping, DiffusionMetadataValue from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate logger = init_logger(__name__) @@ -230,11 +231,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): def _get_diffusion_extra_output_params( self, - custom_output: dict[str, Any] | None, - ) -> dict[str, Any] | None: - """Pick model-specific extra output keys from *custom_output*.""" - if not custom_output: - return None + output: object, + ) -> dict[str, DiffusionMetadataValue] | None: + """Pick model-specific extra output keys from diffusion metadata.""" + metadata: DiffusionMetadataMapping = {} + mm_output = getattr(output, "multimodal_output", None) + if isinstance(mm_output, dict): + raw_metadata = mm_output.get("metadata") + if isinstance(raw_metadata, dict): + metadata = raw_metadata if self._diffusion_extra_output_params is None: params: frozenset[str] = frozenset() @@ -248,9 +253,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if not self._diffusion_extra_output_params: return None - out = {k: custom_output[k] for k in self._diffusion_extra_output_params if k in custom_output} + flat_metadata: dict[str, DiffusionMetadataValue] = {} + for section in metadata.values(): + if isinstance(section, dict): + flat_metadata.update(section) + out = {k: flat_metadata[k] for k in self._diffusion_extra_output_params if k in flat_metadata} return out or None + @staticmethod + def _get_diffusion_text_output(output: object) -> str: + mm_output = getattr(output, "multimodal_output", None) + if isinstance(mm_output, dict) and mm_output.get("text") is not None: + return str(mm_output["text"]) + return "" + def _get_supported_speakers(self) -> set[str]: """Load supported speakers from model config (cached).""" if self._supported_speakers is not None: @@ -2167,8 +2183,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): else: # Diffusion pipeline text output (e.g. single-stage # img2text / text2text) — no AR request_output, so build - # a simple text choice from custom_output. - text_body = (omni_outputs.custom_output or {}).get("text_output", "") + # a simple text choice from diffusion multimodal output. + text_body = self._get_diffusion_text_output(omni_outputs) message = ChatMessage(role=role, content=text_body) choices_data = [ ChatCompletionResponseChoice( @@ -2194,7 +2210,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): response_metrics = {} response_metrics.setdefault("stage_durations", omni_outputs.stage_durations or {}) response_metrics.setdefault("peak_memory_mb", float(omni_outputs.peak_memory_mb or 0.0)) - extra = self._get_diffusion_extra_output_params(omni_outputs.custom_output) + extra = self._get_diffusion_extra_output_params(omni_outputs) if extra: response_metrics.update(extra) choices.extend(choices_data) @@ -3428,7 +3444,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Text output path (img2text / text2text) if is_text_request and result.final_output_type == "text": - text_body = (result.custom_output or {}).get("text_output", "") + text_body = self._get_diffusion_text_output(result) message = ChatMessage(role="assistant", content=text_body) choice = ChatCompletionResponseChoice( index=0, @@ -3447,7 +3463,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): completion_tokens=len(text_body.split()), total_tokens=len(prompt.split()) + len(text_body.split()), ), - metrics=self._get_diffusion_extra_output_params(result.custom_output), + metrics=self._get_diffusion_extra_output_params(result), ) logger.info( "Diffusion chat completed for request %s: text output (%d chars)", @@ -3580,7 +3596,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): completion_tokens=1, total_tokens=len(prompt.split()) + 1, ), - metrics=self._get_diffusion_extra_output_params(result.custom_output), + metrics=self._get_diffusion_extra_output_params(result), ) logger.info( diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 3e2927a27..ed846e7e8 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -29,6 +29,11 @@ from vllm_omni.entrypoints.openai.stage_params import ( from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.outputs.output_metadata import ( + DiffusionMetadataMapping, + DiffusionMultimodalOutput, + DiffusionPayloadValue, +) logger = init_logger(__name__) @@ -58,8 +63,8 @@ class ReferenceAudio: class VideoGenerationArtifacts: """Normalized outputs and profiler metadata extracted from one request.""" - videos: list[Any] - audios: list[Any | None] + videos: list[DiffusionPayloadValue] + audios: list[DiffusionPayloadValue | None] actions: list[VideoAction | None] audio_sample_rate: int output_fps: int @@ -225,8 +230,10 @@ class OmniOpenAIServingVideo: ) result = await self._run_generation(prompt, gen_params, reference_id) - custom_output = self._extract_custom_output(result) - action_only = isinstance(custom_output, dict) and bool(custom_output.get("action_only_output")) + multimodal_output = self._extract_multimodal_output(result) + metadata = multimodal_output.get("metadata") if isinstance(multimodal_output, dict) else {} + common_metadata = metadata.get("common") if isinstance(metadata, dict) else {} + action_only = bool(isinstance(common_metadata, dict) and common_metadata.get("action_only_output")) videos = [{"action_only_output": True}] if action_only else self._extract_video_outputs(result) audios = self._extract_audio_outputs(result, expected_count=len(videos)) actions = self._extract_action_outputs(result, expected_count=len(videos)) @@ -340,10 +347,12 @@ class OmniOpenAIServingVideo: return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb, artifacts.actions[0] @staticmethod - def _resolve_video_fps_multiplier(result: Any) -> int: - custom_output = OmniOpenAIServingVideo._extract_custom_output(result) - if isinstance(custom_output, dict): - multiplier = custom_output.get("video_fps_multiplier") + def _resolve_video_fps_multiplier(result: object) -> int: + multimodal_output = OmniOpenAIServingVideo._extract_multimodal_output(result) + metadata = multimodal_output.get("metadata") + video_metadata = metadata.get("video") if isinstance(metadata, Mapping) else None + if isinstance(video_metadata, Mapping): + multiplier = video_metadata.get("video_fps_multiplier") if multiplier is not None: return int(multiplier) return 1 @@ -390,7 +399,7 @@ class OmniOpenAIServingVideo: prompt: OmniTextPrompt, gen_params: OmniDiffusionSamplingParams, request_id: str, - ) -> Any: + ) -> object: stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None) if not stage_configs: @@ -433,7 +442,7 @@ class OmniOpenAIServingVideo: return result @staticmethod - def _normalize_video_outputs(videos: Any) -> list[Any]: + def _normalize_video_outputs(videos: DiffusionPayloadValue | None) -> list[DiffusionPayloadValue]: if videos is None: return [] if hasattr(videos, "ndim") and videos.ndim == 5: @@ -443,7 +452,7 @@ class OmniOpenAIServingVideo: return [] first = videos[0] if hasattr(first, "ndim") and first.ndim == 5: - flattened: list[Any] = [] + flattened: list[DiffusionPayloadValue] = [] for item in videos: if hasattr(item, "ndim") and item.ndim == 5: flattened.extend([item[i] for i in range(item.shape[0])]) @@ -459,7 +468,7 @@ class OmniOpenAIServingVideo: return videos return [videos] - def _extract_video_outputs(self, result: Any) -> list[Any]: + def _extract_video_outputs(self, result: object) -> list[DiffusionPayloadValue]: videos = None if hasattr(result, "images") and result.images: videos = result.images @@ -469,10 +478,14 @@ class OmniOpenAIServingVideo: videos = request_output["images"] elif hasattr(request_output, "images") and request_output.images: videos = request_output.images - elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output: - videos = request_output.multimodal_output.get("video") - if videos is None and hasattr(result, "multimodal_output") and result.multimodal_output: - videos = result.multimodal_output.get("video") + 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): + videos = multimodal_output.get("video") normalized = self._normalize_video_outputs(videos) if not normalized: @@ -483,17 +496,21 @@ class OmniOpenAIServingVideo: return normalized @staticmethod - def _extract_audio_outputs(result: Any, expected_count: int) -> list[Any | None]: + def _extract_audio_outputs(result: object, expected_count: int) -> list[DiffusionPayloadValue | None]: audio = None - if hasattr(result, "multimodal_output") and result.multimodal_output: - audio = result.multimodal_output.get("audio") + 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 {} - audio = mm_output.get("audio") - elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output: - audio = request_output.multimodal_output.get("audio") + 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 @@ -514,7 +531,7 @@ class OmniOpenAIServingVideo: return [audio] + [None] * max(expected_count - 1, 0) - def _resolve_audio_sample_rate(self, result: Any) -> int: + def _resolve_audio_sample_rate(self, result: object) -> int: result_sample_rate = self._extract_audio_sample_rate_from_result(result) if result_sample_rate is not None: return result_sample_rate @@ -528,38 +545,45 @@ class OmniOpenAIServingVideo: return 24000 @classmethod - def _extract_action_outputs(cls, result: Any, expected_count: int) -> list[VideoAction | None]: - custom_output = cls._extract_custom_output(result) - if not custom_output or "action" not in custom_output: - return [None] * expected_count + def _extract_action_outputs(cls, result: object, expected_count: int) -> list[VideoAction | None]: + multimodal_output = cls._extract_multimodal_output(result) + if "actions" in multimodal_output: + action_payload = multimodal_output["actions"] + metadata = multimodal_output.get("metadata") + action_metadata = metadata.get("actions") if isinstance(metadata, Mapping) else {} + action_metadata = action_metadata if isinstance(action_metadata, Mapping) else {} + action_items = cls._split_action_payload(action_payload, expected_count) + return [ + cls._make_video_action(action_item, action_metadata) if action_item is not None else None + for action_item in action_items + ] - action_payload = custom_output.get("actions", custom_output["action"]) - action_items = cls._split_action_payload(action_payload, expected_count) - return [ - cls._make_video_action(action_item, custom_output) if action_item is not None else None - for action_item in action_items - ] + return [None] * expected_count @staticmethod - def _extract_custom_output(result: Any) -> dict[str, Any]: - custom_output = getattr(result, "custom_output", None) - if isinstance(custom_output, dict): - return custom_output + def _extract_multimodal_output(result: object) -> DiffusionMultimodalOutput: + multimodal_output = getattr(result, "multimodal_output", None) + if isinstance(multimodal_output, Mapping): + return dict(multimodal_output) request_output = getattr(result, "request_output", None) if isinstance(request_output, dict): - custom_output = request_output.get("custom_output") - if custom_output is None: - custom_output = request_output.get("_custom_output") + multimodal_output = request_output.get("multimodal_output") + if multimodal_output is None: + multimodal_output = request_output.get("_multimodal_output") elif request_output is not None: - custom_output = getattr(request_output, "custom_output", None) - if custom_output is None: - custom_output = getattr(request_output, "_custom_output", None) + multimodal_output = getattr(request_output, "multimodal_output", None) + if multimodal_output is None: + multimodal_output = getattr(request_output, "_multimodal_output", None) - return custom_output if isinstance(custom_output, dict) else {} + return dict(multimodal_output) if isinstance(multimodal_output, Mapping) else {} @classmethod - def _split_action_payload(cls, action: Any, expected_count: int) -> list[Any | None]: + def _split_action_payload( + cls, + action: DiffusionPayloadValue, + expected_count: int, + ) -> list[DiffusionPayloadValue | None]: if expected_count <= 0: return [] @@ -573,30 +597,34 @@ class OmniOpenAIServingVideo: return [action] + [None] * (expected_count - 1) @classmethod - def _make_video_action(cls, action: Any, custom_output: dict[str, Any]) -> VideoAction: + def _make_video_action( + cls, + action: DiffusionPayloadValue, + action_metadata: DiffusionMetadataMapping, + ) -> VideoAction: data = cls._to_jsonable(action) if not isinstance(data, list): data = [data] - action_mode = custom_output.get("action_mode") + action_mode = action_metadata.get("action_mode") return VideoAction( data=data, shape=cls._shape_of(action), dtype=cls._dtype_of(action), - raw_action_dim=cls._coerce_optional_int(custom_output.get("raw_action_dim")), + raw_action_dim=cls._coerce_optional_int(action_metadata.get("raw_action_dim")), action_mode=str(action_mode) if action_mode is not None else None, - domain_id=cls._coerce_optional_int(custom_output.get("domain_id")), + domain_id=cls._coerce_optional_int(action_metadata.get("domain_id")), ) @staticmethod - def _index_action(action: Any, index: int) -> Any: + def _index_action(action: DiffusionPayloadValue, index: int) -> DiffusionPayloadValue | None: try: return action[index] except (IndexError, KeyError, TypeError): return None @classmethod - def _to_jsonable(cls, value: Any) -> Any: + def _to_jsonable(cls, value: DiffusionPayloadValue) -> DiffusionPayloadValue: if hasattr(value, "detach"): value = value.detach() if hasattr(value, "cpu"): @@ -613,7 +641,7 @@ class OmniOpenAIServingVideo: return value @classmethod - def _shape_of(cls, value: Any) -> list[int]: + def _shape_of(cls, value: DiffusionPayloadValue) -> list[int]: shape = getattr(value, "shape", None) if shape is not None: try: @@ -627,12 +655,12 @@ class OmniOpenAIServingVideo: return [] @staticmethod - def _dtype_of(value: Any) -> str | None: + def _dtype_of(value: DiffusionPayloadValue) -> str | None: dtype = getattr(value, "dtype", None) return str(dtype) if dtype is not None else None @staticmethod - def _coerce_optional_int(value: Any) -> int | None: + def _coerce_optional_int(value: object) -> int | None: if value is None: return None try: @@ -642,7 +670,7 @@ class OmniOpenAIServingVideo: return None @staticmethod - def _resolve_fps(result: Any) -> int | None: + def _resolve_fps(result: object) -> int | None: """Extract fps from multimodal_output if the model reported it.""" multimodal_output = getattr(result, "multimodal_output", None) if isinstance(multimodal_output, Mapping): @@ -682,7 +710,7 @@ class OmniOpenAIServingVideo: return None @classmethod - def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None: + def _extract_audio_sample_rate_from_result(cls, result: object) -> int | None: multimodal_output = getattr(result, "multimodal_output", None) if isinstance(multimodal_output, Mapping): sample_rate = cls._coerce_audio_sample_rate( @@ -717,7 +745,7 @@ class OmniOpenAIServingVideo: return None @classmethod - def _extract_audio_sample_rate_from_config(cls, config: Any) -> int | None: + def _extract_audio_sample_rate_from_config(cls, config: object) -> int | None: if config is None: return None @@ -748,7 +776,7 @@ class OmniOpenAIServingVideo: return None @staticmethod - def _coerce_audio_sample_rate(value: Any) -> int | None: + def _coerce_audio_sample_rate(value: object) -> int | None: if value is None: return None @@ -761,12 +789,12 @@ class OmniOpenAIServingVideo: return sample_rate if sample_rate > 0 else None @staticmethod - def _extract_stage_durations(result: Any) -> dict[str, float]: + def _extract_stage_durations(result: object) -> dict[str, float]: stage_durations = getattr(result, "stage_durations", None) return stage_durations if isinstance(stage_durations, dict) else {} @staticmethod - def _extract_peak_memory_mb(result: Any) -> float: + def _extract_peak_memory_mb(result: object) -> float: peak_memory_mb = getattr(result, "peak_memory_mb", 0.0) try: return float(peak_memory_mb or 0.0) diff --git a/vllm_omni/model_extras/audiox.py b/vllm_omni/model_extras/audiox.py index 2ab7bc4d8..220972148 100644 --- a/vllm_omni/model_extras/audiox.py +++ b/vllm_omni/model_extras/audiox.py @@ -19,7 +19,7 @@ AUDIOX_EXTRA_BODY_PARAMS = frozenset( "audio_path", # optional reference-audio conditioning } ) -# Echoed back in ``DiffusionOutput.custom_output`` (pipeline_audiox.py). +# Echoed back in diffusion output metadata (pipeline_audiox.py). AUDIOX_EXTRA_OUTPUT_PARAMS = frozenset( { "audiox_task", diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs/__init__.py similarity index 100% rename from vllm_omni/outputs.py rename to vllm_omni/outputs/__init__.py diff --git a/vllm_omni/outputs/output_metadata.py b/vllm_omni/outputs/output_metadata.py new file mode 100644 index 000000000..f3e364a34 --- /dev/null +++ b/vllm_omni/outputs/output_metadata.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any, Literal, TypeAlias, TypedDict + +import numpy as np +import PIL.Image +import torch + +# Final user-facing output modality returned by pipeline stages. +FinalOutputModalityType: TypeAlias = Literal["text", "image", "audio", "video"] + +# Scalar values that may appear in either payload or metadata. +# Examples: text_output="done", fps=16, action_only_output=True. +DiffusionScalar: TypeAlias = str | int | float | bool | None + +# Shape-like metadata. Examples: [1, 16, 3, 480, 832], (16, 480, 832). +DiffusionShape: TypeAlias = list[int] | tuple[int, ...] + +# Named tensor controls used by transfer-style metadata. +# Example: {"pose": pose_tensor, "depth": depth_tensor}. +DiffusionTensorMap: TypeAlias = Mapping[str, torch.Tensor] + +# Values allowed under metadata groups such as video/audio/actions/transfer. +# Examples: {"fps": 16}, {"shape": [1, 16, 3, 480, 832]}, +# {"controls": {"pose": pose_tensor}}, {"hints": ["pose"]}. +DiffusionMetadataValue: TypeAlias = ( + DiffusionScalar + | DiffusionShape + | DiffusionTensorMap + | torch.Tensor + | list["DiffusionMetadataValue"] + | tuple["DiffusionMetadataValue", ...] + | Mapping[str, "DiffusionMetadataValue"] +) + +# Mutable metadata envelope side. Example: +# {"video": {"fps": 16}, "actions": {"raw_action_dim": 14}}. +DiffusionMetadata: TypeAlias = dict[str, DiffusionMetadataValue] + +# Read-only metadata view for validators and serving readers. +DiffusionMetadataMapping: TypeAlias = Mapping[str, DiffusionMetadataValue] + + +class DiffusionTrajectoryPayload(TypedDict, total=False): + # Optional denoising trajectory payload for RL/debug consumers. + latents: torch.Tensor | dict[str, Any] | None + timesteps: torch.Tensor | dict[str, Any] | None + log_probs: torch.Tensor | dict[str, Any] | None + decoded: list[PIL.Image.Image] | None + + +# Concrete payload values produced by post_process. +# Examples: PIL image, video tensor, audio bytes/tensor, action tensor. +DiffusionPayloadValue: TypeAlias = ( + DiffusionScalar + | bytes + | torch.Tensor + | np.ndarray + | PIL.Image.Image + | DiffusionTrajectoryPayload + | list["DiffusionPayloadValue"] + | tuple["DiffusionPayloadValue", ...] + | Mapping[str, "DiffusionPayloadValue"] +) + +# Mutable payload side. Example: {"video": video_tensor, "action": action_tensor}. +DiffusionPayload: TypeAlias = dict[str, DiffusionPayloadValue] + +# Read-only payload view for code that consumes but does not mutate payload. +DiffusionPayloadMapping: TypeAlias = Mapping[str, DiffusionPayloadValue] + +# Serving-facing multimodal output. Non-metadata keys hold payload values. +# Example: {"video": video_tensor, "metadata": {"video": {"fps": 16}}}. +DiffusionMultimodalOutput: TypeAlias = dict[str, DiffusionPayloadValue | DiffusionMetadata] + + +class DiffusionOutputEnvelope(TypedDict, total=False): + # Canonical post_process return shape: + # {"payload": {"video": video_tensor}, "metadata": {"video": {"fps": 16}}} + payload: DiffusionPayload + metadata: DiffusionMetadata + + +# Raw post_process return before formatter normalization. This keeps backward +# compatibility with direct payload returns while accepting the new envelope. +# Examples: image, {"video": video_tensor}, or +# {"payload": {"video": video_tensor}, "metadata": {"video": {"fps": 16}}}. +DiffusionPostprocessRawOutput: TypeAlias = DiffusionPayloadValue | DiffusionPayload | DiffusionOutputEnvelope + + +def strip_internal_metadata(metadata: DiffusionMetadataMapping) -> DiffusionMetadata: + """Return public metadata, dropping model-private postprocess state.""" + return {key: value for key, value in metadata.items() if key != "internal"} + + +def validate_public_diffusion_metadata(metadata: DiffusionMetadataMapping) -> None: + """Validate metadata that is safe to expose outside model postprocess.""" + if "internal" in metadata: + raise ValueError("Diffusion metadata field 'internal' must not escape public formatting.") + validate_diffusion_metadata(metadata) + + +def validate_diffusion_metadata(metadata: DiffusionMetadataMapping) -> None: + """Validate canonical diffusion output metadata fields. + + Unknown model-specific groups and keys are intentionally tolerated. This + validator only guards fields that are part of the shared output protocol or + consumed by serving code. + """ + _validate_video_metadata(_optional_mapping(metadata, "video")) + _validate_audio_metadata(_optional_mapping(metadata, "audio")) + _validate_actions_metadata(_optional_mapping(metadata, "actions")) + _validate_transfer_metadata(_optional_mapping(metadata, "transfer")) + _validate_text_metadata(_optional_mapping(metadata, "text")) + _validate_common_metadata(_optional_mapping(metadata, "common")) + + +def _validate_video_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + _check_optional_positive_number(metadata, "fps", group="video") + _check_optional_int(metadata, "video_fps_multiplier", group="video", min_value=1) + _check_optional_shape(metadata, "shape", group="video") + + +def _validate_audio_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + _check_optional_int(metadata, "sample_rate", group="audio", min_value=1) + _check_optional_non_empty_str(metadata, "audiox_task", group="audio") + + +def _validate_actions_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + _check_optional_int(metadata, "raw_action_dim", group="actions", min_value=1) + _check_optional_non_empty_str(metadata, "action_mode", group="actions") + _check_optional_int(metadata, "domain_id", group="actions") + + +def _validate_transfer_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + controls = metadata.get("controls") + if controls is not None and not _is_tensor_mapping(controls): + raise TypeError("Diffusion metadata field 'transfer.controls' must be a mapping of str to torch.Tensor.") + hints = metadata.get("hints") + if hints is not None and not _is_str_list(hints): + raise TypeError("Diffusion metadata field 'transfer.hints' must be a list of strings.") + + +def _validate_text_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + _check_optional_str(metadata, "text_output", group="text") + _check_optional_str(metadata, "think_text", group="text") + + +def _validate_common_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + value = metadata.get("action_only_output") + if value is not None and not isinstance(value, bool): + raise TypeError("Diffusion metadata field 'common.action_only_output' must be bool.") + + +def _optional_mapping(metadata: DiffusionMetadataMapping, group: str) -> DiffusionMetadataMapping | None: + value = metadata.get(group) + if value is None: + return None + if not isinstance(value, Mapping): + raise TypeError(f"Diffusion metadata group '{group}' must be a mapping.") + return value + + +def _check_optional_positive_number(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None: + value = metadata.get(key) + if value is None: + return + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be a positive number.") + if value <= 0: + raise ValueError(f"Diffusion metadata field '{group}.{key}' must be > 0.") + + +def _check_optional_int( + metadata: DiffusionMetadataMapping, + key: str, + *, + group: str, + min_value: int | None = None, +) -> None: + value = metadata.get(key) + if value is None: + return + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be int.") + if min_value is not None and value < min_value: + raise ValueError(f"Diffusion metadata field '{group}.{key}' must be >= {min_value}.") + + +def _check_optional_str(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None: + value = metadata.get(key) + if value is not None and not isinstance(value, str): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be str.") + + +def _check_optional_non_empty_str(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None: + value = metadata.get(key) + if value is None: + return + if not isinstance(value, str): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be str.") + if not value: + raise ValueError(f"Diffusion metadata field '{group}.{key}' must be non-empty.") + + +def _check_optional_shape(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None: + value = metadata.get(key) + if value is None: + return + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be a list or tuple of positive ints.") + if any(isinstance(dim, bool) or not isinstance(dim, int) or dim <= 0 for dim in value): + raise ValueError(f"Diffusion metadata field '{group}.{key}' must contain positive ints.") + + +def _is_tensor_mapping(value: object) -> bool: + return isinstance(value, Mapping) and all( + isinstance(key, str) and isinstance(tensor, torch.Tensor) for key, tensor in value.items() + ) + + +def _is_str_list(value: object) -> bool: + return isinstance(value, list) and all(isinstance(item, str) for item in value)