[diffusion] feat: support FLUX.2-small-decoder (#22414)

docs/diffusion/compatibility_matrix.md

@@ -17,7 +17,7 @@ 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                 |    ⭕     |         ⭕         |     ✅     |              ⭕               |               ❌               |                   ❌                    |             ❌             |
@@ -37,6 +37,7 @@ default parameters when initializing and generating videos.
 | LTX-2.3                      | `Lightricks/LTX-2.3`                              | 768×512<br>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
+`--<component>-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.<component>`.
+
+### 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.
@@ -69,7 +146,7 @@ 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`<br>`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`<br>`wcde/Z-Image-Turbo-DeJPEG-Lora`                                                                    |

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

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

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,