[feat] Streaming diffusion video generation output (#3737)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
42 files changed
+4322
-65
No files matched your search
@@ -0,0 +1,446 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for `/v1/realtime/video` WebSocket video output streaming."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from fastapi import FastAPI, WebSocket
|
||||
from pytest_mock import MockerFixture
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _fake_video_frames(num_frames: int = 2) -> list[np.ndarray]:
|
||||
return [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(num_frames)]
|
||||
|
||||
|
||||
def _build_test_app(
|
||||
*,
|
||||
mocker: MockerFixture,
|
||||
streaming_chunks: list[tuple[bytes, bool]],
|
||||
mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]],
|
||||
final_chunk: bytes = b"",
|
||||
stall_timeout: float = 5.0,
|
||||
) -> tuple[FastAPI, OmniStreamingVideoOutputHandler, MagicMock]:
|
||||
engine_client = mocker.MagicMock()
|
||||
engine_client.abort = mocker.AsyncMock()
|
||||
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
|
||||
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
|
||||
engine_client.generate = mock_generate
|
||||
|
||||
class FakeStreamingVideoEncoder:
|
||||
encode_calls: list[int] = []
|
||||
|
||||
def encode(self, video):
|
||||
del video
|
||||
idx = len(self.encode_calls)
|
||||
self.encode_calls.append(idx)
|
||||
return streaming_chunks[idx][0]
|
||||
|
||||
def close(self):
|
||||
return final_chunk
|
||||
|
||||
def _make_encoder(*, output_format, fps, video_codec_options=None):
|
||||
del output_format, fps, video_codec_options
|
||||
return FakeStreamingVideoEncoder()
|
||||
|
||||
encoder_factory = mocker.MagicMock(side_effect=_make_encoder)
|
||||
engine_client.streaming_encoder_factory = encoder_factory
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
|
||||
encoder_factory,
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
|
||||
return_value="diffusion",
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
|
||||
handler = OmniStreamingVideoOutputHandler(
|
||||
engine_client=engine_client,
|
||||
model_name="test-model",
|
||||
stage_configs=engine_client.stage_configs,
|
||||
stall_timeout=stall_timeout,
|
||||
start_timeout=5.0,
|
||||
)
|
||||
app = FastAPI()
|
||||
|
||||
@app.websocket("/v1/realtime/video")
|
||||
async def ws_endpoint(websocket: WebSocket):
|
||||
await handler.handle_session(websocket)
|
||||
|
||||
return app, handler, engine_client
|
||||
|
||||
|
||||
class TestStreamingVideoOutputWebSocket:
|
||||
"""WebSocket protocol tests for streaming generated video output."""
|
||||
|
||||
def test_streaming_session_emits_video_start_binary_chunks_and_done(self, mocker: MockerFixture):
|
||||
"""A full session delivers video.start, binary chunks, then session.done."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
for frames, finished in [
|
||||
(_fake_video_frames(2), False),
|
||||
(_fake_video_frames(3), True),
|
||||
]:
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[frames],
|
||||
final_output_type="image",
|
||||
finished=finished,
|
||||
)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
|
||||
mock_generate=mock_generate,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"})
|
||||
|
||||
start = ws.receive_json()
|
||||
assert start["type"] == "video.start"
|
||||
assert start["format"] == "m4s"
|
||||
assert "format" not in start["config"]
|
||||
assert "request_id" in start
|
||||
|
||||
assert ws.receive_bytes() == b"mp4-chunk-0"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-1"
|
||||
|
||||
done = ws.receive_json()
|
||||
assert done["type"] == "session.done"
|
||||
assert done["chunks"] == 2
|
||||
assert done["stopped"] is False
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
engine_client.streaming_encoder_factory.assert_called_once()
|
||||
assert engine_client.streaming_encoder_factory.call_args.kwargs["output_format"] == "m4s"
|
||||
|
||||
def test_streaming_session_rejects_unknown_format(self, mocker: MockerFixture):
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
if False:
|
||||
yield None
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[],
|
||||
mock_generate=mock_generate,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "hello", "format": "webm"})
|
||||
err = ws.receive_json()
|
||||
assert err["type"] == "error"
|
||||
assert "format" in err["message"]
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
|
||||
def test_streaming_session_emits_final_encoder_delta_before_done(self, mocker: MockerFixture):
|
||||
"""A non-empty encoder close delta is delivered as the final binary chunk."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", True)],
|
||||
mock_generate=mock_generate,
|
||||
final_chunk=b"mp4-trailer",
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-0"
|
||||
assert ws.receive_bytes() == b"mp4-trailer"
|
||||
|
||||
done = ws.receive_json()
|
||||
assert done["type"] == "session.done"
|
||||
assert done["chunks"] == 2
|
||||
assert done["stopped"] is False
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
|
||||
def test_generation_error_output_returns_websocket_error(self, mocker: MockerFixture):
|
||||
"""An erroneous OmniRequestOutput from generate is sent to the client as a WebSocket error."""
|
||||
error_message = "diffusion engine exploded"
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=False,
|
||||
)
|
||||
yield OmniRequestOutput.from_error("req-ws", error_message)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", False)],
|
||||
mock_generate=mock_generate,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "hello"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-0"
|
||||
err = ws.receive_json()
|
||||
assert err["type"] == "error"
|
||||
assert error_message in err["message"]
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
|
||||
|
||||
class _MockVideoOutputWebSocket:
|
||||
"""Feeds scripted client messages and records server sends."""
|
||||
|
||||
def __init__(self, messages: list[str]) -> None:
|
||||
self._messages: asyncio.Queue[str] = asyncio.Queue()
|
||||
for message in messages:
|
||||
self._messages.put_nowait(message)
|
||||
self.sent: list[dict | bytes] = []
|
||||
|
||||
def enqueue(self, message: str) -> None:
|
||||
self._messages.put_nowait(message)
|
||||
|
||||
async def accept(self) -> None:
|
||||
return None
|
||||
|
||||
async def receive_text(self) -> str:
|
||||
return await self._messages.get()
|
||||
|
||||
async def send_json(self, data: dict) -> None:
|
||||
self.sent.append(data)
|
||||
|
||||
async def send_bytes(self, data: bytes) -> None:
|
||||
self.sent.append(data)
|
||||
|
||||
@property
|
||||
def sent_json(self) -> list[dict]:
|
||||
return [m for m in self.sent if isinstance(m, dict)]
|
||||
|
||||
|
||||
def _build_handler_for_async_tests(
|
||||
mocker: MockerFixture,
|
||||
*,
|
||||
mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]],
|
||||
streaming_chunks: list[tuple[bytes, bool]] | None = None,
|
||||
stall_timeout: float = 5.0,
|
||||
) -> tuple[OmniStreamingVideoOutputHandler, MagicMock, _MockVideoOutputWebSocket]:
|
||||
engine_client = mocker.MagicMock()
|
||||
engine_client.abort = mocker.AsyncMock()
|
||||
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
|
||||
engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
|
||||
engine_client.generate = mock_generate
|
||||
|
||||
chunks = streaming_chunks or [(b"mp4-chunk-0", True)]
|
||||
|
||||
class FakeStreamingVideoEncoder:
|
||||
def __init__(self) -> None:
|
||||
self.encode_calls = 0
|
||||
|
||||
def encode(self, video):
|
||||
del video
|
||||
idx = self.encode_calls
|
||||
self.encode_calls += 1
|
||||
return chunks[idx][0]
|
||||
|
||||
def close(self):
|
||||
return b""
|
||||
|
||||
encoder_factory = mocker.MagicMock(side_effect=lambda **_: FakeStreamingVideoEncoder())
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
|
||||
encoder_factory,
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
|
||||
return_value="diffusion",
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
|
||||
return_value=[OmniDiffusionSamplingParams()],
|
||||
)
|
||||
|
||||
handler = OmniStreamingVideoOutputHandler(
|
||||
engine_client=engine_client,
|
||||
model_name="test-model",
|
||||
stage_configs=engine_client.stage_configs,
|
||||
stall_timeout=stall_timeout,
|
||||
start_timeout=5.0,
|
||||
)
|
||||
start_msg = json.dumps({"type": "session.start", "prompt": "async test"})
|
||||
ws = _MockVideoOutputWebSocket([start_msg])
|
||||
return handler, engine_client, ws
|
||||
|
||||
|
||||
class TestStreamingVideoOutputStallTimeout:
|
||||
"""Stall-timeout and session.ping behavior."""
|
||||
|
||||
def test_long_generation_without_client_messages_succeeds(self, mocker: MockerFixture):
|
||||
"""Engine progress between yields keeps the session alive without client pings."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=False,
|
||||
)
|
||||
await asyncio.sleep(0.3)
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
|
||||
mock_generate=mock_generate,
|
||||
stall_timeout=0.5,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "slow generation"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-0"
|
||||
assert ws.receive_bytes() == b"mp4-chunk-1"
|
||||
done = ws.receive_json()
|
||||
assert done["type"] == "session.done"
|
||||
assert done["stopped"] is False
|
||||
|
||||
engine_client.abort.assert_not_called()
|
||||
|
||||
def test_stall_timeout_aborts_when_engine_silent(self, mocker: MockerFixture):
|
||||
"""No engine output for longer than stall_timeout aborts the request."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
await asyncio.sleep(1.0)
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
app, _handler, engine_client = _build_test_app(
|
||||
mocker=mocker,
|
||||
streaming_chunks=[(b"mp4-chunk-0", True)],
|
||||
mock_generate=mock_generate,
|
||||
stall_timeout=0.2,
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/realtime/video") as ws:
|
||||
ws.send_json({"type": "session.start", "prompt": "stalled"})
|
||||
assert ws.receive_json()["type"] == "video.start"
|
||||
err = ws.receive_json()
|
||||
assert err["type"] == "error"
|
||||
assert "Stall timeout" in err["message"]
|
||||
|
||||
engine_client.abort.assert_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_ping_returns_pong(self, mocker: MockerFixture):
|
||||
"""session.ping during generation receives session.pong."""
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
await asyncio.sleep(0.5)
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
handler, _engine_client, ws = _build_handler_for_async_tests(
|
||||
mocker,
|
||||
mock_generate=mock_generate,
|
||||
stall_timeout=5.0,
|
||||
)
|
||||
|
||||
async def _send_ping_after_delay() -> None:
|
||||
await asyncio.sleep(0.05)
|
||||
ws.enqueue(json.dumps({"type": "session.ping"}))
|
||||
|
||||
ping_task = asyncio.create_task(_send_ping_after_delay())
|
||||
await handler.handle_session(ws) # type: ignore[arg-type]
|
||||
await ping_task
|
||||
|
||||
assert any(m.get("type") == "session.pong" for m in ws.sent_json)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ping_prevents_stall_during_silent_engine(self, mocker: MockerFixture):
|
||||
"""Client pings refresh the stall clock while the engine produces no output."""
|
||||
|
||||
release = asyncio.Event()
|
||||
|
||||
async def mock_generate(*_args, **_kwargs):
|
||||
await release.wait()
|
||||
yield OmniRequestOutput.from_diffusion(
|
||||
request_id="req-ws",
|
||||
images=[_fake_video_frames(2)],
|
||||
final_output_type="image",
|
||||
finished=True,
|
||||
)
|
||||
|
||||
handler, engine_client, ws = _build_handler_for_async_tests(
|
||||
mocker,
|
||||
mock_generate=mock_generate,
|
||||
stall_timeout=0.3,
|
||||
)
|
||||
|
||||
async def _pinger() -> None:
|
||||
while not release.is_set():
|
||||
ws.enqueue(json.dumps({"type": "session.ping"}))
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
pinger_task = asyncio.create_task(_pinger())
|
||||
session_task = asyncio.create_task(handler.handle_session(ws)) # type: ignore[arg-type]
|
||||
await asyncio.sleep(0.8)
|
||||
release.set()
|
||||
await session_task
|
||||
await pinger_task
|
||||
|
||||
errors = [m for m in ws.sent_json if m.get("type") == "error"]
|
||||
assert not any("Stall timeout" in m.get("message", "") for m in errors)
|
||||
assert any(m.get("type") == "session.pong" for m in ws.sent_json)
|
||||
assert any(m.get("type") == "session.done" for m in ws.sent_json)
|
||||
engine_client.abort.assert_not_called()
|
||||
@@ -48,6 +48,107 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
|
||||
assert mux_calls[0]["audio"] is None
|
||||
|
||||
|
||||
def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch):
|
||||
muxers = []
|
||||
|
||||
class FakeFragmentedMP4Muxer:
|
||||
def __init__(self, *, width, height, fps, video_codec_options=None):
|
||||
self.calls = []
|
||||
muxers.append(
|
||||
{
|
||||
"width": width,
|
||||
"height": height,
|
||||
"fps": fps,
|
||||
"video_codec_options": video_codec_options,
|
||||
"instance": self,
|
||||
}
|
||||
)
|
||||
|
||||
def mux_video_frames(self, frames):
|
||||
self.calls.append(frames)
|
||||
return f"fragment-{len(self.calls)}".encode()
|
||||
|
||||
def close(self):
|
||||
return b"trailer"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.utils.media_utils.FragmentedMP4Muxer",
|
||||
FakeFragmentedMP4Muxer,
|
||||
)
|
||||
|
||||
encoder = video_api_utils.FragmentedMP4VideoEncoder(
|
||||
fps=12,
|
||||
video_codec_options={"preset": "ultrafast"},
|
||||
)
|
||||
assert encoder.encode([np.zeros((4, 6, 3), dtype=np.float32)]) == b"fragment-1"
|
||||
assert encoder.encode([np.ones((4, 6, 3), dtype=np.float32)]) == b"fragment-2"
|
||||
assert encoder.close() == b"trailer"
|
||||
|
||||
assert len(muxers) == 1
|
||||
assert muxers[0]["width"] == 6
|
||||
assert muxers[0]["height"] == 4
|
||||
assert muxers[0]["fps"] == 12.0
|
||||
assert muxers[0]["video_codec_options"] == {"preset": "ultrafast"}
|
||||
assert len(muxers[0]["instance"].calls) == 2
|
||||
|
||||
|
||||
def test_create_streaming_video_encoder_selects_requested_format():
|
||||
assert isinstance(
|
||||
video_api_utils.create_streaming_video_encoder(output_format="m4s", fps=12),
|
||||
video_api_utils.FragmentedMP4VideoEncoder,
|
||||
)
|
||||
|
||||
|
||||
def test_finalize_streaming_mp4_bytes_produces_progressive_mp4():
|
||||
"""Fragment MP4 chunks are remuxed into decodable progressive MP4 bytes."""
|
||||
import numpy as np
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer, finalize_streaming_video_bytes
|
||||
|
||||
def _read_mp4_video_info(mp4_bytes: bytes) -> tuple[int, float, float | None, int, int]:
|
||||
import io
|
||||
|
||||
import av
|
||||
|
||||
with av.open(io.BytesIO(mp4_bytes), mode="r", format="mp4") as container:
|
||||
stream = container.streams.video[0]
|
||||
frame_count = sum(1 for _ in container.decode(stream))
|
||||
assert stream.average_rate is not None
|
||||
fps = float(stream.average_rate)
|
||||
duration_sec = None
|
||||
if stream.duration is not None:
|
||||
assert stream.time_base is not None
|
||||
duration_sec = float(stream.duration * stream.time_base)
|
||||
return frame_count, fps, duration_sec, stream.width, stream.height
|
||||
|
||||
width = 32
|
||||
height = 32
|
||||
fps = 16.0
|
||||
input_frame_count = 2
|
||||
|
||||
muxer = FragmentedMP4Muxer(width=width, height=height, fps=fps)
|
||||
frames = np.zeros((input_frame_count, height, width, 3), dtype=np.uint8)
|
||||
streamed = muxer.mux_video_frames(frames) + muxer.close()
|
||||
|
||||
finalized = finalize_streaming_video_bytes(streamed, input_format="m4s", fps=fps)
|
||||
assert finalized
|
||||
assert finalized != streamed
|
||||
|
||||
streamed_info = _read_mp4_video_info(streamed)
|
||||
final_info = _read_mp4_video_info(finalized)
|
||||
expected_duration = input_frame_count / fps
|
||||
|
||||
assert streamed_info[0] == input_frame_count
|
||||
assert final_info[0] == input_frame_count
|
||||
assert streamed_info[1] == pytest.approx(fps)
|
||||
assert final_info[1] == pytest.approx(fps)
|
||||
assert streamed_info[3:] == (width, height)
|
||||
assert final_info[3:] == (width, height)
|
||||
|
||||
assert final_info[2] == pytest.approx(expected_duration, rel=0.05)
|
||||
assert streamed_info[2] == pytest.approx(expected_duration, rel=0.05)
|
||||
|
||||
|
||||
def test_rife_model_inference_runs_on_dummy_tensors():
|
||||
model = rife_interpolator.Model().eval()
|
||||
img0 = torch.rand(1, 3, 32, 32)
|
||||
|
||||
@@ -1334,7 +1334,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
|
||||
extra_params = {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"use_cfg_zero_star": True,
|
||||
}
|
||||
response = test_client.post(
|
||||
@@ -1352,7 +1352,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
|
||||
captured = engine.captured_sampling_params_list[0]
|
||||
assert captured.extra_args["is_enable_stage2"] is True
|
||||
assert captured.extra_args["pyramid_num_stages"] == 3
|
||||
assert captured.extra_args["pyramid_num_inference_steps_list"] == [20, 20, 20]
|
||||
assert captured.extra_args["pyramid_num_inference_steps_list"] == [1, 1, 1]
|
||||
assert captured.extra_args["use_cfg_zero_star"] is True
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user