From ebb3b20576736ec3ad264788b8721ec61b28aa8d Mon Sep 17 00:00:00 2001 From: Khairul Kabir <250204754+khairulkabir1661@users.noreply.github.com> Date: Fri, 28 Aug 2026 08:05:41 -0700 Subject: [PATCH] [Diffusion] Make Flux2Pipeline text_encoder_out_layers configurable (#6390) Co-authored-by: khairulkabir1661 Co-authored-by: Claude Opus 4.6 --- .../diffusion/diff_model_builders.py | 17 ++++++++--------- .../diffusion/models/flux2/pipeline_flux2.py | 13 +++++++------ 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index c5eadecdb..7c1c9df67 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -242,19 +242,18 @@ def tiny_flux_kontext_builder() -> str: def tiny_flux2_builder() -> str: def shrink_text_encoder(config: dict) -> dict: - # The real checkpoint ships language_model.lm_head.weight as its own - # tensor even though tie_word_embeddings defaults to True; keep it - # untied here too so transformers doesn't drop it on save. config["tie_word_embeddings"] = False - config["text_config"]["num_hidden_layers"] = 31 + config["text_config"]["num_hidden_layers"] = 3 + config["text_config"]["hidden_size"] = 32 config["text_config"]["intermediate_size"] = 64 - config["text_config"]["num_attention_heads"] = 4 - config["text_config"]["head_dim"] = 32 + config["text_config"]["num_attention_heads"] = 2 + config["text_config"]["head_dim"] = 16 config["text_config"]["num_key_value_heads"] = 2 + config["text_config"]["text_encoder_out_layers"] = [0, 1, 2] config["vision_config"]["num_hidden_layers"] = 2 config["vision_config"]["intermediate_size"] = 64 - config["vision_config"]["num_attention_heads"] = 4 - config["vision_config"]["head_dim"] = 32 + config["vision_config"]["num_attention_heads"] = 2 + config["vision_config"]["head_dim"] = 16 return config return build_tiny_from_configs( @@ -262,6 +261,6 @@ def tiny_flux2_builder() -> str: "black-forest-labs/FLUX.2-dev", transform={ "text_encoder": shrink_text_encoder, - "transformer": partial(_shrink_dit_rope_config, num_single_layers=2), + "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, joint_attention_dim=96), }, ) diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 973f8b7ff..139b5275a 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -438,6 +438,12 @@ class Flux2Pipeline( self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) + + if hasattr(self.text_encoder.config, "text_encoder_out_layers"): + self.text_encoder_out_layers = tuple(self.text_encoder.config.text_encoder_out_layers) + else: + self.text_encoder_out_layers = (10, 20, 30) + self._current_timestep = None self._interrupt = False @@ -677,7 +683,6 @@ class Flux2Pipeline( num_images_per_prompt: int = 1, prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - text_encoder_out_layers: tuple[int, ...] = (10, 20, 30), ): device = device or self._execution_device @@ -694,7 +699,7 @@ class Flux2Pipeline( device=device, max_sequence_length=max_sequence_length, system_message=self.system_message, - hidden_states_layers=text_encoder_out_layers, + hidden_states_layers=self.text_encoder_out_layers, ) batch_size, seq_len, _ = prompt_embeds.shape @@ -886,7 +891,6 @@ class Flux2Pipeline( callback_on_step_end: Callable[[int, int, dict], None] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, - text_encoder_out_layers: tuple[int, ...] = (10, 20, 30), caption_upsample_temperature: float = None, ) -> DiffusionOutput: if len(req.prompts) > 1: @@ -922,7 +926,6 @@ class Flux2Pipeline( else num_images_per_prompt ) max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers) caption_upsample_temperature = req.sampling_params.extra_args.get( "caption_upsample_temperature", caption_upsample_temperature ) @@ -977,7 +980,6 @@ class Flux2Pipeline( device=device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, - text_encoder_out_layers=text_encoder_out_layers, ) has_neg_prompt = negative_prompt_embeds is not None or any(req_negative_prompt) @@ -993,7 +995,6 @@ class Flux2Pipeline( device=device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, - text_encoder_out_layers=text_encoder_out_layers, ) # 4. process images