[Diffusion] Add VAE patch parallel support for FLUX.2-dev (#5292)

Signed-off-by: dpeng123 <diyapeng123@gmail.com>

docs/design/feature/vae_parallel.md

@@ -445,6 +445,7 @@ Complete examples in the codebase:
 | **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
 | **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
 | **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
+| **FLUX.2-dev** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_flux2.py` | ✅ | ✅ |
 
 ---
 

docs/user_guide/diffusion_features.md

@@ -117,7 +117,7 @@ The following tables show which models support each feature:
 | **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.1-Kontext-dev**   |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **FLUX.2-dev**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **FLUX.2-dev**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ✅           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |

tests/diffusion/distributed/test_flux2_vae_patch_parallel.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""FLUX.2-dev VAE patch-parallel correctness tests.
+
+Compares vae_patch_parallel_size=1 vs 2 on TP=2 with tiling enabled:
+- Text-to-image decode path
+- Image-to-image encode + decode path (conditioning image)
+"""
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+PROMPT = "a photo of a cat sitting on a laptop keyboard"
+NEGATIVE_PROMPT = "blurry, low quality"
+
+VAE_PP_MEAN_THRESHOLD = 3e-2
+VAE_PP_P99_THRESHOLD = 1e-1
+
+
+def _get_flux2_dev_model() -> str:
+    return os.environ.get("VLLM_TEST_FLUX2_DEV_MODEL", "black-forest-labs/FLUX.2-dev")
+
+
+def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
+    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(arr)
+
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
+    ta = _pil_to_float_rgb_tensor(a)
+    tb = _pil_to_float_rgb_tensor(b)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+    p99_abs_diff = torch.quantile(abs_diff.flatten(), 0.99).item()
+    return abs_diff.mean().item(), p99_abs_diff
+
+
+def _extract_single_image(outputs) -> Image.Image:
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    if not hasattr(first_output, "request_output") or not first_output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")
+
+    req_out = first_output.request_output
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images' key")
+
+    images = req_out.images
+    if images is None or len(images) != 1:
+        raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
+    return images[0]
+
+
+def _run_flux2_generate(
+    *,
+    tp_size: int,
+    height: int,
+    width: int,
+    num_inference_steps: int,
+    seed: int,
+    vae_patch_parallel_size: int,
+    conditioning_image: Image.Image | None = None,
+) -> Image.Image:
+    current_omni_platform.empty_cache()
+
+    request: dict = {"prompt": PROMPT, "negative_prompt": NEGATIVE_PROMPT}
+    if conditioning_image is not None:
+        request["multi_modal_data"] = {"image": conditioning_image}
+
+    with OmniRunner(
+        _get_flux2_dev_model(),
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=tp_size,
+            vae_patch_parallel_size=vae_patch_parallel_size,
+        ),
+        enable_cpu_offload=True,
+        vae_use_tiling=True,
+    ) as runner:
+        outputs = list(
+            runner.omni.generate(
+                request,
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=4.0,
+                    seed=seed,
+                    num_outputs_per_prompt=1,
+                ),
+            )
+        )
+
+    return _extract_single_image(outputs)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_flux2_vae_patch_parallel_decode_tp2(tmp_path: Path):
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("FLUX.2-dev VAE patch parallel requires >= 2 devices.")
+
+    height = 1152
+    width = 1152
+    num_inference_steps = 2
+    seed = 42
+
+    baseline_img = _run_flux2_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        vae_patch_parallel_size=1,
+    )
+    pp2_img = _run_flux2_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        vae_patch_parallel_size=2,
+    )
+
+    baseline_path = tmp_path / "flux2_tp2_vae_pp1.png"
+    pp2_path = tmp_path / "flux2_tp2_vae_pp2.png"
+    baseline_img.save(baseline_path)
+    pp2_img.save(pp2_path)
+
+    assert baseline_img.width == width and baseline_img.height == height
+    assert pp2_img.width == width and pp2_img.height == height
+
+    mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
+    print(
+        "FLUX.2-dev VAE patch parallel decode diff (TP=2, pp=1 vs pp=2): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={VAE_PP_MEAN_THRESHOLD:.6e}, p99<={VAE_PP_P99_THRESHOLD:.6e}; "
+        f"pp1_img={baseline_path}, pp2_img={pp2_path}"
+    )
+    assert mean_abs_diff <= VAE_PP_MEAN_THRESHOLD and p99_abs_diff <= VAE_PP_P99_THRESHOLD, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={VAE_PP_MEAN_THRESHOLD:.6e}, p99<={VAE_PP_P99_THRESHOLD:.6e})"
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_flux2_vae_patch_parallel_i2i_tp2(tmp_path: Path):
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("FLUX.2-dev VAE patch parallel I2I requires >= 2 devices.")
+
+    height = 1152
+    width = 1152
+    num_inference_steps = 2
+    seed = 42
+    conditioning_image = Image.new("RGB", (width, height), (120, 80, 40))
+
+    baseline_img = _run_flux2_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        vae_patch_parallel_size=1,
+        conditioning_image=conditioning_image,
+    )
+    pp2_img = _run_flux2_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        vae_patch_parallel_size=2,
+        conditioning_image=conditioning_image,
+    )
+
+    baseline_path = tmp_path / "flux2_i2i_tp2_vae_pp1.png"
+    pp2_path = tmp_path / "flux2_i2i_tp2_vae_pp2.png"
+    baseline_img.save(baseline_path)
+    pp2_img.save(pp2_path)
+
+    mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
+    print(
+        "FLUX.2-dev VAE patch parallel I2I diff (TP=2, pp=1 vs pp=2): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={VAE_PP_MEAN_THRESHOLD:.6e}, p99<={VAE_PP_P99_THRESHOLD:.6e}; "
+        f"pp1_img={baseline_path}, pp2_img={pp2_path}"
+    )
+    assert mean_abs_diff <= VAE_PP_MEAN_THRESHOLD and p99_abs_diff <= VAE_PP_P99_THRESHOLD, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={VAE_PP_MEAN_THRESHOLD:.6e}, p99<={VAE_PP_P99_THRESHOLD:.6e})"
+    )

tests/e2e/online_serving/test_flux_2_dev_expansion.py

@@ -6,6 +6,7 @@ Coverage:
 - Layerwise CPU offload
 - Ulysses sequence parallelism
 - Ring sequence parallelism
+- VAE patch parallel encode/decode
 
 This test verifies that FLUX.2-dev can be launched with CPU offload enabled,
 accepts text-to-image requests through the OpenAI-compatible API, and returns
@@ -90,6 +91,21 @@ def _get_flux_2_dev_feature_cases(model: str):
             id="ring_2",
             marks=PARALLEL_FEATURE_MARKS,
         ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--tensor-parallel-size",
+                    "2",
+                    "--vae-patch-parallel-size",
+                    "2",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="vae_patch_parallel_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
     ]
 
 

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_flux2.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL_base
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    GridSpec,
+    TileTask,
+)
+
+logger = init_logger(__name__)
+
+
+class DistributedAutoencoderKLFlux2(DistributedAutoencoderKL_base, AutoencoderKLFlux2):
+    def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, height, width = x.shape
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        tiletask_list = []
+        for i in range(0, height, overlap_size):
+            for j in range(0, width, overlap_size):
+                tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // overlap_size, j // overlap_size),
+                        tile,
+                        workload=tile.shape[2] * tile.shape[3],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(2, 3),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "blend_extent": blend_extent,
+                "row_limit": row_limit,
+            },
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+        tile = self.encoder(task.tensor)
+        if self.quant_conv is not None:
+            tile = self.quant_conv(tile)
+        return tile
+
+    def encode_tile_merge(
+        self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+    ) -> torch.Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_extent"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_extent"])
+                result_row.append(tile[:, :, : grid_spec.tile_spec["row_limit"], : grid_spec.tile_spec["row_limit"]])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, height, width = x.shape
+        if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size):
+            if self.is_distributed_enabled():
+                logger.debug("Encode running with distributed executor")
+                return self.distributed_executor.execute(
+                    x,
+                    DistributedOperator(
+                        split=self.encode_tile_split,
+                        exec=self.encode_tile_exec,
+                        merge=self.encode_tile_merge,
+                    ),
+                    broadcast_result=True,
+                )
+            return self._tiled_encode(x)
+
+        enc = self.encoder(x)
+        if self.quant_conv is not None:
+            enc = self.quant_conv(enc)
+        return enc
+
+    def decode(self, z: torch.Tensor, return_dict: bool = True, *args: Any, **kwargs: Any):
+        if not self.is_distributed_enabled():
+            return super().decode(z, return_dict=return_dict, *args, **kwargs)
+
+        split, exec_fn, merge = self._strategy_select(z)
+
+        if split is not None:
+            strategy = "tile" if split == self.tile_split else "patch"
+            logger.info(f"Decode run with distributed executor, split strategy is {strategy}")
+            result = self.distributed_executor.execute(
+                z,
+                DistributedOperator(split=split, exec=exec_fn, merge=merge),
+                broadcast_result=True,
+            )
+            if not return_dict:
+                return (result,)
+
+            from diffusers.models.autoencoders.vae import DecoderOutput
+
+            return DecoderOutput(sample=result)
+
+        return super().decode(z, return_dict=return_dict, *args, **kwargs)

vllm_omni/diffusion/models/flux2/pipeline_flux2.py

@@ -11,7 +11,7 @@ from typing import Any, ClassVar, cast
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler
+from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.pipelines.flux2.pipeline_flux2 import UPSAMPLING_MAX_IMAGE_SIZE
 from diffusers.pipelines.flux2.system_messages import (
@@ -25,6 +25,7 @@ from transformers import AutoConfig, AutoProcessor, PixtralProcessor
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_flux2 import DistributedAutoencoderKLFlux2
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -416,9 +417,9 @@ class Flux2Pipeline(
             system_message_t2i=SYSTEM_MESSAGE_UPSAMPLING_T2I,
             system_message_i2i=SYSTEM_MESSAGE_UPSAMPLING_I2I,
         )
-        self.vae = AutoencoderKLFlux2.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self._execution_device
-        )
+        self.vae = DistributedAutoencoderKLFlux2.from_pretrained(
+            model, subfolder="vae", local_files_only=local_files_only
+        ).to(self._execution_device)
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel)
         transformer_quant_config = _resolve_component_quant_config(od_config.quantization_config, "transformer")
         self.transformer = Flux2Transformer2DModel(