diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py index 5fa4eb2a1..17d4b0263 100644 --- a/tests/e2e/online_serving/test_flux_kontext_expansion.py +++ b/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, ), ] diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index 83ce73b16..55130e0cf 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/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, ...], diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py index 0b3b80cbd..15f0e4e0c 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py +++ b/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: