[diffusion] fix images/edit with 2 images (#17520)

Signed-off-by: Chi McIsaac <chixie.mcisaac@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

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)
 

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

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,

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")