[diffusion] fix: plumb max_sequence_length via diffusers_kwargs (#20930)

Co-authored-by: jiangqc <jqc1569978990@gmail.com>
Co-authored-by: jiangqianchen <jiangqianchen@xiaohongshu.com>

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -349,6 +349,13 @@ class PipelineConfig:
     def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict:
         return tokenizer(prompt, **tok_kwargs)
 
+    def is_flux_v1(self) -> bool:
+        """True if this pipeline is FLUX v1 (dual CLIP + T5 text encoders).
+
+        Used by text encoding (e.g. fixed CLIP context). Other pipelines return False.
+        """
+        return False
+
     def prepare_latent_shape(self, batch, batch_size, num_frames):
         height = batch.height // self.vae_config.arch_config.spatial_compression_ratio
         width = batch.width // self.vae_config.arch_config.spatial_compression_ratio

python/sglang/multimodal_gen/configs/pipeline_configs/flux.py

@@ -87,6 +87,9 @@ class FluxPipelineConfig(ImagePipelineConfig):
         ]
     )
 
+    def is_flux_v1(self) -> bool:
+        return True
+
     def get_text_encoder_attention_mask(self, text_inputs, encoder_index):
         # Flux v1 does not use attention masks for text encoders.
         return None
@@ -472,6 +475,9 @@ class Flux2PipelineConfig(FluxPipelineConfig):
         ]
     )
 
+    def is_flux_v1(self) -> bool:
+        return False
+
     def get_text_encoder_attention_mask(self, text_inputs, encoder_index):
         # Flux2 uses standard attention masks (unlike Flux v1).
         return text_inputs.get("attention_mask")
@@ -490,6 +496,7 @@ class Flux2PipelineConfig(FluxPipelineConfig):
 
     def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
         messages = build_flux2_text_messages(prompts)
+        effective_max_length = tok_kwargs.pop("max_length", 512)
         inputs = tokenizer.apply_chat_template(
             messages,
             add_generation_prompt=False,
@@ -499,7 +506,7 @@ class Flux2PipelineConfig(FluxPipelineConfig):
             padding="max_length",
             truncation=True,
             # 2048 from official github repo, 512 from diffusers
-            max_length=512,
+            max_length=effective_max_length,
         )
 
         return inputs

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -178,6 +178,16 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
         ]
     )
 
+    def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
+        tok_kwargs.setdefault("truncation", True)
+
+        if tok_kwargs.get("max_length") is not None:
+            tok_kwargs["padding"] = "max_length"
+        else:
+            tok_kwargs.setdefault("max_length", 1024)
+            tok_kwargs["padding"] = True
+        return tokenizer(prompts, **tok_kwargs)
+
     def prepare_sigmas(self, sigmas, num_inference_steps):
         return self._prepare_sigmas(sigmas, num_inference_steps)
 

python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py

@@ -90,10 +90,12 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
             )
             for prompt in prompts
         ]
+
+        effective_max_length = tok_kwargs.pop("max_length", 512)
         return tokenizer(
             rendered_prompts,
             padding="max_length",
-            max_length=512,  # TODO (yhyang201): set max length according to config
+            max_length=effective_max_length,
             truncation=True,
             return_tensors="pt",
         )

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -218,6 +218,7 @@ class SamplingParams:
 
     return_file_paths_only: bool = True
     enable_sequence_shard: bool | None = None
+    diffusers_kwargs: dict | None = None
 
     # Prompt enhancement (ErnieImage)
     use_pe: bool | None = None

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -39,11 +39,9 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
 from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import get_global_server_args
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.observability.trace import extract_trace_headers
 
 router = APIRouter(prefix="/v1/images", tags=["images"])
-logger = init_logger(__name__)
 
 
 def _get_extra_field(request, field_name):
@@ -153,6 +151,7 @@ async def generations(
             enable_teacache=request.enable_teacache,
             output_compression=request.output_compression,
             output_quality=request.output_quality,
+            diffusers_kwargs=request.diffusers_kwargs,
             enable_upscaling=request.enable_upscaling,
             upscaling_model_path=request.upscaling_model_path,
             upscaling_scale=request.upscaling_scale,

python/sglang/multimodal_gen/runtime/entrypoints/utils.py

@@ -418,6 +418,9 @@ def prepare_request(
         VSA_sparsity=server_args.attention_backend_config.VSA_sparsity,
     )
     sampling_params.apply_request_extra(req)
+    diffusers_kwargs = getattr(sampling_params, "diffusers_kwargs", None)
+    if diffusers_kwargs and "max_sequence_length" in diffusers_kwargs:
+        req.max_sequence_length = diffusers_kwargs["max_sequence_length"]
 
     req.adjust_size(server_args)
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py

@@ -185,6 +185,9 @@ class TextEncodingStage(PipelineStage):
 
         all_indices: list[int] = list(range(len(self.text_encoders)))
 
+        # Get max_sequence_length from batch if available
+        max_seq_length = getattr(batch, "max_sequence_length", None)
+
         (
             prompt_embeds_list,
             prompt_masks_list,
@@ -196,6 +199,7 @@ class TextEncodingStage(PipelineStage):
             server_args,
             encoder_index=all_indices,
             return_attention_mask=True,
+            max_length=max_seq_length,
         )
 
         for pe in prompt_embeds_list:
@@ -475,6 +479,11 @@ class TextEncodingStage(PipelineStage):
                 encoder_config.tokenizer_kwargs,
                 **text_encoder_extra_arg,
             )
+            # Pass max_length to tokenizer if specified in the request. Flux v1 encoder 0
+            # is CLIP with a fixed 77-token context; overriding breaks tokenization.
+            is_flux_v1 = server_args.pipeline_config.is_flux_v1()
+            if max_length is not None and not (is_flux_v1 and i == 0):
+                tok_kwargs["max_length"] = max_length
 
             text_inputs: dict = server_args.pipeline_config.tokenize_prompt(
                 processed_text_list, tokenizer, tok_kwargs