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

2026-05-22 11:24:46 +08:00
parent 8b473aa0bc
commit ae7c4226eb
10 changed files with 147 additions and 3 deletions
@@ -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,
+            )
+        }
@@ -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 = ""
@@ -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(
@@ -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",
@@ -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
@@ -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,