anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feature] Flux2 klein inpaint (#1180)

Browse Source 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
LancerandHongsheng Liu authored and GitHub committed 2026-04-17 07:55:00 +08:00
1 parent 817e32d548
commit bf64a6b841
6 files changed
+583 -7
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/e2e/offline_inference/test_flux2_klein_inpaint.py
+216
View File
@@ -0,0 +1,216 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end test for Flux2 Klein inpainting.
"""
# ruff: noqa: E402
import os
import sys
from pathlib import Path
import pytest
import torch
from PIL import Image, ImageDraw
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni import Omni
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
MODEL = "black-forest-labs/FLUX.2-klein-4B"
_HEIGHT = 512
_WIDTH = 512
_NUM_INFERENCE_STEPS = 4
def _create_test_image(width: int = _WIDTH, height: int = _HEIGHT, color: tuple = (128, 128, 128)) -> Image.Image:
return Image.new("RGB", (width, height), color)
def _create_test_mask(width: int = _WIDTH, height: int = _HEIGHT) -> Image.Image:
mask = Image.new("L", (width, height), 0)
draw = ImageDraw.Draw(mask)
draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
return mask
def _create_test_inputs(color: tuple = (100, 150, 200)):
return _create_test_image(_WIDTH, _HEIGHT, color), _create_test_mask(_WIDTH, _HEIGHT)
def _extract_images_from_output(outputs: list) -> list[Image.Image]:
images = []
for req_output in outputs:
if hasattr(req_output, "images") and req_output.images:
images.extend(req_output.images)
elif hasattr(req_output, "request_output") and req_output.request_output:
stage_out = req_output.request_output
if isinstance(stage_out, OmniRequestOutput) and hasattr(stage_out, "images"):
images.extend(stage_out.images)
elif isinstance(stage_out, list):
for s in stage_out:
if hasattr(s, "images") and s.images:
images.extend(s.images)
return images
@pytest.mark.core_model
@pytest.mark.diffusion
def test_flux2_klein_inpaint_basic():
m = None
try:
m = Omni(model=MODEL)
input_image, mask_image = _create_test_inputs()
outputs = m.generate(
prompts=[
{
"prompt": "Fill in the masked area with a beautiful garden",
"multi_modal_data": {"image": input_image, "mask_image": mask_image},
}
],
sampling_params_list=OmniDiffusionSamplingParams(
height=_HEIGHT,
width=_WIDTH,
num_inference_steps=_NUM_INFERENCE_STEPS,
guidance_scale=0.0,
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
num_outputs_per_prompt=1,
),
)
images = _extract_images_from_output(list(outputs))
assert len(images) == 1
assert images[0].size == (_WIDTH, _HEIGHT)
finally:
if m is not None and hasattr(m, "close"):
m.close()
@pytest.mark.diffusion
def test_flux2_klein_inpaint_deterministic():
m = None
try:
m = Omni(model=MODEL)
input_image, mask_image = _create_test_inputs()
seed = 12345
gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(seed)
gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(seed)
outputs1 = m.generate(
prompts=[
{
"prompt": "A red flower in a field",
"multi_modal_data": {"image": input_image, "mask_image": mask_image},
}
],
sampling_params_list=OmniDiffusionSamplingParams(
height=_HEIGHT,
width=_WIDTH,
num_inference_steps=_NUM_INFERENCE_STEPS,
guidance_scale=0.0,
generator=gen1,
num_outputs_per_prompt=1,
),
)
outputs2 = m.generate(
prompts=[
{
"prompt": "A red flower in a field",
"multi_modal_data": {"image": input_image, "mask_image": mask_image},
}
],
sampling_params_list=OmniDiffusionSamplingParams(
height=_HEIGHT,
width=_WIDTH,
num_inference_steps=_NUM_INFERENCE_STEPS,
guidance_scale=0.0,
generator=gen2,
num_outputs_per_prompt=1,
),
)
images1 = _extract_images_from_output(list(outputs1))
images2 = _extract_images_from_output(list(outputs2))
assert len(images1) == 1
assert len(images2) == 1
assert list(images1[0].getdata()) == list(images2[0].getdata()), (
"Same input with same seed should produce identical output. "
"This is critical for offline/online consistency."
)
finally:
if m is not None and hasattr(m, "close"):
m.close()
@pytest.mark.diffusion
def test_flux2_klein_inpaint_different_seeds_different_output():
m = None
try:
m = Omni(model=MODEL)
input_image, mask_image = _create_test_inputs()
gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(42)
gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(99999)
outputs1 = m.generate(
prompts=[
{
"prompt": "A beautiful landscape",
"multi_modal_data": {"image": input_image, "mask_image": mask_image},
}
],
sampling_params_list=OmniDiffusionSamplingParams(
height=_HEIGHT,
width=_WIDTH,
num_inference_steps=_NUM_INFERENCE_STEPS,
guidance_scale=0.0,
generator=gen1,
num_outputs_per_prompt=1,
),
)
outputs2 = m.generate(
prompts=[
{
"prompt": "A beautiful landscape",
"multi_modal_data": {"image": input_image, "mask_image": mask_image},
}
],
sampling_params_list=OmniDiffusionSamplingParams(
height=_HEIGHT,
width=_WIDTH,
num_inference_steps=_NUM_INFERENCE_STEPS,
guidance_scale=0.0,
generator=gen2,
num_outputs_per_prompt=1,
),
)
images1 = _extract_images_from_output(list(outputs1))
images2 = _extract_images_from_output(list(outputs2))
assert len(images1) == 1
assert len(images2) == 1
different_pixel_count = sum(1 for p1, p2 in zip(images1[0].getdata(), images2[0].getdata()) if p1 != p2)
assert different_pixel_count > 0, "Different seeds should produce different outputs"
finally:
if m is not None and hasattr(m, "close"):
m.close()
tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py
+160
View File
@@ -0,0 +1,160 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end tests for Flux2 Klein inpainting in online serving mode.
Uses /v1/images/edits endpoint which is the correct API for image inpainting.
"""
import base64
from io import BytesIO
import httpx
import pytest
from PIL import Image, ImageDraw
from tests.conftest import OmniServer, OmniServerParams
MODEL = "black-forest-labs/FLUX.2-klein-4B"
_HEIGHT = 512
_WIDTH = 512
_NUM_INFERENCE_STEPS = 4
def _get_diffusion_feature_cases(model: str):
return [
pytest.param(
OmniServerParams(
model=model,
server_args=["--tensor-parallel-size", "2"],
),
id="tp2_basic",
),
]
def _image_to_base64_jpeg(image: Image.Image) -> str:
buffer = BytesIO()
image.save(buffer, format="JPEG")
buffer.seek(0)
return base64.b64encode(buffer.read()).decode("utf-8")
def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str:
mask = Image.new("L", (width, height), 0)
draw = ImageDraw.Draw(mask)
draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
return _image_to_base64_jpeg(mask)
def _compare_images(img1: Image.Image, img2: Image.Image) -> bool:
return list(img1.getdata()) == list(img2.getdata())
def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs):
url = f"http://{host}:{port}/v1/images/edits"
files = {
"image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"),
"mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"),
}
data = {"prompt": prompt, "model": model, **kwargs}
with httpx.Client(timeout=60.0) as client:
response = client.post(url, files=files, data=data)
response.raise_for_status()
return response.json()
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_flux2_klein_inpaint_basic(omni_server: OmniServer):
input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
mask_b64 = _create_test_mask_base64()
result = _send_edit_request(
host=omni_server.host,
port=omni_server.port,
model=MODEL,
image_b64=input_image_b64,
mask_b64=mask_b64,
prompt="Fill in the masked area with a beautiful garden",
guidance_scale=1.0,
num_inference_steps=_NUM_INFERENCE_STEPS,
n=1,
seed=42,
)
assert "data" in result and len(result["data"]) == 1
img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1]
img = Image.open(BytesIO(base64.b64decode(img_data)))
assert img.size == (_WIDTH, _HEIGHT)
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer):
input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
mask_b64 = _create_test_mask_base64()
prompt = "A red flower in a field"
result1 = _send_edit_request(
host=omni_server.host,
port=omni_server.port,
model=MODEL,
image_b64=input_image_b64,
mask_b64=mask_b64,
prompt=prompt,
guidance_scale=1.0,
num_inference_steps=_NUM_INFERENCE_STEPS,
n=1,
seed=12345,
)
result2 = _send_edit_request(
host=omni_server.host,
port=omni_server.port,
model=MODEL,
image_b64=input_image_b64,
mask_b64=mask_b64,
prompt=prompt,
guidance_scale=1.0,
num_inference_steps=_NUM_INFERENCE_STEPS,
n=1,
seed=12345,
)
img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1]
img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1]
img1 = Image.open(BytesIO(base64.b64decode(img1_data)))
img2 = Image.open(BytesIO(base64.b64decode(img2_data)))
assert _compare_images(img1, img2), (
"Same input with same seed should produce identical output. This is critical for offline/online consistency."
)
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer):
input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
mask_b64 = _create_test_mask_base64()
result = _send_edit_request(
host=omni_server.host,
port=omni_server.port,
model=MODEL,
image_b64=input_image_b64,
mask_b64=mask_b64,
prompt="A beautiful landscape",
guidance_scale=1.0,
num_inference_steps=_NUM_INFERENCE_STEPS,
n=2,
seed=42,
)
assert "data" in result and len(result["data"]) == 2