[Feature] Flux2 klein inpaint (#1180)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/serving/image_edit_api.md

@@ -104,6 +104,8 @@ Content-Type: multipart/form-data
 | `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
 | `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
 | `seed` | integer | null | Random seed for reproducibility |
+| `reference_image` | string or array | null | Reference image for inpainting |
+| `mask_image` | string or array | null | Mask for inpainting (white areas will be inpainted) |
 
 ### Response Format
 

tests/e2e/offline_inference/test_flux2_klein_inpaint.py

@@ -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

@@ -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

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -234,7 +234,15 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         self.transformer = Flux2Transformer2DModel(quant_config=od_config.quantization_config, **transformer_kwargs)
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        self.latent_channels = self.vae.config.latent_channels if hasattr(self.vae, "config") else 16
         self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.mask_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor * 2,
+            vae_latent_channels=self.latent_channels,
+            do_normalize=False,
+            do_binarize=True,
+            do_convert_grayscale=True,
+        )
         self.tokenizer_max_length = 512
         self.default_sample_size = 128
 
@@ -247,6 +255,14 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
         )
 
+    def get_timesteps(self, num_inference_steps, strength, device):
+        init_timestep = min(num_inference_steps * strength, num_inference_steps)
+        t_start = int(max(num_inference_steps - init_timestep, 0))
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+        if hasattr(self.scheduler, "set_begin_index"):
+            self.scheduler.set_begin_index(t_start * self.scheduler.order)
+        return timesteps, num_inference_steps - t_start
+
     @staticmethod
     def _get_qwen3_prompt_embeds(
         text_encoder: Qwen3ForCausalLM,
@@ -582,6 +598,54 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
 
         return image_latents, image_latent_ids
 
+    def prepare_mask_latents(
+        self,
+        mask,
+        masked_image,
+        batch_size,
+        num_channels_latents,
+        num_images_per_prompt,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+    ):
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+        mask = torch.nn.functional.interpolate(mask, size=(height, width))
+        mask = mask.to(device=device, dtype=dtype)
+
+        batch_size = batch_size * num_images_per_prompt
+
+        if masked_image is not None:
+            masked_image = masked_image.to(device=device, dtype=dtype)
+            if masked_image.shape[1] != num_channels_latents:
+                masked_image_latents = self._encode_vae_image(image=masked_image, generator=generator)
+            else:
+                masked_image_latents = masked_image
+        else:
+            masked_image_latents = None
+
+        if mask.shape[0] < batch_size:
+            if not batch_size % mask.shape[0] == 0:
+                raise ValueError("The passed mask and the required batch size don't match.")
+            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
+        if masked_image_latents is not None and masked_image_latents.shape[0] < batch_size:
+            if not batch_size % masked_image_latents.shape[0] == 0:
+                raise ValueError("The passed mask and the required batch size don't match.")
+            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)
+
+        if masked_image_latents is not None:
+            masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
+            masked_image_latents = self._pack_latents(masked_image_latents)
+
+        mask = mask.repeat(1, self.latent_channels, 1, 1)
+        mask = self._patchify_latents(mask)
+        mask = self._pack_latents(mask)
+
+        return mask, masked_image_latents
+
     def check_inputs(
         self,
         prompt,
@@ -590,7 +654,15 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         prompt_embeds=None,
         callback_on_step_end_tensor_inputs=None,
         guidance_scale=None,
+        strength=None,
+        num_inference_steps=None,
     ):
