66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniRunner
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
|
|
|
MODEL_ID = "HiDream-ai/HiDream-O1-Image"
|
|
PROMPT = "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'"
|
|
BASELINE_PATH = Path(__file__).resolve().parents[2] / "assets" / "hidream_o1" / "hidream_o1_baseline.png"
|
|
SSIM_THRESHOLD = 0.8
|
|
PSNR_THRESHOLD = 19.0
|
|
|
|
|
|
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
|
def test_hidream_o1_image_tp2_matches_baseline(
|
|
accuracy_artifact_root: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
|
|
model = os.environ.get("HIDREAM_O1_IMAGE_MODEL", MODEL_ID)
|
|
with OmniRunner(
|
|
model,
|
|
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
|
|
enforce_eager=True,
|
|
) as runner:
|
|
outputs = runner.omni.generate(
|
|
{"prompt": PROMPT},
|
|
OmniDiffusionSamplingParams(
|
|
height=2048,
|
|
width=2048,
|
|
num_inference_steps=50,
|
|
guidance_scale=5.0,
|
|
seed=42,
|
|
),
|
|
)
|
|
images = extract_images_from_outputs(outputs)
|
|
assert images
|
|
image = images[0].convert("RGB")
|
|
|
|
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
|
|
image.save(output_dir / "vllm_omni_tp2.png")
|
|
assert BASELINE_PATH.exists()
|
|
baseline = Image.open(BASELINE_PATH).convert("RGB")
|
|
resized = image.resize(baseline.size, Image.Resampling.LANCZOS)
|
|
assert_similarity(
|
|
model_name=MODEL_ID,
|
|
vllm_image=resized,
|
|
diffusers_image=baseline,
|
|
ssim_threshold=SSIM_THRESHOLD,
|
|
psnr_threshold=PSNR_THRESHOLD,
|
|
)
|