diff --git a/tests/diffusion/models/flux2/test_flux2_klein_validation.py b/tests/diffusion/models/flux2/test_flux2_klein_validation.py new file mode 100644 index 000000000..de9ff6f39 --- /dev/null +++ b/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)) 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 5543bc208..2ea267be1 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/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)}")