minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[diffusion] model: support FLUX.2-klein-base (#25661)

Browse Source
This commit is contained in:
Alex O. P.
2026-05-22 11:24:46 +08:00
committed by GitHub
parent 8b473aa0bc
commit ae7c4226eb
10 changed files with 147 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/diffusion-ci-gt-gen.yml
+2
View File
@@ -85,6 +85,7 @@ jobs:
"diffusers": [
"flux_2_image_t2i",
"flux_2_klein_image_t2i",
"flux_2_klein_base_image_t2i",
"flux_2_ti2i",
"flux_image_t2i",
"qwen_image_edit_2509_ti2i",
@@ -120,6 +121,7 @@ jobs:
h200_cases = {
"flux_2_image_t2i",
"flux_2_klein_image_t2i",
"flux_2_klein_base_image_t2i",
"flux_2_ti2i",
}
include = []
docs/diffusion/compatibility_matrix.md
+2
View File
@@ -66,6 +66,8 @@ default parameters when initializing and generating videos.
| FLUX.2-dev-NVFP4 | `black-forest-labs/FLUX.2-dev-NVFP4` |
| FLUX.2-Klein-4B | `black-forest-labs/FLUX.2-klein-4B` |
| FLUX.2-Klein-9B | `black-forest-labs/FLUX.2-klein-9B` |
| FLUX.2-Klein-Base-4B | `black-forest-labs/FLUX.2-klein-base-4B` |
| FLUX.2-Klein-Base-9B | `black-forest-labs/FLUX.2-klein-base-9B` |
| Z-Image | `Tongyi-MAI/Z-Image` |
| Z-Image-Turbo | `Tongyi-MAI/Z-Image-Turbo` |
| GLM-Image | `zai-org/GLM-Image` |
docs_new/docs/sglang-diffusion/compatibility_matrix.mdx
+8
View File
@@ -435,6 +435,14 @@ Optimization columns are abbreviated to keep the matrix readable:
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-9B</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>black-forest-labs/FLUX.2-klein-9B</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-Base-4B</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>black-forest-labs/FLUX.2-klein-base-4B</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-Base-9B</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>black-forest-labs/FLUX.2-klein-base-9B</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Z-Image</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>Tongyi-MAI/Z-Image</code></td>
docs_new/docs/sglang-diffusion/dynamic_batching.mdx
+2
View File
@@ -65,6 +65,8 @@ An initial implementation of dynamic batching for T2I and T2V models can be foun
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-dev-NVFP4</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-4B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-9B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-Base-4B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FLUX.2-Klein-Base-9B</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>?</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Z-Image</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>?</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Z-Image-Turbo</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
<tr><td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>GLM-Image</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td><td style={{textAlign: "center", padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>-</td></tr>
python/sglang/multimodal_gen/configs/pipeline_configs/flux.py
+25
View File
@@ -798,3 +798,28 @@ class Flux2KleinPipelineConfig(Flux2PipelineConfig):
return_tensors=return_tensors,
**tok_kwargs,
)
@dataclass
class Flux2KleinBasePipelineConfig(Flux2KleinPipelineConfig):
# Undistilled Klein base model, with guidance embeddings
should_use_guidance: bool = True
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
txt_seq_lens = self.require_text_seq_lens(
batch,
0,
negative=True,
expected_batch_size=batch.negative_prompt_embeds[0].shape[0],
)
return {
"freqs_cis": self.get_freqs_cis(
batch.negative_prompt_embeds[0],
batch.width,
batch.height,
device,
rotary_emb,
batch,
txt_seq_lens,
)
}
python/sglang/multimodal_gen/configs/sample/flux.py
+8
View File
@@ -29,3 +29,11 @@ class Flux2KleinSamplingParams(Flux2SamplingParams):
# Klein is step-distilled, so default to 4 steps
guidance_scale: float = 1.0
num_inference_steps: int = 4
@dataclass
class Flux2KleinBaseSamplingParams(FluxSamplingParams):
# Klein-base is undistilled
num_inference_steps: int = 50
guidance_scale: float = 4.0
negative_prompt: str = ""
python/sglang/multimodal_gen/registry.py
+20 -2
View File
@@ -46,6 +46,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.ernie_image import (
ErnieImagePipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.flux import (
Flux2KleinBasePipelineConfig,
Flux2KleinPipelineConfig,
Flux2PipelineConfig,
)
@@ -85,6 +86,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
)
from sglang.multimodal_gen.configs.sample.ernie_image import ErnieImageSamplingParams
from sglang.multimodal_gen.configs.sample.flux import (
Flux2KleinBaseSamplingParams,
Flux2KleinSamplingParams,
Flux2SamplingParams,
FluxSamplingParams,
@@ -788,8 +790,24 @@ def _register_configs():
"black-forest-labs/FLUX.2-klein-9B",
],
model_detectors=[
lambda hf_id: "flux.2-klein" in hf_id.lower()
or "flux2-klein" in hf_id.lower()
lambda hf_id: (
"flux.2-klein" in hf_id.lower() or "flux2-klein" in hf_id.lower()
)
and "base" not in hf_id.lower()
],
)
register_configs(
sampling_param_cls=Flux2KleinBaseSamplingParams,
pipeline_config_cls=Flux2KleinBasePipelineConfig,
hf_model_paths=[
"black-forest-labs/FLUX.2-klein-base-4B",
"black-forest-labs/FLUX.2-klein-base-9B",
],
model_detectors=[
lambda hf_id: (
"flux.2-klein" in hf_id.lower() or "flux2-klein" in hf_id.lower()
)
and "base" in hf_id.lower()
],
)
register_configs(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py
+3 -1
View File
@@ -396,7 +396,9 @@ class TextEncodingStage(PipelineStage):
result.add_check(
"negative_prompt",
batch.negative_prompt,
lambda x: not batch.do_classifier_free_guidance or V.string_not_none(x),
lambda x: not batch.do_classifier_free_guidance
or V.string_not_none(x)
or isinstance(x, str),
)
result.add_check(
"do_classifier_free_guidance",
python/sglang/multimodal_gen/test/server/gpu_cases.py
+10
View File
@@ -36,6 +36,7 @@ from sglang.multimodal_gen.test.test_utils import (
DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST,
DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST,
DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST,
DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST,
DEFAULT_JOYAI_IMAGE_EDIT_MODEL_NAME_FOR_TEST,
DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST,
@@ -123,6 +124,15 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
),
T2I_sampling_params,
),
DiffusionTestCase(
"flux_2_klein_base_image_t2i",
DiffusionServerArgs(
model_path=DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST,
),
T2I_sampling_params,
run_consistency_check=False,
run_component_accuracy_check=False,
),
# TODO: replace with a faster model to test the --dit-layerwise-offload
# TODO: currently, we don't support sending more than one request in test, and setting `num_outputs_per_prompt` to 2 doesn't guarantee the denoising be executed twice,
# so we do one warmup and send one request instead
python/sglang/multimodal_gen/test/server/perf_baselines.json
+67
View File
@@ -321,6 +321,73 @@
"expected_median_denoise_ms": 39.47,
"estimated_full_test_time_s": 120.5
},
"flux_2_klein_base_image_t2i": {
"stages_ms": {
"InputValidationStage": 0.06,
"TextEncodingStage": 505.51,
"ImageVAEEncodingStage": 0.01,
"LatentPreparationStage": 1.14,
"TimestepPreparationStage": 53.76,
"DenoisingStage": 12365.93,
"DecodingStage": 11.74
},
"denoise_step_ms": {
"0": 72.0,
"1": 192.12,
"2": 217.22,
"3": 217.5,
"4": 257.12,
"5": 262.19,
"6": 246.18,
"7": 224.94,
"8": 252.71,
"9": 262.82,
"10": 246.37,
"11": 231.87,
"12": 253.34,
"13": 260.94,
"14": 242.16,
"15": 235.92,
"16": 254.31,
"17": 262.02,
"18": 243.92,
"19": 239.61,
"20": 253.47,
"21": 259.75,
"22": 245.48,
"23": 240.91,
"24": 253.84,
"25": 255.36,
"26": 248.31,
"27": 243.75,
"28": 250.5,
"29": 251.34,
"30": 246.87,
"31": 243.8,
"32": 249.8,
"33": 254.02,
"34": 247.4,
"35": 244.77,
"36": 252.03,
"37": 248.85,
"38": 249.74,
"39": 248.6,
"40": 252.29,
"41": 249.25,
"42": 249.85,
"43": 249.68,
"44": 252.65,
"45": 249.2,
"46": 249.32,
"47": 248.62,
"48": 252.87,
"49": 249.19
},
"expected_e2e_ms": 13075.51,
"expected_avg_denoise_ms": 243.34,
"expected_median_denoise_ms": 249.22,
"estimated_full_test_time_s": 124.4
},
"layerwise_offload": {
"stages_ms": {
"TextEncodingStage": 176.59,