vllm-project--vllm-omni
wehub-resource-sync
201 行
7.5 KiB
Bash
201 行
7.5 KiB
Bash
#!/bin/bash
|
|
# End-to-end test script for unified quantization framework (PR #1764).
|
|
#
|
|
# Tests FP8 quantization for:
|
|
# 1. Z-Image-Turbo (single-stage, ~20GB VRAM)
|
|
# 2. Qwen-Image (single-stage, ~20GB VRAM)
|
|
# 3. FLUX.1-dev (single-stage, ~25GB VRAM with fp8)
|
|
# 4. BAGEL (multi-stage LLM+DiT, ~55GB VRAM)
|
|
#
|
|
# Usage:
|
|
# bash tests/e2e/offline_inference/run_quantization_e2e.sh [--skip-flux] [--skip-bagel]
|
|
#
|
|
# Expected: all runs produce images without errors.
|
|
# Key check for BAGEL: FP8 only applies to diffusion stage, NOT the LLM stage.
|
|
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
|
REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
|
|
OUTPUT_DIR="${REPO_ROOT}/test_quant_outputs"
|
|
mkdir -p "$OUTPUT_DIR"
|
|
|
|
SKIP_FLUX=false
|
|
SKIP_BAGEL=false
|
|
for arg in "$@"; do
|
|
case "$arg" in
|
|
--skip-flux) SKIP_FLUX=true ;;
|
|
--skip-bagel) SKIP_BAGEL=true ;;
|
|
esac
|
|
done
|
|
|
|
PASS=0
|
|
FAIL=0
|
|
SKIP=0
|
|
|
|
run_test() {
|
|
local name="$1"
|
|
shift
|
|
echo ""
|
|
echo "============================================================"
|
|
echo " TEST: $name"
|
|
echo "============================================================"
|
|
if "$@"; then
|
|
echo " PASS: $name"
|
|
PASS=$((PASS + 1))
|
|
else
|
|
echo " FAIL: $name (exit code $?)"
|
|
FAIL=$((FAIL + 1))
|
|
fi
|
|
}
|
|
|
|
# ─── 1. Z-Image-Turbo ────────────────────────────────────────────────────────
|
|
# Defaults from README: 1024x1024, 50 steps, cfg-scale 4.0, seed 42
|
|
|
|
run_test "Z-Image-Turbo BF16 (baseline)" \
|
|
python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
|
|
--model Tongyi-MAI/Z-Image-Turbo \
|
|
--prompt "a cup of coffee on the table" \
|
|
--seed 42 --num-inference-steps 50 \
|
|
--height 1024 --width 1024 \
|
|
--cfg-scale 4.0 \
|
|
--output "$OUTPUT_DIR/zimage_bf16.png"
|
|
|
|
run_test "Z-Image-Turbo FP8" \
|
|
python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
|
|
--model Tongyi-MAI/Z-Image-Turbo \
|
|
--prompt "a cup of coffee on the table" \
|
|
--seed 42 --num-inference-steps 50 \
|
|
--height 1024 --width 1024 \
|
|
--cfg-scale 4.0 \
|
|
--quantization fp8 \
|
|
--output "$OUTPUT_DIR/zimage_fp8.png"
|
|
|
|
# ─── 2. Qwen-Image ───────────────────────────────────────────────────────────
|
|
# Defaults from README: 1024x1024, 50 steps, cfg-scale 4.0, seed 142
|
|
|
|
run_test "Qwen-Image BF16 (baseline)" \
|
|
python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
|
|
--model Qwen/Qwen-Image \
|
|
--prompt "a cup of coffee on the table" \
|
|
--seed 142 --num-inference-steps 50 \
|
|
--height 1024 --width 1024 \
|
|
--cfg-scale 4.0 \
|
|
--output "$OUTPUT_DIR/qwen_bf16.png"
|
|
|
|
run_test "Qwen-Image FP8" \
|
|
python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
|
|
--model Qwen/Qwen-Image \
|
|
--prompt "a cup of coffee on the table" \
|
|
--seed 142 --num-inference-steps 50 \
|
|
--height 1024 --width 1024 \
|
|
--cfg-scale 4.0 \
|
|
--quantization fp8 \
|
|
--output "$OUTPUT_DIR/qwen_fp8.png"
|
|
|
|
# ─── 3. FLUX.1-dev ───────────────────────────────────────────────────────────
|
|
# Defaults from first test run: 1024x1024, 20 steps, guidance-scale 3.5
|
|
|
|
if [ "$SKIP_FLUX" = true ]; then
|
|
echo ""
|
|
echo " SKIP: FLUX.1-dev (--skip-flux)"
|
|
SKIP=$((SKIP + 1))
|
|
SKIP=$((SKIP + 1))
|
|
else
|
|
run_test "FLUX.1-dev BF16 (baseline)" \
|
|
python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
|
|
--model black-forest-labs/FLUX.1-dev \
|
|
--prompt "a cup of coffee on the table" \
|
|
--seed 42 --num-inference-steps 20 \
|
|
--height 1024 --width 1024 \
|
|
--guidance-scale 3.5 --cfg-scale 1.0 \
|
|
--output "$OUTPUT_DIR/flux_bf16.png"
|
|
|
|
run_test "FLUX.1-dev FP8" \
|
|
python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
|
|
--model black-forest-labs/FLUX.1-dev \
|
|
--prompt "a cup of coffee on the table" \
|
|
--seed 42 --num-inference-steps 20 \
|
|
--height 1024 --width 1024 \
|
|
--guidance-scale 3.5 --cfg-scale 1.0 \
|
|
--quantization fp8 \
|
|
--output "$OUTPUT_DIR/flux_fp8.png"
|
|
fi
|
|
|
|
# ─── 4. BAGEL (multi-stage) ──────────────────────────────────────────────────
|
|
# Defaults from README: 50 steps, cfg-text-scale 4.0, cfg-img-scale 1.5
|
|
# BAGEL end2end.py saves images as output_1_stage_None_0.png in cwd.
|
|
# We copy them to OUTPUT_DIR with proper names for LPIPS comparison.
|
|
|
|
if [ "$SKIP_BAGEL" = true ]; then
|
|
echo ""
|
|
echo " SKIP: BAGEL (--skip-bagel)"
|
|
SKIP=$((SKIP + 1))
|
|
SKIP=$((SKIP + 1))
|
|
else
|
|
BAGEL_WORKDIR=$(mktemp -d)
|
|
|
|
run_test "BAGEL BF16 (baseline)" \
|
|
bash -c "cd '$BAGEL_WORKDIR' && python '$REPO_ROOT/examples/offline_inference/bagel/end2end.py' \
|
|
--model ByteDance-Seed/BAGEL-7B-MoT \
|
|
--modality text2img \
|
|
--prompts 'A cute cat' \
|
|
--steps 50 \
|
|
--cfg-text-scale 4.0 \
|
|
--cfg-img-scale 1.5 \
|
|
--seed 52 \
|
|
&& cp -f output_1_stage_None_0.png '$OUTPUT_DIR/bagel_bf16.png'"
|
|
|
|
run_test "BAGEL FP8 (diffusion-only quantization)" \
|
|
bash -c "cd '$BAGEL_WORKDIR' && python '$REPO_ROOT/examples/offline_inference/bagel/end2end.py' \
|
|
--model ByteDance-Seed/BAGEL-7B-MoT \
|
|
--modality text2img \
|
|
--prompts 'A cute cat' \
|
|
--steps 50 \
|
|
--cfg-text-scale 4.0 \
|
|
--cfg-img-scale 1.5 \
|
|
--seed 52 \
|
|
--quantization fp8 \
|
|
&& cp -f output_1_stage_None_0.png '$OUTPUT_DIR/bagel_fp8.png'"
|
|
|
|
rm -rf "$BAGEL_WORKDIR"
|
|
fi
|
|
|
|
# ─── 5. LPIPS quality comparison ─────────────────────────────────────────────
|
|
# Compares each BF16/FP8 image pair using perceptual LPIPS distance.
|
|
# Requires: pip install lpips torchvision
|
|
|
|
echo ""
|
|
echo "============================================================"
|
|
echo " LPIPS Quality Comparison"
|
|
echo "============================================================"
|
|
|
|
if python -c "import lpips" 2>/dev/null; then
|
|
run_test "LPIPS BF16-vs-FP8 quality check" \
|
|
python "$SCRIPT_DIR/compute_lpips.py" \
|
|
--image-dir "$OUTPUT_DIR" \
|
|
--threshold 0.1
|
|
else
|
|
echo " SKIP: lpips not installed (pip install lpips torchvision)"
|
|
SKIP=$((SKIP + 1))
|
|
fi
|
|
|
|
# ─── Summary ─────────────────────────────────────────────────────────────────
|
|
|
|
echo ""
|
|
echo "============================================================"
|
|
echo " SUMMARY"
|
|
echo "============================================================"
|
|
echo " PASS: $PASS"
|
|
echo " FAIL: $FAIL"
|
|
echo " SKIP: $SKIP"
|
|
echo " Output images: $OUTPUT_DIR/"
|
|
if [ -f "$OUTPUT_DIR/lpips_results.md" ]; then
|
|
echo " LPIPS report: $OUTPUT_DIR/lpips_results.md"
|
|
fi
|
|
echo "============================================================"
|
|
|
|
if [ "$FAIL" -gt 0 ]; then
|
|
exit 1
|
|
fi
|