diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 7005ac70d..aa102d606 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index b6dffd597..198ff7d0f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index d2b9c5a62..3fccaa8b9 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/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) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index d26276863..3f67bfede 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/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", ) diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 5c8ea67ac..0eef02b01 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index d19e9648c..848d4f151 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index d65a1287f..b3f9b6549 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/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) 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 4d1395022..8396f3448 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 @@ -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