From c560e142169ea816f1441e18772a2cbcde9761b4 Mon Sep 17 00:00:00 2001 From: Ratish P <114130421+Ratish1@users.noreply.github.com> Date: Tue, 20 Jan 2026 09:33:33 +0530 Subject: [PATCH] [diffusion] fix: enforce 16-pixel alignment for flux2 to prevent shape mismatch crash (#17302) --- .../configs/pipeline_configs/flux.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index a5c58e684..9aab2badd 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -473,13 +473,23 @@ class Flux2PipelineConfig(FluxPipelineConfig): def calculate_condition_image_size( self, image, width, height ) -> Optional[tuple[int, int]]: + vae_scale_factor = self.vae_config.arch_config.vae_scale_factor + multiple_of = vae_scale_factor * 2 + target_area: int = 1024 * 1024 if width is not None and height is not None: + new_width, new_height = width, height if width * height > target_area: scale = math.sqrt(target_area / (width * height)) - width = int(width * scale) - height = int(height * scale) - return width, height + new_width = int(width * scale) + new_height = int(height * scale) + + # Flux requires multiples of (VAE scale 8 * Patch size 2) + new_width = (new_width // multiple_of) * multiple_of + new_height = (new_height // multiple_of) * multiple_of + + if new_width != width or new_height != height: + return new_width, new_height return None @@ -628,7 +638,6 @@ class Flux2PipelineConfig(FluxPipelineConfig): class Flux2KleinPipelineConfig(Flux2PipelineConfig): # Klein is distilled, so no guidance embeddings should_use_guidance: bool = False - task_type: ModelTaskType = ModelTaskType.T2I text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))