[Feat] support CFG Parallel for FLUX kontext (#2281)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

tests/e2e/online_serving/test_flux_kontext_expansion.py

@@ -5,6 +5,8 @@ and are supported by the FluxKontext model.
 
 import pytest
 
+from tests.helpers import skip_if_gated_repo_inaccessible
+from tests.helpers.mark import hardware_marks
 from tests.helpers.media import generate_synthetic_image
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
 
@@ -13,6 +15,8 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
 NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
 MODEL = "black-forest-labs/FLUX.1-Kontext-dev"
+skip_if_gated_repo_inaccessible(MODEL)
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2)
 
 
 def _get_diffusion_feature_cases(model: str):
@@ -26,7 +30,20 @@ def _get_diffusion_feature_cases(model: str):
                     "--enable-cpu-offload",
                 ],
             ),
-            id="parallel_001",
+            id="parallel_tp_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--cfg-parallel-size",
+                    "2",
+                ],
+            ),
+            id="parallel_cfg_2",
+            marks=PARALLEL_FEATURE_MARKS,
         ),
     ]
 

vllm_omni/diffusion/distributed/cfg_parallel.py

@@ -463,6 +463,31 @@ class CFGParallelMixin(metaclass=ABCMeta):
         """
         raise NotImplementedError("Subclasses must implement diffuse")
 
+    def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool) -> bool:
+        """
+        Check if CFG parallel configuration is valid.
+
+        Args:
+            true_cfg_scale: The classifier-free guidance scale value
+            has_neg_prompt: Whether a negative prompt is provided
+
+        Returns:
+            True if CFG parallel configuration is valid, False otherwise
+        """
+        if get_classifier_free_guidance_world_size() == 1:
+            return True
+
+        if true_cfg_scale <= 1:
+            logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.")
+            return False
+
+        if not has_neg_prompt:
+            logger.warning(
+                "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings."
+            )
+            return False
+        return True
+
     def scheduler_step(
         self,
         noise_pred: torch.Tensor | tuple[torch.Tensor, ...],

vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py

@@ -24,6 +24,7 @@ from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokeniz
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.flux import (
@@ -72,6 +73,7 @@ def get_flux_kontext_post_process_func(od_config: OmniDiffusionConfig) -> Callab
 class FluxKontextPipeline(
     nn.Module,
     FluxPipelineMixin,
+    CFGParallelMixin,
     SupportImageInput,
     ProgressBarMixin,
     DiffusionPipelineProfilerMixin,
@@ -157,6 +159,7 @@ class FluxKontextPipeline(
         self._attention_kwargs = None
         self._joint_attention_kwargs = None
         self._num_timesteps = None
+        self._current_timestep = None
         self._interrupt = False
         self._callback_tensor_inputs = ["latents", "prompt_embeds"]
         self.latent_channels = self.vae.config.latent_channels if hasattr(self.vae, "config") else 16
@@ -465,6 +468,113 @@ class FluxKontextPipeline(
     def num_timesteps(self):
         return self._num_timesteps
 
+    @property
+    def current_timestep(self) -> torch.Tensor | None:
+        return self._current_timestep
+
+    def predict_noise(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        guidance: torch.Tensor | None,
+        pooled_projections: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        txt_ids: torch.Tensor,
+        img_ids: torch.Tensor,
+        joint_attention_kwargs: dict[str, Any] | None,
+        return_dict: bool = False,
+    ) -> torch.Tensor:
+        return self.transformer(
+            hidden_states=hidden_states,
+            timestep=timestep,
+            guidance=guidance,
+            pooled_projections=pooled_projections,
+            encoder_hidden_states=encoder_hidden_states,
+            txt_ids=txt_ids,
+            img_ids=img_ids,
+            joint_attention_kwargs=joint_attention_kwargs,
+            return_dict=return_dict,
+        )[0]
+
+    def diffuse(
+        self,
+        prompt_embeds: torch.Tensor,
+        pooled_prompt_embeds: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_pooled_prompt_embeds: torch.Tensor | None,
+        latents: torch.Tensor,
+        image_latents: torch.Tensor | None,
+        latent_ids: torch.Tensor,
+        text_ids: torch.Tensor,
+        negative_text_ids: torch.Tensor | None,
+        timesteps: torch.Tensor,
+        do_true_cfg: bool,
+        guidance: torch.Tensor | None,
+        true_cfg_scale: float,
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+    ) -> torch.Tensor:
+        self.scheduler.set_begin_index(0)
+        self.transformer.do_true_cfg = do_true_cfg
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+                latent_model_input = latents
+                if image_latents is not None:
+                    latent_model_input = torch.cat([latents, image_latents], dim=1)
+                timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
+
+                positive_kwargs = {
+                    "hidden_states": latent_model_input,
+                    "timestep": timestep / 1000,
+                    "guidance": guidance,
+                    "pooled_projections": pooled_prompt_embeds,
+                    "encoder_hidden_states": prompt_embeds,
+                    "txt_ids": text_ids,
+                    "img_ids": latent_ids,
+                    "joint_attention_kwargs": self.joint_attention_kwargs,
+                    "return_dict": False,
+                }
+                negative_kwargs = None
+                if do_true_cfg:
+                    negative_kwargs = {
+                        "hidden_states": latent_model_input,
+                        "timestep": timestep / 1000,
+                        "guidance": guidance,
+                        "pooled_projections": negative_pooled_prompt_embeds,
+                        "encoder_hidden_states": negative_prompt_embeds,
+                        "txt_ids": negative_text_ids,
+                        "img_ids": latent_ids,
+                        "joint_attention_kwargs": self.joint_attention_kwargs,
+                        "return_dict": False,
+                    }
+
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg,
+                    true_cfg_scale,
+                    positive_kwargs,
+                    negative_kwargs,
+                    cfg_normalize=False,
+                    output_slice=latents.size(1),
+                )
+
+                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                pbar.update()
+        return latents
+
     @torch.no_grad()
     def forward(
         self,
@@ -503,8 +613,13 @@ class FluxKontextPipeline(
 
         if isinstance(first_prompt, str):
             prompt = first_prompt
+            # Keep negative_prompt from function parameter (passed via API)
         elif first_prompt:
             prompt = first_prompt.get("prompt") or ""
+            # Prefer negative_prompt from prompt dict, fallback to function parameter
+            np_from_dict = first_prompt.get("negative_prompt")
+            if np_from_dict is not None:
+                negative_prompt = np_from_dict
         # else: keep original prompt unchanged
 
         # Handle image from prompt data
@@ -524,8 +639,17 @@ class FluxKontextPipeline(
         height = req.sampling_params.height or height
         width = req.sampling_params.width or width
         num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas or sigmas
         guidance_scale = req.sampling_params.guidance_scale or guidance_scale
+        true_cfg_scale = (
+            req.sampling_params.true_cfg_scale if req.sampling_params.true_cfg_scale is not None else true_cfg_scale
+        )
         generator = req.sampling_params.generator or generator
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
         latents = (
             req.sampling_params.extra_args.get("latents")
             if req.sampling_params.extra_args.get("latents") is not None
@@ -568,6 +692,8 @@ class FluxKontextPipeline(
         )
         do_true_cfg = true_cfg_scale > 1.0 and has_neg_prompt
 
+        do_true_cfg = do_true_cfg and self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
+
         # 1. Prepare text embeddings
         prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
             prompt=prompt,
@@ -651,57 +777,23 @@ class FluxKontextPipeline(
             self._joint_attention_kwargs = {}
 
         # 5. Denoising loop
-        self.scheduler.set_begin_index(0)
-        with self.progress_bar(total=len(timesteps)) as pbar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                latent_model_input = latents
-                if image_latents is not None:
-                    latent_model_input = torch.cat([latents, image_latents], dim=1)
-                timestep = t.expand(latents.shape[0]).to(latents.dtype)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    timestep=timestep / 1000,
-                    guidance=guidance,
-                    pooled_projections=pooled_prompt_embeds,
-                    encoder_hidden_states=prompt_embeds,
-                    txt_ids=text_ids,
-                    img_ids=latent_ids,
-                    joint_attention_kwargs=self.joint_attention_kwargs,
-                    return_dict=False,
-                )[0]
-                noise_pred = noise_pred[:, : latents.size(1)]
-
-                if do_true_cfg:
-                    neg_noise_pred = self.transformer(
-                        hidden_states=latent_model_input,
-                        timestep=timestep / 1000,
-                        guidance=guidance,
-                        pooled_projections=negative_pooled_prompt_embeds,
-                        encoder_hidden_states=negative_prompt_embeds,
-                        txt_ids=negative_text_ids,
-                        img_ids=latent_ids,
-                        joint_attention_kwargs=self.joint_attention_kwargs,
-                        return_dict=False,
-                    )[0]
-                    neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
-                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
-
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-
-                pbar.update()
+        latents = self.diffuse(
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            latents=latents,
+            image_latents=image_latents,
+            latent_ids=latent_ids,
+            text_ids=text_ids,
+            negative_text_ids=negative_text_ids if do_true_cfg else None,
+            timesteps=timesteps,
+            do_true_cfg=do_true_cfg,
+            guidance=guidance,
+            true_cfg_scale=true_cfg_scale,
+            callback_on_step_end=callback_on_step_end,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
         if output_type == "latent":
             image = latents
         else: