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

.github/workflows/diffusion-ci-gt-gen.yml

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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()
+            lambda hf_id: (
-            or "flux2-klein" in hf_id.lower()
+                "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

@@ -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

@@ -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

@@ -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,