diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py new file mode 100644 index 000000000..e73195017 --- /dev/null +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +import base64 +import gc +import io +import os +from pathlib import Path + +import pytest +import requests +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from PIL import Image + +from tests.conftest import ( + OmniServer, + _run_post_test_cleanup, + _run_pre_test_cleanup, +) +from tests.e2e.accuracy.utils import assert_similarity, model_output_dir +from tests.utils import hardware_test + +MODEL_ID = "Qwen/Qwen-Image" +MODEL_ENV_VAR = "QWEN_IMAGE_MODEL" +PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style." +NEGATIVE_PROMPT = "blurry, low quality" +WIDTH = 512 +HEIGHT = 512 +NUM_INFERENCE_STEPS = 20 +TRUE_CFG_SCALE = 4.0 +SEED = 42 +SSIM_THRESHOLD = 0.97 +PSNR_THRESHOLD = 30.0 + + +def _model_name() -> str: + return os.environ.get(MODEL_ENV_VAR, MODEL_ID) + + +def _local_files_only(model: str) -> bool: + return Path(model).exists() + + +def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image: + server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] + with OmniServer(model, server_args, use_omni=True) as omni_server: + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/images/generations", + json={ + "model": omni_server.model, + "prompt": PROMPT, + "size": f"{WIDTH}x{HEIGHT}", + "n": 1, + "response_format": "b64_json", + "negative_prompt": NEGATIVE_PROMPT, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": TRUE_CFG_SCALE, + "seed": SEED, + }, + timeout=600, + ) + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == 1 + image_bytes = base64.b64decode(payload["data"][0]["b64_json"]) + image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + image.load() + image.save(output_path) + return image + + +def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: + _run_pre_test_cleanup(enable_force=True) + pipe: DiffusionPipeline | None = None + try: + pipe = DiffusionPipeline.from_pretrained( + model, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + local_files_only=_local_files_only(model), + ).to("cuda") + generator = torch.Generator(device="cuda").manual_seed(SEED) + result = pipe( # pyright: ignore[reportCallIssue] + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + width=WIDTH, + height=HEIGHT, + num_inference_steps=NUM_INFERENCE_STEPS, + true_cfg_scale=TRUE_CFG_SCALE, + generator=generator, + ) + output_image = result.images[0].convert("RGB") + output_image.save(output_path) + return output_image + finally: + if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"): + pipe.maybe_free_model_hooks() + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + _run_post_test_cleanup(enable_force=True) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: + model = _model_name() + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + + vllm_output = _run_vllm_omni_qwen_image(model=model, output_path=output_dir / "vllm_omni.png") + diffusers_output = _run_diffusers_qwen_image(model=model, output_path=output_dir / "diffusers.png") + + assert_similarity( + model_name=MODEL_ID, + vllm_image=vllm_output, + diffusers_image=diffusers_output, + width=WIDTH, + height=HEIGHT, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + ) diff --git a/tests/e2e/accuracy/test_qwen_image_layered.py b/tests/e2e/accuracy/test_qwen_image_layered.py new file mode 100644 index 000000000..04b13df3b --- /dev/null +++ b/tests/e2e/accuracy/test_qwen_image_layered.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import base64 +import gc +import io +import os +from pathlib import Path + +import pytest +import requests +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from PIL import Image + +from tests.conftest import ( + OmniServer, + _run_post_test_cleanup, + _run_pre_test_cleanup, +) +from tests.e2e.accuracy.utils import assert_image_sequence_similarity, model_output_dir +from tests.utils import hardware_test + +MODEL_ID = "Qwen/Qwen-Image-Layered" +MODEL_ENV_VAR = "QWEN_IMAGE_LAYERED_MODEL" +PROMPT = "decompose into layers" +NEGATIVE_PROMPT = " " +NUM_INFERENCE_STEPS = 20 +TRUE_CFG_SCALE = 4.0 +SEED = 777 +LAYERS = 3 +RESOLUTION = 640 +SSIM_THRESHOLD = 0.97 +PSNR_THRESHOLD = 30.0 + + +def _model_name() -> str: + return os.environ.get(MODEL_ENV_VAR, MODEL_ID) + + +def _local_files_only(model: str) -> bool: + return Path(model).exists() + + +def _normalize_layered_images(images: object) -> list[Image.Image]: + if not isinstance(images, list) or not images: + raise AssertionError(f"Unexpected layered output container: {type(images).__name__}") + + first_item = images[0] + if isinstance(first_item, Image.Image): + return [image.convert("RGBA") for image in images if isinstance(image, Image.Image)] + if isinstance(first_item, (list, tuple)): + return [image.convert("RGBA") for image in first_item if isinstance(image, Image.Image)] + raise AssertionError(f"Unexpected layered image element type: {type(first_item).__name__}") + + +def _run_vllm_omni_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]: + input_image.save(output_dir / "input.png") + server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] + with OmniServer(model, server_args, use_omni=True) as omni_server: + buffer = io.BytesIO() + input_image.save(buffer, format="PNG") + buffer.seek(0) + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/images/edits", + data={ + "model": omni_server.model, + "prompt": PROMPT, + "size": "auto", + "n": 1, + "response_format": "b64_json", + "negative_prompt": NEGATIVE_PROMPT, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": TRUE_CFG_SCALE, + "seed": SEED, + "layers": LAYERS, + "resolution": RESOLUTION, + }, + files=[("image", ("input.png", buffer, "image/png"))], + timeout=600, + ) + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == LAYERS + output_images = [] + for item in payload["data"]: + image_bytes = base64.b64decode(item["b64_json"]) + image = Image.open(io.BytesIO(image_bytes)).convert("RGBA") + image.load() + output_images.append(image) + for index, image in enumerate(output_images, start=1): + image.save(output_dir / f"vllm_omni_layer_{index}.png") + return output_images + + +def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]: + _run_pre_test_cleanup(enable_force=True) + pipe: DiffusionPipeline | None = None + try: + pipe = DiffusionPipeline.from_pretrained( + model, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + local_files_only=_local_files_only(model), + ).to("cuda") + generator = torch.Generator(device="cuda").manual_seed(SEED) + result = pipe( # pyright: ignore[reportCallIssue] + image=input_image, + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + num_inference_steps=NUM_INFERENCE_STEPS, + true_cfg_scale=TRUE_CFG_SCALE, + generator=generator, + num_images_per_prompt=1, + layers=LAYERS, + resolution=RESOLUTION, + ) + output_images = _normalize_layered_images(result.images) + assert len(output_images) == LAYERS, f"Expected {LAYERS} diffusers layers, got {len(output_images)}" + for index, image in enumerate(output_images, start=1): + image.save(output_dir / f"diffusers_layer_{index}.png") + return output_images + finally: + if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"): + pipe.maybe_free_model_hooks() + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + _run_post_test_cleanup(enable_force=True) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_qwen_image_layered_matches_diffusers(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> None: + model = _model_name() + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + input_image = qwen_bear_image.convert("RGBA") + + vllm_outputs = _run_vllm_omni_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir) + diffusers_outputs = _run_diffusers_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir) + + assert_image_sequence_similarity( + model_name=MODEL_ID, + vllm_images=vllm_outputs, + diffusers_images=diffusers_outputs, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + compare_mode="RGBA", + ) diff --git a/tests/e2e/accuracy/utils.py b/tests/e2e/accuracy/utils.py index eb0eea757..d722b69b0 100644 --- a/tests/e2e/accuracy/utils.py +++ b/tests/e2e/accuracy/utils.py @@ -21,13 +21,14 @@ def assert_similarity( model_name: str, vllm_image: Image.Image, diffusers_image: Image.Image, - width: int, - height: int, ssim_threshold: float, psnr_threshold: float, + width: int | None = None, + height: int | None = None, + compare_mode: str = "RGB", ) -> None: - requested_size = (width, height) - if diffusers_image.size != requested_size: + requested_size = (width, height) if width is not None and height is not None else None + if requested_size is not None and diffusers_image.size != requested_size: pytest.skip( "Skipping as diffusers baseline output is corrupt and not comparable: " f"dimensions do not match requested size; requested={requested_size}, got={diffusers_image.size}." @@ -37,7 +38,11 @@ def assert_similarity( f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}" ) - ssim_score, psnr_score = compute_image_ssim_psnr(prediction=vllm_image, reference=diffusers_image) + ssim_score, psnr_score = compute_image_ssim_psnr( + prediction=vllm_image, + reference=diffusers_image, + compare_mode=compare_mode, + ) print(f"{model_name} similarity metrics:") print(f" SSIM: value={ssim_score:.6f}, threshold>={ssim_threshold:.6f}, range=[-1, 1], higher_is_better=True") print( @@ -52,13 +57,37 @@ def assert_similarity( ) +def assert_image_sequence_similarity( + *, + model_name: str, + vllm_images: list[Image.Image], + diffusers_images: list[Image.Image], + ssim_threshold: float, + psnr_threshold: float, + compare_mode: str = "RGB", +) -> None: + assert len(vllm_images) == len(diffusers_images), ( + f"Output image count mismatch for {model_name}: online={len(vllm_images)}, diffusers={len(diffusers_images)}" + ) + for index, (vllm_image, diffusers_image) in enumerate(zip(vllm_images, diffusers_images, strict=True), start=1): + assert_similarity( + model_name=f"{model_name}[layer={index}]", + vllm_image=vllm_image, + diffusers_image=diffusers_image, + ssim_threshold=ssim_threshold, + psnr_threshold=psnr_threshold, + compare_mode=compare_mode, + ) + + def compute_image_ssim_psnr( *, prediction: Image.Image, reference: Image.Image, + compare_mode: str = "RGB", ) -> tuple[float, float]: - pred_tensor = _pil_to_batched_tensor(prediction) - ref_tensor = _pil_to_batched_tensor(reference) + pred_tensor = _pil_to_batched_tensor(prediction, compare_mode=compare_mode) + ref_tensor = _pil_to_batched_tensor(reference, compare_mode=compare_mode) ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0) psnr_metric = PeakSignalNoiseRatio(data_range=1.0) @@ -68,7 +97,7 @@ def compute_image_ssim_psnr( return ssim_value, psnr_value -def _pil_to_batched_tensor(image: Image.Image) -> torch.Tensor: - array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 +def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Tensor: + array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0 tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) return tensor