[Bugfix] Reject empty prompts in Flux2 Klein diffusion pipeline (#3711)

Signed-off-by: MmMaiIIi <2360655509@qq.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>

tests/diffusion/models/flux2/test_flux2_klein_validation.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import (
+    Flux2KleinPipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_pipeline():
+    pipeline = object.__new__(Flux2KleinPipeline)
+    pipeline.vae_scale_factor = 8
+    pipeline.is_distilled = True
+    pipeline._guidance_scale = 0.0
+    return pipeline
+
+
+def _check(pipeline, prompt, prompt_embeds=None):
+    pipeline.check_inputs(
+        prompt=prompt,
+        height=512,
+        width=512,
+        prompt_embeds=prompt_embeds,
+        num_inference_steps=4,
+        guidance_scale=0.0,
+    )
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        "",
+        "   ",
+    ],
+)
+def test_rejects_empty_or_whitespace_prompt(prompt):
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError, match="empty or whitespace-only"):
+        _check(pipe, prompt)
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        ["valid prompt", ""],
+        ["valid prompt", "   "],
+        ["   "],
+    ],
+)
+def test_rejects_list_with_empty_or_whitespace_element(prompt):
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError, match="empty or whitespace-only"):
+        _check(pipe, prompt)
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        "valid prompt",
+        ["valid prompt", "another valid prompt"],
+    ],
+)
+def test_accepts_valid_prompts(prompt):
+    pipe = _make_pipeline()
+    _check(pipe, prompt)
+
+
+def test_rejects_none_without_embeds():
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError):
+        _check(pipe, None)
+
+
+def test_accepts_none_with_prompt_embeds():
+    pipe = _make_pipeline()
+    _check(pipe, None, prompt_embeds=torch.randn(1, 256, 4096))

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -701,6 +701,10 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             raise ValueError(
                 "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
             )
+        elif isinstance(prompt, str) and not prompt.strip():
+            raise ValueError("`prompt` cannot be empty or whitespace-only.")
+        elif isinstance(prompt, list) and any(isinstance(p, str) and not p.strip() for p in prompt):
+            raise ValueError("`prompt` cannot contain empty or whitespace-only strings.")
         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)}")