207 lines
6.9 KiB
Python
207 lines
6.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""FLUX.2-dev VAE patch-parallel correctness tests.
|
|
|
|
Compares vae_patch_parallel_size=1 vs 2 on TP=2 with tiling enabled:
|
|
- Text-to-image decode path
|
|
- Image-to-image encode + decode path (conditioning image)
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniRunner
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
|
|
PROMPT = "a photo of a cat sitting on a laptop keyboard"
|
|
NEGATIVE_PROMPT = "blurry, low quality"
|
|
|
|
VAE_PP_MEAN_THRESHOLD = 3e-2
|
|
VAE_PP_P99_THRESHOLD = 1e-1
|
|
|
|
|
|
def _get_flux2_dev_model() -> str:
|
|
return os.environ.get("VLLM_TEST_FLUX2_DEV_MODEL", "black-forest-labs/FLUX.2-dev")
|
|
|
|
|
|
def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
|
|
arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
|
|
return torch.from_numpy(arr)
|
|
|
|
|
|
def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
|
|
ta = _pil_to_float_rgb_tensor(a)
|
|
tb = _pil_to_float_rgb_tensor(b)
|
|
assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
|
|
abs_diff = torch.abs(ta - tb)
|
|
p99_abs_diff = torch.quantile(abs_diff.flatten(), 0.99).item()
|
|
return abs_diff.mean().item(), p99_abs_diff
|
|
|
|
|
|
def _extract_single_image(outputs) -> Image.Image:
|
|
first_output = outputs[0]
|
|
assert first_output.final_output_type == "image"
|
|
if not isinstance(first_output, OmniRequestOutput):
|
|
raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}")
|
|
|
|
images = first_output.images
|
|
if images is None or len(images) != 1:
|
|
raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
|
|
return images[0]
|
|
|
|
|
|
def _run_flux2_generate(
|
|
*,
|
|
tp_size: int,
|
|
height: int,
|
|
width: int,
|
|
num_inference_steps: int,
|
|
seed: int,
|
|
vae_patch_parallel_size: int,
|
|
conditioning_image: Image.Image | None = None,
|
|
) -> Image.Image:
|
|
current_omni_platform.empty_cache()
|
|
|
|
request: dict = {"prompt": PROMPT, "negative_prompt": NEGATIVE_PROMPT}
|
|
if conditioning_image is not None:
|
|
request["multi_modal_data"] = {"image": conditioning_image}
|
|
|
|
with OmniRunner(
|
|
_get_flux2_dev_model(),
|
|
parallel_config=DiffusionParallelConfig(
|
|
tensor_parallel_size=tp_size,
|
|
vae_patch_parallel_size=vae_patch_parallel_size,
|
|
),
|
|
enable_cpu_offload=True,
|
|
vae_use_tiling=True,
|
|
) as runner:
|
|
outputs = list(
|
|
runner.omni.generate(
|
|
request,
|
|
OmniDiffusionSamplingParams(
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
guidance_scale=4.0,
|
|
seed=seed,
|
|
num_outputs_per_prompt=1,
|
|
),
|
|
)
|
|
)
|
|
|
|
return _extract_single_image(outputs)
|
|
|
|
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.parallel
|
|
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
|
def test_flux2_vae_patch_parallel_decode_tp2(tmp_path: Path):
|
|
if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
|
|
pytest.skip("FLUX.2-dev VAE patch parallel requires >= 2 devices.")
|
|
|
|
height = 1152
|
|
width = 1152
|
|
num_inference_steps = 2
|
|
seed = 42
|
|
|
|
baseline_img = _run_flux2_generate(
|
|
tp_size=2,
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
seed=seed,
|
|
vae_patch_parallel_size=1,
|
|
)
|
|
pp2_img = _run_flux2_generate(
|
|
tp_size=2,
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
seed=seed,
|
|
vae_patch_parallel_size=2,
|
|
)
|
|
|
|
baseline_path = tmp_path / "flux2_tp2_vae_pp1.png"
|
|
pp2_path = tmp_path / "flux2_tp2_vae_pp2.png"
|
|
baseline_img.save(baseline_path)
|
|
pp2_img.save(pp2_path)
|
|
|
|
assert baseline_img.width == width and baseline_img.height == height
|
|
assert pp2_img.width == width and pp2_img.height == height
|
|
|
|
mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
|
|
print(
|
|
"FLUX.2-dev VAE patch parallel decode diff (TP=2, pp=1 vs pp=2): "
|
|
f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
|
|
f"thresholds: mean<={VAE_PP_MEAN_THRESHOLD:.6e}, p99<={VAE_PP_P99_THRESHOLD:.6e}; "
|
|
f"pp1_img={baseline_path}, pp2_img={pp2_path}"
|
|
)
|
|
assert mean_abs_diff <= VAE_PP_MEAN_THRESHOLD and p99_abs_diff <= VAE_PP_P99_THRESHOLD, (
|
|
f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
|
|
f"(thresholds: mean<={VAE_PP_MEAN_THRESHOLD:.6e}, p99<={VAE_PP_P99_THRESHOLD:.6e})"
|
|
)
|
|
|
|
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.parallel
|
|
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
|
def test_flux2_vae_patch_parallel_i2i_tp2(tmp_path: Path):
|
|
if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
|
|
pytest.skip("FLUX.2-dev VAE patch parallel I2I requires >= 2 devices.")
|
|
|
|
height = 1152
|
|
width = 1152
|
|
num_inference_steps = 2
|
|
seed = 42
|
|
conditioning_image = Image.new("RGB", (width, height), (120, 80, 40))
|
|
|
|
baseline_img = _run_flux2_generate(
|
|
tp_size=2,
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
seed=seed,
|
|
vae_patch_parallel_size=1,
|
|
conditioning_image=conditioning_image,
|
|
)
|
|
pp2_img = _run_flux2_generate(
|
|
tp_size=2,
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
seed=seed,
|
|
vae_patch_parallel_size=2,
|
|
conditioning_image=conditioning_image,
|
|
)
|
|
|
|
baseline_path = tmp_path / "flux2_i2i_tp2_vae_pp1.png"
|
|
pp2_path = tmp_path / "flux2_i2i_tp2_vae_pp2.png"
|
|
baseline_img.save(baseline_path)
|
|
pp2_img.save(pp2_path)
|
|
|
|
mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
|
|
print(
|
|
"FLUX.2-dev VAE patch parallel I2I diff (TP=2, pp=1 vs pp=2): "
|
|
f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
|
|
f"thresholds: mean<={VAE_PP_MEAN_THRESHOLD:.6e}, p99<={VAE_PP_P99_THRESHOLD:.6e}; "
|
|
f"pp1_img={baseline_path}, pp2_img={pp2_path}"
|
|
)
|
|
assert mean_abs_diff <= VAE_PP_MEAN_THRESHOLD and p99_abs_diff <= VAE_PP_P99_THRESHOLD, (
|
|
f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
|
|
f"(thresholds: mean<={VAE_PP_MEAN_THRESHOLD:.6e}, p99<={VAE_PP_P99_THRESHOLD:.6e})"
|
|
)
|