From 9709192ce96d9a22e84a77e577f6c1217aee5076 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 9 Apr 2026 15:53:14 +0800 Subject: [PATCH] [diffusion] feat: support FLUX.2-small-decoder (#22414) --- docs/diffusion/compatibility_matrix.md | 131 ++++++++++++++---- docs/diffusion/index.md | 2 +- .../configs/models/vaes/flux.py | 2 +- .../models/vaes/autoencoder_kl_flux2.py | 5 +- 4 files changed, 110 insertions(+), 30 deletions(-) diff --git a/docs/diffusion/compatibility_matrix.md b/docs/diffusion/compatibility_matrix.md index 038e43892..9f5f13978 100644 --- a/docs/diffusion/compatibility_matrix.md +++ b/docs/diffusion/compatibility_matrix.md @@ -16,27 +16,28 @@ default parameters when initializing and generating videos. ### Video Generation Models -| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) | Sparse Linear Attention (SLA) | Sage Sparse Linear Attention (SageSLA) | Sparse Video Gen 2 (SVG2) | -|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|:----------------------------:|:-----------------------------------------------:|:----------------------------------:| -| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ❌ | ❌ | -| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ❌ | ❌ | -| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ | ❌ | ❌ | ❌ | -| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ | ❌ | -| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ | ❌ | -| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | -| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | -| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | -| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | -| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | -| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | -| TurboWan2.1 T2V 1.3B | `IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | -| TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | -| TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | -| TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | -| LTX-2 | `Lightricks/LTX-2` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| LTX-2.3 | `Lightricks/LTX-2.3` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) | Sparse Linear Attention (SLA) | Sage Sparse Linear Attention (SageSLA) | Sparse Video Gen 2 (SVG2) | +|:-----------------------------|:--------------------------------------------------|:---------------------|:--------:|:-----------------:|:---------:|:----------------------------:|:-----------------------------:|:--------------------------------------:|:-------------------------:| +| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ❌ | ❌ | +| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ❌ | ❌ | +| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ | ❌ | ❌ | ❌ | +| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ | ❌ | +| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ | ❌ | +| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | +| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | +| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | +| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | +| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | +| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | +| TurboWan2.1 T2V 1.3B | `IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | +| TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | +| TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | +| TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | +| LTX-2 | `Lightricks/LTX-2` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| LTX-2.3 | `Lightricks/LTX-2.3` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | **Note**: + 1. Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. 2. SageSLA is based on SpargeAttn. Install it first with `pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation` 3. LTX-2 two-stage generation uses `--pipeline-class-name LTX2TwoStagePipeline`. The spatial upsampler and distilled LoRA are auto-resolved from the model snapshot by default, and can still be overridden with `--spatial-upsampler-path` and `--distilled-lora-path`. @@ -57,6 +58,82 @@ default parameters when initializing and generating videos. | Qwen Image Edit | `Qwen/Qwen-Image-Edit` | | Qwen Image Edit 2511 | `Qwen/Qwen-Image-Edit-2511` | +## Supported Components + +SGLang Diffusion supports overriding individual pipeline components with +`---path`. The value can be either a Hugging Face repo ID or a local +component directory. + +The same overrides can also be provided in config files through +`component_paths.`. + +### Common Syntax + +CLI: + +```bash +sglang generate \ + --model-path black-forest-labs/FLUX.2-dev \ + --vae-path black-forest-labs/FLUX.2-small-decoder \ + --transformer-path /models/flux2/transformer +``` + +Config file: + +```yaml +model_path: black-forest-labs/FLUX.2-dev +component_paths: + vae: black-forest-labs/FLUX.2-small-decoder + transformer: /models/flux2/transformer +``` + +Use the component name from the pipeline's `model_index.json` or the native pipeline's registered module name: + +| Component Type | Supported Keys | Notes | +|:------------------|:---------------------------------------------------------------------------------------------------------------------------|:--------------------------------------------------------------| +| VAE | `vae`, `video_vae`, `audio_vae` | `vae` is the common image-generation override | +| Transformer / DiT | `transformer`, `video_dit`, `audio_dit` | `transformer` is the standard override for the main denoiser | +| Text / Preprocess | `text_encoder`, `text_encoder_2`, `tokenizer`, `processor`, `image_processor` | Replacement encoders often need matching preprocessing assets | +| Auxiliary | `scheduler`, `spatial_upsampler`, `vocoder`, `connectors`, `dual_tower_bridge`, `image_encoder`, `vision_language_encoder` | Only valid for pipelines that expose these components | + +### Known Component Repos + +The table below lists concrete Hugging Face component repos that are already used in SGLang Diffusion docs or tests. It is not an exhaustive catalog of all compatible component repos. + +| Base Model | Override Key | Example Repo | Notes | +|:-------------------------------|:--------------|:-----------------------------------------|:------------------------------------------| +| `black-forest-labs/FLUX.2-dev` | `vae` | `black-forest-labs/FLUX.2-small-decoder` | Decoder-only FLUX.2 VAE override | +| `black-forest-labs/FLUX.2-dev` | `vae` | `fal/FLUX.2-Tiny-AutoEncoder` | Existing tested custom VAE path | + +### VAE + +- `--vae-path` is the common image-generation override. +- `--video-vae-path` and `--audio-vae-path` are only relevant for pipelines with separate video or audio VAEs. + +### Transformer / DiT + +- `--transformer-path` is the standard override for the main denoising transformer. +- For quantized transformers, prefer `--transformer-path` or `--transformer-weights-path`; see `quantization.md`. +- `--video-dit-path` and `--audio-dit-path` are only for pipelines that split denoisers by modality. + +### Text Encoders and Preprocessors + +- `--text-encoder-path` and `--text-encoder-2-path` override primary and secondary text encoders. +- `--tokenizer-path`, `--processor-path`, and `--image-processor-path` are useful when the replacement encoder requires matching preprocessing assets. + +### Auxiliary Components + +- `--scheduler-path` is only relevant when the pipeline exposes a scheduler component. +- `--spatial-upsampler-path` is mainly for two-stage pipelines such as `LTX2TwoStagePipeline`. +- `--vocoder-path`, `--connectors-path`, `--dual-tower-bridge-path`, `--image-encoder-path`, and `--vision-language-encoder-path` are only valid for pipelines that expose those components. + +### Notes + +1. Component overrides are only valid when the target pipeline actually uses + that component. +2. The override key should match the component name in the pipeline's + `model_index.json` or the native pipeline's registered module name. + ## Verified LoRA Examples This section lists example LoRAs that have been explicitly tested and verified with each base model in the **SGLang Diffusion** pipeline. @@ -68,14 +145,14 @@ This section lists example LoRAs that have been explicitly tested and verified w ### Verified LoRAs by Base Model -| Base Model | Supported LoRAs | -|:-----------------|:----------------| -| Wan2.2 | `lightx2v/Wan2.2-Distill-Loras`
`Cseti/wan2.2-14B-Arcane_Jinx-lora-v1` | -| Wan2.1 | `lightx2v/Wan2.1-Distill-Loras` | -| Z-Image-Turbo | `tarn59/pixel_art_style_lora_z_image_turbo`
`wcde/Z-Image-Turbo-DeJPEG-Lora` | -| Qwen-Image | `lightx2v/Qwen-Image-Lightning`
`flymy-ai/qwen-image-realism-lora`
`prithivMLmods/Qwen-Image-HeadshotX`
`starsfriday/Qwen-Image-EVA-LoRA` | -| Qwen-Image-Edit | `ostris/qwen_image_edit_inpainting`
`lightx2v/Qwen-Image-Edit-2511-Lightning` | -| Flux | `dvyio/flux-lora-simple-illustration`
`XLabs-AI/flux-furry-lora`
`XLabs-AI/flux-RealismLora` | +| Base Model | Supported LoRAs | +|:----------------|:---------------------------------------------------------------------------------------------------------------------------------------------------| +| Wan2.2 | `lightx2v/Wan2.2-Distill-Loras`
`Cseti/wan2.2-14B-Arcane_Jinx-lora-v1` | +| Wan2.1 | `lightx2v/Wan2.1-Distill-Loras` | +| Z-Image-Turbo | `tarn59/pixel_art_style_lora_z_image_turbo`
`wcde/Z-Image-Turbo-DeJPEG-Lora` | +| Qwen-Image | `lightx2v/Qwen-Image-Lightning`
`flymy-ai/qwen-image-realism-lora`
`prithivMLmods/Qwen-Image-HeadshotX`
`starsfriday/Qwen-Image-EVA-LoRA` | +| Qwen-Image-Edit | `ostris/qwen_image_edit_inpainting`
`lightx2v/Qwen-Image-Edit-2511-Lightning` | +| Flux | `dvyio/flux-lora-simple-illustration`
`XLabs-AI/flux-furry-lora`
`XLabs-AI/flux-RealismLora` | ## Special requirements diff --git a/docs/diffusion/index.md b/docs/diffusion/index.md index db384c090..5c467522c 100644 --- a/docs/diffusion/index.md +++ b/docs/diffusion/index.md @@ -28,7 +28,7 @@ sglang serve --model-path Qwen/Qwen-Image --port 30010 ## Start Here - [Installation](installation.md): install SGLang Diffusion and platform dependencies -- [Compatibility Matrix](compatibility_matrix.md): check model and optimization support +- [Compatibility Matrix](compatibility_matrix.md): check model, optimization, and component override support - [CLI](api/cli.md): run one-off generation jobs or launch a persistent server - [OpenAI-Compatible API](api/openai_api.md): send image and video requests to the HTTP server - [Attention Backends](performance/attention_backends.md): choose the best backend for your model and hardware diff --git a/python/sglang/multimodal_gen/configs/models/vaes/flux.py b/python/sglang/multimodal_gen/configs/models/vaes/flux.py index 33308640a..4543ec9b4 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/flux.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/flux.py @@ -28,7 +28,7 @@ class FluxVAEArchConfig(VAEArchConfig): @dataclass class Flux2VAEArchConfig(FluxVAEArchConfig): - pass + decoder_block_out_channels: tuple[int, ...] | None = None @dataclass diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py index 7410358fb..5c1bf65ea 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py @@ -49,6 +49,9 @@ class AutoencoderKLFlux2(ParallelTiledVAE): down_block_types: Tuple[str, ...] = arch_config.down_block_types up_block_types: Tuple[str, ...] = arch_config.up_block_types block_out_channels: Tuple[int, ...] = arch_config.block_out_channels + decoder_block_out_channels: Optional[Tuple[int, ...]] = getattr( + arch_config, "decoder_block_out_channels", None + ) layers_per_block: int = arch_config.layers_per_block act_fn: str = arch_config.act_fn latent_channels: int = arch_config.latent_channels @@ -79,7 +82,7 @@ class AutoencoderKLFlux2(ParallelTiledVAE): in_channels=latent_channels, out_channels=out_channels, up_block_types=up_block_types, - block_out_channels=block_out_channels, + block_out_channels=decoder_block_out_channels or block_out_channels, layers_per_block=layers_per_block, norm_num_groups=norm_num_groups, act_fn=act_fn,