95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
|
# 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,
|
||
|
|
},
|
||
|
|
"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),
|
||
|
|
({"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
|