229 lines
7.7 KiB
Python
229 lines
7.7 KiB
Python
from __future__ import annotations
|
|
|
|
import gc
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import requests
|
|
import torch
|
|
from diffusers import QwenImageEditPipeline, QwenImageEditPlusPipeline
|
|
from PIL import Image
|
|
|
|
from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
|
|
from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
|
|
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniServer
|
|
|
|
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
|
|
|
|
|
SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
|
|
MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509"
|
|
WIDTH = 512
|
|
HEIGHT = 512
|
|
NUM_INFERENCE_STEPS = 20
|
|
TRUE_CFG_SCALE = 4.0
|
|
SEED = 42
|
|
SSIM_THRESHOLD = 0.94
|
|
PSNR_THRESHOLD = 28.0
|
|
|
|
PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure."
|
|
PROMPT_MULTIPLE_IMAGE = "Put the cartoon bear mascot and the furry rabbit into one coherent scene with a painterly oil artwork style and consistent lighting."
|
|
NEGATIVE_PROMPT = "low quality, blurry, artifacts, distortion"
|
|
SERVER_ARGS = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
|
|
|
|
|
|
def _run_vllm_omni_image_edit(
|
|
*,
|
|
omni_server: OmniServer,
|
|
prompt: str,
|
|
input_images: list[Image.Image],
|
|
output_path: Path,
|
|
) -> Image.Image:
|
|
response = requests.post(
|
|
f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
|
|
data={
|
|
"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,
|
|
},
|
|
files=[
|
|
("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png"))
|
|
for index, image in enumerate(input_images)
|
|
],
|
|
timeout=600,
|
|
)
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
assert len(payload["data"]) == 1
|
|
image = decode_base64_image(payload["data"][0]["b64_json"])
|
|
image.load()
|
|
image.save(output_path)
|
|
return image
|
|
|
|
|
|
def _run_diffusers_image_edit(
|
|
*,
|
|
model: str,
|
|
pipeline_class: type[QwenImageEditPipeline] | type[QwenImageEditPlusPipeline],
|
|
prompt: str,
|
|
input_images: list[Image.Image],
|
|
output_path: Path,
|
|
) -> Image.Image:
|
|
run_pre_test_cleanup(enable_force=True)
|
|
pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None
|
|
device = torch.device("cuda:0")
|
|
torch.cuda.set_device(device)
|
|
try:
|
|
images = input_images[0] if len(input_images) == 1 else input_images
|
|
pipe = pipeline_class.from_pretrained(
|
|
model,
|
|
torch_dtype=torch.bfloat16,
|
|
trust_remote_code=True,
|
|
).to(device)
|
|
pipe.set_progress_bar_config(disable=False)
|
|
generator = torch.Generator(device=device).manual_seed(SEED)
|
|
result = pipe( # pyright: ignore[reportCallIssue]
|
|
prompt=prompt,
|
|
image=images,
|
|
negative_prompt=NEGATIVE_PROMPT,
|
|
num_inference_steps=NUM_INFERENCE_STEPS,
|
|
true_cfg_scale=TRUE_CFG_SCALE,
|
|
width=WIDTH,
|
|
height=HEIGHT,
|
|
generator=generator,
|
|
)
|
|
output_image = result.images[0].convert("RGB") # pyright: ignore[reportAttributeAccessIssue]
|
|
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.accelerator.empty_cache()
|
|
run_post_test_cleanup(enable_force=True)
|
|
|
|
|
|
def _vllm_omni_output_single_image(
|
|
accuracy_artifact_root: Path,
|
|
qwen_bear_image: Image.Image,
|
|
) -> Image.Image:
|
|
output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL)
|
|
output_path = output_dir / "vllm_omni_single.png"
|
|
with OmniServer(model=SINGLE_MODEL, serve_args=SERVER_ARGS) as server:
|
|
output = _run_vllm_omni_image_edit(
|
|
omni_server=server,
|
|
prompt=PROMPT_SINGLE_IMAGE,
|
|
input_images=[qwen_bear_image],
|
|
output_path=output_path,
|
|
)
|
|
return output
|
|
|
|
|
|
def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> Image.Image:
|
|
output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL)
|
|
output_path = output_dir / "diffusers_single.png"
|
|
return _run_diffusers_image_edit(
|
|
model=SINGLE_MODEL,
|
|
pipeline_class=QwenImageEditPipeline,
|
|
prompt=PROMPT_SINGLE_IMAGE,
|
|
input_images=[qwen_bear_image],
|
|
output_path=output_path,
|
|
)
|
|
|
|
|
|
def _vllm_omni_output_multiple_image(
|
|
accuracy_artifact_root: Path,
|
|
qwen_bear_image: Image.Image,
|
|
rabbit_image: Image.Image,
|
|
) -> Image.Image:
|
|
output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
|
|
output_path = output_dir / "vllm_omni_multiple.png"
|
|
with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
|
|
output = _run_vllm_omni_image_edit(
|
|
omni_server=server,
|
|
prompt=PROMPT_MULTIPLE_IMAGE,
|
|
input_images=[qwen_bear_image, rabbit_image],
|
|
output_path=output_path,
|
|
)
|
|
return output
|
|
|
|
|
|
def _diffusers_output_multiple_image(
|
|
accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
|
|
) -> Image.Image:
|
|
output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
|
|
output_path = output_dir / "diffusers_multiple.png"
|
|
return _run_diffusers_image_edit(
|
|
model=MULTIPLE_MODEL,
|
|
pipeline_class=QwenImageEditPlusPipeline,
|
|
prompt=PROMPT_MULTIPLE_IMAGE,
|
|
input_images=[qwen_bear_image, rabbit_image],
|
|
output_path=output_path,
|
|
)
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
|
def test_qwen_image_edit_single_matches_diffusers(
|
|
accuracy_artifact_root: Path,
|
|
qwen_bear_image: Image.Image,
|
|
) -> None:
|
|
vllm_image = _vllm_omni_output_single_image(
|
|
accuracy_artifact_root=accuracy_artifact_root,
|
|
qwen_bear_image=qwen_bear_image,
|
|
)
|
|
diffusers_image = _diffusers_output_single_image(
|
|
accuracy_artifact_root=accuracy_artifact_root,
|
|
qwen_bear_image=qwen_bear_image,
|
|
)
|
|
assert_similarity(
|
|
model_name=SINGLE_MODEL,
|
|
vllm_image=vllm_image,
|
|
diffusers_image=diffusers_image,
|
|
width=WIDTH,
|
|
height=HEIGHT,
|
|
ssim_threshold=SSIM_THRESHOLD,
|
|
psnr_threshold=PSNR_THRESHOLD,
|
|
)
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
|
@pytest.mark.skip(
|
|
reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
|
|
)
|
|
def test_qwen_image_edit_multiple_matches_diffusers(
|
|
accuracy_artifact_root: Path,
|
|
qwen_bear_image: Image.Image,
|
|
rabbit_image: Image.Image,
|
|
) -> None:
|
|
vllm_image = _vllm_omni_output_multiple_image(
|
|
accuracy_artifact_root=accuracy_artifact_root,
|
|
qwen_bear_image=qwen_bear_image,
|
|
rabbit_image=rabbit_image,
|
|
)
|
|
diffusers_image = _diffusers_output_multiple_image(
|
|
accuracy_artifact_root=accuracy_artifact_root,
|
|
qwen_bear_image=qwen_bear_image,
|
|
rabbit_image=rabbit_image,
|
|
)
|
|
assert_similarity(
|
|
model_name=MULTIPLE_MODEL,
|
|
vllm_image=vllm_image,
|
|
diffusers_image=diffusers_image,
|
|
width=WIDTH,
|
|
height=HEIGHT,
|
|
ssim_threshold=SSIM_THRESHOLD,
|
|
psnr_threshold=PSNR_THRESHOLD,
|
|
)
|