Skip to content

[Model] [WIP] Add Flux2 Model support #210

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
LawJarp-A wants to merge 29 commits into
base: main
Choose a base branch
from
Draft

[Model] [WIP] Add Flux2 Model support #210

LawJarp-A wants to merge 29 commits into from

Conversation

@LawJarp-A
Copy link

@LawJarp-A LawJarp-A commented Dec 5, 2025 •
edited
Loading

Add Flux 2.0 diffusion model support

Refer issue: #153

Purpose

Add support for Flux 2.0 text-to-image diffusion model with dual-stream (8 blocks) + single-stream (48 blocks) transformer architecture.

Key components:

  • Flux2Transformer2DModel with 4D RoPE and vLLM-optimized layers
  • Flux2Pipeline with Mistral3-Small-24B text encoder
  • 128-channel latent space with 2x2 patch packing
  • Registry integration and VLLM_OMNI_DIFFUSERS_PATH environment variable support

Files added:

  • vllm_omni/diffusion/models/flux2/ (transformer and pipeline)
  • Updated registry and diffusion __init__.py

Test Plan

# Setup
export VLLM_OMNI_DIFFUSERS_PATH=/path/to/diffusers/src

# Test imports
python -c "from vllm_omni.diffusion.models.flux2 import Flux2Pipeline; print('✅')"

# Test generation
python -c "
from vllm_omni import Omni
model = Omni('black-forest-labs/FLUX.2-dev')
outputs = model.generate(prompt='a sunset', height=1024, width=1024)
"

Test Result

WIP

Checklist

  • Purpose stated (Flux 2.0 support)
  • Test plan provided
  • Test results documented
  • Update supported_models.md for Flux 2.0
  • Release notes: "Added Flux 2.0 diffusion model with dual+single stream transformer architecture"

ZJY0516 and others added 18 commits November 30, 2025 19:35
@ZJY0516
init attention
308210c 
Signed-off-by: zjy0516 <[email protected]>
@ZJY0516
add
882a88b 
Signed-off-by: zjy0516 <[email protected]>
@ZJY0516
update
c4e7161 
Signed-off-by: zjy0516 <[email protected]>
@ZJY0516
update
9a35310 
Signed-off-by: zjy0516 <[email protected]>
@SamitHuang
Merge branch 'main' into diffusion-attention
dc215ca
@SamitHuang @ZJY0516
fix fa bugs
f82a6e2 
Signed-off-by: samithuang <[email protected]>
@SamitHuang @ZJY0516
support setting attention backend via VLLM_ATTENTION_BACKEND
8d01e7c 
Signed-off-by: samithuang <[email protected]>
@SamitHuang @ZJY0516
linting
dfb4a6a 
Signed-off-by: samithuang <[email protected]>
@SamitHuang @ZJY0516
rm func
c200035 
Signed-off-by: samithuang <[email protected]>
@SamitHuang @ZJY0516
rename to use DIFFUSION_ATTENTION_BACKEND
e7a1ce2 
Signed-off-by: samithuang <[email protected]>
@SamitHuang @ZJY0516
rm SPDA key, not limited to fa3
4986cf1 
Signed-off-by: samithuang <[email protected]>
@ZJY0516
fix pre-commit
304dda6 
Signed-off-by: zjy0516 <[email protected]>
@ZJY0516
Merge branch 'main' into diffusion-attention
b1dfd0b
@ZJY0516
update
68e4f3c 
Signed-off-by: zjy0516 <[email protected]>
@ZJY0516
update
c018fc2 
Signed-off-by: zjy0516 <[email protected]>
@ZJY0516
update
e730266 
Signed-off-by: zjy0516 <[email protected]>
@LawJarp-A
add flux 2.0 diffusion model support with dual+single stream transfor…
ee49dd0 
…mer architecture
@LawJarp-A
Merge branch 'main' into feat/flux2-support
627f85b 
Signed-off-by: Prajwal A <[email protected]>
@LawJarp-A LawJarp-A changed the title Add Flux2 Model support [Model] Add Flux2 Model support Dec 5, 2025
@LawJarp-A LawJarp-A changed the title [Model] Add Flux2 Model support [Model] [WIP] Add Flux2 Model support Dec 5, 2025
@LawJarp-A LawJarp-A mentioned this pull request Dec 5, 2025
Open
1 task
@LawJarp-A
Copy link
Author

