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

[Frontend] MiniMax-H3: Add latent-mask editing to the ComfyUI extension (#7575)

Browse Source 
Signed-off-by: chen hongwei <1792043268@qq.com>
Signed-off-by: Wang Zhipeng <wangzhipeng628@gmail.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: Wang Zhipeng <wangzhipeng628@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
chen hongweiWang ZhipengClaude Opus 5.5
 authored and GitHub committed 2026-10-06 16:32:40 +00:00
1 parent 7be49aa0ce
commit bad88bca25
15 files changed
+325 -1878
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/e2e/features/comfyui/mock_videos_server.py
-115
View File
@@ -1,115 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Minimal mock of vLLM-Omni's ``/v1/videos`` API for e2e tests.
Records every received multipart field to ``MOCK_STATE_FILE`` (JSON) so a test
can assert exactly what the client serialized, and returns a synthetic MP4.
"""
import asyncio
import json
import os
import subprocess
import uuid
from fastapi import FastAPI, Request
from fastapi.responses import Response
from starlette.datastructures import UploadFile
app = FastAPI()
_JOBS: dict[str, dict] = {}
_STATE_FILE = os.environ.get("MOCK_STATE_FILE")
def _record(fields: dict) -> None:
if _STATE_FILE:
with open(_STATE_FILE, "w") as f:
json.dump(fields, f)
def _synthetic_mp4() -> bytes:
return subprocess.run(
[
"ffmpeg",
"-hide_banner",
"-loglevel",
"error",
"-y",
"-f",
"lavfi",
"-i",
"testsrc2=size=320x240:rate=24:duration=1",
"-f",
"lavfi",
"-i",
"sine=frequency=440:sample_rate=32000:duration=1",
"-map",
"0:v",
"-map",
"1:a",
"-shortest",
"-c:v",
"libx264",
"-pix_fmt",
"yuv420p",
"-c:a",
"aac",
"-ar",
"32000",
"-ac",
"2",
"-movflags",
"frag_keyframe+empty_moov",
"-f",
"mp4",
"pipe:1",
],
check=True,
capture_output=True,
).stdout
@app.post("/v1/videos")
async def create_video(request: Request):
form = await request.form()
fields = {}
for key, value in form.items():
if isinstance(value, UploadFile):
data = await value.read()
fields[key] = {
"file": True,
"size": len(data),
"filename": value.filename,
"content_type": value.content_type,
}
if value.content_type == "application/json":
fields[key]["json"] = json.loads(data)
else:
fields[key] = str(value)
_record(fields)
job_id = str(uuid.uuid4())
_JOBS[job_id] = {"status": "queued"}
return {"id": job_id, "status": "queued"}
@app.get("/v1/videos/{job_id}")
async def get_video(job_id: str):
if job_id not in _JOBS:
return {"status": "failed", "error": "unknown job"}
return {"id": job_id, "status": "completed"}
@app.get("/v1/videos/{job_id}/content")
async def get_video_content(job_id: str):
if job_id not in _JOBS:
return Response(status_code=404)
mp4 = await asyncio.to_thread(_synthetic_mp4)
return Response(content=mp4, media_type="video/mp4")
@app.delete("/v1/videos/{job_id}")
async def delete_video(job_id: str):
_JOBS.pop(job_id, None)
return {"deleted": True}
tests/e2e/features/comfyui/test_comfyui_integration.py
+50
View File
@@ -30,6 +30,7 @@ from comfyui_vllm_omni.nodes import (
VLLMOmniGenerateImage,
VLLMOmniGenerateMusic,
VLLMOmniGenerateVideo,
VLLMOmniLatentMaskEditing,
VLLMOmniTTS,
VLLMOmniUnderstanding,
VLLMOmniVideoReferences,
@@ -1344,3 +1345,52 @@ async def test_music_generation_node_minimax_music3(api_server: str, sampling_ca
assert len(result) == 1
assert result[0]["sample_rate"] == 24000
assert result[0]["waveform"].shape == (1, 1, 24000)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"server_case",
[
pytest.param(
ServerCase(
served_model="MiniMaxAI/MiniMax-H3",
stage_list=["diffusion"],
stage_configs=[H3_STAGE_CONFIG],
outputs=[_build_diffusion_video_output()],
),
id="minimax-h3",
),
],
indirect=True,
)
@pytest.mark.parametrize(
"sampling_case",
[pytest.param(SamplingCase(kind=SamplingKind.VIDEO_NONE, sampling_params=None), id="latent-mask")],
indirect=True,
)
async def test_video_generation_node_minimax_h3_latent_mask_editing(api_server: str, sampling_case: SamplingCase):
"""Latent-mask editing: the raw video/audio masks must be accepted by the real /v1/videos form parser."""
mask_node = VLLMOmniLatentMaskEditing()
(latent_edit,) = mask_node.get_latent_edit(
source_video=VideoInput(b"mock-source-video"),
video_mask=torch.zeros((VIDEO_HEIGHT, VIDEO_WIDTH), dtype=torch.float32),
audio_mask=0.5,
)
node = VLLMOmniGenerateVideo()
result = await node.generate(
url=api_server,
model="MiniMaxAI/MiniMax-H3",
prompt="Restyle the clip.",
negative_prompt="",
width=VIDEO_WIDTH,
height=VIDEO_HEIGHT,
fps=VIDEO_FPS,
duration=VIDEO_DURATION,
model_params=H3_MODEL_PARAMS,
latent_edit=latent_edit,
)
assert isinstance(result, tuple)
assert len(result) == 1
assert isinstance(result[0], VideoInput)
tests/e2e/features/comfyui/test_latent_mask.py
+50 -49
View File
@@ -1,13 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for latent-mask serialization helpers."""
"""Unit tests for latent-mask serialization helpers and the temporal-mask node."""
import json
import pytest
import torch
from comfyui_vllm_omni.utils.latent_mask import scalar_mask_to_json, video_mask_to_grid_json
from comfyui_vllm_omni.nodes import VLLMOmniMiniMaxH3TemporalMask
from comfyui_vllm_omni.utils.latent_mask import scalar_mask_to_json, video_mask_to_json
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -21,60 +22,60 @@ def test_scalar_mask_to_json():
scalar_mask_to_json(-0.1)
def test_video_mask_grid_shape():
mask = torch.zeros(256, 448)
grid = json.loads(video_mask_to_grid_json(mask, width=448, height=256, num_frames=107))
assert len(grid) == 2 + 5 * ((107 - 5) // 17)
assert len(grid[0]) == 256 // 16
assert len(grid[0][0]) == 448 // 16
def test_video_mask_to_json_downsamples_spatial_mask():
# A 2D mask keeps its rank; only the spatial axes shrink by the VAE stride.
grid = json.loads(video_mask_to_json(torch.full((120, 160), 0.5)))
assert len(grid) == 8 # ceil(120 / 16)
assert len(grid[0]) == 10
assert grid[0][0] == 0.5
def test_video_mask_grid_shape_non_aligned_frames():
# 96 snaps up to 107 (17n+5), so Tv == 32, matching the server.
grid = json.loads(video_mask_to_grid_json(torch.zeros(256, 448), width=448, height=256, num_frames=96))
assert len(grid) == 32
def test_video_mask_to_json_keeps_frame_axis():
# A 3D mask keeps one slice per frame; the server maps frames to latents.
grid = json.loads(video_mask_to_json(torch.full((22, 120, 160), 0.5)))
assert len(grid) == 22
assert len(grid[0]) == 8
assert len(grid[0][0]) == 10
def test_video_mask_grid_shape_short_clip():
# num_frames <= 5 -> Tv == 2.
grid = json.loads(video_mask_to_grid_json(torch.zeros(64, 64), width=64, height=64, num_frames=1))
assert len(grid) == 2
def test_video_mask_to_json_area_averages_and_rounds():
mask = torch.zeros(16, 32)
mask[:, 16:] = 1.0
mask[0, :16] = 1.0 / 3.0
assert json.loads(video_mask_to_json(mask)) == [[0.0208, 1.0]]
def test_video_mask_grid_shape_unaligned_width():
# width=500 floors to 480 (multiple of 32) -> gw == 30, not 31.
grid = json.loads(video_mask_to_grid_json(torch.zeros(64, 64), width=500, height=256, num_frames=22))
assert len(grid[0][0]) == 30
def test_video_mask_to_json_fits_upload_limit():
# A soft full-resolution 1344x768 mask serialized raw would exceed 8 MiB.
payload = video_mask_to_json(torch.rand(768, 1344))
assert len(payload.encode()) < 8 * 1024 * 1024 // 100
def test_video_mask_3d_input():
grid = json.loads(video_mask_to_grid_json(torch.ones(1, 256, 448), width=448, height=256, num_frames=107))
assert grid[0][0][0] == 1.0
def test_video_mask_to_json_rejects_wrong_rank():
with pytest.raises(ValueError):
video_mask_to_json(torch.zeros(4))
with pytest.raises(ValueError):
video_mask_to_json(torch.zeros(1, 1, 1, 1))
def test_video_mask_temporal():
# 3D mask: first half 0 (preserve), second half 1 (regenerate).
mask = torch.zeros(7, 64, 64)
mask[4:] = 1.0
grid = json.loads(video_mask_to_grid_json(mask, width=64, height=64, num_frames=22))
assert len(grid) == 7
assert all(v == 0.0 for row in grid[0] for v in row)
assert all(v == 1.0 for row in grid[-1] for v in row)
def test_video_mask_temporal_resize():
# 2 temporal slices resize to tv == 7.
mask = torch.zeros(2, 64, 64)
mask[1] = 1.0
grid = json.loads(video_mask_to_grid_json(mask, width=64, height=64, num_frames=22))
assert len(grid) == 7
def test_video_mask_uniform_preserved():
grid = json.loads(video_mask_to_grid_json(torch.full((64, 64), 0.5), width=64, height=64, num_frames=107))
assert all(v == 0.5 for row in grid for cell in row for v in cell)
def test_video_mask_values_in_range():
grid = json.loads(video_mask_to_grid_json(torch.rand(128, 128), width=128, height=128, num_frames=22))
assert all(0.0 <= v <= 1.0 for row in grid for cell in row for v in cell)
@pytest.mark.parametrize(
("mode", "duration", "preserve_fraction", "frames", "prefix"),
[
# 5 s source extended to 10 s: 240 frames align to 243; the 120 source
# frames snap down to a 107-frame prefix (32 latents on the server).
("extension", 10.0, 0.5, 243, 107),
# 5 s output keeps 20% of the source: 24 frames snap down to 22.
("continuation", 5.0, 0.2, 124, 22),
],
)
def test_temporal_mask_is_frame_space(mode, duration, preserve_fraction, frames, prefix):
images = torch.zeros(120, 8, 8, 3)
mask, preview_fps, preview_images, preview_mask = VLLMOmniMiniMaxH3TemporalMask().build(
images, 24.0, duration, mode, preserve_fraction
)
assert mask.shape == (frames, 1, 1)
assert bool((mask[:prefix] == 0).all())
assert bool((mask[prefix:] == 1).all())
assert preview_fps == 24.0
assert preview_images.shape[0] == frames
assert torch.equal(preview_mask, mask)
tests/e2e/features/comfyui/test_latent_mask_editing.py
+57 -74
View File
@@ -3,97 +3,80 @@
"""End-to-end test for latent-mask editing serialization.
Exercises ``VLLMOmniClient.generate_video`` with a ``latent_edit`` payload
against a mock ``/v1/videos`` server and asserts the multipart fields the
client produced. Runs on CPU: the ComfyUI ``comfy_api`` / ``comfy_extras``
modules are mocked by this directory's ``conftest.py``.
Exercises ``VLLMOmniClient.generate_video`` with a ``latent_edit`` payload and
asserts the exact multipart fields the client produces — in particular that the
masks are uploaded as JSON *file* parts (the server's ``_parse_video_form``
declares ``video_noise_mask``/``audio_noise_mask`` as ``UploadFile``, so a plain
string field would be rejected with a 422). Runs on CPU: ``comfy_api`` /
``comfy_extras`` are mocked by this directory's ``conftest.py``.
"""
import asyncio
import json
import os
import socket
import subprocess
import sys
import time
from unittest.mock import AsyncMock
import pytest
import torch
from comfy_api.input import VideoInput
from comfyui_vllm_omni.utils.api_client import VLLMOmniClient
from comfyui_vllm_omni.utils import api_client
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_TESTS_DIR = os.path.dirname(os.path.abspath(__file__))
def _free_port() -> int:
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
@pytest.fixture
def mock_server(tmp_path):
port = _free_port()
state_file = tmp_path / "state.json"
env = dict(os.environ, MOCK_STATE_FILE=str(state_file))
proc = subprocess.Popen(
[sys.executable, "-m", "uvicorn", "mock_videos_server:app", "--host", "127.0.0.1", "--port", str(port)],
cwd=_TESTS_DIR,
env=env,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
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
def _fields_by_name(calls):
fields = calls.call_args_list[0].kwargs["data"]._fields
return {options["name"]: (options, headers, value) for options, headers, value in fields}
async def _generate(**kwargs):
return await api_client.VLLMOmniClient("http://localhost/v1").generate_video(
model="MiniMaxAI/MiniMax-H3",
prompt="restyle the clip",
width=160,
height=120,
num_frames=22,
fps=24,
**kwargs,
)
try:
deadline = time.time() + 30
while time.time() < deadline:
try:
with socket.create_connection(("127.0.0.1", port), timeout=1):
break
except OSError:
time.sleep(0.2)
yield f"http://127.0.0.1:{port}/v1", state_file
finally:
proc.terminate()
proc.wait(timeout=10)
def test_latent_edit_serialization(mock_server):
base_url, state_file = mock_server
async def test_latent_edit_serialization(api_calls):
latent_edit = {
"source_video": VideoInput(b"mock_source_video"),
"video_mask": torch.zeros(1, 120, 160),
"audio_mask": 0.5,
}
# A non-trivial video mask requires a source video. The mocked VideoInput
# only needs to provide ``save_to`` for the client's multipart upload.
source_video = VideoInput(b"mock_source_video")
mask = torch.zeros(1, 120, 160)
latent_edit = {"source_video": source_video, "video_mask": mask, "audio_mask": 0.5}
out = await _generate(latent_edit=latent_edit)
assert out == b"generated-video"
async def run():
client = VLLMOmniClient(base_url)
return await client.generate_video(
model="MiniMaxAI/MiniMax-H3",
prompt="restyle the clip",
width=160,
height=120,
num_frames=22,
fps=24,
latent_edit=latent_edit,
)
fields = _fields_by_name(api_calls)
out = asyncio.run(run())
assert out is not None
source = fields["source_video"]
assert source[0]["filename"] == "source.mp4"
assert source[1]["Content-Type"] == "video/mp4"
with open(state_file) as f:
fields = json.load(f)
# video_noise_mask must be a JSON file part, not a string field. The mask
# keeps its frame axis and is downsampled by the VAE stride (the server
# resolves the latent grid).
video = fields["video_noise_mask"]
assert video[0]["filename"] == "video-mask.json"
assert video[1]["Content-Type"] == "application/json"
grid = json.loads(video[2].decode())
assert len(grid) == 1 # one temporal slice
assert len(grid[0]) == 8 # ceil(120 / 16)
assert len(grid[0][0]) == 10 # 160 / 16
assert fields["source_video"]["file"] is True
assert fields["source_video"]["content_type"] == "video/mp4"
for name in ("video_noise_mask", "audio_noise_mask"):
assert fields[name]["file"] is True
assert fields[name]["content_type"] == "application/json"
assert fields["audio_noise_mask"]["json"] == 0.5
video_mask = fields["video_noise_mask"]["json"]
assert len(video_mask) == 7 # num_frames=22 aligns to 22 (17n+5) -> Tv = 7
assert len(video_mask[0]) == 6 # height 120 floors to 96 (multiple of 32) -> gh 6
assert len(video_mask[0][0]) == 10 # width 160 -> gw 10
# A scalar audio_mask is sent as a JSON file part carrying the scalar value.
audio = fields["audio_noise_mask"]
assert audio[0]["filename"] == "audio-mask.json"
assert audio[1]["Content-Type"] == "application/json"
assert audio[2].decode() == "0.5"
tests/e2e/features/comfyui/test_latent_mask_workflows.py
+97
View File
@@ -0,0 +1,97 @@
# 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 Latent Mask Editing.json"
)
@pytest.fixture
def workflow():
return json.loads(WORKFLOW.read_text())
def _source_of(workflow: dict, node: dict, input_name: str) -> dict:
nodes = {n["id"]: n for n in workflow["nodes"]}
links = {link[0]: link for link in workflow["links"]}
link_id = next(i["link"] for i in node["inputs"] if i["name"] == input_name)
assert link_id is not None, f"{node['type']}.{input_name} is not connected"
return nodes[links[link_id][1]]
def _cases(workflow: dict) -> dict[str, dict]:
return {
node["title"].removeprefix("Generate Video(").removesuffix(")"): node
for node in workflow["nodes"]
if node["type"] == "VLLMOmniGenerateVideo"
}
def test_latent_mask_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)
def test_latent_mask_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"]:
kind = inputs[input_["name"]][0]
assert input_["type"] == ("COMBO" if isinstance(kind, list) else kind)
assert tuple(output["type"] for output in node["outputs"]) == cls.RETURN_TYPES
def test_latent_mask_workflow_defaults(workflow):
cases = _cases(workflow)
assert set(cases) == {"Object Removal", "Inpainting", "Continuation", "Extension"}
for generate in cases.values():
assert generate["widgets_values"][0] == "http://127.0.0.1:8000/v1"
assert generate["widgets_values"][4:6] == [1344, 768]
@pytest.mark.parametrize("case", ["Object Removal", "Inpainting"])
def test_spatial_cases_use_a_composited_mask(workflow, case):
edit = _source_of(workflow, _cases(workflow)[case], "latent_edit")
assert _source_of(workflow, edit, "source_video")["type"] == "LoadVideo"
assert _source_of(workflow, edit, "video_mask")["type"] == "MaskComposite"
@pytest.mark.parametrize("case", ["Continuation", "Extension"])
def test_temporal_cases_build_a_per_frame_mask(workflow, case):
# The server pads a short [T, H, W] mask with its last slice, so the mask must
# come from the Temporal Mask node (one slice per output frame), not a
# hand-batched stack of a few slices.
generate = _cases(workflow)[case]
edit = _source_of(workflow, generate, "latent_edit")
temporal = _source_of(workflow, edit, "video_mask")
assert temporal["type"] == "VLLMOmniMiniMaxH3TemporalMask"
assert _source_of(workflow, temporal, "images")["type"] == "GetVideoComponents"
assert _source_of(workflow, edit, "source_video")["type"] == "LoadVideo"
_, duration, mode, _ = temporal["widgets_values"]
assert mode == case.lower()
assert duration == generate["widgets_values"][7]