anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: bjf-frz <frz123db@gmail.com>
This commit is contained in:
bjf-frz authored and GitHub committed 2026-07-14 17:24:03 +08:00
1 parent 4b9d09cb87
commit 1a4a042a83
44 files changed
+1373 -478
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+57 -25
View File
@@ -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
+50
View File
@@ -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
+144 -32
View File
@@ -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
+97
View File
@@ -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
+13 -6
View File
@@ -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
+38 -13
View File
@@ -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)