LawJarp-A commented Dec 5, 2025 •
edited
Loading

Pending tasks

  • Test imports work - from vllm_omni.diffusion.models.flux2 import Flux2Pipeline
  • Verify weight loading - Run pipeline.load_weights() with real Flux2 checkpoint
  • End-to-end generation test - Generate one image to validate full pipeline
  • Update supported_models.md - Add Flux2 documentation entry

@david6666666 david6666666 mentioned this pull request Dec 8, 2025
Open
41 tasks
@LawJarp-A
Copy link
Author

LawJarp-A commented Dec 9, 2025

cc: @hsliuustc0106 @ZJY0516
Hey, could not test the original weights due to GPU constraints.
Accd to: huggingface.co/blog/flux-2 (will need ~64GB)
Additionally, they have provided quantized weights as well which I can try. But do suggest how I can proceed. Thanks!

@ZJY0516 ZJY0516 self-requested a review December 9, 2025 11:42
ZJY0516
ZJY0516 reviewed Dec 9, 2025
View reviewed changes
vllm_omni/diffusion/__init__.py Outdated
@@ -1,0 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Diffusion model support for vllm-omni."""
Copy link
Collaborator

@ZJY0516 ZJY0516 Dec 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need this?

vllm_omni/diffusion/models/flux2/pipeline_flux2.py
"""Load transformer weights."""
self.load_transformer()

def load_transformer(self):
Copy link
Collaborator

@ZJY0516 ZJY0516 Dec 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please use weight loader introduced in #157

vllm_omni/diffusion/models/flux2/flux2_transformer.py Outdated

import torch
import torch.nn as nn
from diffusers.models.attention_dispatch import dispatch_attention_fn
Copy link
Collaborator

@ZJY0516 ZJY0516 Dec 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use attention layer introduced in #115

vllm_omni/diffusion/models/flux2/flux2_transformer.py
rope_theta: int = 2000,
eps: float = 1e-6,
):
super().__init__()
Copy link
Collaborator

@ZJY0516 ZJY0516 Dec 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please try to read some properties from od_config.tf_model_config like

vllm-omni/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

Lines 502 to 503 in c45cdea

model_config = od_config.tf_model_config
num_layers = model_config.num_layers

Because we may use a little model with less layers to test in ci

@LawJarp-A
Merge branch 'pr-115' into feat/flux2-support
48be98c
@LawJarp-A
use implemented attention layers
78078ca 
Signed-off-by: Prajwal A <[email protected]>
@LawJarp-A
remove test code
b858c84 
Signed-off-by: Prajwal A <[email protected]>
@LawJarp-A
Merge upstream/main into feat/flux2-support
aba9fe7 
- Resolved conflicts in registry.py by keeping both Flux2Pipeline and QwenImageEditPipeline
- Resolved conflicts in z_image_transformer.py by removing duplicate Attention import
- Includes latest changes from main: CI tests, new examples, model loader refactor
@LawJarp-A
use tf_model_config
ca246ce 
Signed-off-by: Prajwal A <[email protected]>
@LawJarp-A
style: Fix ruff linting issues in Flux2 models
8f41ed6 
- Replace deprecated typing.Tuple/Dict/List with tuple/dict/list
- Fix ambiguous variable name 'l' to 'layer_dim'
- Remove unused variables batch_size and img_seq_len
- Add TYPE_CHECKING import for OmniDiffusionConfig

Signed-off-by: Prajwal A <[email protected]>
@LawJarp-A
chore: Remove outdated local diffusers path comments
e72a20b 
- Remove obsolete comments about local diffusers repo
- Remove reference to deleted __init__.py file
- All imports are now using standard diffusers package

Signed-off-by: Prajwal A <[email protected]>
@SamitHuang SamitHuang mentioned this pull request Dec 10, 2025
Open
5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ZJY0516 ZJY0516 ZJY0516 left review comments

@hsliuustc0106 hsliuustc0106 Awaiting requested review from hsliuustc0106 hsliuustc0106 will be requested when the pull request is marked ready for review hsliuustc0106 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@LawJarp-A @ZJY0516 @SamitHuang