[REFACT]Refactor diffusion outputs to payload metadata (#4922)

Signed-off-by: bjf-frz <frz123db@gmail.com>

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.
 

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

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))

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}")

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"),

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}},
+    }

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(

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

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."""

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)

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:

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),
-            },
         )

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,

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
 

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":

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)
 

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,
         )

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,
         )
 

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",

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

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,
+                        },
+                    },
                 },
             )
 

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,
         )

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),
         )

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,
         )
 

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

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

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,
         )
 

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 {},
         )

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,
         )
 

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),
         )

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"
 
 

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

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

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

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

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

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",

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):

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__)

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(

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)

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",

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)