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

[diffusion] CI: minor refactor CI for less code duplication (#13905)

Browse Source
This commit is contained in:
Mick
2025-11-25 18:44:11 +08:00
committed by GitHub
parent 59464dbf15
commit 6f094deff0
11 changed files with 516 additions and 597 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
View File

@@ -105,7 +105,7 @@ def shard_rotary_emb_for_sp(emb):
class PipelineConfig:
"""The base configuration class for a generation pipeline."""
task_type: ModelTaskType
task_type: ModelTaskType = ModelTaskType.I2I
model_path: str = ""
pipeline_config_path: str | None = None

6
python/sglang/multimodal_gen/configs/sample/base.py
View File

@@ -318,6 +318,12 @@ class SamplingParams:
return sampling_params
def output_size_str(self) -> str:
return f"{self.width}x{self.height}"
def seconds(self) -> float:
return self.num_frames / self.fps
@staticmethod
def add_cli_args(parser: Any) -> Any:
"""Add CLI arguments for SamplingParam fields"""

37
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -8,7 +8,6 @@ This module provides a consolidated interface for generating videos using
diffusion models.
"""
import logging
import multiprocessing as mp
import os
import time
@@ -21,23 +20,18 @@ import torch
import torchvision
from einops import rearrange
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
# Suppress verbose logging from imageio, which is triggered when saving images.
logging.getLogger("imageio").setLevel(logging.WARNING)
logging.getLogger("imageio_ffmpeg").setLevel(logging.WARNING)
# Suppress Pillow plugin import logs when app log level is DEBUG
logging.getLogger("PIL").setLevel(logging.WARNING)
logging.getLogger("PIL.Image").setLevel(logging.WARNING)
from sglang.multimodal_gen.configs.sample.base import DataType, SamplingParams
from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
from sglang.multimodal_gen.runtime.launch_server import launch_server
from sglang.multimodal_gen.runtime.managers.schedulerbase import SchedulerBase
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.logging_utils import (
init_logger,
suppress_other_loggers,
)
logger = init_logger(__name__)
@@ -185,15 +179,16 @@ class DiffGenerator:
if save_output:
if save_file_path:
os.makedirs(os.path.dirname(save_file_path), exist_ok=True)
if data_type == DataType.VIDEO:
imageio.mimsave(
save_file_path,
frames,
fps=fps,
format=data_type.get_default_extension(),
)
else:
imageio.imwrite(save_file_path, frames[0])
with suppress_other_loggers():
if data_type == DataType.VIDEO:
imageio.mimsave(
save_file_path,
frames,
fps=fps,
format=data_type.get_default_extension(),
)
else:
imageio.imwrite(save_file_path, frames[0])
logger.info("Saved output to %s", save_file_path)
else:
logger.warning("No output path provided, output not saved")

5
python/sglang/multimodal_gen/runtime/entrypoints/utils.py
View File

@@ -8,11 +8,6 @@ This module provides a consolidated interface for generating videos using
diffusion models.
"""
import logging
# Suppress verbose logging from imageio, which is triggered when saving images.
logging.getLogger("imageio").setLevel(logging.WARNING)
logging.getLogger("imageio_ffmpeg").setLevel(logging.WARNING)
from sglang.multimodal_gen.configs.sample.base import SamplingParams
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req

2
python/sglang/multimodal_gen/runtime/platforms/cuda.py
View File

@@ -276,7 +276,7 @@ class CudaPlatformBase(Platform):
return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
logger.info("Using FlashAttention (FA3 for hopper, FA4 for blackwell) backend.")
logger.info("Using FlashAttention (FA3 for hopper, FA4 for blackwell) backend")
return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend"

7
python/sglang/multimodal_gen/runtime/utils/logging_utils.py
View File

@@ -409,6 +409,13 @@ def suppress_other_loggers(not_suppress_on_main_rank: bool = False):
original_levels[logger_name] = logger.level
logger.setLevel(logging.WARNING)
# Suppress verbose logging from imageio, which is triggered when saving images.
logging.getLogger("imageio").setLevel(logging.WARNING)
logging.getLogger("imageio_ffmpeg").setLevel(logging.WARNING)
# Suppress Pillow plugin import logs when app log level is DEBUG
logging.getLogger("PIL").setLevel(logging.WARNING)
logging.getLogger("PIL.Image").setLevel(logging.WARNING)
try:
yield
finally:

301
python/sglang/multimodal_gen/test/server/test_server_common.py
View File

@@ -7,12 +7,11 @@ If the actual run is significantly better than the baseline, the improved cases
from __future__ import annotations
import base64
import os
import time
from pathlib import Path
from typing import Any, Callable
import openai
import pytest
from openai import OpenAI
@@ -26,6 +25,7 @@ from sglang.multimodal_gen.test.server.test_server_utils import (
ServerManager,
WarmupRunner,
download_image_from_url,
get_generate_fn,
)
from sglang.multimodal_gen.test.server.testcase_configs import (
BASELINE_CONFIG,
@@ -33,12 +33,10 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
PerformanceSummary,
ScenarioConfig,
)
from sglang.multimodal_gen.test.slack_utils import upload_file_to_slack
from sglang.multimodal_gen.test.test_utils import (
get_dynamic_server_port,
is_image_url,
read_perf_logs,
validate_image,
validate_openai_video,
wait_for_req_perf_record,
)
@@ -50,13 +48,16 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
"""Start a diffusion server for a single case and tear it down afterwards."""
default_port = get_dynamic_server_port()
port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port))
server_args = case.server_args
sampling_params = case.sampling_params
extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "")
extra_args += f" --num-gpus {case.num_gpus} --ulysses-degree {case.num_gpus}"
extra_args += (
f" --num-gpus {server_args.num_gpus} --ulysses-degree {server_args.num_gpus}"
)
# start server
manager = ServerManager(
model=case.model_path,
model=server_args.model_path,
port=port,
wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")),
extra_args=extra_args,
@@ -64,25 +65,31 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
ctx = manager.start()
try:
# Reconstruct output size for OpenAI API
output_size = sampling_params.output_size
warmup = WarmupRunner(
port=ctx.port,
model=case.model_path,
prompt=case.prompt or "A colorful raccoon icon",
output_size=case.output_size,
model=server_args.model_path,
prompt=sampling_params.prompt or "A colorful raccoon icon",
output_size=output_size,
)
warmup.run_text_warmups(case.warmup_text)
warmup.run_text_warmups(case.server_args.warmup_text)
if case.warmup_edit > 0 and case.edit_prompt and case.image_path:
if (
case.server_args.warmup_edit > 0
and case.sampling_params.prompt
and sampling_params.image_path
):
# Handle URL or local path
image_path = case.image_path
if case.is_image_url():
image_path = download_image_from_url(str(case.image_path))
image_path = sampling_params.image_path
if is_image_url(sampling_params.image_path):
image_path = download_image_from_url(str(sampling_params.image_path))
else:
image_path = Path(case.image_path)
image_path = Path(sampling_params.image_path)
warmup.run_edit_warmups(
count=case.warmup_edit,
edit_prompt=case.edit_prompt,
count=case.server_args.warmup_edit,
edit_prompt=case.sampling_params.prompt,
image_path=image_path,
)
except Exception as exc:
@@ -141,14 +148,16 @@ Consider updating perf_baselines.json with the snippets below:
def run_and_collect(
self,
ctx: ServerContext,
generate_fn: Callable[[], str],
case_id: str,
generate_fn: Callable[[str, openai.Client], str],
) -> RequestPerfRecord:
"""Run generation and collect performance records."""
log_path = ctx.perf_log_path
prev_len = len(read_perf_logs(log_path))
log_wait_timeout = 30
rid = generate_fn()
client = self._client(ctx)
rid = generate_fn(case_id, client)
req_perf_record, _ = wait_for_req_perf_record(
rid,
@@ -159,241 +168,6 @@ Consider updating perf_baselines.json with the snippets below:
return req_perf_record
def get_generate_fn(
self,
ctx: ServerContext,
case: DiffusionTestCase,
) -> Callable[[], str]:
"""Return appropriate generation function for the case."""
client = self._client(ctx)
def _create_and_download_video(
*,
model: str,
size: str,
prompt: str | None = None,
seconds: int | None = None,
input_reference: Any | None = None,
) -> str:
"""
Create a video job via /v1/videos, poll until completion,
then download the binary content and validate it.
"""
create_kwargs: dict[str, Any] = {
"model": model,
"size": size,
}
if prompt is not None:
create_kwargs["prompt"] = prompt
if seconds is not None:
create_kwargs["seconds"] = seconds
if input_reference is not None:
create_kwargs["input_reference"] = input_reference # triggers multipart
# create video job
job = client.videos.create(**create_kwargs) # type: ignore[attr-defined]
video_id = job.id
job_completed = False
is_baseline_generation_mode = (
os.environ.get("SGLANG_GEN_BASELINE", "0") == "1"
)
timeout = 3600.0 if is_baseline_generation_mode else 1200.0
deadline = time.time() + timeout
while True:
page = client.videos.list() # type: ignore[attr-defined]
item = next((v for v in page.data if v.id == video_id), None)
if item and getattr(item, "status", None) == "completed":
job_completed = True
break
if time.time() > deadline:
break
time.sleep(1)
if not job_completed:
if is_baseline_generation_mode:
logger.warning(
f"{case.id}: video job {video_id} timed out during baseline generation. "
"Attempting to collect performance data anyway."
)
return video_id
pytest.fail(f"{case.id}: video job {video_id} did not complete in time")
# download video
resp = client.videos.download_content(video_id=video_id) # type: ignore[attr-defined]
content = resp.read()
validate_openai_video(content)
tmp_path = f"{video_id}.mp4"
with open(tmp_path, "wb") as f:
f.write(content)
upload_file_to_slack(
case_id=case.id,
model=case.model_path,
prompt=case.prompt,
file_path=tmp_path,
origin_file_path=case.image_path,
)
os.remove(tmp_path)
return video_id
# for all tests, seconds = case.seconds or fallback 4 seconds
video_seconds = case.seconds or 4
# -------------------------
# IMAGE MODE
# -------------------------
def generate_image() -> str:
"""T2I: Text to Image generation."""
if not case.prompt:
pytest.skip(f"{case.id}: no text prompt configured")
response = client.images.with_raw_response.generate(
model=case.model_path,
prompt=case.prompt,
n=1,
size=case.output_size,
response_format="b64_json",
)
result = response.parse()
validate_image(result.data[0].b64_json)
img_data = base64.b64decode(result.data[0].b64_json)
tmp_path = f"{result.created}.png"
with open(tmp_path, "wb") as f:
f.write(img_data)
upload_file_to_slack(
case_id=case.id,
model=case.model_path,
prompt=case.prompt,
file_path=tmp_path,
)
os.remove(tmp_path)
return str(result.created)
def generate_image_edit() -> str:
"""TI2I: Text + Image ? Image edit."""
if not case.edit_prompt or not case.image_path:
pytest.skip(f"{case.id}: no edit config")
# Handle URL or local path
if case.is_image_url():
image_path = download_image_from_url(str(case.image_path))
else:
image_path = Path(case.image_path)
if not image_path.exists():
pytest.skip(f"{case.id}: file missing: {image_path}")
with image_path.open("rb") as fh:
response = client.images.with_raw_response.edit(
model=case.model_path,
image=fh,
prompt=case.edit_prompt,
n=1,
size=case.output_size,
response_format="b64_json",
)
rid = response.headers.get("x-request-id", "")
result = response.parse()
validate_image(result.data[0].b64_json)
img_data = base64.b64decode(result.data[0].b64_json)
tmp_path = f"{rid}.png"
with open(tmp_path, "wb") as f:
f.write(img_data)
upload_file_to_slack(
case_id=case.id,
model=case.model_path,
prompt=case.edit_prompt,
file_path=tmp_path,
origin_file_path=case.image_path,
)
os.remove(tmp_path)
return rid
# -------------------------
# VIDEO MODE
# -------------------------
def generate_video() -> str:
"""T2V: Text ? Video."""
if not case.prompt:
pytest.skip(f"{case.id}: no text prompt configured")
return _create_and_download_video(
model=case.model_path,
prompt=case.prompt,
size=case.output_size,
seconds=video_seconds,
)
def generate_image_to_video() -> str:
"""I2V: Image ? Video (optional prompt)."""
if not case.image_path:
pytest.skip(f"{case.id}: no input image configured")
# Handle URL or local path
if case.is_image_url():
image_path = download_image_from_url(str(case.image_path))
else:
image_path = Path(case.image_path)
if not image_path.exists():
pytest.skip(f"{case.id}: file missing: {image_path}")
with image_path.open("rb") as fh:
return _create_and_download_video(
model=case.model_path,
prompt=case.edit_prompt,
size=case.output_size,
seconds=video_seconds,
input_reference=fh,
)
def generate_text_image_to_video() -> str:
"""TI2V: Text + Image ? Video."""
if not case.edit_prompt or not case.image_path:
pytest.skip(f"{case.id}: no edit config")
# Handle URL or local path
if case.is_image_url():
image_path = download_image_from_url(str(case.image_path))
else:
image_path = Path(case.image_path)
if not image_path.exists():
pytest.skip(f"{case.id}: file missing: {image_path}")
with image_path.open("rb") as fh:
return _create_and_download_video(
model=case.model_path,
prompt=case.edit_prompt,
size=case.output_size,
seconds=video_seconds,
input_reference=fh,
)
if case.modality == "video":
if case.image_path and case.edit_prompt:
return generate_text_image_to_video
elif case.image_path:
return generate_image_to_video
else:
return generate_video
# Image modality
if case.edit_prompt and case.image_path:
return generate_image_edit
return generate_image
def _validate_and_record(
self,
case: DiffusionTestCase,
@@ -420,7 +194,7 @@ Consider updating perf_baselines.json with the snippets below:
if not is_baseline_generation_mode:
missing_scenario = True
validator_name = case.custom_validator or "default"
validator_name = case.server_args.custom_validator or "default"
validator_class = VALIDATOR_REGISTRY.get(validator_name, PerformanceValidator)
validator = validator_class(
@@ -440,7 +214,7 @@ Consider updating perf_baselines.json with the snippets below:
self._check_for_improvement(case, summary, scenario)
try:
validator.validate(perf_record, case.num_frames)
validator.validate(perf_record, case.sampling_params.num_frames)
except AssertionError as e:
logger.error(f"Performance validation failed for {case.id}:\n{e}")
self._dump_baseline_for_testcase(case, summary, missing_scenario)
@@ -448,7 +222,7 @@ Consider updating perf_baselines.json with the snippets below:
result = {
"test_name": case.id,
"modality": case.modality,
"modality": case.server_args.modality,
"e2e_ms": summary.e2e_ms,
"avg_denoise_ms": summary.avg_denoise_ms,
"median_denoise_ms": summary.median_denoise_ms,
@@ -569,7 +343,7 @@ Consider updating perf_baselines.json with the snippets below:
}
# Video-specific metrics
if case.modality == "video":
if case.server_args.modality == "video":
if "per_frame_generation" not in baseline["stages_ms"]:
baseline["stages_ms"]["per_frame_generation"] = (
round(summary.avg_frame_time_ms, 2)
@@ -598,9 +372,14 @@ Consider updating perf_baselines.json with the snippets below:
- test_diffusion_perf[qwen_image_edit]
- etc.
"""
generate_fn = self.get_generate_fn(diffusion_server, case)
generate_fn = get_generate_fn(
model_path=case.server_args.model_path,
modality=case.server_args.modality,
sampling_params=case.sampling_params,
)
perf_record = self.run_and_collect(
diffusion_server,
case.id,
generate_fn,
)
self._validate_and_record(case, perf_record)

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

@@ -4,6 +4,7 @@ Server management and performance validation for diffusion tests.
from __future__ import annotations
import base64
import os
import shlex
import subprocess
@@ -13,21 +14,29 @@ import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Sequence
from typing import Any, Callable, Sequence
from urllib.request import urlopen
from openai import OpenAI
import pytest
from openai import Client, OpenAI
from sglang.multimodal_gen.benchmarks.compare_perf import calculate_upper_bound
from sglang.multimodal_gen.runtime.utils.common import kill_process_tree
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.server.testcase_configs import (
DiffusionSamplingParams,
PerformanceSummary,
ScenarioConfig,
ToleranceConfig,
)
from sglang.multimodal_gen.test.test_utils import prepare_perf_log, validate_image
from sglang.multimodal_gen.test.slack_utils import upload_file_to_slack
from sglang.multimodal_gen.test.test_utils import (
is_image_url,
prepare_perf_log,
validate_image,
validate_openai_video,
)
logger = init_logger(__name__)
@@ -456,3 +465,232 @@ VALIDATOR_REGISTRY = {
"default": PerformanceValidator,
"video": VideoPerformanceValidator,
}
def get_generate_fn(
model_path: str,
modality: str,
sampling_params: DiffusionSamplingParams,
) -> Callable[[str, Client], str]:
"""Return appropriate generation function for the case."""
def _create_and_download_video(
client,
case_id,
*,
model: str,
size: str,
prompt: str | None = None,
seconds: int | None = None,
input_reference: Any | None = None,
) -> str:
"""
Create a video job via /v1/videos, poll until completion,
then download the binary content and validate it.
"""
create_kwargs: dict[str, Any] = {
"model": model,
"size": size,
}
if prompt is not None:
create_kwargs["prompt"] = prompt
if seconds is not None:
create_kwargs["seconds"] = seconds
if input_reference is not None:
create_kwargs["input_reference"] = input_reference # triggers multipart
job = client.videos.create(**create_kwargs) # type: ignore[attr-defined]
video_id = job.id
job_completed = False
is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1"
timeout = 3600.0 if is_baseline_generation_mode else 1200.0
deadline = time.time() + timeout
while True:
page = client.videos.list() # type: ignore[attr-defined]
item = next((v for v in page.data if v.id == video_id), None)
if item and getattr(item, "status", None) == "completed":
job_completed = True
break
if time.time() > deadline:
break
time.sleep(1)
if not job_completed:
if is_baseline_generation_mode:
logger.warning(
f"{id}: video job {video_id} timed out during baseline generation. "
"Attempting to collect performance data anyway."
)
return video_id
pytest.fail(f"{id}: video job {video_id} did not complete in time")
# download video
resp = client.videos.download_content(video_id=video_id) # type: ignore[attr-defined]
content = resp.read()
validate_openai_video(content)
tmp_path = f"{video_id}.mp4"
with open(tmp_path, "wb") as f:
f.write(content)
upload_file_to_slack(
case_id=case_id,
model=model_path,
prompt=sampling_params.prompt,
file_path=tmp_path,
origin_file_path=sampling_params.image_path,
)
os.remove(tmp_path)
return video_id
video_seconds = sampling_params.seconds or 4
def generate_image(case_id, client) -> str:
"""T2I: Text to Image generation."""
if not sampling_params.prompt:
pytest.skip(f"{id}: no text prompt configured")
response = client.images.with_raw_response.generate(
model=model_path,
prompt=sampling_params.prompt,
n=1,
size=sampling_params.output_size,
response_format="b64_json",
)
result = response.parse()
validate_image(result.data[0].b64_json)
img_data = base64.b64decode(result.data[0].b64_json)
tmp_path = f"{result.created}.png"
with open(tmp_path, "wb") as f:
f.write(img_data)
upload_file_to_slack(
case_id=case_id,
model=model_path,
prompt=sampling_params.prompt,
file_path=tmp_path,
)
os.remove(tmp_path)
return str(result.created)
def generate_image_edit(case_id, client) -> str:
"""TI2I: Text + Image ? Image edit."""
if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config")
if is_image_url(sampling_params.image_path):
image_path = download_image_from_url(str(sampling_params.image_path))
else:
image_path = Path(sampling_params.image_path)
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
with image_path.open("rb") as fh:
response = client.images.with_raw_response.edit(
model=model_path,
image=fh,
prompt=sampling_params.prompt,
n=1,
size=sampling_params.output_size,
response_format="b64_json",
)
rid = response.headers.get("x-request-id", "")
result = response.parse()
validate_image(result.data[0].b64_json)
img_data = base64.b64decode(result.data[0].b64_json)
tmp_path = f"{rid}.png"
with open(tmp_path, "wb") as f:
f.write(img_data)
upload_file_to_slack(
case_id=case_id,
model=model_path,
prompt=sampling_params.prompt,
file_path=tmp_path,
origin_file_path=sampling_params.image_path,
)
os.remove(tmp_path)
return rid
def generate_video(case_id, client) -> str:
"""T2V: Text ? Video."""
if not sampling_params.prompt:
pytest.skip(f"{id}: no text prompt configured")
return _create_and_download_video(
client,
case_id,
model=model_path,
prompt=sampling_params.prompt,
size=sampling_params.output_size,
seconds=video_seconds,
)
def generate_image_to_video(case_id, client) -> str:
"""I2V: Image ? Video (optional prompt)."""
if not sampling_params.image_path:
pytest.skip(f"{id}: no input image configured")
if is_image_url(sampling_params.image_path):
image_path = download_image_from_url(str(sampling_params.image_path))
else:
image_path = Path(sampling_params.image_path)
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
with image_path.open("rb") as fh:
return _create_and_download_video(
client,
case_id,
model=model_path,
prompt=sampling_params.prompt,
size=sampling_params.output_size,
seconds=video_seconds,
input_reference=fh,
)
def generate_text_image_to_video(case_id, client) -> str:
"""TI2V: Text + Image ? Video."""
if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config")
if is_image_url(sampling_params.image_path):
image_path = download_image_from_url(str(sampling_params.image_path))
else:
image_path = Path(sampling_params.image_path)
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
with image_path.open("rb") as fh:
return _create_and_download_video(
client,
case_id,
model=model_path,
prompt=sampling_params.prompt,
size=sampling_params.output_size,
seconds=video_seconds,
input_reference=fh,
)
if modality == "video":
if sampling_params.image_path and sampling_params.prompt:
fn = generate_text_image_to_video
elif sampling_params.image_path:
fn = generate_image_to_video
else:
fn = generate_video
elif sampling_params.prompt and sampling_params.image_path:
fn = generate_image_edit
else:
fn = generate_image
return fn

338
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -108,17 +108,27 @@ class BaselineConfig:
@dataclass(frozen=True)
class DiffusionTestCase:
class DiffusionServerArgs:
"""Configuration for a single model/scenario test case."""
id: str # pytest test id and scenario name
model_path: str # HF repo or local path
modality: str = "image" # "image" or "video" or "3d"
warmup_text: int = 1 # number of text-to-image/video warmups
warmup_edit: int = 0 # number of image/video-edit warmups
custom_validator: str | None = None # optional custom validator name
# resources
num_gpus: int = 1
@dataclass(frozen=True)
class DiffusionSamplingParams:
"""Configuration for a single model/scenario test case."""
output_size: str = "1024x1024" # output image dimensions (or video resolution)
# inputs and conditioning
prompt: str | None = None # text prompt for generation
edit_prompt: str | None = None # prompt for editing
image_path: Path | str | None = None # input image/video for editing (Path or URL)
# duration
@@ -126,21 +136,14 @@ class DiffusionTestCase:
num_frames: int | None = None # for video: number of frames
fps: int | None = None # for video: frames per second
warmup_text: int = 1 # number of text-to-image/video warmups
warmup_edit: int = 0 # number of image/video-edit warmups
custom_validator: str | None = None # optional custom validator name
# resources
num_gpus: int = 1
@dataclass(frozen=True)
class DiffusionTestCase:
"""Configuration for a single model/scenario test case."""
def is_image_url(self) -> bool:
"""Check if image_edit_path is a URL."""
if self.image_path is None:
return False
return isinstance(self.image_path, str) and (
self.image_path.startswith("http://")
or self.image_path.startswith("https://")
)
id: str # pytest test id and scenario name
server_args: DiffusionServerArgs
sampling_params: DiffusionSamplingParams
def sample_step_indices(
@@ -214,49 +217,63 @@ class PerformanceSummary:
ONE_GPU_CASES_A: list[DiffusionTestCase] = [
# === Text to Image (T2I) ===
DiffusionTestCase(
id="qwen_image_t2i",
model_path="Qwen/Qwen-Image",
modality="image",
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
warmup_text=1,
warmup_edit=0,
"qwen_image_t2i",
DiffusionServerArgs(
model_path="Qwen/Qwen-Image",
modality="image",
warmup_text=1,
warmup_edit=0,
),
DiffusionSamplingParams(
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
),
),
DiffusionTestCase(
id="flux_image_t2i",
model_path="black-forest-labs/FLUX.1-dev",
modality="image",
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
warmup_text=1,
warmup_edit=0,
"flux_image_t2i",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.1-dev",
modality="image",
warmup_text=1,
warmup_edit=0,
),
DiffusionSamplingParams(
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
),
),
# === Text and Image to Image (TI2I) ===
DiffusionTestCase(
id="qwen_image_edit_ti2i",
model_path="Qwen/Qwen-Image-Edit",
modality="image",
prompt=None, # not used for editing
output_size="1024x1536",
warmup_text=0,
warmup_edit=1,
edit_prompt="Convert 2D style to 3D style",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
"qwen_image_edit_ti2i",
DiffusionServerArgs(
model_path="Qwen/Qwen-Image-Edit",
warmup_text=0,
warmup_edit=1,
modality="image",
),
DiffusionSamplingParams(
prompt="Convert 2D style to 3D style",
output_size="1024x1536",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
),
),
]
ONE_GPU_CASES_B: list[DiffusionTestCase] = [
# === Text to Video (T2V) ===
DiffusionTestCase(
id="wan2_1_t2v_1.3b",
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
modality="video",
prompt="A curious raccoon",
output_size="848x480",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
"wan2_1_t2v_1.3b",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
modality="video",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
),
DiffusionSamplingParams(
prompt="A curious raccoon",
output_size="848x480",
),
),
# NOTE(mick): flaky
# DiffusionTestCase(
@@ -270,126 +287,161 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
# custom_validator="video",
# ),
DiffusionTestCase(
id="fast_hunyuan_video",
model_path="FastVideo/FastHunyuan-diffusers",
modality="video",
prompt="A curious raccoon",
output_size="720x480",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
"fast_hunyuan_video",
DiffusionServerArgs(
model_path="FastVideo/FastHunyuan-diffusers",
modality="video",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
),
DiffusionSamplingParams(
prompt="A curious raccoon",
output_size="720x480",
),
),
# === Text and Image to Video (TI2V) ===
DiffusionTestCase(
id="wan2_2_ti2v_5b",
model_path="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
modality="video",
output_size="832x1104",
prompt="Animate this image",
edit_prompt="Add dynamic motion to the scene",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
"wan2_2_ti2v_5b",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
modality="video",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
),
DiffusionSamplingParams(
output_size="832x1104",
prompt="Add dynamic motion to the scene",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
),
),
DiffusionTestCase(
id="fastwan2_2_ti2v_5b",
model_path="FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
modality="video",
output_size="832x1104",
prompt="Animate this image",
edit_prompt="Add dynamic motion to the scene",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
"fastwan2_2_ti2v_5b",
DiffusionServerArgs(
model_path="FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
modality="video",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
),
DiffusionSamplingParams(
output_size="832x1104",
prompt="Add dynamic motion to the scene",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
),
),
]
TWO_GPU_CASES_A = [
DiffusionTestCase(
id="wan2_2_i2v_a14b_2gpu",
model_path="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
modality="video",
prompt="generate",
warmup_text=0,
warmup_edit=0,
output_size="832x1104",
edit_prompt="generate",
image_path="https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true",
custom_validator="video",
num_gpus=2,
num_frames=1,
"wan2_2_i2v_a14b_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
modality="video",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
num_gpus=2,
),
DiffusionSamplingParams(
prompt="generate",
output_size="832x1104",
image_path="https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true",
num_frames=1,
),
),
DiffusionTestCase(
id="wan2_2_t2v_a14b_2gpu",
model_path="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
modality="video",
prompt="A curious raccoon",
output_size="720x480",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
num_gpus=2,
"wan2_2_t2v_a14b_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
modality="video",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
num_gpus=2,
),
DiffusionSamplingParams(
prompt="A curious raccoon",
output_size="720x480",
),
),
DiffusionTestCase(
id="wan2_1_t2v_14b_2gpu",
model_path="Wan-AI/Wan2.1-T2V-14B-Diffusers",
modality="video",
prompt="A curious raccoon",
output_size="720x480",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
num_gpus=2,
"wan2_1_t2v_14b_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-14B-Diffusers",
warmup_text=0,
warmup_edit=0,
modality="video",
num_gpus=2,
custom_validator="video",
),
DiffusionSamplingParams(
prompt="A curious raccoon",
output_size="720x480",
),
),
]
TWO_GPU_CASES_B = [
DiffusionTestCase(
id="wan2_1_i2v_14b_480P_2gpu",
model_path="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
output_size="832x1104",
modality="video",
prompt="Animate this image",
edit_prompt="Add dynamic motion to the scene",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
num_gpus=2,
"wan2_1_i2v_14b_480P_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
warmup_text=0,
warmup_edit=0,
modality="video",
custom_validator="video",
num_gpus=2,
),
DiffusionSamplingParams(
output_size="832x1104",
prompt="Add dynamic motion to the scene",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
),
),
DiffusionTestCase(
id="wan2_1_i2v_14b_720P_2gpu",
model_path="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
modality="video",
prompt="Animate this image",
edit_prompt="Add dynamic motion to the scene",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
output_size="832x1104",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
num_gpus=2,
"wan2_1_i2v_14b_720P_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
modality="video",
warmup_text=0,
warmup_edit=0,
custom_validator="video",
num_gpus=2,
),
DiffusionSamplingParams(
prompt="Add dynamic motion to the scene",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
output_size="832x1104",
),
),
DiffusionTestCase(
id="qwen_image_t2i_2_gpus",
model_path="Qwen/Qwen-Image",
modality="image",
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
warmup_text=1,
warmup_edit=0,
num_gpus=2,
"qwen_image_t2i_2_gpus",
DiffusionServerArgs(
model_path="Qwen/Qwen-Image",
modality="image",
warmup_text=1,
warmup_edit=0,
num_gpus=2,
),
DiffusionSamplingParams(
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
),
),
DiffusionTestCase(
id="flux_image_t2i_2_gpus",
model_path="black-forest-labs/FLUX.1-dev",
modality="image",
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
warmup_text=1,
warmup_edit=0,
"flux_image_t2i_2_gpus",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.1-dev",
modality="image",
warmup_text=1,
warmup_edit=0,
),
DiffusionSamplingParams(
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
),
),
]

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

@@ -25,6 +25,15 @@ from sglang.multimodal_gen.runtime.utils.perf_logger import (
logger = init_logger(__name__)
def is_image_url(image_path: str | Path | None) -> bool:
"""Check if image_path is a URL."""
if image_path is None:
return False
return isinstance(image_path, str) and (
image_path.startswith("http://") or image_path.startswith("https://")
)
def run_command(command) -> Optional[float]:
"""Runs a command and returns the execution time and status."""
print(f"Running command: {shlex.join(command)}")

162
python/sglang/multimodal_gen/test/utils.py
View File

@@ -1,162 +0,0 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
import json
import os
import numpy as np
import torch
from pytorch_msssim import ms_ssim, ssim
from torchvision.io import read_video
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
"""
Compute SSIM between two videos.
Args:
video1_path: Path to the first video.
video2_path: Path to the second video.
use_ms_ssim: Whether to use Multi-Scale Structural Similarity(MS-SSIM) instead of SSIM.
"""
print(f"Computing SSIM between {video1_path} and {video2_path}...")
if not os.path.exists(video1_path):
raise FileNotFoundError(f"Video1 not found: {video1_path}")
if not os.path.exists(video2_path):
raise FileNotFoundError(f"Video2 not found: {video2_path}")
frames1, _, _ = read_video(video1_path, pts_unit="sec", output_format="TCHW")
frames2, _, _ = read_video(video2_path, pts_unit="sec", output_format="TCHW")
# Ensure same number of frames
min_frames = min(frames1.shape[0], frames2.shape[0])
frames1 = frames1[:min_frames]
frames2 = frames2[:min_frames]
frames1 = frames1.float() / 255.0
frames2 = frames2.float() / 255.0
if torch.cuda.is_available():
frames1 = frames1.cuda()
frames2 = frames2.cuda()
ssim_values = []
# Process each frame individually
for i in range(min_frames):
img1 = frames1[i : i + 1]
img2 = frames2[i : i + 1]
with torch.no_grad():
if use_ms_ssim:
value = ms_ssim(img1, img2, data_range=1.0)
else:
value = ssim(img1, img2, data_range=1.0)
ssim_values.append(value.item())
if ssim_values:
mean_ssim = np.mean(ssim_values)
min_ssim = np.min(ssim_values)
max_ssim = np.max(ssim_values)
min_frame_idx = np.argmin(ssim_values)
max_frame_idx = np.argmax(ssim_values)
print(f"Mean SSIM: {mean_ssim:.4f}")
print(f"Min SSIM: {min_ssim:.4f} (at frame {min_frame_idx})")
print(f"Max SSIM: {max_ssim:.4f} (at frame {max_frame_idx})")
return mean_ssim, min_ssim, max_ssim
else:
print("No SSIM values calculated")
return 0, 0, 0
def compare_folders(reference_folder, generated_folder, use_ms_ssim=True):
"""
Compare videos with the same filename between reference_folder and generated_folder
Example usage:
results = compare_folders(reference_folder, generated_folder,
args.use_ms_ssim)
for video_name, ssim_value in results.items():
if ssim_value is not None:
print(
f"{video_name}: {ssim_value[0]:.4f}, Min SSIM: {ssim_value[1]:.4f}, Max SSIM: {ssim_value[2]:.4f}"
)
else:
print(f"{video_name}: Error during comparison")
valid_ssims = [v for v in results.values() if v is not None]
if valid_ssims:
avg_ssim = np.mean([v[0] for v in valid_ssims])
print(f"\nAverage SSIM across all videos: {avg_ssim:.4f}")
else:
print("\nNo valid SSIM values to average")
"""
reference_videos = [f for f in os.listdir(reference_folder) if f.endswith(".mp4")]
results = {}
for video_name in reference_videos:
ref_path = os.path.join(reference_folder, video_name)
gen_path = os.path.join(generated_folder, video_name)
if os.path.exists(gen_path):
print(f"\nComparing {video_name}...")
try:
ssim_value = compute_video_ssim_torchvision(
ref_path, gen_path, use_ms_ssim
)
results[video_name] = ssim_value
except Exception as e:
print(f"Error comparing {video_name}: {e}")
results[video_name] = None
else:
print(f"\nSkipping {video_name} - no matching file in generated folder")
return results
def write_ssim_results(
output_dir, ssim_values, reference_path, generated_path, num_inference_steps, prompt
):
"""
Write SSIM results to a JSON file in the same directory as the generated videos.
"""
try:
logger.info(f"Attempting to write SSIM results to directory: {output_dir}")
if not os.path.exists(output_dir):
os.makedirs(output_dir, exist_ok=True)
mean_ssim, min_ssim, max_ssim = ssim_values
result = {
"mean_ssim": mean_ssim,
"min_ssim": min_ssim,
"max_ssim": max_ssim,
"reference_video": reference_path,
"generated_video": generated_path,
"parameters": {
"num_inference_steps": num_inference_steps,
"prompt": prompt,
},
}
test_name = f"steps{num_inference_steps}_{prompt[:100]}"
result_file = os.path.join(output_dir, f"{test_name}_ssim.json")
logger.info(f"Writing JSON results to: {result_file}")
with open(result_file, "w") as f:
json.dump(result, f, indent=2)
logger.info(f"SSIM results written to {result_file}")
return True
except Exception as e:
logger.error(f"ERROR writing SSIM results: {str(e)}")
return False