From bf64a6b841d5ddb705ac1d6f0c48b484a333ccd9 Mon Sep 17 00:00:00 2001 From: Lancer Date: Fri, 17 Apr 2026 07:55:00 +0800 Subject: [PATCH] [Feature] Flux2 klein inpaint (#1180) Signed-off-by: Lancer Co-authored-by: Hongsheng Liu --- docs/serving/image_edit_api.md | 2 + .../test_flux2_klein_inpaint.py | 216 ++++++++++++++++++ .../test_flux2_klein_inpaint_expansion.py | 160 +++++++++++++ .../flux2_klein/pipeline_flux2_klein.py | 199 +++++++++++++++- vllm_omni/entrypoints/openai/api_server.py | 11 + vllm_omni/entrypoints/openai/serving_chat.py | 2 +- 6 files changed, 583 insertions(+), 7 deletions(-) create mode 100644 tests/e2e/offline_inference/test_flux2_klein_inpaint.py create mode 100644 tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md index d254ac06a..79303e1a6 100644 --- a/docs/serving/image_edit_api.md +++ b/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 diff --git a/tests/e2e/offline_inference/test_flux2_klein_inpaint.py b/tests/e2e/offline_inference/test_flux2_klein_inpaint.py new file mode 100644 index 000000000..ac7169361 --- /dev/null +++ b/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() diff --git a/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py b/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py new file mode 100644 index 000000000..f59a0e783 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index 437dd58d0..021d28c2a 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/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( diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 11ba59e43..1e4575836 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 34ddbbd30..28d6ef277 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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"] = {}