chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] CI: add tests to validate the size, extension, and format of output images/videos. (#15736)

Browse Source
This commit is contained in:
Prozac614
2025-12-25 14:00:22 +08:00
committed by GitHub
parent 49ab72f837
commit 10a9573efd
2 changed files with 306 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

119
python/sglang/multimodal_gen/test/server/test_server_utils.py
View File

@@ -35,10 +35,13 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
)
from sglang.multimodal_gen.test.slack_utils import upload_file_to_slack
from sglang.multimodal_gen.test.test_utils import (
get_expected_image_format,
is_image_url,
prepare_perf_log,
validate_image,
validate_image_file,
validate_openai_video,
validate_video_file,
)
logger = init_logger(__name__)
@@ -77,6 +80,49 @@ def download_image_from_url(url: str) -> Path:
raise
def parse_dimensions(size_string: str | None) -> tuple[int | None, int | None]:
"""Parse a size string in "widthxheight" format to (width, height) tuple.
Args:
size_string: Size string in "widthxheight" format (e.g., "1024x1024") or None.
Spaces are automatically stripped.
Returns:
Tuple of (width, height) as integers if parsing succeeds, (None, None) otherwise.
"""
if not size_string:
return (None, None)
# Strip spaces from the entire string
size_string = size_string.strip()
if not size_string:
return (None, None)
# Split by "x"
parts = size_string.split("x")
if len(parts) != 2:
return (None, None)
# Strip spaces from each part and try to convert to int
try:
width_str = parts[0].strip()
height_str = parts[1].strip()
if not width_str or not height_str:
return (None, None)
width = int(width_str)
height = int(height_str)
# Validate that both are positive
if width <= 0 or height <= 0:
return (None, None)
return (width, height)
except ValueError:
return (None, None)
@dataclass
class ServerContext:
"""Context for a running diffusion server."""
@@ -752,9 +798,17 @@ def get_generate_fn(
content = resp.read()
validate_openai_video(content)
tmp_path = f"{video_id}.mp4"
expected_filename = f"{video_id}.mp4"
tmp_path = expected_filename
with open(tmp_path, "wb") as f:
f.write(content)
# Validate output file
expected_width, expected_height = parse_dimensions(size)
validate_video_file(
tmp_path, expected_filename, expected_width, expected_height
)
upload_file_to_slack(
case_id=case_id,
model=model_path,
@@ -773,6 +827,10 @@ def get_generate_fn(
if not sampling_params.prompt:
pytest.skip(f"{id}: no text prompt configured")
# Request parameters that affect output format
req_output_format = None # Not specified in current request
req_background = None # Not specified in current request
response = client.images.with_raw_response.generate(
model=model_path,
prompt=sampling_params.prompt,
@@ -784,9 +842,24 @@ def get_generate_fn(
validate_image(result.data[0].b64_json)
img_data = base64.b64decode(result.data[0].b64_json)
tmp_path = f"{result.created}.png"
# Infer expected format from request parameters
expected_ext = get_expected_image_format(req_output_format, req_background)
expected_filename = f"{result.created}.{expected_ext}"
tmp_path = expected_filename
with open(tmp_path, "wb") as f:
f.write(img_data)
# Validate output file
expected_width, expected_height = parse_dimensions(output_size)
validate_image_file(
tmp_path,
expected_filename,
expected_width,
expected_height,
output_format=req_output_format,
background=req_background,
)
upload_file_to_slack(
case_id=case_id,
model=model_path,
@@ -818,6 +891,10 @@ def get_generate_fn(
image_paths = new_image_paths
# Request parameters that affect output format
req_output_format = None # Not specified in current request
req_background = None # Not specified in current request
images = [open(image_path, "rb") for image_path in image_paths]
try:
response = client.images.with_raw_response.edit(
@@ -838,9 +915,24 @@ def get_generate_fn(
validate_image(result.data[0].b64_json)
img_data = base64.b64decode(result.data[0].b64_json)
tmp_path = f"{rid}.png"
# Infer expected format from request parameters
expected_ext = get_expected_image_format(req_output_format, req_background)
expected_filename = f"{rid}.{expected_ext}"
tmp_path = expected_filename
with open(tmp_path, "wb") as f:
f.write(img_data)
# Validate output file
expected_width, expected_height = parse_dimensions(output_size)
validate_image_file(
tmp_path,
expected_filename,
expected_width,
expected_height,
output_format=req_output_format,
background=req_background,
)
upload_file_to_slack(
case_id=case_id,
model=model_path,
@@ -869,6 +961,10 @@ def get_generate_fn(
f"{id}: image_path must be a URL for URL direct test: {url}"
)
# Request parameters that affect output format
req_output_format = None # Not specified in current request
req_background = None # Not specified in current request
response = client.images.with_raw_response.edit(
model=model_path,
prompt=sampling_params.prompt,
@@ -885,9 +981,24 @@ def get_generate_fn(
# Save and upload result for verification
img_data = base64.b64decode(result.data[0].b64_json)
tmp_path = f"{rid}.png"
# Infer expected format from request parameters
expected_ext = get_expected_image_format(req_output_format, req_background)
expected_filename = f"{rid}.{expected_ext}"
tmp_path = expected_filename
with open(tmp_path, "wb") as f:
f.write(img_data)
# Validate output file
expected_width, expected_height = parse_dimensions(sampling_params.output_size)
validate_image_file(
tmp_path,
expected_filename,
expected_width,
expected_height,
output_format=req_output_format,
background=req_background,
)
upload_file_to_slack(
case_id=case_id,
model=model_path,

204
python/sglang/multimodal_gen/test/test_utils.py
View File

@@ -12,6 +12,7 @@ import unittest
from pathlib import Path
from typing import Optional
import cv2
from PIL import Image
from sglang.multimodal_gen.configs.sample.sampling_params import DataType
@@ -90,9 +91,11 @@ def get_dynamic_server_port() -> int:
return base_port + 1000
def is_mp4(data):
idx = data.find(b"ftyp")
return 0 <= idx <= 32
def is_mp4(data: bytes) -> bool:
"""Check if data represents a valid MP4 file by magic bytes."""
if len(data) < 8:
return False
return data[4:8] == b"ftyp"
def is_jpeg(data: bytes) -> bool:
@@ -105,6 +108,30 @@ def is_png(data):
return data.startswith(b"\x89PNG\r\n\x1a\n")
def is_webp(data: bytes) -> bool:
# WebP files start with: RIFF....WEBP
return data[:4] == b"RIFF" and data[8:12] == b"WEBP"
def get_expected_image_format(
output_format: str | None = None,
background: str | None = None,
) -> str:
"""Infer expected image format based on request parameters.
Args:
output_format: The output_format parameter from the request (png/jpeg/webp/jpg)
background: The background parameter from the request (transparent/opaque/auto)
Returns:
Expected file extension: "jpg", "png", or "webp"
"""
fmt = (output_format or "").lower()
if fmt in {"png", "webp", "jpeg", "jpg"}:
return "jpg" if fmt == "jpeg" else fmt
if (background or "auto").lower() == "transparent":
return "png"
return "jpg" # Default
def wait_for_port(host="127.0.0.1", port=30010, deadline=300.0, interval=0.5):
end = time.time() + deadline
last_err = None
@@ -206,22 +233,173 @@ def validate_image(b64_json: str) -> None:
def validate_video(b64_json: str) -> None:
"""Decode and validate that video is a valid format."""
video_bytes = base64.b64decode(b64_json)
is_mp4 = (
video_bytes[:4] == b"\x00\x00\x00\x18" or video_bytes[:4] == b"\x00\x00\x00\x1c"
)
is_webm = video_bytes[:4] == b"\x1a\x45\xdf\xa3"
assert is_mp4 or is_webm, "Video must be MP4 or WebM"
assert is_mp4(video_bytes) or is_webm, "Video must be MP4 or WebM"
def validate_openai_video(video_bytes: bytes) -> None:
"""Validate that video is MP4 or WebM by magic bytes."""
is_mp4 = (
video_bytes.startswith(b"\x00\x00\x00\x18")
or video_bytes.startswith(b"\x00\x00\x00\x1c")
or video_bytes[4:8] == b"ftyp"
)
is_webm = video_bytes.startswith(b"\x1a\x45\xdf\xa3")
assert is_mp4 or is_webm, "Video must be MP4 or WebM"
assert is_mp4(video_bytes) or is_webm, "Video must be MP4 or WebM"
def validate_image_file(
file_path: str,
expected_filename: str,
expected_width: int | None = None,
expected_height: int | None = None,
output_format: str | None = None,
background: str | None = None,
) -> None:
"""Validate image output file: existence, extension, size, filename, format, dimensions."""
# Infer expected format from request parameters
expected_ext = get_expected_image_format(output_format, background)
# 1. File existence
assert os.path.exists(file_path), f"Image file does not exist: {file_path}"
# 2. Extension check
assert file_path.endswith(
f".{expected_ext}"
), f"Expected .{expected_ext} extension, got: {file_path}"
# 3. File size > 0
file_size = os.path.getsize(file_path)
assert file_size > 0, f"Image file is empty: {file_path}"
# 4. Filename validation
actual_filename = os.path.basename(file_path)
assert (
actual_filename == expected_filename
), f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'"
# 5. Image format validation (magic bytes check based on expected format)
with open(file_path, "rb") as f:
header = f.read(12) # Read enough bytes for webp detection
if expected_ext == "png":
assert is_png(header), f"File is not a valid PNG: {file_path}"
elif expected_ext == "jpg":
assert is_jpeg(header), f"File is not a valid JPEG: {file_path}"
elif expected_ext == "webp":
assert is_webp(header), f"File is not a valid WebP: {file_path}"
# 6. Image dimension validation (reuse PIL)
if expected_width is not None and expected_height is not None:
with Image.open(file_path) as img:
width, height = img.size
assert (
width == expected_width
), f"Width mismatch: expected {expected_width}, got {width}"
assert (
height == expected_height
), f"Height mismatch: expected {expected_height}, got {height}"
def _get_video_dimensions_from_metadata(
cap: cv2.VideoCapture,
) -> tuple[int, int] | None:
"""Get video dimensions from metadata properties.
Args:
cap: OpenCV VideoCapture object
Returns:
Tuple of (width, height) if successful, None if metadata is invalid
"""
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
if width == 0 or height == 0:
return None
return (int(width), int(height))
def _get_video_dimensions_from_frame(cap: cv2.VideoCapture) -> tuple[int, int]:
"""Get video dimensions by reading the first frame.
Args:
cap: OpenCV VideoCapture object
Returns:
Tuple of (width, height)
Raises:
ValueError: If unable to read a frame from the video
"""
ret, frame = cap.read()
if not ret or frame is None:
raise ValueError("Unable to read video frame to get dimensions")
# frame.shape is (height, width, channels)
height, width = frame.shape[:2]
return (int(width), int(height))
def get_video_dimensions(file_path: str) -> tuple[int, int]:
"""Get video dimensions (width, height) from a video file.
Tries to get dimensions from metadata first, falls back to reading first frame.
Args:
file_path: Path to the video file
Returns:
Tuple of (width, height)
Raises:
ValueError: If unable to get video dimensions
"""
cap = cv2.VideoCapture(file_path)
try:
# Try to get dimensions from metadata first
dimensions = _get_video_dimensions_from_metadata(cap)
if dimensions is not None:
return dimensions
# Fall back to reading first frame
return _get_video_dimensions_from_frame(cap)
finally:
cap.release()
def validate_video_file(
file_path: str,
expected_filename: str,
expected_width: int | None = None,
expected_height: int | None = None,
) -> None:
"""Validate video output file: existence, extension, size, filename, format, dimensions."""
# 1. File existence
assert os.path.exists(file_path), f"Video file does not exist: {file_path}"
# 2. Extension check
assert file_path.endswith(".mp4"), f"Expected .mp4 extension, got: {file_path}"
# 3. File size > 0
file_size = os.path.getsize(file_path)
assert file_size > 0, f"Video file is empty: {file_path}"
# 4. Filename validation
actual_filename = os.path.basename(file_path)
assert (
actual_filename == expected_filename
), f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'"
# 5. Video format validation (reuse is_mp4)
with open(file_path, "rb") as f:
header = f.read(32)
assert is_mp4(header), f"File is not a valid MP4: {file_path}"
# 6. Video dimension validation (using OpenCV)
if expected_width is not None and expected_height is not None:
actual_width, actual_height = get_video_dimensions(file_path)
assert (
actual_width == expected_width
), f"Video width mismatch: expected {expected_width}, got {actual_width}"
assert (
actual_height == expected_height
), f"Video height mismatch: expected {expected_height}, got {actual_height}"
@dataclasses.dataclass