From acf941da4da8bdca36cf4e8e8ea16d1e2df38e95 Mon Sep 17 00:00:00 2001 From: Diya Peng <50388704+dpeng123@users.noreply.github.com> Date: Fri, 24 Jul 2026 00:00:57 -0700 Subject: [PATCH] [Diffusion] Add VAE patch parallel support for FLUX.2-dev (#5292) Signed-off-by: dpeng123 --- docs/design/feature/vae_parallel.md | 1 + docs/user_guide/diffusion_features.md | 2 +- .../test_flux2_vae_patch_parallel.py | 210 ++++++++++++++++++ .../test_flux_2_dev_expansion.py | 16 ++ .../autoencoders/autoencoder_kl_flux2.py | 116 ++++++++++ .../diffusion/models/flux2/pipeline_flux2.py | 9 +- 6 files changed, 349 insertions(+), 5 deletions(-) create mode 100644 tests/diffusion/distributed/test_flux2_vae_patch_parallel.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_flux2.py diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md index 00f41378a..c465449f8 100644 --- a/docs/design/feature/vae_parallel.md +++ b/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` | ✅ | ✅ | --- diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index bf714b3a0..005a0e899 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | diff --git a/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py b/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py new file mode 100644 index 000000000..e861bd61a --- /dev/null +++ b/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})" + ) diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index 868a1e3bd..9bc3308b8 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/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, + ), ] diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_flux2.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_flux2.py new file mode 100644 index 000000000..7c0cee0f5 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 351d94f7e..0dbfe0f8c 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/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(