diff --git a/.github/workflows/diffusion-ci-gt-gen.yml b/.github/workflows/diffusion-ci-gt-gen.yml new file mode 100644 index 000000000..ef039180b --- /dev/null +++ b/.github/workflows/diffusion-ci-gt-gen.yml @@ -0,0 +1,129 @@ +name: Diffusion CI Ground Truth Generation + +on: + workflow_dispatch: + inputs: + ref: + description: 'Git ref to checkout' + required: false + default: '' + type: string + case_ids: + description: 'Specific case IDs to run (space-separated, optional)' + required: false + default: '' + type: string + +concurrency: + group: diffusion-ci-gt-gen-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: write + actions: read + +jobs: + multimodal-diffusion-gen-1gpu: + if: github.repository == 'sgl-project/sglang' + runs-on: 1-gpu-runner + strategy: + matrix: + part: [0, 1] + timeout-minutes: 60 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Install dependencies + run: bash scripts/ci/ci_install_dependency.sh diffusion + + - name: Generate outputs + run: | + cd python + python -m sglang.multimodal_gen.test.scripts.gen_diffusion_ci_outputs \ + --suite 1-gpu \ + --partition-id ${{ matrix.part }} \ + --total-partitions 2 \ + --out-dir ./diffusion-ci-outputs \ + --continue-on-error \ + ${{ inputs.case_ids != '' && format('--case-ids {0}', inputs.case_ids) || '' }} + + - name: Upload artifact + uses: actions/upload-artifact@v4 + with: + name: diffusion-gen-1gpu-part${{ matrix.part }} + path: python/diffusion-ci-outputs + retention-days: 7 + + multimodal-diffusion-gen-2gpu: + if: github.repository == 'sgl-project/sglang' + runs-on: 2-gpu-runner + strategy: + matrix: + part: [0, 1] + timeout-minutes: 60 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Install dependencies + run: bash scripts/ci/ci_install_dependency.sh diffusion + + - name: Generate outputs + run: | + cd python + python -m sglang.multimodal_gen.test.scripts.gen_diffusion_ci_outputs \ + --suite 2-gpu \ + --partition-id ${{ matrix.part }} \ + --total-partitions 2 \ + --out-dir ./diffusion-ci-outputs \ + --continue-on-error \ + ${{ inputs.case_ids != '' && format('--case-ids {0}', inputs.case_ids) || '' }} + + - name: Upload artifact + uses: actions/upload-artifact@v4 + with: + name: diffusion-gen-2gpu-part${{ matrix.part }} + path: python/diffusion-ci-outputs + retention-days: 7 + + diffusion-ci-push: + needs: [multimodal-diffusion-gen-1gpu, multimodal-diffusion-gen-2gpu] + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-latest + steps: + - name: Checkout sgl-test-files + uses: actions/checkout@v4 + with: + repository: sgl-project/sgl-test-files + path: sgl-test-files + ref: main + token: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + pattern: diffusion-gen-* + path: combined + merge-multiple: true + + - name: Copy image and video frame files + run: | + mkdir -p sgl-test-files/diffusion-ci/consistency_gt + find combined \( -name "*.png" -o -name "*.jpg" -o -name "*.jpeg" -o -name "*.webp" \) -type f -exec cp -f {} sgl-test-files/diffusion-ci/consistency_gt/ \; + + - name: Git commit and push + env: + GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} + run: | + cd sgl-test-files + git config user.email "github-actions[bot]@users.noreply.github.com" + git config user.name "github-actions[bot]" + git remote set-url origin "https://x-access-token:${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }}@github.com/sgl-project/sgl-test-files.git" + git add diffusion-ci/consistency_gt/ + git diff --staged --quiet || git commit -m "diffusion-ci: update consistency_gt images [automated]" + git push origin main diff --git a/python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py b/python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py new file mode 100755 index 000000000..bcc82fab9 --- /dev/null +++ b/python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +""" +Generate diffusion CI outputs for consistency testing. + +This script reuses the CI test code by calling run_suite.py with SGLANG_GEN_GT=1, +ensuring that GT generation uses exactly the same code path as CI tests. + +Usage: + python gen_diffusion_ci_outputs.py --suite 1-gpu --partition-id 0 --total-partitions 2 --out-dir ./output + python gen_diffusion_ci_outputs.py --suite 1-gpu --case-ids qwen_image_t2i flux_image_t2i --out-dir ./output +""" + +import argparse +import os +import sys +from pathlib import Path + +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.test.run_suite import SUITES, collect_test_items, run_pytest + +logger = init_logger(__name__) + + +def main(): + """Main entry point.""" + parser = argparse.ArgumentParser(description="Generate diffusion CI outputs") + parser.add_argument( + "--suite", + type=str, + choices=["1-gpu", "2-gpu"], + required=True, + help="Test suite to run (1-gpu or 2-gpu)", + ) + parser.add_argument( + "--partition-id", + type=int, + required=False, + help="Partition ID for matrix partitioning (0-based)", + ) + parser.add_argument( + "--total-partitions", + type=int, + required=False, + help="Total number of partitions", + ) + parser.add_argument( + "--out-dir", + type=str, + required=True, + help="Output directory for generated files", + ) + parser.add_argument( + "--continue-on-error", + action="store_true", + help="Continue processing other cases if one fails", + ) + parser.add_argument( + "--case-ids", + type=str, + nargs="*", + required=False, + help="Specific case IDs to run (space-separated). If provided, only these cases will be run.", + ) + + args = parser.parse_args() + + # Validate partition arguments + if args.partition_id is not None and args.total_partitions is not None: + if args.partition_id < 0 or args.partition_id >= args.total_partitions: + parser.error(f"partition-id must be in range [0, {args.total_partitions})") + elif args.partition_id is not None or args.total_partitions is not None: + parser.error( + "Both --partition-id and --total-partitions must be provided together" + ) + + # Create output directory + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + # Set environment variables for GT generation mode + os.environ["SGLANG_GEN_GT"] = "1" + os.environ["SGLANG_GT_OUTPUT_DIR"] = str(out_dir.absolute()) + os.environ["SGLANG_SKIP_CONSISTENCY"] = ( + "1" # Skip consistency checks in GT gen mode + ) + + logger.info(f"GT generation mode enabled") + logger.info(f"Output directory: {out_dir}") + + # Resolve test files path (same as run_suite.py) + current_file_path = Path(__file__).resolve() + test_root_dir = current_file_path.parent.parent # scripts -> test + target_dir = test_root_dir / "server" + + # Get files from suite (same as run_suite.py) + suite_files_rel = SUITES[args.suite] + suite_files_abs = [] + for f_rel in suite_files_rel: + f_abs = target_dir / f_rel + if not f_abs.exists(): + logger.warning(f"Test file {f_rel} not found in {target_dir}. Skipping.") + continue + suite_files_abs.append(str(f_abs)) + + if not suite_files_abs: + logger.error(f"No valid test files found for suite '{args.suite}'.") + sys.exit(1) + + # Build pytest filter for case_ids if provided + filter_expr = None + if args.case_ids: + # pytest parametrized test format: test_diffusion_generation[case_id] + filters = [f"test_diffusion_generation[{case_id}]" for case_id in args.case_ids] + filter_expr = " or ".join(filters) + logger.info(f"Filtering by case IDs: {args.case_ids}") + + # Collect all test items (same as run_suite.py) + all_test_items = collect_test_items(suite_files_abs, filter_expr=filter_expr) + + if not all_test_items: + logger.warning(f"No test items found for suite '{args.suite}'.") + sys.exit(0) + + # Partition by test items (same as run_suite.py) + partition_id = args.partition_id if args.partition_id is not None else 0 + total_partitions = args.total_partitions if args.total_partitions is not None else 1 + + my_items = [ + item + for i, item in enumerate(all_test_items) + if i % total_partitions == partition_id + ] + + logger.info( + f"Partition {partition_id}/{total_partitions}: " + f"running {len(my_items)} of {len(all_test_items)} test items" + ) + + if not my_items: + logger.warning("No items assigned to this partition. Exiting success.") + sys.exit(0) + + # Run pytest with the specific test items (same as run_suite.py) + exit_code = run_pytest(my_items) + + if exit_code != 0: + if args.continue_on_error: + logger.warning(f"pytest exited with code {exit_code}") + else: + sys.exit(exit_code) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py b/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py index 182223486..982e51374 100644 --- a/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py +++ b/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py @@ -117,7 +117,7 @@ def _run_case(case: DiffusionTestCase) -> dict: modality=case.server_args.modality, sampling_params=sp, ) - rid = gen(case.id, client) + rid, _ = gen(case.id, client) rec = wait_for_req_perf_record( rid, ctx.perf_log_path, 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 6efeba304..3491a3c3b 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -1,5 +1,5 @@ """ -Config-driven diffusion performance test with pytest parametrization. +Config-driven diffusion generation test with pytest parametrization. If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed @@ -8,6 +8,7 @@ If the actual run is significantly better than the baseline, the improved cases from __future__ import annotations import os +from pathlib import Path from typing import Any, Callable import openai @@ -33,6 +34,8 @@ from sglang.multimodal_gen.test.server.testcase_configs import ( ScenarioConfig, ) from sglang.multimodal_gen.test.test_utils import ( + _consistency_gt_filenames, + extract_key_frames_from_video, get_dynamic_server_port, wait_for_req_perf_record, ) @@ -61,6 +64,12 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port)) sampling_params = case.sampling_params extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "") + + # In GT generation mode, force --backend diffusers + if os.environ.get("SGLANG_GEN_GT", "0") == "1": + if "--backend" not in extra_args: + extra_args = "--backend diffusers " + extra_args.strip() + extra_args += f" --num-gpus {server_args.num_gpus}" if server_args.tp_size is not None: @@ -189,14 +198,18 @@ Consider updating perf_baselines.json with the snippets below: self, ctx: ServerContext, case_id: str, - generate_fn: Callable[[str, openai.Client], str], - ) -> RequestPerfRecord: - """Run generation and collect performance records.""" + generate_fn: Callable[[str, openai.Client], tuple[str, bytes]], + ) -> tuple[RequestPerfRecord, bytes]: + """Run generation and collect performance records. + + Returns: + Tuple of (performance_record, content_bytes) + """ log_path = ctx.perf_log_path log_wait_timeout = 30 client = self._client(ctx) - rid = generate_fn(case_id, client) + rid, content = generate_fn(case_id, client) req_perf_record = wait_for_req_perf_record( rid, @@ -204,7 +217,7 @@ Consider updating perf_baselines.json with the snippets below: timeout=log_wait_timeout, ) - return req_perf_record + return (req_perf_record, content) def _validate_and_record( self, @@ -404,11 +417,65 @@ Consider updating perf_baselines.json with the snippets below: """ logger.error(output) + def _save_gt_output( + self, + case: DiffusionTestCase, + content: bytes, + ) -> None: + """Save generated content as ground truth files. + + Args: + case: Test case configuration + content: Generated content bytes (image or video) + """ + gt_output_dir = os.environ.get("SGLANG_GT_OUTPUT_DIR") + if not gt_output_dir: + logger.error("SGLANG_GT_OUTPUT_DIR not set, cannot save GT output") + return + + out_dir = Path(gt_output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + num_gpus = case.server_args.num_gpus + is_video = case.server_args.modality == "video" + + if is_video: + # Extract key frames from video + frames = extract_key_frames_from_video( + content, num_frames=case.sampling_params.num_frames + ) + + if len(frames) != 3: + logger.warning( + f"{case.id}: expected 3 frames, got {len(frames)}, skipping frame save" + ) + return + + # Save frames (reuse naming from _consistency_gt_filenames) + filenames = _consistency_gt_filenames(case.id, num_gpus, is_video=True) + from PIL import Image + + for frame, fn in zip(frames, filenames): + frame_path = out_dir / fn + Image.fromarray(frame).save(frame_path) + logger.info(f"Saved GT frame: {frame_path}") + else: + # Save image + from sglang.multimodal_gen.test.test_utils import detect_image_format + + detected_format = detect_image_format(content) + filenames = _consistency_gt_filenames( + case.id, num_gpus, is_video=False, output_format=detected_format + ) + output_path = out_dir / filenames[0] + output_path.write_bytes(content) + logger.info(f"Saved GT image: {output_path} (format: {detected_format})") + def _test_lora_api_functionality( self, ctx: ServerContext, case: DiffusionTestCase, - generate_fn: Callable[[str, openai.Client], str], + generate_fn: Callable[[str, openai.Client], tuple[str, bytes]], ) -> None: """ Test LoRA API functionality with end-to-end validation: merge, unmerge, and set_lora. @@ -423,8 +490,8 @@ Consider updating perf_baselines.json with the snippets below: assert resp.status_code == 200, f"unmerge_lora_weights failed: {resp.text}" logger.info("[LoRA E2E] Verifying generation after unmerge for %s", case.id) - output_after_unmerge = generate_fn(case.id, client) - assert output_after_unmerge is not None, "Generation after unmerge failed" + rid_after_unmerge, _ = generate_fn(case.id, client) + assert rid_after_unmerge is not None, "Generation after unmerge failed" logger.info("[LoRA E2E] Generation after unmerge succeeded") # Test 2: merge_lora_weights - API should succeed and generation should work @@ -433,8 +500,8 @@ Consider updating perf_baselines.json with the snippets below: assert resp.status_code == 200, f"merge_lora_weights failed: {resp.text}" logger.info("[LoRA E2E] Verifying generation after re-merge for %s", case.id) - output_after_merge = generate_fn(case.id, client) - assert output_after_merge is not None, "Generation after merge failed" + rid_after_merge, _ = generate_fn(case.id, client) + assert rid_after_merge is not None, "Generation after merge failed" logger.info("[LoRA E2E] Generation after merge succeeded") # Test 3: set_lora (re-set the same adapter) - API should succeed and generation should work @@ -443,8 +510,8 @@ Consider updating perf_baselines.json with the snippets below: assert resp.status_code == 200, f"set_lora failed: {resp.text}" logger.info("[LoRA E2E] Verifying generation after set_lora for %s", case.id) - output_after_set = generate_fn(case.id, client) - assert output_after_set is not None, "Generation after set_lora failed" + rid_after_set, _ = generate_fn(case.id, client) + assert rid_after_set is not None, "Generation after set_lora failed" logger.info("[LoRA E2E] Generation after set_lora succeeded") # Test 4: list_loras - API should return the expected list of LoRA adapters @@ -468,7 +535,7 @@ Consider updating perf_baselines.json with the snippets below: self, ctx: ServerContext, case: DiffusionTestCase, - generate_fn: Callable[[str, openai.Client], str], + generate_fn: Callable[[str, openai.Client], tuple[str, bytes]], second_lora_path: str, ) -> None: """ @@ -483,8 +550,8 @@ Consider updating perf_baselines.json with the snippets below: logger.info( "[LoRA Switch E2E] Testing generation with initial LoRA for %s", case.id ) - output_initial = generate_fn(case.id, client) - assert output_initial is not None, "Generation with initial LoRA failed" + rid_initial, _ = generate_fn(case.id, client) + assert rid_initial is not None, "Generation with initial LoRA failed" logger.info("[LoRA Switch E2E] Generation with initial LoRA succeeded") # Test 2: Switch to second LoRA and generate @@ -502,8 +569,8 @@ Consider updating perf_baselines.json with the snippets below: logger.info( "[LoRA Switch E2E] Verifying generation with second LoRA for %s", case.id ) - output_second = generate_fn(case.id, client) - assert output_second is not None, "Generation with second LoRA failed" + rid_second, _ = generate_fn(case.id, client) + assert rid_second is not None, "Generation with second LoRA failed" logger.info("[LoRA Switch E2E] Generation with second LoRA succeeded") # Test 3: Switch back to original LoRA and generate @@ -515,10 +582,8 @@ Consider updating perf_baselines.json with the snippets below: "[LoRA Switch E2E] Verifying generation after switching back for %s", case.id, ) - output_switched_back = generate_fn(case.id, client) - assert ( - output_switched_back is not None - ), "Generation after switching back failed" + rid_switched_back, _ = generate_fn(case.id, client) + assert rid_switched_back is not None, "Generation after switching back failed" logger.info("[LoRA Switch E2E] Generation after switching back succeeded") logger.info( @@ -557,7 +622,7 @@ Consider updating perf_baselines.json with the snippets below: self, ctx: ServerContext, case: DiffusionTestCase, - generate_fn: Callable[[str, openai.Client], str], + generate_fn: Callable[[str, openai.Client], tuple[str, bytes]], first_lora_path: str, second_lora_path: str, ) -> None: @@ -581,7 +646,8 @@ Consider updating perf_baselines.json with the snippets below: assert ( resp.status_code == 200 ), f"set_lora with multiple adapters failed: {resp.text}" - assert generate_fn(case.id, client) is not None + rid, _ = generate_fn(case.id, client) + assert rid is not None # Test 2: Different strengths resp = requests.post( @@ -596,7 +662,8 @@ Consider updating perf_baselines.json with the snippets below: assert ( resp.status_code == 200 ), f"set_lora with different strengths failed: {resp.text}" - assert generate_fn(case.id, client) is not None + rid, _ = generate_fn(case.id, client) + assert rid is not None # Test 3: Different targets requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"}) @@ -612,14 +679,16 @@ Consider updating perf_baselines.json with the snippets below: assert ( resp.status_code == 200 ), f"set_lora with cached adapters failed: {resp.text}" - assert generate_fn(case.id, client) is not None + rid, _ = generate_fn(case.id, client) + assert rid is not None # Test 4: Switch back to single LoRA resp = requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"}) assert ( resp.status_code == 200 ), f"set_lora back to single adapter failed: {resp.text}" - assert generate_fn(case.id, client) is not None + rid, _ = generate_fn(case.id, client) + assert rid is not None logger.info("[Multi-LoRA] All multi-LoRA tests passed for %s", case.id) @@ -742,22 +811,30 @@ Consider updating perf_baselines.json with the snippets below: "input_reference is not supported" in detail ), f"Unexpected error detail for T2V input_reference: {detail}" - def test_diffusion_perf( + def test_diffusion_generation( self, case: DiffusionTestCase, diffusion_server: ServerContext, ): """Single parametrized test that runs for all cases. + This test performs: + 1. Generation + 2. Performance validation against baselines + 3. Consistency validation against ground truth + Pytest will execute this test once per case in ONE_GPU_CASES, with test IDs like: - - test_diffusion_perf[qwen_image_text] - - test_diffusion_perf[qwen_image_edit] + - test_diffusion_generation[qwen_image_text] + - test_diffusion_generation[qwen_image_edit] - etc. """ + # Check if we're in GT generation mode + is_gt_gen_mode = os.environ.get("SGLANG_GEN_GT", "0") == "1" + # Dynamic LoRA loading test - tests LayerwiseOffload + set_lora interaction # Server starts WITHOUT lora_path, then set_lora is called after startup - if case.server_args.dynamic_lora_path: + if case.server_args.dynamic_lora_path and not is_gt_gen_mode: self._test_dynamic_lora_loading(diffusion_server, case) generate_fn = get_generate_fn( @@ -765,12 +842,20 @@ Consider updating perf_baselines.json with the snippets below: modality=case.server_args.modality, sampling_params=case.sampling_params, ) - perf_record = self.run_and_collect( + + # Single generation - output is reused for both validations + perf_record, content = self.run_and_collect( diffusion_server, case.id, generate_fn, ) + if is_gt_gen_mode: + # GT generation mode: save output and skip all validations/tests + self._save_gt_output(case, content) + return + + # Validation 1: Performance self._validate_and_record(case, perf_record) # Test /v1/models endpoint for router compatibility 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 5cbb364d6..c137e0b76 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -647,7 +647,7 @@ def get_generate_fn( model_path: str, modality: str, sampling_params: DiffusionSamplingParams, -) -> Callable[[str, Client], str]: +) -> Callable[[str, Client], tuple[str, bytes]]: """Return appropriate generation function for the case.""" # Allow override via environment variable (useful for AMD where large resolutions cause slow VAE) output_size = os.environ.get("SGLANG_TEST_OUTPUT_SIZE", sampling_params.output_size) @@ -717,7 +717,7 @@ def get_generate_fn( f"{case_id}: video job {video_id} timed out during baseline generation. " "Attempting to collect performance data anyway." ) - return video_id + return (video_id, b"") if is_amd: logger.warning( @@ -755,14 +755,14 @@ def get_generate_fn( ) os.remove(tmp_path) - return video_id + return (video_id, content) video_seconds = sampling_params.seconds or 4 - def generate_image(case_id, client) -> str: + def generate_image(case_id, client) -> tuple[str, bytes]: """T2I: Text to Image generation.""" if not sampling_params.prompt: - pytest.skip(f"{id}: no text prompt configured") + pytest.skip(f"{case_id}: no text prompt configured") # Request parameters that affect output format req_output_format = None # Not specified in current request @@ -813,12 +813,12 @@ def get_generate_fn( ) os.remove(tmp_path) - return rid + return (rid, img_data) - def generate_image_edit(case_id, client) -> str: - """TI2I: Text + Image ? Image edit.""" + def generate_image_edit(case_id, client) -> tuple[str, bytes]: + """TI2I: Text + Image -> Image edit.""" if not sampling_params.prompt or not sampling_params.image_path: - pytest.skip(f"{id}: no edit config") + pytest.skip(f"{case_id}: no edit config") image_paths = sampling_params.image_path @@ -832,7 +832,7 @@ def get_generate_fn( else: new_image_paths.append(Path(image_path)) if not image_path.exists(): - pytest.skip(f"{id}: file missing: {image_path}") + pytest.skip(f"{case_id}: file missing: {image_path}") image_paths = new_image_paths @@ -896,12 +896,12 @@ def get_generate_fn( ) os.remove(tmp_path) - return rid + return (rid, img_data) - def generate_image_edit_url(case_id, client) -> str: + def generate_image_edit_url(case_id, client) -> tuple[str, bytes]: """TI2I: Text + Image ? Image edit using direct URL transfer (no pre-download).""" if not sampling_params.prompt or not sampling_params.image_path: - pytest.skip(f"{id}: no edit config") + pytest.skip(f"{case_id}: no edit config") # Handle both single URL and list of URLs image_urls = sampling_params.image_path if not isinstance(image_urls, list): @@ -911,7 +911,7 @@ def get_generate_fn( for url in image_urls: if not is_image_url(url): pytest.skip( - f"{id}: image_path must be a URL for URL direct test: {url}" + f"{case_id}: image_path must be a URL for URL direct test: {url}" ) # Request parameters that affect output format @@ -965,12 +965,12 @@ def get_generate_fn( ) os.remove(tmp_path) - return rid + return (rid, img_data) - def generate_video(case_id, client) -> str: + def generate_video(case_id, client) -> tuple[str, bytes]: """T2V: Text ? Video.""" if not sampling_params.prompt: - pytest.skip(f"{id}: no text prompt configured") + pytest.skip(f"{case_id}: no text prompt configured") # Build extra_body for optional features extra_body = {} @@ -987,17 +987,17 @@ def get_generate_fn( extra_body=extra_body if extra_body else None, ) - def generate_image_to_video(case_id, client) -> str: - """I2V: Image ? Video (optional prompt).""" + def generate_image_to_video(case_id, client) -> tuple[str, bytes]: + """I2V: Image -> Video (optional prompt).""" if not sampling_params.image_path: - pytest.skip(f"{id}: no input image configured") + pytest.skip(f"{case_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}") + pytest.skip(f"{case_id}: file missing: {image_path}") # Build extra_body for optional features extra_body = {} @@ -1016,9 +1016,9 @@ def get_generate_fn( extra_body=extra_body if extra_body else None, ) - def generate_text_url_image_to_video(case_id, client) -> str: + def generate_text_url_image_to_video(case_id, client) -> tuple[str, bytes]: if not sampling_params.prompt or not sampling_params.image_path: - pytest.skip(f"{id}: no edit config") + pytest.skip(f"{case_id}: no edit config") # Build extra_body for optional features extra_body = {"reference_url": sampling_params.image_path} @@ -1039,17 +1039,17 @@ def get_generate_fn( }, ) - def generate_text_image_to_video(case_id, client) -> str: - """TI2V: Text + Image ? Video.""" + def generate_text_image_to_video(case_id, client) -> tuple[str, bytes]: + """TI2V: Text + Image -> Video.""" if not sampling_params.prompt or not sampling_params.image_path: - pytest.skip(f"{id}: no edit config") + pytest.skip(f"{case_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}") + pytest.skip(f"{case_id}: file missing: {image_path}") # Build extra_body for optional features extra_body = {} diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index f517f11c8..921ec5519 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -1,12 +1,15 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import base64 +import io import json import os import socket +import tempfile import time from pathlib import Path import cv2 +import numpy as np from PIL import Image from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var @@ -81,6 +84,19 @@ def is_webp(data: bytes) -> bool: return data[:4] == b"RIFF" and data[8:12] == b"WEBP" +def detect_image_format(data: bytes) -> str: + """Detect image format from bytes (magic). Returns 'png'|'jpeg'|'webp'; default 'png'.""" + if len(data) < 12: + return "png" + if is_png(data): + return "png" + if is_jpeg(data): + return "jpeg" + if is_webp(data): + return "webp" + return "png" + + def get_expected_image_format( output_format: str | None = None, background: str | None = None, @@ -358,3 +374,88 @@ def validate_video_file( assert ( actual_height == expected_height ), f"Video height mismatch: expected {expected_height}, got {actual_height}" + + +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" + of = output_format.lower() + if of == "jpeg": + return "jpg" + if of in ("png", "webp", "jpg"): + return of + return "png" + + +def _consistency_gt_filenames( + case_id: str, num_gpus: int, is_video: bool, output_format: str | None = None +) -> list[str]: + """Return the list of GT image filenames for a case. Reused by GT generation and consistency check.""" + 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", + ] + ext = output_format_to_ext(output_format) + return [f"{case_id}_{n}gpu.{ext}"] + + +def extract_key_frames_from_video( + video_bytes: bytes, + num_frames: int | None = None, +) -> list[np.ndarray]: + """ + Extract key frames (first, middle, last) from video bytes. + + Args: + video_bytes: Raw video bytes (MP4 format) + num_frames: Total number of frames (if known), used for validation + + Returns: + List of numpy arrays [first_frame, middle_frame, last_frame]. + """ + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + tmp.write(video_bytes) + tmp_path = tmp.name + + try: + cap = cv2.VideoCapture(tmp_path) + if not cap.isOpened(): + raise ValueError("Failed to open video file") + + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + if total_frames < 1: + raise ValueError("Video has no frames") + + first_idx = 0 + mid_idx = total_frames // 2 + last_idx = total_frames - 1 + key_indices = [first_idx, mid_idx, last_idx] + + frames = [] + for idx in key_indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, frame = cap.read() + if not ret: + raise ValueError(f"Failed to read frame at index {idx}") + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames.append(frame_rgb) + + cap.release() + logger.info( + f"Extracted {len(frames)} key frames from video " + f"(total: {total_frames}, indices: {key_indices})" + ) + return frames + + finally: + os.unlink(tmp_path) + + +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)