[Frontend] Add MiniMax-H3 references and Ref2VA workflow (#7483)

Signed-off-by: Linze-Shi <linzeshi0@gmail.com>

tests/e2e/features/comfyui/test_comfyui_integration.py

@@ -16,9 +16,11 @@ import time
 import traceback
 from collections.abc import Iterable, Sequence
 from enum import Enum
+from pathlib import Path
 from types import SimpleNamespace
 from typing import Any, NamedTuple
 
+import av
 import pytest
 import requests
 import torch
@@ -82,6 +84,7 @@ class SamplingKind(str, Enum):
     VIDEO_FASTH3 = "video_fasth3"
     VIDEO_REF2VA_IMAGE_AUDIO = "video_ref2va_image_audio"
     VIDEO_REF2VA_MULTI_VIDEO = "video_ref2va_multi_video"
+    VIDEO_REF2VA_MIXED = "video_ref2va_mixed"
 
 
 # Pre-defined arguments to be used in function calls during the tests
@@ -153,6 +156,13 @@ H3_MODEL_PARAMS = MiniMaxH3ModelSpecificParams(
     }
 )
 
+H3_STAGE_CONFIG = {
+    "stage_type": "diffusion",
+    "final_output": True,
+    "final_output_type": "video",
+    "engine_args": {"model_class_name": "MiniMaxH3Pipeline"},
+}
+
 LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10}
 
 
@@ -442,17 +452,30 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
                     "audio_flow_shift": 3.0,
                 },
             )
-        elif sampling_case.kind is SamplingKind.VIDEO_REF2VA_MULTI_VIDEO:
+        elif sampling_case.kind in (SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, SamplingKind.VIDEO_REF2VA_MIXED):
             assert len(received_sampling_params_list) == 1
             assert isinstance(prompt, dict)
             multi_modal_data = prompt.get("multi_modal_data")
             assert isinstance(multi_modal_data, dict)
             input_videos = multi_modal_data.get("video")
             assert isinstance(input_videos, list)
-            assert len(input_videos) == 2
             assert all(isinstance(path, str) for path in input_videos)
-            assert "image" not in multi_modal_data
-            assert "audio" not in multi_modal_data
+            if sampling_case.kind is SamplingKind.VIDEO_REF2VA_MULTI_VIDEO:
+                assert len(input_videos) == 2
+                assert "image" not in multi_modal_data
+                assert "audio" not in multi_modal_data
+            else:
+                assert [Path(path).read_bytes() for path in input_videos] == [b"video-1", b"video-2", b"video-3"]
+                images = multi_modal_data.get("image")
+                assert isinstance(images, list)
+                assert len(images) == 6
+                assert [image.getpixel((0, 0)) for image in images] == [(i * 40,) * 3 for i in range(1, 7)]
+                audios = multi_modal_data.get("audio")
+                assert isinstance(audios, list)
+                assert len(audios) == 3
+                for path, sample_rate in zip(audios, (24000, 32000, 48000)):
+                    with av.open(path) as container:
+                        assert container.streams.audio[0].sample_rate == sample_rate
             _assert_sampling_param_values(
                 received_sampling_params_list[0],
                 {
@@ -970,7 +993,7 @@ async def test_fast_h3_deployment_node(api_server: str, sampling_case: SamplingC
             ServerCase(
                 served_model="MiniMaxAI/MiniMax-H3",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
+                stage_configs=[H3_STAGE_CONFIG],
                 outputs=[_build_diffusion_video_output()],
             ),
             SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
@@ -981,13 +1004,24 @@ async def test_fast_h3_deployment_node(api_server: str, sampling_case: SamplingC
             ServerCase(
                 served_model="MiniMaxAI/MiniMax-H3",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
+                stage_configs=[H3_STAGE_CONFIG],
                 outputs=[_build_diffusion_video_output()],
             ),
             SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),
             "multi_video",
             id="ref2va-multi-video",
         ),
+        pytest.param(
+            ServerCase(
+                served_model="MiniMaxAI/MiniMax-H3",
+                stage_list=["diffusion"],
+                stage_configs=[H3_STAGE_CONFIG],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MIXED, sampling_params=None),
+            "mixed",
+            id="ref2va-mixed-12-references",
+        ),
     ],
     indirect=["server_case", "sampling_case"],
 )
