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>
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Unit tests for latent-mask serialization helpers."""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
import torch
|
|
from comfyui_vllm_omni.utils.latent_mask import scalar_mask_to_json, video_mask_to_grid_json
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_scalar_mask_to_json():
|
|
assert scalar_mask_to_json(0.5) == "0.5"
|
|
assert scalar_mask_to_json(1.0) == "1.0"
|
|
with pytest.raises(ValueError):
|
|
scalar_mask_to_json(1.5)
|
|
with pytest.raises(ValueError):
|
|
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_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_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_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_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_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)
|