[Bugfix] Fix Flux2klein Text Input Processing (#3098)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/e2e/offline_inference/test_flux2_klein.py

@@ -16,6 +16,7 @@ import pytest
 import torch
 from PIL import Image, ImageDraw
 
+from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
@@ -24,8 +25,6 @@ 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"
@@ -66,6 +65,18 @@ def _extract_images_from_output(outputs: list) -> list[Image.Image]:
     return images
 
 
+# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097
+@pytest.mark.core_model
+@pytest.mark.diffusion
+def test_flux2_klein_can_accept_text_inputs():
+    model = Omni(model=MODEL)
+    outputs = model.generate(
+        "a cup of coffee on the table",
+        OmniDiffusionSamplingParams(num_inference_steps=2, seed=42),
+    )
+    assert len(outputs[0].images) == 1
+
+
 @pytest.mark.core_model
 @pytest.mark.diffusion
 def test_flux2_klein_inpaint_basic():

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -829,14 +829,21 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
                 """Taking only the first image for now.""",
             )
         first_prompt = req.prompts[0]
-        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+        if isinstance(first_prompt, str):
+            multi_modal_data = {}
+            prompt = first_prompt
+            raw_image = None
+            mask_image = None
+            reference_image = None
+        else:
+            multi_modal_data = first_prompt.get("multi_modal_data", {})
+            prompt = first_prompt.get("prompt") or ""
+            raw_image = multi_modal_data.get("image")
+            mask_image = multi_modal_data.get("mask_image")
+            reference_image = multi_modal_data.get("reference_image")
 
-        if (
-            raw_image := None
-            if isinstance(first_prompt, str)
-            else first_prompt.get("multi_modal_data", {}).get("image")
-        ) is None:
-            pass  # use image from param list
+        if raw_image is None:
+            image = None
         elif isinstance(raw_image, list):
             image = [PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image, im) for im in raw_image]
         else:
@@ -956,11 +963,6 @@ 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")