[Model] Add Flux2 support

[LawJarp-A]

Add FLUX 2 diffusion model support

Ref: vllm-project/vllm-omni#153

Summary

Adds support for FLUX 2 text-to-image diffusion with a dual-stream + single-stream transformer architecture.

What’s included

• Flux2Transformer2DModel: dual-stream (8 blocks) + single-stream (48 blocks), 4D RoPE, vLLM-optimized linear/norm layers.
• Flux2Pipeline: Mistral3-based prompt embeddings, 128-channel latents with 2×2 patch packing, FlowMatch Euler scheduler, VAE decode handling.
• Registry integration: Flux2Pipeline added to DiffusionModelRegistry and post-process function registry.
• Dependency update: pin diffusers==0.36.0 (includes pipelines.flux2 + AutoencoderKLFlux2) — no custom diffusers path / monkey patch.
• Docs:
  • Add FLUX 2 offline inference page under User Guide
  • Add FLUX 2 to supported models list

Files changed

• Added: vllm_omni/diffusion/models/flux2/ (pipeline_flux2.py, flux2_transformer.py, __init__.py)
• Updated: vllm_omni/diffusion/registry.py
• Updated: pyproject.toml (diffusers pin)
• Docs:
  • docs/user_guide/examples/offline_inference/flux2.md
  • docs/.nav.yml
  • docs/models/supported_models.md

Test plan

import torch
from vllm_omni.entrypoints.omni import Omni

omni = Omni(model="black-forest-labs/FLUX.2-dev", num_gpus=1)
gen = torch.Generator(device="cuda").manual_seed(0)
imgs = omni.generate(
    "a serene mountain landscape at sunset",
    height=1024,
    width=1024,
    num_inference_steps=20,
    guidance_scale=4.0,
    generator=gen,
    num_images_per_prompt=1,
    num_outputs_per_prompt=1,
)
imgs[0].save("flux2_smoke.png")
print("saved flux2_smoke.png")
PY

Test results

• Import + registry sanity check (local .venv)
• Full generation run (hardware-dependent)

Checklist

• Purpose stated
• Test plan provided
• Update supported_models.md
• Docs added + nav updated

[LawJarp-A]

cc: @ZJY0516

[hsliuustc0106]

can this model use text_to_image.md as an example? #274

[hsliuustc0106]

I think is no longer necessary ofter text_to_image.md

[hsliuustc0106]

@ZJY0516 please help check whether we need to change diffusers version, another PR #289 needs to rely on diffusers/main

[ZJY0516]

ValueError: There is no module or parameter named 'transformer.transformer_blocks.0.attn.add_k_proj' in Flux2Pipeline
Loading safetensors checkpoint shards:  14% Completed | 1/7 [00:01<00:11,  1.98s/it]