bad88bca25
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>
83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""End-to-end test for latent-mask editing serialization.
|
|
|
|
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 json
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
import torch
|
|
from comfy_api.input import VideoInput
|
|
from comfyui_vllm_omni.utils import api_client
|
|
|
|
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
|
|
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
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,
|
|
}
|
|
|
|
out = await _generate(latent_edit=latent_edit)
|
|
assert out == b"generated-video"
|
|
|
|
fields = _fields_by_name(api_calls)
|
|
|
|
source = fields["source_video"]
|
|
assert source[0]["filename"] == "source.mp4"
|
|
assert source[1]["Content-Type"] == "video/mp4"
|
|
|
|
# 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
|
|
|
|
# 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"
|