diff --git a/.github/workflows/diffusion-ci-gt-gen.yml b/.github/workflows/diffusion-ci-gt-gen.yml index 909cfff54..39b267a0e 100644 --- a/.github/workflows/diffusion-ci-gt-gen.yml +++ b/.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 = [] diff --git a/docs/diffusion/compatibility_matrix.md b/docs/diffusion/compatibility_matrix.md index cd7f2d5b5..435a83c36 100644 --- a/docs/diffusion/compatibility_matrix.md +++ b/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` | diff --git a/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx b/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx index 9f234bfa2..5091df562 100644 --- a/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx +++ b/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx @@ -435,6 +435,14 @@ Optimization columns are abbreviated to keep the matrix readable: 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 diff --git a/docs_new/docs/sglang-diffusion/dynamic_batching.mdx b/docs_new/docs/sglang-diffusion/dynamic_batching.mdx index b05a6eb89..795536a49 100644 --- a/docs_new/docs/sglang-diffusion/dynamic_batching.mdx +++ b/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 FLUX.2-dev-NVFP4?? FLUX.2-Klein-4B✅❌ FLUX.2-Klein-9B?? + FLUX.2-Klein-Base-4B?? + FLUX.2-Klein-Base-9B?? Z-Image?- Z-Image-Turbo✅- GLM-Image❌- diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index 876c233ed..bbb2d4023 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/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, + ) + } diff --git a/python/sglang/multimodal_gen/configs/sample/flux.py b/python/sglang/multimodal_gen/configs/sample/flux.py index 0b094957b..5b4d81632 100644 --- a/python/sglang/multimodal_gen/configs/sample/flux.py +++ b/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 = "" diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 484219238..9e7612162 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/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() - 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( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index acb9387d0..c5e619561 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/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", diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index dab086cfc..daa4e541d 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 9aaf37698..80981acf9 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/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,