38 lines
1007 B
Python
38 lines
1007 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Tests for Flux1 Schnell."""
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from vllm_omni.entrypoints.omni import Omni
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
MODEL = "black-forest-labs/FLUX.1-schnell"
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.diffusion
|
|
def test_flux_schnell_text_to_image():
|
|
"""Test FLUX.1-schnell text-to-image generation."""
|
|
omni = Omni(model=MODEL)
|
|
|
|
omni_outputs = list(
|
|
omni.generate(
|
|
prompts=["A photo of a cat sitting on a laptop"],
|
|
sampling_params_list=OmniDiffusionSamplingParams(
|
|
height=512,
|
|
width=512,
|
|
num_inference_steps=2,
|
|
seed=42,
|
|
),
|
|
)
|
|
)
|
|
|
|
assert len(omni_outputs) > 0
|
|
images = omni_outputs[0].images
|
|
assert len(images) == 1
|
|
assert isinstance(images[0], Image.Image)
|
|
assert images[0].size == (512, 512)
|