From 6f094deff035579a3984b9950e8ba0dad779215a Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 25 Nov 2025 18:44:11 +0800 Subject: [PATCH] [diffusion] CI: minor refactor CI for less code duplication (#13905) --- .../configs/pipeline_configs/base.py | 2 +- .../multimodal_gen/configs/sample/base.py | 6 + .../entrypoints/diffusion_generator.py | 37 +- .../runtime/entrypoints/utils.py | 5 - .../multimodal_gen/runtime/platforms/cuda.py | 2 +- .../runtime/utils/logging_utils.py | 7 + .../test/server/test_server_common.py | 301 +++------------- .../test/server/test_server_utils.py | 244 ++++++++++++- .../test/server/testcase_configs.py | 338 ++++++++++-------- .../sglang/multimodal_gen/test/test_utils.py | 9 + python/sglang/multimodal_gen/test/utils.py | 162 --------- 11 files changed, 516 insertions(+), 597 deletions(-) delete mode 100644 python/sglang/multimodal_gen/test/utils.py diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 55b3db3bc..929201877 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/sample/base.py b/python/sglang/multimodal_gen/configs/sample/base.py index 18b4ea276..71e9caa26 100644 --- a/python/sglang/multimodal_gen/configs/sample/base.py +++ b/python/sglang/multimodal_gen/configs/sample/base.py @@ -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""" diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 945cbe81a..96862c36e 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -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") diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index b36f51450..9f7cfe032 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 2df5495b1..90320daed 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -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" diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index 01281405b..fc7c41c61 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -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: diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 0c3590e54..f3f1e8ef2 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 3c1b419cf..4e3e788a9 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index d02b4dd1d..723f6191c 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -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", + ), ), ] diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 9446680a1..56c74246e 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -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)}") diff --git a/python/sglang/multimodal_gen/test/utils.py b/python/sglang/multimodal_gen/test/utils.py deleted file mode 100644 index b1d7620c9..000000000 --- a/python/sglang/multimodal_gen/test/utils.py +++ /dev/null @@ -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