[diffusion] fix: fix Flux.2 condition image resize (#14232)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/multimodal_gen/configs/models/vaes/base.py

@@ -143,6 +143,9 @@ class VAEConfig(ModelConfig):
     def get_vae_scale_factor(self):
         return 2 ** (len(self.arch_config.block_out_channels) - 1)
 
+    def encode_sample_mode(self):
+        return "argmax"
+
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig":
         kwargs = {}

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

@@ -190,8 +190,15 @@ class PipelineConfig:
         return sigmas
 
     ## For ImageVAEEncodingStage
-    def resize_condition_image(self, image, target_width, target_height):
-        return image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS)
+    def preprocess_condition_image(
+        self, image, target_width, target_height, _vae_image_processor
+    ):
+        """
+        preprocess the condition image, returns (image, final_image_width, final_image_height)
+        """
+        return image.resize(
+            (target_width, target_height), PIL.Image.Resampling.LANCZOS
+        ), (target_width, target_height)
 
     def prepare_image_processor_kwargs(self, batch):
         return {}

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

@@ -4,6 +4,7 @@ from typing import Callable, List, Optional
 
 import PIL
 import torch
+from diffusers.image_processor import VaeImageProcessor
 
 from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
@@ -465,8 +466,19 @@ class Flux2PipelineConfig(FluxPipelineConfig):
 
         return None
 
-    def resize_condition_image(self, image, target_width, target_height):
-        return image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS)
+    def preprocess_condition_image(
+        self, image, target_width, target_height, vae_image_processor: VaeImageProcessor
+    ):
+        img = image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS)
+        image_width, image_height = img.size
+        vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
+        multiple_of = vae_scale_factor * 2
+        image_width = (image_width // multiple_of) * multiple_of
+        image_height = (image_height // multiple_of) * multiple_of
+        img = vae_image_processor.preprocess(
+            img, height=image_height, width=image_width, resize_mode="crop"
+        )
+        return img, (image_width, image_height)
 
     def postprocess_image_latent(self, latent_condition, batch):
         batch_size = batch.batch_size

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

@@ -273,8 +273,13 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
             "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
         }
 
-    def resize_condition_image(self, image, target_width, target_height):
-        return resize(image, target_height, target_width, resize_mode="default")
+    def preprocess_condition_image(
+        self, image, target_width, target_height, _vae_image_processor
+    ):
+        return resize(image, target_height, target_width, resize_mode="default"), (
+            target_width,
+            target_height,
+        )
 
     def postprocess_image_latent(self, latent_condition, batch):
         batch_size = batch.batch_size

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

@@ -299,7 +299,6 @@ class SamplingParams:
         from sglang.multimodal_gen.registry import get_model_info
 
         model_info = get_model_info(model_path)
-        logger.debug(f"Found model info: {model_info}")
         if model_info is not None:
             sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs)
         else:
@@ -317,8 +316,6 @@ class SamplingParams:
         user_sampling_params = SamplingParams(*args, **kwargs)
         # TODO: refactor
         sampling_params._merge_with_user_params(user_sampling_params)
-        sampling_params.width_not_provided = user_sampling_params.width is None
-        sampling_params.height_not_provided = user_sampling_params.height is None
         sampling_params._adjust(server_args)
 
         return sampling_params