@@ -1017,6 +1051,17 @@ async def test_video_generation_node_minimax_h3_ref2va(
         prompt = (
             "Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2."
         )
+    elif ref_mode == "mixed":
+        inputs = {f"image_{i}": torch.full((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), i * 40 / 255) for i in range(1, 7)}
+        inputs.update({f"video_{i}": VideoInput(f"video-{i}".encode()) for i in range(1, 4)})
+        inputs.update(
+            {
+                f"audio_{i}": {"waveform": torch.zeros(1, 1, sample_rate), "sample_rate": sample_rate}
+                for i, sample_rate in enumerate((24000, 32000, 48000), start=1)
+            }
+        )
+        (references,) = refs_node.get_references(**inputs)
+        prompt = "Combine the subjects from the images with the reference videos and audio clips."
     else:
         raise AssertionError(f"Unknown ref_mode: {ref_mode}")
 

tests/e2e/features/comfyui/test_h3_reference_workflow.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import json
+from graphlib import TopologicalSorter
+from pathlib import Path
+
+import pytest
+from comfyui_vllm_omni import nodes as omni_nodes
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+WORKFLOW = (
+    Path(__file__).resolve().parents[4]
+    / "apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax-H3 Reference to Video.json"
+)
+
+
+@pytest.fixture
+def workflow():
+    return json.loads(WORKFLOW.read_text())
+
+
+def test_reference_workflow_connections(workflow):
+    nodes = {node["id"]: node for node in workflow["nodes"]}
+    links = {link[0]: link for link in workflow["links"]}
+    assert len(nodes) == len(workflow["nodes"])
+    assert len(links) == len(workflow["links"])
+    graph: dict[int, set[int]] = {node_id: set() for node_id in nodes}
+    for link_id, source, output_slot, target, input_slot, kind in links.values():
+        output = nodes[source]["outputs"][output_slot]
+        input_ = nodes[target]["inputs"][input_slot]
+        assert output["type"] == input_["type"] == kind
+        assert link_id in output["links"]
+        assert input_["link"] == link_id
+        graph[target].add(source)
+    assert len(tuple(TopologicalSorter(graph).static_order())) == len(nodes)
+    for node in nodes.values():
+        for slot, input_ in enumerate(node.get("inputs", [])):
+            if input_["link"] is not None:
+                assert links[input_["link"]][3:5] == [node["id"], slot]
+        for slot, output in enumerate(node.get("outputs", [])):
+            for link_id in output.get("links") or []:
+                assert links[link_id][1:3] == [node["id"], slot]
+
+
+def test_reference_workflow_matches_omni_node_interfaces(workflow):
+    for node in workflow["nodes"]:
+        if not node["type"].startswith("VLLMOmni"):
+            continue
+        cls = getattr(omni_nodes, node["type"])
+        schema = cls.INPUT_TYPES()
+        inputs = {**schema.get("required", {}), **schema.get("optional", {})}
+        for input_ in node["inputs"]:
+            assert input_["type"] == inputs[input_["name"]][0]
+        assert tuple(output["type"] for output in node["outputs"]) == cls.RETURN_TYPES
+        expected_widgets = len(schema.get("required", {}))
+        if "seed" in schema.get("required", {}):
+            expected_widgets += 1
+        assert len(node["widgets_values"]) == expected_widgets
+    refs = next(node for node in workflow["nodes"] if node["type"] == "VLLMOmniVideoReferences")
+    assert [input_["name"] for input_ in refs["inputs"]] == list(
+        omni_nodes.VLLMOmniVideoReferences.INPUT_TYPES()["optional"]
+    )
+
+
+def test_reference_workflow_default_route_and_output(workflow):
+    nodes = {node["id"]: node for node in workflow["nodes"]}
+    links = {link[0]: link for link in workflow["links"]}
+    generate = next(node for node in nodes.values() if node["type"] == "VLLMOmniGenerateVideo")
+    inputs = {input_["name"]: input_["link"] for input_ in generate["inputs"]}
+    assert inputs["frame"] is None
+    assert inputs["lora"] is None
+    for name, kind in (
+        ("references", "VLLMOmniVideoReferences"),
+        ("sampling_params", "VLLMOmniDiffusionSampling"),
+        ("model_params", "VLLMOmniMiniMaxH3Params"),
+    ):
+        assert nodes[links[inputs[name]][1]]["type"] == kind
+    save = next(node for node in nodes.values() if node["type"] == "SaveVideo")
+    assert links[save["inputs"][0]["link"]][1] == generate["id"]
+    assert save["widgets_values"] == ["video/MiniMax-H3-Ref2VA", "mp4", "h264"]
+
+
+def test_reference_workflow_defaults_and_portable_assets(workflow):
+    by_type = {node["type"]: node for node in workflow["nodes"]}
+    values = by_type["VLLMOmniGenerateVideo"]["widgets_values"]
+    assert values[0] == "http://localhost:8000/v1"
+    assert values[1] == "MiniMaxAI/MiniMax-H3"
+    assert "<Picture 1>" in values[2]
+    assert values[4:8] == [1344, 768, 24, 5.167]
+    num_frames = round(values[7] * values[6])
+    assert num_frames == 124
+    assert values[3] == ""
+    assert (num_frames - 5) % 17 == 0
+    assert 4 <= values[7] <= 15
+    assert by_type["VLLMOmniDiffusionSampling"]["widgets_values"] == [1, 50, 1.0, 1.0, False, False, 42, "fixed"]
+    assert by_type["VLLMOmniMiniMaxH3Params"]["widgets_values"] == [3.0, 12.0]
+    assert by_type["VLLMOmniRemoteLoRA"]["widgets_values"][0] == ""
+    for kind in ("LoadImage", "LoadVideo", "LoadAudio"):
+        filename = by_type[kind]["widgets_values"][0]
+        assert filename and Path(filename).name == filename
+        assert ":" not in filename and "\\" not in filename
+    assert set(by_type) == {
+        "LoadImage",
+        "LoadVideo",
+        "LoadAudio",
+        "SaveVideo",
+        "MarkdownNote",
+        "VLLMOmniVideoReferences",
+        "VLLMOmniGenerateVideo",
+        "VLLMOmniDiffusionSampling",
+        "VLLMOmniMiniMaxH3Params",
+        "VLLMOmniRemoteLoRA",
+    }

tests/e2e/features/comfyui/test_video_references.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from io import BytesIO
+from unittest.mock import AsyncMock
+
+import av
+import pytest
+import torch
+from comfy_api.input import VideoInput
+from comfyui_vllm_omni.nodes import VLLMOmniVideoReferences
+from comfyui_vllm_omni.utils import api_client
+from PIL import Image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def api_calls(monkeypatch):
+    calls = AsyncMock(return_value={"id": "test-video", "status": "completed"})
+    monkeypatch.setattr(api_client, "url_json", calls)
+    monkeypatch.setattr(api_client, "url_bytes", AsyncMock(return_value=b"generated-video"))
+    monkeypatch.setattr(api_client, "bytes_to_video", lambda data: data)
+    return calls
+
+
+async def _generate(references, **kwargs):
+    return await api_client.VLLMOmniClient("http://localhost/v1").generate_video(
+        model="MiniMaxAI/MiniMax-H3",
+        prompt="Use the reference subjects and sounds.",
+        width=896,
+        height=512,
+        num_frames=90,
+        fps=24,
+        references=references,
+        **kwargs,
+    )
+
+
+def _references(image_count, video_count, audio_count):
+    refs = {f"image_{i}": torch.full((1, 16, 16, 3), i / 10) for i in range(1, image_count + 1)}
+    refs.update({f"video_{i}": VideoInput(f"video-{i}".encode()) for i in range(1, video_count + 1)})
+    refs.update(
+        {f"audio_{i}": {"waveform": torch.zeros(1, 1, 24000), "sample_rate": 24000} for i in range(1, audio_count + 1)}
+    )
+    return refs
+
+
+def test_reference_node_preserves_existing_ports_and_exposes_h3_limits():
+    ports = VLLMOmniVideoReferences.INPUT_TYPES()["optional"]
+    assert list(ports)[:6] == ["image_1", "image_2", "audio_1", "audio_2", "video_1", "video_2"]
+    assert {name for name, kind in ports.items() if kind == ("IMAGE",)} == {f"image_{i}" for i in range(1, 10)}
+    assert {name for name, kind in ports.items() if kind == ("VIDEO",)} == {"video_1", "video_2", "video_3"}
+    assert {name for name, kind in ports.items() if kind == ("AUDIO",)} == {"audio_1", "audio_2", "audio_3"}
+    inputs = _references(6, 3, 3)
+    (refs,) = VLLMOmniVideoReferences().get_references(**inputs)
+    assert refs.keys() == inputs.keys()
+    assert all(refs[name] is value for name, value in inputs.items())
+
+
+@pytest.mark.parametrize(
+    "counts", [(1, 0, 0), (9, 0, 0), (0, 3, 0), (1, 0, 3), (0, 1, 3), (6, 3, 3), (9, 3, 0), (9, 0, 3)]
+)
+async def test_reference_uploads_keep_numeric_order_and_media_types(api_calls, counts):
+    refs = _references(*counts)
+    result = await _generate(dict(reversed(list(refs.items()))))
+    assert result == b"generated-video"
+    fields = api_calls.call_args_list[0].kwargs["data"]._fields
+    uploads = [
+        (options["filename"], headers["Content-Type"], value)
+        for options, headers, value in fields
+        if options["name"] == "input_references"
+    ]
+    expected = [
+        (f"{kind}_{i}.{extension}", content_type)
+        for kind, count, extension, content_type in (
+            ("image", counts[0], "png", "image/png"),
+            ("video", counts[1], "mp4", "video/mp4"),
+            ("audio", counts[2], "mp3", "audio/mpeg"),
+        )
+        for i in range(1, count + 1)
+    ]
+    assert [(name, content_type) for name, content_type, _ in uploads] == expected
+    for name, content_type, payload in uploads:
+        data = payload.getvalue()
+        if content_type == "image/png":
+            index = int(name.removeprefix("image_").removesuffix(".png"))
+            assert Image.open(BytesIO(data)).getpixel((0, 0)) == (int(index / 10 * 255),) * 3
+        elif content_type == "video/mp4":
+            assert data == name.removesuffix(".mp4").replace("_", "-").encode()
+        else:
+            with av.open(BytesIO(data)) as container:
+                assert container.streams.audio[0].sample_rate == 24000
+                assert sum(frame.samples for frame in container.decode(audio=0)) > 0
+
+
+async def test_sparse_reference_slots_are_not_dropped(api_calls):
+    inputs = _references(9, 3, 3)
+    inputs = {name: inputs[name] for name in ("audio_3", "image_9", "video_3", "image_2")}
+    (refs,) = VLLMOmniVideoReferences().get_references(image_1=None, **inputs)
+    await _generate(refs)
+    fields = api_calls.call_args_list[0].kwargs["data"]._fields
+    assert [options["filename"] for options, _, _ in fields if options["name"] == "input_references"] == [
+        "image_2.png",
+        "image_9.png",
+        "video_3.mp4",
+        "audio_3.mp3",
+    ]
+
+
+@pytest.mark.parametrize(
+    "counts, message",
+    [
+        ((0, 0, 0), "at least one image or video"),
+        ((0, 0, 3), "at least one image or video"),
+        ((9, 3, 1), "at most 12"),
+        ((9, 3, 3), "at most 12"),
+    ],
+)
+async def test_invalid_reference_combinations_fail_before_upload(api_calls, counts, message):
+    with pytest.raises(ValueError, match=message):
+        await _generate(_references(*counts))
+    api_calls.assert_not_called()
+
+
+@pytest.mark.parametrize("name", ["image_10", "video_4", "audio_4", "image_0"])
+async def test_unsupported_reference_slots_are_not_silently_ignored(api_calls, name):
+    refs = _references(1, 0, 0)
+    refs[name] = object()
+    with pytest.raises(ValueError, match="Unsupported reference input"):
+        await _generate(refs)
+    api_calls.assert_not_called()
+
+
+async def test_frame_and_references_are_mutually_exclusive(api_calls):
+    with pytest.raises(ValueError, match="only one of frame or references"):
+        await _generate(_references(1, 0, 0), frame=torch.zeros(1, 16, 16, 3))
+    api_calls.assert_not_called()