[diffusion] CI: GT generation flow for diffusion CI (#19236)
Co-authored-by: Prozac614 <dwt614707404@163.com>
This commit is contained in:
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -117,7 +117,7 @@ def _run_case(case: DiffusionTestCase) -> dict:
|
|||||||
modality=case.server_args.modality,
|
modality=case.server_args.modality,
|
||||||
sampling_params=sp,
|
sampling_params=sp,
|
||||||
)
|
)
|
||||||
rid = gen(case.id, client)
|
rid, _ = gen(case.id, client)
|
||||||
rec = wait_for_req_perf_record(
|
rec = wait_for_req_perf_record(
|
||||||
rid,
|
rid,
|
||||||
ctx.perf_log_path,
|
ctx.perf_log_path,
|
||||||
|
|||||||
@@ -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
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
import os
|
import os
|
||||||
|
from pathlib import Path
|
||||||
from typing import Any, Callable
|
from typing import Any, Callable
|
||||||
|
|
||||||
import openai
|
import openai
|
||||||
@@ -33,6 +34,8 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
|
|||||||
ScenarioConfig,
|
ScenarioConfig,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.test.test_utils import (
|
from sglang.multimodal_gen.test.test_utils import (
|
||||||
|
_consistency_gt_filenames,
|
||||||
|
extract_key_frames_from_video,
|
||||||
get_dynamic_server_port,
|
get_dynamic_server_port,
|
||||||
wait_for_req_perf_record,
|
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))
|
port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port))
|
||||||
sampling_params = case.sampling_params
|
sampling_params = case.sampling_params
|
||||||
extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "")
|
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}"
|
extra_args += f" --num-gpus {server_args.num_gpus}"
|
||||||
|
|
||||||
if server_args.tp_size is not None:
|
if server_args.tp_size is not None:
|
||||||
@@ -189,14 +198,18 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
self,
|
self,
|
||||||
ctx: ServerContext,
|
ctx: ServerContext,
|
||||||
case_id: str,
|
case_id: str,
|
||||||
generate_fn: Callable[[str, openai.Client], str],
|
generate_fn: Callable[[str, openai.Client], tuple[str, bytes]],
|
||||||
) -> RequestPerfRecord:
|
) -> tuple[RequestPerfRecord, bytes]:
|
||||||
"""Run generation and collect performance records."""
|
"""Run generation and collect performance records.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tuple of (performance_record, content_bytes)
|
||||||
|
"""
|
||||||
log_path = ctx.perf_log_path
|
log_path = ctx.perf_log_path
|
||||||
log_wait_timeout = 30
|
log_wait_timeout = 30
|
||||||
|
|
||||||
client = self._client(ctx)
|
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(
|
req_perf_record = wait_for_req_perf_record(
|
||||||
rid,
|
rid,
|
||||||
@@ -204,7 +217,7 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
timeout=log_wait_timeout,
|
timeout=log_wait_timeout,
|
||||||
)
|
)
|
||||||
|
|
||||||
return req_perf_record
|
return (req_perf_record, content)
|
||||||
|
|
||||||
def _validate_and_record(
|
def _validate_and_record(
|
||||||
self,
|
self,
|
||||||
@@ -404,11 +417,65 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
"""
|
"""
|
||||||
logger.error(output)
|
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(
|
def _test_lora_api_functionality(
|
||||||
self,
|
self,
|
||||||
ctx: ServerContext,
|
ctx: ServerContext,
|
||||||
case: DiffusionTestCase,
|
case: DiffusionTestCase,
|
||||||
generate_fn: Callable[[str, openai.Client], str],
|
generate_fn: Callable[[str, openai.Client], tuple[str, bytes]],
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
Test LoRA API functionality with end-to-end validation: merge, unmerge, and set_lora.
|
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}"
|
assert resp.status_code == 200, f"unmerge_lora_weights failed: {resp.text}"
|
||||||
|
|
||||||
logger.info("[LoRA E2E] Verifying generation after unmerge for %s", case.id)
|
logger.info("[LoRA E2E] Verifying generation after unmerge for %s", case.id)
|
||||||
output_after_unmerge = generate_fn(case.id, client)
|
rid_after_unmerge, _ = generate_fn(case.id, client)
|
||||||
assert output_after_unmerge is not None, "Generation after unmerge failed"
|
assert rid_after_unmerge is not None, "Generation after unmerge failed"
|
||||||
logger.info("[LoRA E2E] Generation after unmerge succeeded")
|
logger.info("[LoRA E2E] Generation after unmerge succeeded")
|
||||||
|
|
||||||
# Test 2: merge_lora_weights - API should succeed and generation should work
|
# 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}"
|
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)
|
logger.info("[LoRA E2E] Verifying generation after re-merge for %s", case.id)
|
||||||
output_after_merge = generate_fn(case.id, client)
|
rid_after_merge, _ = generate_fn(case.id, client)
|
||||||
assert output_after_merge is not None, "Generation after merge failed"
|
assert rid_after_merge is not None, "Generation after merge failed"
|
||||||
logger.info("[LoRA E2E] Generation after merge succeeded")
|
logger.info("[LoRA E2E] Generation after merge succeeded")
|
||||||
|
|
||||||
# Test 3: set_lora (re-set the same adapter) - API should succeed and generation should work
|
# 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}"
|
assert resp.status_code == 200, f"set_lora failed: {resp.text}"
|
||||||
|
|
||||||
logger.info("[LoRA E2E] Verifying generation after set_lora for %s", case.id)
|
logger.info("[LoRA E2E] Verifying generation after set_lora for %s", case.id)
|
||||||
output_after_set = generate_fn(case.id, client)
|
rid_after_set, _ = generate_fn(case.id, client)
|
||||||
assert output_after_set is not None, "Generation after set_lora failed"
|
assert rid_after_set is not None, "Generation after set_lora failed"
|
||||||
logger.info("[LoRA E2E] Generation after set_lora succeeded")
|
logger.info("[LoRA E2E] Generation after set_lora succeeded")
|
||||||
|
|
||||||
# Test 4: list_loras - API should return the expected list of LoRA adapters
|
# 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,
|
self,
|
||||||
ctx: ServerContext,
|
ctx: ServerContext,
|
||||||
case: DiffusionTestCase,
|
case: DiffusionTestCase,
|
||||||
generate_fn: Callable[[str, openai.Client], str],
|
generate_fn: Callable[[str, openai.Client], tuple[str, bytes]],
|
||||||
second_lora_path: str,
|
second_lora_path: str,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
@@ -483,8 +550,8 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
logger.info(
|
logger.info(
|
||||||
"[LoRA Switch E2E] Testing generation with initial LoRA for %s", case.id
|
"[LoRA Switch E2E] Testing generation with initial LoRA for %s", case.id
|
||||||
)
|
)
|
||||||
output_initial = generate_fn(case.id, client)
|
rid_initial, _ = generate_fn(case.id, client)
|
||||||
assert output_initial is not None, "Generation with initial LoRA failed"
|
assert rid_initial is not None, "Generation with initial LoRA failed"
|
||||||
logger.info("[LoRA Switch E2E] Generation with initial LoRA succeeded")
|
logger.info("[LoRA Switch E2E] Generation with initial LoRA succeeded")
|
||||||
|
|
||||||
# Test 2: Switch to second LoRA and generate
|
# Test 2: Switch to second LoRA and generate
|
||||||
@@ -502,8 +569,8 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
logger.info(
|
logger.info(
|
||||||
"[LoRA Switch E2E] Verifying generation with second LoRA for %s", case.id
|
"[LoRA Switch E2E] Verifying generation with second LoRA for %s", case.id
|
||||||
)
|
)
|
||||||
output_second = generate_fn(case.id, client)
|
rid_second, _ = generate_fn(case.id, client)
|
||||||
assert output_second is not None, "Generation with second LoRA failed"
|
assert rid_second is not None, "Generation with second LoRA failed"
|
||||||
logger.info("[LoRA Switch E2E] Generation with second LoRA succeeded")
|
logger.info("[LoRA Switch E2E] Generation with second LoRA succeeded")
|
||||||
|
|
||||||
# Test 3: Switch back to original LoRA and generate
|
# 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",
|
"[LoRA Switch E2E] Verifying generation after switching back for %s",
|
||||||
case.id,
|
case.id,
|
||||||
)
|
)
|
||||||
output_switched_back = generate_fn(case.id, client)
|
rid_switched_back, _ = generate_fn(case.id, client)
|
||||||
assert (
|
assert rid_switched_back is not None, "Generation after switching back failed"
|
||||||
output_switched_back is not None
|
|
||||||
), "Generation after switching back failed"
|
|
||||||
logger.info("[LoRA Switch E2E] Generation after switching back succeeded")
|
logger.info("[LoRA Switch E2E] Generation after switching back succeeded")
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -557,7 +622,7 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
self,
|
self,
|
||||||
ctx: ServerContext,
|
ctx: ServerContext,
|
||||||
case: DiffusionTestCase,
|
case: DiffusionTestCase,
|
||||||
generate_fn: Callable[[str, openai.Client], str],
|
generate_fn: Callable[[str, openai.Client], tuple[str, bytes]],
|
||||||
first_lora_path: str,
|
first_lora_path: str,
|
||||||
second_lora_path: str,
|
second_lora_path: str,
|
||||||
) -> None:
|
) -> None:
|
||||||
@@ -581,7 +646,8 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
assert (
|
assert (
|
||||||
resp.status_code == 200
|
resp.status_code == 200
|
||||||
), f"set_lora with multiple adapters failed: {resp.text}"
|
), 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
|
# Test 2: Different strengths
|
||||||
resp = requests.post(
|
resp = requests.post(
|
||||||
@@ -596,7 +662,8 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
assert (
|
assert (
|
||||||
resp.status_code == 200
|
resp.status_code == 200
|
||||||
), f"set_lora with different strengths failed: {resp.text}"
|
), 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
|
# Test 3: Different targets
|
||||||
requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"})
|
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 (
|
assert (
|
||||||
resp.status_code == 200
|
resp.status_code == 200
|
||||||
), f"set_lora with cached adapters failed: {resp.text}"
|
), 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
|
# Test 4: Switch back to single LoRA
|
||||||
resp = requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"})
|
resp = requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"})
|
||||||
assert (
|
assert (
|
||||||
resp.status_code == 200
|
resp.status_code == 200
|
||||||
), f"set_lora back to single adapter failed: {resp.text}"
|
), 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)
|
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
|
"input_reference is not supported" in detail
|
||||||
), f"Unexpected error detail for T2V input_reference: {detail}"
|
), f"Unexpected error detail for T2V input_reference: {detail}"
|
||||||
|
|
||||||
def test_diffusion_perf(
|
def test_diffusion_generation(
|
||||||
self,
|
self,
|
||||||
case: DiffusionTestCase,
|
case: DiffusionTestCase,
|
||||||
diffusion_server: ServerContext,
|
diffusion_server: ServerContext,
|
||||||
):
|
):
|
||||||
"""Single parametrized test that runs for all cases.
|
"""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,
|
Pytest will execute this test once per case in ONE_GPU_CASES,
|
||||||
with test IDs like:
|
with test IDs like:
|
||||||
- test_diffusion_perf[qwen_image_text]
|
- test_diffusion_generation[qwen_image_text]
|
||||||
- test_diffusion_perf[qwen_image_edit]
|
- test_diffusion_generation[qwen_image_edit]
|
||||||
- etc.
|
- 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
|
# Dynamic LoRA loading test - tests LayerwiseOffload + set_lora interaction
|
||||||
# Server starts WITHOUT lora_path, then set_lora is called after startup
|
# 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)
|
self._test_dynamic_lora_loading(diffusion_server, case)
|
||||||
|
|
||||||
generate_fn = get_generate_fn(
|
generate_fn = get_generate_fn(
|
||||||
@@ -765,12 +842,20 @@ Consider updating perf_baselines.json with the snippets below:
|
|||||||
modality=case.server_args.modality,
|
modality=case.server_args.modality,
|
||||||
sampling_params=case.sampling_params,
|
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,
|
diffusion_server,
|
||||||
case.id,
|
case.id,
|
||||||
generate_fn,
|
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)
|
self._validate_and_record(case, perf_record)
|
||||||
|
|
||||||
# Test /v1/models endpoint for router compatibility
|
# Test /v1/models endpoint for router compatibility
|
||||||
|
|||||||
@@ -647,7 +647,7 @@ def get_generate_fn(
|
|||||||
model_path: str,
|
model_path: str,
|
||||||
modality: str,
|
modality: str,
|
||||||
sampling_params: DiffusionSamplingParams,
|
sampling_params: DiffusionSamplingParams,
|
||||||
) -> Callable[[str, Client], str]:
|
) -> Callable[[str, Client], tuple[str, bytes]]:
|
||||||
"""Return appropriate generation function for the case."""
|
"""Return appropriate generation function for the case."""
|
||||||
# Allow override via environment variable (useful for AMD where large resolutions cause slow VAE)
|
# 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)
|
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. "
|
f"{case_id}: video job {video_id} timed out during baseline generation. "
|
||||||
"Attempting to collect performance data anyway."
|
"Attempting to collect performance data anyway."
|
||||||
)
|
)
|
||||||
return video_id
|
return (video_id, b"")
|
||||||
|
|
||||||
if is_amd:
|
if is_amd:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
@@ -755,14 +755,14 @@ def get_generate_fn(
|
|||||||
)
|
)
|
||||||
os.remove(tmp_path)
|
os.remove(tmp_path)
|
||||||
|
|
||||||
return video_id
|
return (video_id, content)
|
||||||
|
|
||||||
video_seconds = sampling_params.seconds or 4
|
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."""
|
"""T2I: Text to Image generation."""
|
||||||
if not sampling_params.prompt:
|
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
|
# Request parameters that affect output format
|
||||||
req_output_format = None # Not specified in current request
|
req_output_format = None # Not specified in current request
|
||||||
@@ -813,12 +813,12 @@ def get_generate_fn(
|
|||||||
)
|
)
|
||||||
os.remove(tmp_path)
|
os.remove(tmp_path)
|
||||||
|
|
||||||
return rid
|
return (rid, img_data)
|
||||||
|
|
||||||
def generate_image_edit(case_id, client) -> str:
|
def generate_image_edit(case_id, client) -> tuple[str, bytes]:
|
||||||
"""TI2I: Text + Image ? Image edit."""
|
"""TI2I: Text + Image -> Image edit."""
|
||||||
if not sampling_params.prompt or not sampling_params.image_path:
|
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
|
image_paths = sampling_params.image_path
|
||||||
|
|
||||||
@@ -832,7 +832,7 @@ def get_generate_fn(
|
|||||||
else:
|
else:
|
||||||
new_image_paths.append(Path(image_path))
|
new_image_paths.append(Path(image_path))
|
||||||
if not image_path.exists():
|
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
|
image_paths = new_image_paths
|
||||||
|
|
||||||
@@ -896,12 +896,12 @@ def get_generate_fn(
|
|||||||
)
|
)
|
||||||
os.remove(tmp_path)
|
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)."""
|
"""TI2I: Text + Image ? Image edit using direct URL transfer (no pre-download)."""
|
||||||
if not sampling_params.prompt or not sampling_params.image_path:
|
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
|
# Handle both single URL and list of URLs
|
||||||
image_urls = sampling_params.image_path
|
image_urls = sampling_params.image_path
|
||||||
if not isinstance(image_urls, list):
|
if not isinstance(image_urls, list):
|
||||||
@@ -911,7 +911,7 @@ def get_generate_fn(
|
|||||||
for url in image_urls:
|
for url in image_urls:
|
||||||
if not is_image_url(url):
|
if not is_image_url(url):
|
||||||
pytest.skip(
|
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
|
# Request parameters that affect output format
|
||||||
@@ -965,12 +965,12 @@ def get_generate_fn(
|
|||||||
)
|
)
|
||||||
os.remove(tmp_path)
|
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."""
|
"""T2V: Text ? Video."""
|
||||||
if not sampling_params.prompt:
|
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
|
# Build extra_body for optional features
|
||||||
extra_body = {}
|
extra_body = {}
|
||||||
@@ -987,17 +987,17 @@ def get_generate_fn(
|
|||||||
extra_body=extra_body if extra_body else None,
|
extra_body=extra_body if extra_body else None,
|
||||||
)
|
)
|
||||||
|
|
||||||
def generate_image_to_video(case_id, client) -> str:
|
def generate_image_to_video(case_id, client) -> tuple[str, bytes]:
|
||||||
"""I2V: Image ? Video (optional prompt)."""
|
"""I2V: Image -> Video (optional prompt)."""
|
||||||
if not sampling_params.image_path:
|
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):
|
if is_image_url(sampling_params.image_path):
|
||||||
image_path = download_image_from_url(str(sampling_params.image_path))
|
image_path = download_image_from_url(str(sampling_params.image_path))
|
||||||
else:
|
else:
|
||||||
image_path = Path(sampling_params.image_path)
|
image_path = Path(sampling_params.image_path)
|
||||||
if not image_path.exists():
|
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
|
# Build extra_body for optional features
|
||||||
extra_body = {}
|
extra_body = {}
|
||||||
@@ -1016,9 +1016,9 @@ def get_generate_fn(
|
|||||||
extra_body=extra_body if extra_body else None,
|
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:
|
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
|
# Build extra_body for optional features
|
||||||
extra_body = {"reference_url": sampling_params.image_path}
|
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:
|
def generate_text_image_to_video(case_id, client) -> tuple[str, bytes]:
|
||||||
"""TI2V: Text + Image ? Video."""
|
"""TI2V: Text + Image -> Video."""
|
||||||
if not sampling_params.prompt or not sampling_params.image_path:
|
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):
|
if is_image_url(sampling_params.image_path):
|
||||||
image_path = download_image_from_url(str(sampling_params.image_path))
|
image_path = download_image_from_url(str(sampling_params.image_path))
|
||||||
else:
|
else:
|
||||||
image_path = Path(sampling_params.image_path)
|
image_path = Path(sampling_params.image_path)
|
||||||
if not image_path.exists():
|
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
|
# Build extra_body for optional features
|
||||||
extra_body = {}
|
extra_body = {}
|
||||||
|
|||||||
@@ -1,12 +1,15 @@
|
|||||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||||
import base64
|
import base64
|
||||||
|
import io
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import socket
|
import socket
|
||||||
|
import tempfile
|
||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
|
import numpy as np
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
|
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"
|
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(
|
def get_expected_image_format(
|
||||||
output_format: str | None = None,
|
output_format: str | None = None,
|
||||||
background: str | None = None,
|
background: str | None = None,
|
||||||
@@ -358,3 +374,88 @@ def validate_video_file(
|
|||||||
assert (
|
assert (
|
||||||
actual_height == expected_height
|
actual_height == expected_height
|
||||||
), f"Video height mismatch: expected {expected_height}, got {actual_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)
|
||||||
|
|||||||
Reference in New Issue
Block a user