From d46397809852cd5599e659a63133cced30e549e5 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Fri, 17 Apr 2026 02:02:30 -0600 Subject: [PATCH] [Model] Support Flux1 Schnell (#2528) Signed-off-by: Alex Brooks --- docs/models/supported_models.md | 1 + docs/user_guide/diffusion_features.md | 1 + .../offline_inference/text_to_image/README.md | 3 +- tests/e2e/offline_inference/test_flux.py | 37 +++++++++++++++++++ .../diffusion/models/flux/flux_transformer.py | 4 +- .../diffusion/models/flux/pipeline_flux.py | 7 +++- 6 files changed, 49 insertions(+), 4 deletions(-) create mode 100644 tests/e2e/offline_inference/test_flux.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0f9c8fff6..b4976b09c 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -47,6 +47,7 @@ th { | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | | `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | +| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 45953b852..4e7003cce 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -109,6 +109,7 @@ The following tables show which models support each feature: |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:| | **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | +| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | | **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index cc295e827..c71773972 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -29,7 +29,8 @@ This folder provides several entrypoints for experimenting with text-to-image di | `AIDC-AI/Ovis-Image-7B` | 1024 x 1024 | 71.8 | 17.1 | | `OmniGen2/OmniGen2` | 1024 x 1024 | 20.1 | 14.7 | | `stabilityai/stable-diffusion-3.5-medium` | 1024 x 1024 | 20.1 | 15.6 | -| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 77.6 | 31.4 | +| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 33.9 | 31.4 | +| `black-forest-labs/FLUX.1-schnell` | 1024 x 1024 | 33.9 | 31.4 | | `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 | | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 | | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) | diff --git a/tests/e2e/offline_inference/test_flux.py b/tests/e2e/offline_inference/test_flux.py new file mode 100644 index 000000000..02c6787be --- /dev/null +++ b/tests/e2e/offline_inference/test_flux.py @@ -0,0 +1,37 @@ +# 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) diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 297c62675..dd88ee76c 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -526,10 +526,10 @@ class FluxTransformer2DModel(nn.Module): def __init__( self, - od_config: OmniDiffusionConfig = None, + od_config: OmniDiffusionConfig | None = None, patch_size: int = 1, in_channels: int = 64, - out_channels: int = None, + out_channels: int | None = None, num_layers: int = 19, num_single_layers: int = 38, attention_head_dim: int = 128, diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index 6f43e8dbb..70d572d9a 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -30,6 +30,7 @@ from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixi from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific logger = logging.getLogger(__name__) @@ -106,7 +107,11 @@ class FluxPipeline(nn.Module, FluxPipelineMixin, CFGParallelMixin, DiffusionPipe self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device ) - self.transformer = FluxTransformer2DModel(od_config=od_config, quant_config=od_config.quantization_config) + + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, FluxTransformer2DModel) + self.transformer = FluxTransformer2DModel( + **transformer_kwargs, od_config=od_config, quant_config=od_config.quantization_config + ) self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer_2 = T5TokenizerFast.from_pretrained(