+        if strength is not None:
+            if strength < 0 or strength > 1:
+                raise ValueError(f"strength must be between 0 and 1, got {strength}")
+        if num_inference_steps is not None and num_inference_steps <= 0:
+            raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}")
+
         if (
             height is not None
             and height % (self.vae_scale_factor * 2) != 0
@@ -622,7 +694,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
             raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
 
-        if guidance_scale > 1.0 and self.is_distilled:
+        if guidance_scale is not None and guidance_scale > 1.0 and self.is_distilled:
             logger.warning(f"Guidance scale {guidance_scale} is ignored for step-wise distilled models.")
 
     @property
@@ -653,11 +725,14 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         self,
         req: OmniDiffusionRequest,
         image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
+        reference_image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
+        mask_image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
         prompt: str | list[str] | None = None,
         height: int | None = None,
         width: int | None = None,
         num_inference_steps: int = 50,
         sigmas: list[float] | None = None,
+        strength: float = 1.0,
         guidance_scale: float | None = 4.0,
         num_images_per_prompt: int = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
@@ -671,6 +746,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         callback_on_step_end_tensor_inputs: list[str] = ["latents"],
         max_sequence_length: int = 512,
         text_encoder_out_layers: tuple[int, ...] = (9, 18, 27),
+        padding_mask_crop: int | None = None,
     ) -> DiffusionOutput:
         r"""
         Function invoked when calling the pipeline for generation.
@@ -804,6 +880,8 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             prompt_embeds=prompt_embeds,
             callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
             guidance_scale=guidance_scale,
+            strength=strength,
+            num_inference_steps=num_inference_steps,
         )
 
         self._guidance_scale = guidance_scale
@@ -848,6 +926,9 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         if image is not None and not isinstance(image, list):
             image = [image]
 
+        multiple_of = self.vae_scale_factor * 2
+        crops_coords = None
+        resize_mode = "crop"
         condition_images = None
         if image is not None:
             for img in image:
@@ -860,10 +941,14 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
                     img = self.image_processor._resize_to_target_area(img, 1024 * 1024)
                     image_width, image_height = img.size
 
-                multiple_of = self.vae_scale_factor * 2
                 image_width = (image_width // multiple_of) * multiple_of
                 image_height = (image_height // multiple_of) * multiple_of
-                img = self.image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop")
+                if padding_mask_crop is not None:
+                    crops_coords = self.mask_processor.get_crop_region(mask_image, width, height, pad=padding_mask_crop)
+                    resize_mode = "fill"
+                img = self.image_processor.preprocess(
+                    img, height=image_height, width=image_width, crops_coords=crops_coords, resize_mode=resize_mode
+                )
                 condition_images.append(img)
                 height = height or image_height
                 width = width or image_width
@@ -871,6 +956,16 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         height = height or self.default_sample_size * self.vae_scale_factor
         width = width or self.default_sample_size * self.vae_scale_factor
 
+        # Get mask_image and reference_image
+        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if req.prompts else {}
+        mask_image = multi_modal_data.get("mask_image")
+        reference_image = multi_modal_data.get("reference_image")
+
+        if mask_image is not None and (image is None or (isinstance(image, list) and len(image) == 0)):
+            raise ValueError("image must be provided when using mask_image for inpainting")
+
+        init_image = condition_images[0] if condition_images else None
+
         # 5. prepare latent variables
         num_channels_latents = self.transformer.config.in_channels // 4
         latents, latent_ids = self.prepare_latents(
@@ -884,6 +979,8 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             latents=latents,
         )
 
+        original_latent_ids = latent_ids
+
         image_latents = None
         image_latent_ids = None
         if condition_images is not None:
@@ -895,6 +992,71 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
                 dtype=self.vae.dtype,
             )
 
+        # Preprocess reference_image
+        if reference_image is not None and not (
+            isinstance(reference_image, torch.Tensor) and reference_image.size(1) == self.latent_channels
+        ):
+            if (
+                isinstance(reference_image, list)
+                and isinstance(reference_image[0], torch.Tensor)
+                and reference_image[0].ndim == 4
+            ):
+                reference_image = torch.cat(reference_image, dim=0)
+            img_reference = reference_image[0] if isinstance(reference_image, list) else reference_image
+            reference_image_height, reference_image_width = self.image_processor.get_default_height_width(img_reference)
+
+            reference_image_width = reference_image_width // multiple_of * multiple_of
+            reference_image_height = reference_image_height // multiple_of * multiple_of
+            reference_image = self.image_processor.resize(
+                reference_image, reference_image_height, reference_image_width
+            )
+            reference_image = self.image_processor.preprocess(
+                reference_image,
+                reference_image_height,
+                reference_image_width,
+                crops_coords=crops_coords,
+                resize_mode=resize_mode,
+            )
+        else:
+            pass
+
+        reference_image_latents = None
+        reference_image_latent_ids = None
+        if reference_image is not None:
+            reference_image_latents, reference_image_latent_ids = self.prepare_image_latents(
+                images=[reference_image],
+                batch_size=batch_size * num_images_per_prompt,
+                generator=generator,
+                device=device,
+                dtype=self.vae.dtype,
+            )
+
+        if reference_image_latent_ids is not None:
+            latent_ids = torch.cat([latent_ids, reference_image_latent_ids], dim=1)
+        elif image_latent_ids is not None:
+            latent_ids = torch.cat([latent_ids, image_latent_ids], dim=1)
+
+        mask = None
+        if mask_image is not None:
+            mask_condition = self.mask_processor.preprocess(
+                mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords
+            )
+            if init_image is not None:
+                masked_image = init_image * (mask_condition < 0.5).to(init_image.dtype)
+            else:
+                masked_image = None
+            mask, _ = self.prepare_mask_latents(
+                mask_condition,
+                masked_image,
+                batch_size,
+                self.latent_channels,
+                num_images_per_prompt,
+                height,
+                width,
+                prompt_embeds.dtype,
+                device,
+                generator,
+            )
         # 6. Prepare timesteps
         sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
         if hasattr(self.scheduler.config, "use_flow_sigmas") and self.scheduler.config.use_flow_sigmas:
@@ -908,6 +1070,13 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             sigmas=sigmas,
             mu=mu,
         )
+        if reference_image is not None or mask_image is not None:
+            timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+            if num_inference_steps < 1:
+                raise ValueError(
+                    f"After adjusting the num_inference_steps by strength parameter: {strength}, the number of pipeline"
+                    f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline."
+                )
         self._num_timesteps = len(timesteps)
 
         # 7. Denoising loop
@@ -924,7 +1093,10 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             latent_model_input = latents.to(self.transformer.dtype)
             latent_image_ids = latent_ids
 
-            if image_latents is not None:
+            if reference_image_latents is not None:
+                latent_model_input = torch.cat([latents, reference_image_latents], dim=1)
+                latent_image_ids = latent_ids
+            elif image_latents is not None:
                 latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype)
                 latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1)
 
@@ -953,7 +1125,9 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
                 negative_kwargs = None
 
             # For editing pipelines, we need to slice the output to remove condition latents
-            output_slice = latents.size(1) if image_latents is not None else None
+            output_slice = (
+                latents.size(1) if (image_latents is not None or reference_image_latents is not None) else None
+            )
 
             noise_pred = self.predict_noise_maybe_with_cfg(
                 do_true_cfg=self.do_classifier_free_guidance,
@@ -964,9 +1138,22 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
                 output_slice=output_slice,
             )
 
+            latents_dtype = latents.dtype
             # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync
             latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, self.do_classifier_free_guidance)
 
+            if mask is not None and image_latents is not None:
+                init_latents_proper = image_latents
+                if i < len(timesteps) - 1:
+                    noise_timestep = timesteps[i + 1]
+                    init_latents_proper = self.scheduler.scale_noise(
+                        init_latents_proper, torch.tensor([noise_timestep], device=device), latents
+                    )
+                latents = (1 - mask) * init_latents_proper + mask * latents
+
+            if latents.dtype != latents_dtype and torch.backends.mps.is_available():
+                latents = latents.to(latents_dtype)
+
             if callback_on_step_end is not None:
                 callback_kwargs = {}
                 for k in callback_on_step_end_tensor_inputs:
@@ -978,7 +1165,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
 
         self._current_timestep = None
 
-        latents = self._unpack_latents_with_ids(latents, latent_ids)
+        latents = self._unpack_latents_with_ids(latents, original_latent_ids)
 
         latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
         latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(

vllm_omni/entrypoints/openai/api_server.py

@@ -1426,6 +1426,9 @@ async def edit_images(
     background: str | None = Form("auto"),
     output_compression: Annotated[int, Form(ge=0, le=100)] = 100,
     user: str | None = Form(None),  # unused now
+    # vllm-omni extensions for image editing
+    mask_image: str | UploadFile | None = None,
+    reference_image: str | UploadFile | None = None,
     # vllm-omni extensions for diffusion control
     negative_prompt: str | None = Form(None),
     num_inference_steps: int | None = Form(None),
@@ -1479,6 +1482,14 @@ async def edit_images(
         prompt["multi_modal_data"] = {}
         prompt["multi_modal_data"]["image"] = pil_images
 
+        if mask_image is not None:
+            loaded = await _load_input_images([mask_image])
+            prompt["multi_modal_data"]["mask_image"] = loaded[0]
+
+        if reference_image is not None:
+            loaded = await _load_input_images([reference_image])
+            prompt["multi_modal_data"]["reference_image"] = loaded[0]
+
         # 3 Build sample params
         gen_params = OmniDiffusionSamplingParams()
         # 3.0 Init with system default values

vllm_omni/entrypoints/openai/serving_chat.py

@@ -2197,7 +2197,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 except Exception as e:  # pragma: no cover - safeguard
                     logger.warning("Failed to parse LoRA request: %s", e)
 
-            # Add reference image if provided
+            # Add reference image if provided (from messages content)
             if pil_images:
                 if len(pil_images) == 1:
                     gen_prompt["multi_modal_data"] = {}