diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index a88baffc9..489223f9a 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -980,7 +980,7 @@ STANDALONE_FILES = { # measured value that must be copied into STANDALONE_FILE_EST_TIMES. STANDALONE_FILE_EST_TIMES = { "bcg-diffusion": { - "../single_test_file/test_diffusion_bcg_zimage_turbo.py": 300.0, + "../single_test_file/test_diffusion_bcg_zimage_turbo.py": 420.0, }, "1-gpu": { "../single_test_file/test_update_weights_from_disk.py": 1200.0, diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_zimage_turbo.py b/python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_zimage_turbo.py index e035359a7..979b388a8 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_zimage_turbo.py +++ b/python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_zimage_turbo.py @@ -5,9 +5,20 @@ import tempfile import unittest from pathlib import Path -from sglang.multimodal_gen.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST +from openai import OpenAI + +from sglang.multimodal_gen.test.server.test_server_utils import ServerManager +from sglang.multimodal_gen.test.test_utils import ( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + get_dynamic_server_port, +) from sglang.test.test_utils import CustomTestCase +IMAGE_SIZE = "512x512" +PROMPT_SWITCH_NUM_INFERENCE_STEPS = 4 +PROMPT_SWITCH_SEED = 0 +BCG_CAPTURE_MARKER = "[Diffusion BCG] captured" + class TestDiffusionBCGZImageTurbo(CustomTestCase): def test_zimage_turbo_true_bcg_generate(self): @@ -94,6 +105,68 @@ class TestDiffusionBCGZImageTurbo(CustomTestCase): self.assertIn("DenoisingStage", stage_names) self.assertGreater(len(perf.get("denoise_steps_ms", [])), 0) + def test_zimage_turbo_bcg_prompt_switch_reuses_captured_graph(self): + port = get_dynamic_server_port() + extra_args = ( + "--model-type diffusion " + "--num-gpus 1 " + "--strict-ports " + "--enable-breakable-cuda-graph " + f"--warmup-resolutions {IMAGE_SIZE} " + "--bcg-text-buckets 256" + ) + manager = ServerManager( + model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + port=port, + wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")), + extra_args=extra_args, + ) + ctx = manager.start() + try: + client = OpenAI( + api_key="sglang-anything", + base_url=f"http://localhost:{ctx.port}/v1", + timeout=float( + os.environ.get("SGLANG_TEST_OPENAI_REQUEST_TIMEOUT_SECS", "600") + ), + max_retries=0, + ) + self._generate_prompt_switch_image( + client, + "a red cube on a white table, centered product photo", + ) + self._generate_prompt_switch_image( + client, + "a blue butterfly flying above green grass, watercolor illustration", + ) + finally: + server_log = ctx.stdout_file.read_text(encoding="utf-8", errors="ignore") + ctx.cleanup() + + capture_count = server_log.count(BCG_CAPTURE_MARKER) + self.assertEqual( + capture_count, + 1, + "BCG should be captured during warmup and then reused by different " + f"prompt requests. Observed {capture_count} capture log line(s).\n" + f"Server log tail:\n{ctx.log_tail(lines=400)}", + ) + + def _generate_prompt_switch_image(self, client: OpenAI, prompt: str): + response = client.images.with_raw_response.generate( + model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + prompt=prompt, + n=1, + size=IMAGE_SIZE, + response_format="b64_json", + extra_body={ + "num_inference_steps": PROMPT_SWITCH_NUM_INFERENCE_STEPS, + "seed": PROMPT_SWITCH_SEED, + }, + ) + parsed = response.parse() + self.assertTrue(parsed.data[0].b64_json) + if __name__ == "__main__": unittest.main()