From c36de62bfcfc5721bcf30825f3a5399975ee010d Mon Sep 17 00:00:00 2001 From: Chi McIsaac <153383231+qimcis@users.noreply.github.com> Date: Thu, 5 Mar 2026 03:56:39 -0500 Subject: [PATCH] [diffusion] fix images/edit with 2 images (#17520) Signed-off-by: Chi McIsaac Co-authored-by: Mick --- .../configs/pipeline_configs/flux.py | 11 ++-- .../pipelines_core/stages/image_encoding.py | 9 +++ .../test/server/perf_baselines.json | 66 +++++++++++++++++++ .../test/server/testcase_configs.py | 65 +++++++++++++----- 4 files changed, 129 insertions(+), 22 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index 884c47bd3..71d0c0128 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -509,12 +509,6 @@ class Flux2PipelineConfig(FluxPipelineConfig): def postprocess_image_latent(self, latent_condition, batch): batch_size = batch.batch_size - # get image_latent_ids right after scale & shift - image_latent_ids = _prepare_image_ids([latent_condition]) - image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1) - image_latent_ids = image_latent_ids.to(get_local_torch_device()) - batch.condition_image_latent_ids = image_latent_ids - # latent: (1, 128, 32, 32) packed = self.maybe_pack_latents( latent_condition, None, batch @@ -526,6 +520,11 @@ class Flux2PipelineConfig(FluxPipelineConfig): image_latents = image_latents.repeat(batch_size, 1, 1) return image_latents + def prepare_condition_image_latent_ids(self, image_latents, batch): + image_latent_ids = _prepare_image_ids(image_latents) + image_latent_ids = image_latent_ids.repeat(batch.batch_size, 1, 1) + batch.condition_image_latent_ids = image_latent_ids.to(get_local_torch_device()) + def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch): txt_ids = _prepare_text_ids(prompt_embeds).to(device=device) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index 0738c87f6..988cad561 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py @@ -225,6 +225,10 @@ class ImageVAEEncodingStage(PipelineStage): images = [images] all_image_latents = [] + prepare_condition_image_latent_ids = getattr( + server_args.pipeline_config, "prepare_condition_image_latent_ids", None + ) + condition_latents = [] if callable(prepare_condition_image_latent_ids) else None for image in images: image = self.preprocess( image, @@ -309,12 +313,17 @@ class ImageVAEEncodingStage(PipelineStage): latent_condition -= shift_factor latent_condition = latent_condition * scaling_factor + if condition_latents is not None: + condition_latents.append(latent_condition) + image_latent = server_args.pipeline_config.postprocess_image_latent( latent_condition, batch ) all_image_latents.append(image_latent) batch.image_latent = torch.cat(all_image_latents, dim=1) + if condition_latents is not None: + prepare_condition_image_latent_ids(condition_latents, batch) self.offload_model() return batch diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 8b275ddb2..757652442 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -447,6 +447,72 @@ "expected_avg_denoise_ms": 1077.26, "expected_median_denoise_ms": 1096.5 }, + "flux_2_ti2i_multi_image_cache_dit": { + "stages_ms": { + "ImageVAEEncodingStage": 282.83, + "DenoisingStage": 26936.93, + "DecodingStage": 129.33, + "TextEncodingStage": 737.01, + "LatentPreparationStage": 0.84, + "TimestepPreparationStage": 20.57, + "InputValidationStage": 84.29 + }, + "denoise_step_ms": { + "0": 1846.55, + "1": 232.6, + "2": 1621.05, + "3": 1606.1, + "4": 1441.78, + "5": 59.2, + "6": 60.32, + "7": 235.36, + "8": 1443.21, + "9": 59.71, + "10": 59.79, + "11": 236.12, + "12": 1439.11, + "13": 59.97, + "14": 60.46, + "15": 239.29, + "16": 1441.32, + "17": 60.76, + "18": 60.68, + "19": 240.16, + "20": 1442.46, + "21": 60.14, + "22": 61.1, + "23": 239.26, + "24": 1443.28, + "25": 59.41, + "26": 60.74, + "27": 238.02, + "28": 1444.38, + "29": 59.09, + "30": 59.12, + "31": 241.69, + "32": 1441.99, + "33": 60.44, + "34": 61.93, + "35": 241.0, + "36": 1443.56, + "37": 60.55, + "38": 61.07, + "39": 238.11, + "40": 1443.33, + "41": 59.08, + "42": 60.43, + "43": 239.23, + "44": 1444.53, + "45": 61.77, + "46": 61.65, + "47": 239.84, + "48": 1443.89, + "49": 59.92 + }, + "expected_e2e_ms": 28591.93, + "expected_avg_denoise_ms": 533.42, + "expected_median_denoise_ms": 235.74 + }, "flux_image_t2i_2_gpus": { "stages_ms": { "InputValidationStage": 0.03, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index b4291a47b..b32dc9dd9 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -321,6 +321,13 @@ MULTI_IMAGE_TI2I_sampling_params = DiffusionSamplingParams( ], direct_url_test=True, ) +MULTI_IMAGE_TI2I_UPLOAD_sampling_params = DiffusionSamplingParams( + prompt="The magician bear is on the left, the alchemist bear is on the right, facing each other in the central park square.", + image_path=[ + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_1.jpg", + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_2.jpg", + ], +) MULTI_FRAME_I2I_sampling_params = DiffusionSamplingParams( prompt="a high quality, cute halloween themed illustration, consistent style and lighting", image_path=[ @@ -699,22 +706,6 @@ TWO_GPU_CASES_A = [ ), ] -# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536. -if not current_platform.is_hip(): - TWO_GPU_CASES_A.append( - DiffusionTestCase( - "turbo_wan2_2_i2v_a14b_2gpu", - DiffusionServerArgs( - model_path="IPostYellow/TurboWan2.2-I2V-A14B-Diffusers", - modality="video", - custom_validator="video", - num_gpus=2, - tp_size=2, - ), - TURBOWAN_I2V_sampling_params, - ) - ) - TWO_GPU_CASES_B = [ DiffusionTestCase( "wan2_1_i2v_14b_480P_2gpu", @@ -800,6 +791,48 @@ TWO_GPU_CASES_B = [ ), ] +if not current_platform.is_hip(): + # Flux2 multi-image edit with cache-dit, regression test + ONE_GPU_CASES_B.append( + DiffusionTestCase( + "flux_2_ti2i_multi_image_cache_dit", + DiffusionServerArgs( + model_path="black-forest-labs/FLUX.2-dev", + modality="image", + enable_cache_dit=True, + ), + MULTI_IMAGE_TI2I_UPLOAD_sampling_params, + ) + ) + # Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536. + ONE_GPU_CASES_B.append( + DiffusionTestCase( + "turbo_wan2_1_t2v_1.3b", + DiffusionServerArgs( + model_path="IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers", + modality="video", + custom_validator="video", + ), + DiffusionSamplingParams( + prompt=T2V_PROMPT, + ), + ) + ) + # Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536. + TWO_GPU_CASES_A.append( + DiffusionTestCase( + "turbo_wan2_2_i2v_a14b_2gpu", + DiffusionServerArgs( + model_path="IPostYellow/TurboWan2.2-I2V-A14B-Diffusers", + modality="video", + custom_validator="video", + num_gpus=2, + tp_size=2, + ), + TURBOWAN_I2V_sampling_params, + ) + ) + # Load global configuration BASELINE_CONFIG = BaselineConfig.load( Path(__file__).with_name("perf_baselines.json")