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

[diffusion] CI: add consistency test (#15236)

Browse Source 
Co-authored-by: daiweitao <dwti614707404@163.com>
This commit is contained in:
Prozac614
2026-04-07 09:50:23 +08:00
committed by GitHub
co-authored by daiweitao
parent e4b1366a46
commit ef2d4013d7
11 changed files with 1026 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test-amd-rocm720.yml
+2
View File
@@ -633,6 +633,7 @@ jobs:
-e SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE=0.6 \
-e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \
-e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \
-e SGLANG_SKIP_CONSISTENCY=1 \
-e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \
-e AITER_JIT_DIR=/sgl-data/aiter-kernels \
-e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \
@@ -762,6 +763,7 @@ jobs:
-e SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE=0.6 \
-e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \
-e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \
-e SGLANG_SKIP_CONSISTENCY=1 \
-e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \
-e AITER_JIT_DIR=/sgl-data/aiter-kernels \
-e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \
.github/workflows/pr-test-amd.yml
+2
View File
@@ -638,6 +638,7 @@ jobs:
-e SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE=0.6 \
-e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \
-e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \
-e SGLANG_SKIP_CONSISTENCY=1 \
-e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \
-e AITER_JIT_DIR=/sgl-data/aiter-kernels \
-e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \
@@ -765,6 +766,7 @@ jobs:
-e SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE=0.6 \
-e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \
-e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \
-e SGLANG_SKIP_CONSISTENCY=1 \
-e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \
-e AITER_JIT_DIR=/sgl-data/aiter-kernels \
-e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \
python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
+12 -3
View File
@@ -506,16 +506,25 @@ class LoRAPipeline(ComposedPipelineBase):
return bool(self.cur_adapter_name)
return target in self.cur_adapter_name
def load_lora_adapter(self, lora_path: str, lora_nickname: str, rank: int):
def load_lora_adapter(
self,
lora_path: str,
lora_nickname: str,
rank: int,
weight_name: str | None = None,
):
"""
Load the LoRA, and setup the lora_adapters for later weight replacement
"""
assert lora_path is not None
if weight_name is None and lora_path == self.server_args.lora_path:
weight_name = self.server_args.lora_weight_name
# Only rank 0 downloads to avoid race conditions where other ranks
# try to load incomplete downloads
if rank == 0:
lora_local_path = maybe_download_lora(lora_path)
lora_local_path = maybe_download_lora(lora_path, weight_name=weight_name)
else:
lora_local_path = None
@@ -525,7 +534,7 @@ class LoRAPipeline(ComposedPipelineBase):
# Non-rank-0 workers now download (will hit cache since rank 0 completed)
if rank != 0:
lora_local_path = maybe_download_lora(lora_path)
lora_local_path = maybe_download_lora(lora_path, weight_name=weight_name)
raw_state_dict = load_file(lora_local_path)
lora_state_dict = normalize_lora_state_dict(raw_state_dict, logger=logger)
python/sglang/multimodal_gen/runtime/server_args.py
+7
View File
@@ -149,6 +149,7 @@ class ServerArgs:
lora_path: str | None = None
lora_nickname: str = "default" # for swapping adapters in the pipeline
lora_scale: float = 1.0 # LoRA scale for merging (e.g., 0.125 for Hyper-SD)
lora_weight_name: str | None = None
# Component path overrides (key = model_index.json component name, value = path)
component_paths: dict[str, str] = field(default_factory=dict)
@@ -849,6 +850,12 @@ class ServerArgs:
default=ServerArgs.lora_scale,
help="LoRA scale for merging (e.g., 0.125 for Hyper-SD). Same as lora_scale in Diffusers",
)
parser.add_argument(
"--lora-weight-name",
type=str,
default=ServerArgs.lora_weight_name,
help="Specific safetensors filename to load from a multi-file LoRA repo",
)
# Add pipeline configuration arguments
PipelineConfig.add_cli_args(parser)
python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
+18 -5
View File
@@ -378,7 +378,10 @@ def check_gguf_file(model: str | os.PathLike) -> bool:
def maybe_download_lora(
model_name_or_path: str, local_dir: str | None = None, download: bool = True
model_name_or_path: str,
local_dir: str | None = None,
download: bool = True,
weight_name: str | None = None,
) -> str:
"""
Check if the model path is a Hugging Face Hub model ID and download it if needed.
@@ -386,6 +389,8 @@ def maybe_download_lora(
model_name_or_path: Local path or Hugging Face Hub model ID
local_dir: Local directory to save the model
download: Whether to download the model from Hugging Face Hub
weight_name: Specific safetensors filename to load (pins deterministic selection
for repos with multiple weight files)
Returns:
Local path to the model
@@ -403,14 +408,22 @@ def maybe_download_lora(
if os.path.isfile(local_path):
return local_path
weight_name = _best_guess_weight_name(local_path, file_extension=".safetensors")
if weight_name is not None:
target = os.path.join(local_path, weight_name)
if not os.path.isfile(target):
raise FileNotFoundError(
f"Specified lora_weight_name '{weight_name}' not found in {local_path}"
)
return target
guessed = _best_guess_weight_name(local_path, file_extension=".safetensors")
# AMD workaround: PR 15813 changed from model_name_or_path to local_path,
# which can return None. Fall back to original behavior on ROCm.
if weight_name is None and current_platform.is_rocm():
weight_name = _best_guess_weight_name(
if guessed is None and current_platform.is_rocm():
guessed = _best_guess_weight_name(
model_name_or_path, file_extension=".safetensors"
)
return os.path.join(local_path, weight_name)
return os.path.join(local_path, guessed)
def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:
python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py
+5
View File
@@ -15,6 +15,7 @@ ONE_NPU_CASES: list[DiffusionTestCase] = [
modality="image",
),
T2I_sampling_params,
run_consistency_check=False,
),
# === Text to Video (T2V) ===
DiffusionTestCase(
@@ -27,6 +28,7 @@ ONE_NPU_CASES: list[DiffusionTestCase] = [
DiffusionSamplingParams(
prompt=T2V_PROMPT,
),
run_consistency_check=False,
),
]
@@ -41,6 +43,7 @@ TWO_NPU_CASES: list[DiffusionTestCase] = [
tp_size=2,
),
T2I_sampling_params,
run_consistency_check=False,
),
DiffusionTestCase(
"qwen_image_t2i_2npu",
@@ -53,6 +56,7 @@ TWO_NPU_CASES: list[DiffusionTestCase] = [
ring_degree=2,
),
T2I_sampling_params,
run_consistency_check=False,
),
]
@@ -70,5 +74,6 @@ EIGHT_NPU_CASES: list[DiffusionTestCase] = [
DiffusionSamplingParams(
prompt=T2V_PROMPT,
),
run_consistency_check=False,
),
]
python/sglang/multimodal_gen/test/server/consistency_threshold.json
+205
View File
@@ -0,0 +1,205 @@
{
"_comment": "Some cases use lower thresholds; raise them if quality/perf improves later.",
"cases": {
"qwen_image_t2i": {
"clip_threshold": 0.97,
"ssim_threshold": 0.84,
"psnr_threshold": 16.0,
"mean_abs_diff_threshold": 11.5
},
"flux_image_t2i": {
"clip_threshold": 0.92,
"ssim_threshold": 0.95,
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 8.0
},
"flux_2_klein_image_t2i": {
"clip_threshold": 0.86,
"ssim_threshold": 0.60,
"psnr_threshold": 10.0,
"mean_abs_diff_threshold": 56.0
},
"zimage_image_t2i": {
"clip_threshold": 0.92,
"ssim_threshold": 0.90,
"psnr_threshold": 22.0,
"mean_abs_diff_threshold": 8.0
},
"zimage_image_t2i_multi_lora": {
"clip_threshold": 0.92,
"ssim_threshold": 0.92,
"psnr_threshold": 22.0,
"mean_abs_diff_threshold": 8.0
},
"qwen_image_t2i_cache_dit_enabled": {
"clip_threshold": 0.92,
"ssim_threshold": 0.86,
"psnr_threshold": 17.0,
"mean_abs_diff_threshold": 10.0
},
"layerwise_offload": {
"clip_threshold": 0.92,
"ssim_threshold": 0.91,
"psnr_threshold": 21.5,
"mean_abs_diff_threshold": 8.0
},
"zimage_image_t2i_fp8": {
"clip_threshold": 0.92,
"ssim_threshold": 0.84,
"psnr_threshold": 19.0,
"mean_abs_diff_threshold": 10.0
},
"sana_image_t2i": {
"clip_threshold": 0.91,
"ssim_threshold": 0.88,
"psnr_threshold": 21.0,
"mean_abs_diff_threshold": 8.4
},
"qwen_image_edit_2509_ti2i": {
"clip_threshold": 0.92,
"ssim_threshold": 0.65,
"psnr_threshold": 13.0,
"mean_abs_diff_threshold": 26.0
},
"qwen_image_layered_i2i": {
"clip_threshold": 0.92,
"ssim_threshold": 0.94,
"psnr_threshold": 28.0,
"mean_abs_diff_threshold": 8.0
},
"mova_360p_1gpu": {
"clip_threshold": 0.90,
"ssim_threshold": 0.87,
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 10.0
},
"wan2_1_t2v_1_3b_lora_1gpu": {
"clip_threshold": 0.54,
"ssim_threshold": 0.40,
"psnr_threshold": 13.2,
"mean_abs_diff_threshold": 32.0
},
"wan2_2_ti2v_5b": {
"clip_threshold": 0.90,
"ssim_threshold": 0.81,
"psnr_threshold": 20.4,
"mean_abs_diff_threshold": 10.0
},
"fastwan2_2_ti2v_5b": {
"clip_threshold": 0.90,
"ssim_threshold": 0.88,
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 10.0
},
"turbo_wan2_1_t2v_1.3b": {
"clip_threshold": 0.90,
"ssim_threshold": 0.52,
"psnr_threshold": 9.5,
"mean_abs_diff_threshold": 46.0
},
"zimage_image_t2i_multi_lora": {
"clip_threshold": 0.90,
"ssim_threshold": 0.52,
"psnr_threshold": 9.5,
"mean_abs_diff_threshold": 46.0
},
"fsdp-inference": {
"clip_threshold": 0.92,
"ssim_threshold": 0.90,
"psnr_threshold": 21.5,
"mean_abs_diff_threshold": 8.0
},
"zimage_image_t2i_2_gpus_non_square": {
"clip_threshold": 0.92,
"ssim_threshold": 0.76,
"psnr_threshold": 14.8,
"mean_abs_diff_threshold": 17.5
},
"flux_2_image_t2i_2_gpus": {
"clip_threshold": 0.54,
"ssim_threshold": 0.9,
"psnr_threshold": 19,
"mean_abs_diff_threshold": 8.0
},
"zimage_image_t2i_2_gpus": {
"clip_threshold": 0.92,
"ssim_threshold": 0.90,
"psnr_threshold": 21.5,
"mean_abs_diff_threshold": 8.0
},
"flux_image_t2i_2_gpus": {
"clip_threshold": 0.92,
"ssim_threshold": 0.90,
"psnr_threshold": 18.7,
"mean_abs_diff_threshold": 8.0
},
"flux_2_klein_ti2i_2_gpus": {
"clip_threshold": 0.92,
"ssim_threshold": 0.77,
"psnr_threshold": 18.4,
"mean_abs_diff_threshold": 18.0
},
"wan2_2_t2v_a14b_teacache_2gpu": {
"clip_threshold": 0.90,
"ssim_threshold": 0.72,
"psnr_threshold": 17.8,
"mean_abs_diff_threshold": 16.0
},
"wan2_1_t2v_14b_2gpu": {
"clip_threshold": 0.90,
"ssim_threshold": 0.84,
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 10.0
},
"mova_360p_ring1_uly2": {
"clip_threshold": 0.90,
"ssim_threshold": 0.91,
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 10.0
},
"wan2_1_i2v_14b_lora_2gpu": {
"clip_threshold": 0.90,
"ssim_threshold": 0.81,
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 10.0
},
"wan2_2_t2v_a14b_2gpu": {
"clip_threshold": 0.90,
"ssim_threshold": 0.72,
"psnr_threshold": 17.8,
"mean_abs_diff_threshold": 16.0
},
"wan2_2_t2v_a14b_lora_2gpu": {
"clip_threshold": 0.90,
"ssim_threshold": 0.81,
"psnr_threshold": 22.2,
"mean_abs_diff_threshold": 10.0
},
"mova_360p_ring2_uly1": {
"clip_threshold": 0.90,
"ssim_threshold": 0.91,
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 10.0
},
"wan2_1_i2v_14b_480P_2gpu": {
"clip_threshold": 0.76,
"ssim_threshold": 0.51,
"psnr_threshold": 14.8,
"mean_abs_diff_threshold": 23.0
},
"wan2_1_i2v_14b_720P_2gpu": {
"clip_threshold": 0.90,
"ssim_threshold": 0.89,
"psnr_threshold": 24.0,
"mean_abs_diff_threshold": 10.0
}
},
"default_clip_threshold_image": 0.92,
"default_clip_threshold_video": 0.90,
"default_ssim_threshold_image": 0.95,
"default_psnr_threshold_image": 28.0,
"default_mean_abs_diff_threshold_image": 8.0,
"default_ssim_threshold_video": 0.92,
"default_psnr_threshold_video": 24.0,
"default_mean_abs_diff_threshold_video": 10.0
}
python/sglang/multimodal_gen/test/server/test_server_common.py
+151
View File
@@ -35,8 +35,16 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
)
from sglang.multimodal_gen.test.test_utils import (
_consistency_gt_filenames,
_get_consistency_gt_dir,
compare_with_gt,
extract_key_frames_from_video,
get_consistency_gt_candidates,
get_consistency_gt_remote_files,
get_consistency_thresholds,
get_dynamic_server_port,
gt_exists,
image_bytes_to_numpy,
load_consistency_gt,
wait_for_req_perf_record,
)
@@ -445,6 +453,146 @@ Consider updating perf_baselines.json with the snippets below:
"""
logger.error(output)
def _validate_consistency(
self,
case: DiffusionTestCase,
content: bytes,
) -> None:
"""Validate output consistency against ground truth using CLIP similarity."""
if os.environ.get("SGLANG_SKIP_CONSISTENCY", "0") == "1":
logger.info(
f"[Consistency] Skipping consistency check for {case.id} (SGLANG_SKIP_CONSISTENCY=1)"
)
return
if not content:
logger.warning(
f"[Consistency] Skipping consistency check for {case.id}: "
"content is empty (generation may have timed out)"
)
return
num_gpus = case.server_args.num_gpus
is_video = case.server_args.modality == "video"
output_format = case.sampling_params.output_format
if not gt_exists(
case.id, num_gpus, is_video=is_video, output_format=output_format
):
if _get_consistency_gt_dir() is not None:
names = ", ".join(
get_consistency_gt_candidates(
case.id, num_gpus, is_video, output_format
)
)
else:
names = ", ".join(
_consistency_gt_filenames(
case.id, num_gpus, is_video, output_format
)
)
logger.error(f"""
--- MISSING GROUND TRUTH DETECTED ---
GT image(s) not found for '{case.id}'.
Add the expected file(s) to sglang-ci-data in diffusion-ci/consistency_gt/ with naming (n=num_gpus).
Image: {case.id}_{{n}}gpu.<ext> (ext from output_format: png, jpg, webp)
Video: {case.id}_{{n}}gpu_frame_0.png, {case.id}_{{n}}gpu_frame_mid.png, {case.id}_{{n}}gpu_frame_last.png
For this case, expected file(s): {names}
Repository: https://github.com/sglang-bot/sglang-ci-data (path: diffusion-ci/consistency_gt/)
(Optional) Per-case override in consistency_threshold.json:
"cases": {{
"{case.id}": {{
"clip_threshold": 0.92,
"ssim_threshold": 0.95,
"psnr_threshold": 28.0,
"mean_abs_diff_threshold": 8.0
}}
}}
""")
pytest.fail(
f"GT not found for {case.id}. See logs for instructions to add GT."
)
gt_data = load_consistency_gt(
case.id, num_gpus, is_video=is_video, output_format=output_format
)
thresholds = get_consistency_thresholds(case.id, is_video=is_video)
if is_video:
output_frames = extract_key_frames_from_video(content)
else:
output_frames = [image_bytes_to_numpy(content)]
result = compare_with_gt(
output_frames=output_frames,
gt_data=gt_data,
thresholds=thresholds,
case_id=case.id,
)
if not result.passed:
failed_frames = []
video_gt_info = ""
if is_video:
gt_remote_files = get_consistency_gt_remote_files(
case.id,
num_gpus,
is_video=True,
output_format=output_format,
)
video_gt_info = "\n".join(
f" - {filename}: {url}" for filename, url in gt_remote_files
)
for metric in result.frame_metrics:
failed_metrics = []
if not metric.clip_passed:
failed_metrics.append("clip")
if not metric.ssim_passed:
failed_metrics.append("ssim")
if not metric.psnr_passed:
failed_metrics.append("psnr")
if not metric.mean_abs_diff_passed:
failed_metrics.append("mean_abs_diff")
if failed_metrics:
failed_frames.append(
f" - f{metric.frame_index} "
f"[{', '.join(failed_metrics)}] "
f"clip={metric.clip_similarity:.4f} "
f"ssim={metric.ssim:.4f} "
f"psnr={metric.psnr:.4f} "
f"mean_abs_diff={metric.mean_abs_diff:.4f}"
)
pytest.fail(
f"Consistency check failed for {case.id}:\n"
f" Metrics: sim={result.min_similarity:.4f}, "
f"ssim={result.min_ssim:.4f}, "
f"psnr={result.min_psnr:.4f}, "
f"mean_abs_diff={result.max_mean_abs_diff:.4f}\n"
f" Thresholds: clip>={result.thresholds.clip_threshold}, "
f"ssim>={result.thresholds.ssim_threshold}, "
f"psnr>={result.thresholds.psnr_threshold}, "
f"mean_abs_diff<={result.thresholds.mean_abs_diff_threshold}\n"
f" Failed frames:\n"
+ "\n".join(failed_frames)
+ (
f"\n Compared GT frame files and links:\n{video_gt_info}"
if video_gt_info
else ""
)
)
logger.info(
f"[Consistency] {case.id}: PASSED "
f"(min_similarity={result.min_similarity:.4f}, "
f"min_ssim={result.min_ssim:.4f}, "
f"min_psnr={result.min_psnr:.4f}, "
f"max_mean_abs_diff={result.max_mean_abs_diff:.4f})"
)
def _save_gt_output(
self,
case: DiffusionTestCase,
@@ -903,6 +1051,9 @@ Consider updating perf_baselines.json with the snippets below:
if case.run_t2v_input_reference_check:
self._test_t2v_rejects_input_reference(diffusion_server, case)
if case.run_consistency_check:
self._validate_consistency(case, content)
# LoRA API functionality test with E2E validation (only for LoRA-enabled cases)
if case.run_lora_basic_api_check:
self._test_lora_api_functionality(diffusion_server, case, generate_fn)
python/sglang/multimodal_gen/test/server/testcase_configs.py
+8
View File
@@ -248,6 +248,7 @@ class DiffusionTestCase:
server_args: DiffusionServerArgs
sampling_params: DiffusionSamplingParams
run_perf_check: bool = True
run_consistency_check: bool = True
run_models_api_check: bool = True
run_t2v_input_reference_check: bool = True
run_lora_basic_api_check: bool = False
@@ -777,6 +778,7 @@ if not current_platform.is_hip():
enable_warmup=False,
),
HUNYUAN3D_SHAPE_sampling_params,
run_consistency_check=False,
),
)
# Skip turbowan on AMD: Triton requires 81920 shared memory, but AMD only has 65536.
@@ -802,6 +804,7 @@ ONE_GPU_CASES_C = [
modality="image",
),
T2I_sampling_params,
run_consistency_check=False,
)
]
@@ -851,6 +854,10 @@ TWO_GPU_CASES_A = [
custom_validator="video",
num_gpus=2,
lora_path="Cseti/wan2.2-14B-Arcane_Jinx-lora-v1",
extras=[
"--lora-weight-name",
"985347-wan22_14B-low-Nfj1nx-e65.safetensors",
],
),
DiffusionSamplingParams(
prompt="Nfj1nx with blue hair, a woman walking in a cyberpunk city at night",
@@ -938,6 +945,7 @@ TWO_GPU_CASES_A = [
extras=["--pipeline-class-name LTX2TwoStagePipeline"],
),
T2V_sampling_params,
run_consistency_check=False,
),
]
python/sglang/multimodal_gen/test/test_consistency_metrics.py
+115
View File
@@ -0,0 +1,115 @@
import math
import numpy as np
from sglang.multimodal_gen.test import test_utils
from sglang.multimodal_gen.test.test_utils import (
ConsistencyThresholds,
LoadedConsistencyGT,
compare_with_gt,
compute_mean_abs_diff,
compute_psnr,
compute_ssim,
)
def _solid_image(value: int, size: int = 32) -> np.ndarray:
return np.full((size, size, 3), value, dtype=np.uint8)
def test_pixel_metrics_identical_image():
image = _solid_image(128)
ssim = compute_ssim(image, image)
psnr = compute_psnr(image, image)
mean_abs_diff = compute_mean_abs_diff(image, image)
assert ssim == 1.0
assert math.isinf(psnr)
assert mean_abs_diff == 0.0
def test_pixel_metrics_detect_different_image():
image = _solid_image(128)
other = _solid_image(0)
ssim = compute_ssim(image, other)
psnr = compute_psnr(image, other)
mean_abs_diff = compute_mean_abs_diff(image, other)
assert ssim < 0.95
assert psnr < 28.0
assert mean_abs_diff > 8.0
def test_compare_with_gt_passes_for_identical_image(monkeypatch):
gt_image = _solid_image(128)
monkeypatch.setattr(
test_utils,
"compute_clip_embedding",
lambda image: np.array([1.0, 0.0], dtype=np.float32),
)
result = compare_with_gt(
output_frames=[gt_image.copy()],
gt_data=LoadedConsistencyGT(
images=[gt_image.copy()],
embeddings=[np.array([1.0, 0.0], dtype=np.float32)],
),
thresholds=ConsistencyThresholds(
clip_threshold=0.92,
ssim_threshold=0.95,
psnr_threshold=28.0,
mean_abs_diff_threshold=8.0,
),
case_id="unit_image_pass",
)
assert result.passed is True
assert result.min_similarity == 1.0
assert result.min_ssim == 1.0
assert math.isinf(result.min_psnr)
assert result.max_mean_abs_diff == 0.0
def test_compare_with_gt_uses_worst_frame_for_video(monkeypatch):
gt_frame_0 = _solid_image(128)
gt_frame_1 = _solid_image(128)
bad_frame = _solid_image(0)
monkeypatch.setattr(
test_utils,
"compute_clip_embedding",
lambda image: np.array([1.0, 0.0], dtype=np.float32),
)
result = compare_with_gt(
output_frames=[gt_frame_0.copy(), bad_frame],
gt_data=LoadedConsistencyGT(
images=[gt_frame_0.copy(), gt_frame_1.copy()],
embeddings=[
np.array([1.0, 0.0], dtype=np.float32),
np.array([1.0, 0.0], dtype=np.float32),
],
),
thresholds=ConsistencyThresholds(
clip_threshold=0.92,
ssim_threshold=0.95,
psnr_threshold=28.0,
mean_abs_diff_threshold=8.0,
),
case_id="unit_video_fail",
)
assert result.passed is False
assert result.min_similarity == 1.0
assert result.min_ssim < 0.95
assert result.min_psnr < 28.0
assert result.max_mean_abs_diff > 8.0
assert any(
not metric.ssim_passed
or not metric.psnr_passed
or not metric.mean_abs_diff_passed
for metric in result.frame_metrics
)
python/sglang/multimodal_gen/test/test_utils.py
+501 -1
View File
@@ -7,12 +7,15 @@ import socket
import subprocess
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
from urllib.parse import urljoin
import cv2
import httpx
import numpy as np
import requests
from PIL import Image
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
@@ -22,8 +25,27 @@ from sglang.multimodal_gen.runtime.utils.perf_logger import (
get_diffusion_perf_log_dir,
)
if TYPE_CHECKING:
from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
logger = init_logger(__name__)
SGL_TEST_FILES_CONSISTENCY_GT_BASE = "https://raw.githubusercontent.com/sglang-bot/sglang-ci-data/main/diffusion-ci/consistency_gt"
CONSISTENCY_THRESHOLD_JSON_PATH = (
Path(__file__).resolve().parent / "server" / "consistency_threshold.json"
)
CLIP_MODEL_NAME = "openai/clip-vit-large-patch14"
DEFAULT_CLIP_THRESHOLD_IMAGE = 0.92
DEFAULT_CLIP_THRESHOLD_VIDEO = 0.90
DEFAULT_SSIM_THRESHOLD_IMAGE = 0.95
DEFAULT_PSNR_THRESHOLD_IMAGE = 28.0
DEFAULT_MEAN_ABS_DIFF_THRESHOLD_IMAGE = 8.0
DEFAULT_SSIM_THRESHOLD_VIDEO = 0.92
DEFAULT_PSNR_THRESHOLD_VIDEO = 24.0
DEFAULT_MEAN_ABS_DIFF_THRESHOLD_VIDEO = 10.0
_clip_model_cache: dict[str, Any] = {}
_consistency_gt_cache: dict[str, Any] = {}
# ---------------------------------------------------------------------------
# Common model IDs for diffusion tests
#
@@ -560,10 +582,232 @@ def validate_video_file(
), f"Video height mismatch: expected {expected_height}, got {actual_height}"
def _load_threshold_json() -> dict[str, Any]:
"""Load consistency_threshold.json; returns {} if missing."""
if not CONSISTENCY_THRESHOLD_JSON_PATH.exists():
return {}
with CONSISTENCY_THRESHOLD_JSON_PATH.open("r", encoding="utf-8") as f:
return json.load(f)
@dataclass
class ConsistencyThresholds:
clip_threshold: float
ssim_threshold: float
psnr_threshold: float
mean_abs_diff_threshold: float
def get_consistency_thresholds(
case_id: str,
is_video: bool,
metadata: dict[str, Any] | None = None,
) -> ConsistencyThresholds:
"""Get all consistency thresholds for a case."""
if metadata is None:
metadata = _load_threshold_json()
case_meta = metadata.get("cases", {}).get(case_id, {})
suffix = "video" if is_video else "image"
defaults = {
"clip_threshold": metadata.get(
f"default_clip_threshold_{suffix}",
DEFAULT_CLIP_THRESHOLD_VIDEO if is_video else DEFAULT_CLIP_THRESHOLD_IMAGE,
),
"ssim_threshold": metadata.get(
f"default_ssim_threshold_{suffix}",
DEFAULT_SSIM_THRESHOLD_VIDEO if is_video else DEFAULT_SSIM_THRESHOLD_IMAGE,
),
"psnr_threshold": metadata.get(
f"default_psnr_threshold_{suffix}",
DEFAULT_PSNR_THRESHOLD_VIDEO if is_video else DEFAULT_PSNR_THRESHOLD_IMAGE,
),
"mean_abs_diff_threshold": metadata.get(
f"default_mean_abs_diff_threshold_{suffix}",
(
DEFAULT_MEAN_ABS_DIFF_THRESHOLD_VIDEO
if is_video
else DEFAULT_MEAN_ABS_DIFF_THRESHOLD_IMAGE
),
),
}
return ConsistencyThresholds(
clip_threshold=float(
case_meta.get("clip_threshold", defaults["clip_threshold"])
),
ssim_threshold=float(
case_meta.get("ssim_threshold", defaults["ssim_threshold"])
),
psnr_threshold=float(
case_meta.get("psnr_threshold", defaults["psnr_threshold"])
),
mean_abs_diff_threshold=float(
case_meta.get(
"mean_abs_diff_threshold", defaults["mean_abs_diff_threshold"]
)
),
)
def get_clip_threshold(
case: "DiffusionTestCase",
metadata: dict[str, Any] | None = None,
) -> float:
"""Get CLIP similarity threshold for a consistency test case."""
return get_consistency_thresholds(
case_id=case.id,
is_video=case.server_args.modality == "video",
metadata=metadata,
).clip_threshold
@dataclass
class FrameConsistencyMetrics:
frame_index: int
clip_similarity: float
ssim: float
psnr: float
mean_abs_diff: float
clip_passed: bool
ssim_passed: bool
psnr_passed: bool
mean_abs_diff_passed: bool
@dataclass
class ConsistencyResult:
"""Result of a consistency comparison."""
case_id: str
passed: bool
similarity_scores: list[float]
min_similarity: float
threshold: float
min_ssim: float
min_psnr: float
max_mean_abs_diff: float
thresholds: ConsistencyThresholds
frame_metrics: list[FrameConsistencyMetrics]
@dataclass
class LoadedConsistencyGT:
images: list[np.ndarray]
embeddings: list[np.ndarray]
def get_clip_model() -> tuple[Any, Any]:
"""Get CLIP model and processor."""
global _clip_model_cache
if "model" not in _clip_model_cache:
try:
import torch
from transformers import CLIPModel, CLIPProcessor
except ImportError as exc:
raise ImportError(
"transformers and torch are required for CLIP consistency check."
) from exc
logger.info(f"Loading CLIP model: {CLIP_MODEL_NAME}")
processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME)
model = CLIPModel.from_pretrained(CLIP_MODEL_NAME)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()
_clip_model_cache["model"] = model
_clip_model_cache["processor"] = processor
_clip_model_cache["device"] = device
logger.info(f"CLIP model loaded on {device}")
return _clip_model_cache["model"], _clip_model_cache["processor"]
def compute_clip_embedding(image: np.ndarray) -> np.ndarray:
"""Compute a normalized CLIP image embedding."""
try:
import torch
except ImportError as exc:
raise ImportError("torch is required for CLIP consistency check.") from exc
model, processor = get_clip_model()
device = _clip_model_cache["device"]
pil_image = Image.fromarray(image)
inputs = processor(images=pil_image, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
image_features = model.get_image_features(**inputs)
if hasattr(image_features, "image_embeds"):
image_features = image_features.image_embeds
elif hasattr(image_features, "pooler_output"):
image_features = image_features.pooler_output
image_features = image_features / image_features.norm(dim=-1, keepdim=True)
return image_features.cpu().numpy().flatten()
def compute_clip_similarity(emb1: np.ndarray, emb2: np.ndarray) -> float:
"""Compute cosine similarity between two CLIP embeddings."""
return float(np.dot(emb1, emb2))
def _ensure_rgb_uint8_image(image: np.ndarray) -> np.ndarray:
"""Normalize image input for pixel-wise consistency metrics."""
if image.ndim != 3 or image.shape[2] != 3:
raise ValueError(f"Expected RGB HWC image, got shape={image.shape}")
if image.dtype == np.uint8:
return image
image = np.clip(image, 0, 255)
return image.astype(np.uint8)
def compute_ssim(image: np.ndarray, gt_image: np.ndarray) -> float:
"""Compute SSIM between two RGB images."""
from skimage.metrics import structural_similarity
image = _ensure_rgb_uint8_image(image)
gt_image = _ensure_rgb_uint8_image(gt_image)
if image.shape != gt_image.shape:
raise ValueError(
f"Image shape mismatch for SSIM: output={image.shape}, gt={gt_image.shape}"
)
return float(structural_similarity(image, gt_image, channel_axis=2, data_range=255))
def compute_psnr(image: np.ndarray, gt_image: np.ndarray) -> float:
"""Compute PSNR between two RGB images."""
from skimage.metrics import peak_signal_noise_ratio
image = _ensure_rgb_uint8_image(image)
gt_image = _ensure_rgb_uint8_image(gt_image)
if image.shape != gt_image.shape:
raise ValueError(
f"Image shape mismatch for PSNR: output={image.shape}, gt={gt_image.shape}"
)
return float(peak_signal_noise_ratio(gt_image, image, data_range=255))
def compute_mean_abs_diff(image: np.ndarray, gt_image: np.ndarray) -> float:
"""Compute mean absolute pixel difference between two RGB images."""
image = _ensure_rgb_uint8_image(image)
gt_image = _ensure_rgb_uint8_image(gt_image)
if image.shape != gt_image.shape:
raise ValueError(
f"Image shape mismatch for mean_abs_diff: output={image.shape}, gt={gt_image.shape}"
)
return float(np.abs(image.astype(np.float32) - gt_image.astype(np.float32)).mean())
def output_format_to_ext(output_format: str | None) -> str:
"""Map output_format to file extension. Used by GT naming and consistency check."""
if not output_format:
return "png"
return "jpg"
of = output_format.lower()
if of == "jpeg":
return "jpg"
@@ -587,6 +831,151 @@ def _consistency_gt_filenames(
return [f"{case_id}_{n}gpu.{ext}"]
def get_consistency_gt_candidates(
case_id: str, num_gpus: int, is_video: bool, output_format: str | None = None
) -> list[str]:
"""Return candidate GT filenames for local consistency data."""
n = num_gpus
if is_video:
return [
f"{case_id}_{n}gpu_frame_0.png",
f"{case_id}_{n}gpu_frame_mid.png",
f"{case_id}_{n}gpu_frame_last.png",
]
base = f"{case_id}_{n}gpu"
preferred = output_format_to_ext(output_format)
exts = [preferred] + [e for e in ("png", "jpg", "webp") if e != preferred]
return [f"{base}.{e}" for e in exts]
def get_consistency_gt_remote_files(
case_id: str, num_gpus: int, is_video: bool, output_format: str | None = None
) -> list[tuple[str, str]]:
"""Return GT filenames with their remote raw URLs."""
filenames = _consistency_gt_filenames(case_id, num_gpus, is_video, output_format)
return [
(filename, f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{filename}")
for filename in filenames
]
def _get_consistency_gt_dir() -> Path | None:
"""Return the local GT directory when configured."""
d = os.environ.get("SGLANG_CONSISTENCY_GT_DIR")
if not d:
return None
return Path(d).resolve()
def _get_consistency_gt_cache_key(
case_id: str,
num_gpus: int,
is_video: bool,
output_format: str | None,
) -> str:
gt_dir = _get_consistency_gt_dir()
source = str(gt_dir) if gt_dir is not None else "remote"
return f"{case_id}:{num_gpus}:{is_video}:{output_format or ''}:{source}"
def load_consistency_gt(
case_id: str,
num_gpus: int,
is_video: bool = False,
output_format: str | None = None,
) -> LoadedConsistencyGT:
"""Load GT images and CLIP embeddings for consistency checks."""
cache_key = _get_consistency_gt_cache_key(
case_id, num_gpus, is_video, output_format
)
cached = _consistency_gt_cache.get(cache_key)
if cached is not None:
return cached
filenames = _consistency_gt_filenames(case_id, num_gpus, is_video, output_format)
images: list[np.ndarray] = []
gt_dir = _get_consistency_gt_dir()
if gt_dir is not None:
candidates = get_consistency_gt_candidates(
case_id, num_gpus, is_video, output_format
)
if is_video:
for fn in candidates:
path = gt_dir / fn
if not path.exists():
raise FileNotFoundError(f"GT image not found: {path}")
arr = np.array(Image.open(path).convert("RGB"))
images.append(arr)
else:
path = None
for fn in candidates:
candidate = gt_dir / fn
if candidate.exists():
path = candidate
break
if path is None:
raise FileNotFoundError(
f"GT image not found in {gt_dir}. Tried: {', '.join(candidates)}"
)
images.append(np.array(Image.open(path).convert("RGB")))
logger.info(f"Loaded {len(images)} GT images for {case_id} from {gt_dir}")
else:
for fn in filenames:
url = f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{fn}"
resp = requests.get(url, timeout=30)
if resp.status_code != 200:
raise FileNotFoundError(f"GT image not found: {url}")
images.append(np.array(Image.open(io.BytesIO(resp.content)).convert("RGB")))
logger.info(f"Loaded {len(images)} GT images for {case_id} from sglang-ci-data")
embeddings = [compute_clip_embedding(arr) for arr in images]
loaded_gt = LoadedConsistencyGT(images=images, embeddings=embeddings)
_consistency_gt_cache[cache_key] = loaded_gt
return loaded_gt
def load_gt_embeddings(
case_id: str,
num_gpus: int,
is_video: bool = False,
output_format: str | None = None,
) -> list[np.ndarray]:
"""Load GT images and convert them into CLIP embeddings."""
return load_consistency_gt(
case_id=case_id,
num_gpus=num_gpus,
is_video=is_video,
output_format=output_format,
).embeddings
def gt_exists(
case_id: str,
num_gpus: int,
is_video: bool = False,
output_format: str | None = None,
) -> bool:
"""Check whether GT image(s) exist."""
gt_dir = _get_consistency_gt_dir()
if gt_dir is not None:
candidates = get_consistency_gt_candidates(
case_id, num_gpus, is_video, output_format
)
if is_video:
return all((gt_dir / c).exists() for c in candidates)
return any((gt_dir / c).exists() for c in candidates)
filenames = _consistency_gt_filenames(case_id, num_gpus, is_video, output_format)
fn = filenames[0]
url = f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{fn}"
try:
r = requests.head(url, timeout=10)
return r.status_code == 200
except Exception:
return False
def extract_key_frames_from_video(
video_bytes: bytes,
num_frames: int | None = None,
@@ -643,3 +1032,114 @@ def image_bytes_to_numpy(image_bytes: bytes) -> np.ndarray:
"""Convert image bytes to numpy array."""
img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
return np.array(img)
def compare_with_gt(
output_frames: list[np.ndarray],
gt_data: LoadedConsistencyGT,
thresholds: ConsistencyThresholds,
case_id: str,
) -> ConsistencyResult:
"""Compare output frames with GT using CLIP and pixel-level metrics."""
if len(output_frames) != len(gt_data.embeddings):
raise ValueError(
f"Frame count mismatch: output={len(output_frames)}, gt={len(gt_data.embeddings)}"
)
similarity_scores = []
frame_metrics: list[FrameConsistencyMetrics] = []
for i, (out_frame, gt_frame, gt_emb) in enumerate(
zip(output_frames, gt_data.images, gt_data.embeddings)
):
out_frame = _ensure_rgb_uint8_image(out_frame)
gt_frame = _ensure_rgb_uint8_image(gt_frame)
if out_frame.shape != gt_frame.shape:
raise ValueError(
f"Frame shape mismatch for case {case_id}, frame {i}: "
f"output={out_frame.shape}, gt={gt_frame.shape}"
)
out_emb = compute_clip_embedding(out_frame)
clip_similarity = compute_clip_similarity(out_emb, gt_emb)
ssim = compute_ssim(out_frame, gt_frame)
psnr = compute_psnr(out_frame, gt_frame)
mean_abs_diff = compute_mean_abs_diff(out_frame, gt_frame)
similarity_scores.append(clip_similarity)
frame_metrics.append(
FrameConsistencyMetrics(
frame_index=i,
clip_similarity=clip_similarity,
ssim=ssim,
psnr=psnr,
mean_abs_diff=mean_abs_diff,
clip_passed=clip_similarity >= thresholds.clip_threshold,
ssim_passed=ssim >= thresholds.ssim_threshold,
psnr_passed=psnr >= thresholds.psnr_threshold,
mean_abs_diff_passed=(
mean_abs_diff <= thresholds.mean_abs_diff_threshold
),
)
)
min_similarity = min(similarity_scores)
min_ssim = min(metric.ssim for metric in frame_metrics)
min_psnr = min(metric.psnr for metric in frame_metrics)
max_mean_abs_diff = max(metric.mean_abs_diff for metric in frame_metrics)
passed = all(
metric.clip_passed
and metric.ssim_passed
and metric.psnr_passed
and metric.mean_abs_diff_passed
for metric in frame_metrics
)
result = ConsistencyResult(
case_id=case_id,
passed=passed,
similarity_scores=similarity_scores,
min_similarity=min_similarity,
threshold=thresholds.clip_threshold,
min_ssim=min_ssim,
min_psnr=min_psnr,
max_mean_abs_diff=max_mean_abs_diff,
thresholds=thresholds,
frame_metrics=frame_metrics,
)
status = "PASSED" if passed else "FAILED"
print(f"\n{'=' * 60}")
print(f"[CLIP Consistency] {case_id}: {status}")
print(
" Thresholds: "
f"clip>={thresholds.clip_threshold}, "
f"ssim>={thresholds.ssim_threshold}, "
f"psnr>={thresholds.psnr_threshold}, "
f"mean_abs_diff<={thresholds.mean_abs_diff_threshold}"
)
print(f" Min similarity: {min_similarity:.4f}")
print(f" Min SSIM: {min_ssim:.4f}")
print(f" Min PSNR: {min_psnr:.4f}")
print(f" Max mean_abs_diff: {max_mean_abs_diff:.4f}")
print(" Frame details:")
for metric in frame_metrics:
frame_status = (
"PASS"
if (
metric.clip_passed
and metric.ssim_passed
and metric.psnr_passed
and metric.mean_abs_diff_passed
)
else "FAIL"
)
print(
f" Frame {metric.frame_index}: "
f"clip={metric.clip_similarity:.4f} "
f"ssim={metric.ssim:.4f} "
f"psnr={metric.psnr:.4f} "
f"mean_abs_diff={metric.mean_abs_diff:.4f} "
f"{frame_status}"
)
print(f"{'=' * 60}\n")
return result