fa1d03ab8d
Signed-off-by: chen hongwei <1792043268@qq.com> Signed-off-by: Weiming Liao <liaowm5@gmail.com> Co-authored-by: chen hongwei <1792043268@qq.com> Co-authored-by: chen hongwei <54873389+avicii-forever@users.noreply.github.com>
100 lines
3.1 KiB
Python
100 lines
3.1 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
|
|
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``.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
import pytest
|
|
import torch
|
|
from comfy_api.input import VideoInput
|
|
from comfyui_vllm_omni.utils.api_client import VLLMOmniClient
|
|
|
|
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,
|
|
)
|
|
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
|
|
|
|
# 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}
|
|
|
|
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,
|
|
)
|
|
|
|
out = asyncio.run(run())
|
|
assert out is not None
|
|
|
|
with open(state_file) as f:
|
|
fields = json.load(f)
|
|
|
|
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
|