From 089bb88706d4e3038f370981a4d3d4b7e2910eec Mon Sep 17 00:00:00 2001
From: Alex Brooks <albrooks@redhat.com>
Date: Thu, 30 Jul 2026 03:55:42 -0600
Subject: [PATCH] [Model] Remove Text Encoder Output Layer Hardcoding in Flux2
 (#5589)

Signed-off-by: Alex Brooks <albrooks@redhat.com>
---
 .../text_encoder/config.json                  | 30 ++-----------------
 .../flux2_klein/pipeline_flux2_klein.py       | 19 +++++++-----
 2 files changed, 14 insertions(+), 35 deletions(-)

diff --git a/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json b/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json
index 004914a7f..19a6b4063 100644
--- a/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json
+++ b/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json
@@ -13,31 +13,6 @@
   "initializer_range": 0.02,
   "intermediate_size": 64,
   "layer_types": [
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
-    "full_attention",
     "full_attention",
     "full_attention",
     "full_attention"
@@ -46,7 +21,7 @@
   "max_window_layers": 36,
   "model_type": "qwen3",
   "num_attention_heads": 2,
-  "num_hidden_layers": 28,
+  "num_hidden_layers": 3,
   "num_key_value_heads": 2,
   "rms_norm_eps": 1e-06,
   "rope_scaling": null,
@@ -56,5 +31,6 @@
   "transformers_version": "4.56.1",
   "use_cache": true,
   "use_sliding_window": false,
-  "vocab_size": 151936
+  "vocab_size": 151936,
+  "text_encoder_out_layers":[0, 1, 2]
 }
diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
index 184b5d402..a4350373b 100644
--- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
+++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
@@ -271,6 +271,15 @@ class Flux2KleinPipeline(
         self.tokenizer_max_length = 512
         self.default_sample_size = 128
 
+        # text_encoder_out_layers is hardcoded to (9, 18, 27) in Diffusers,
+        # but this won't necessarily work for all text encoders. We allow
+        # setting the text_encoder_out_layers in the config to allow
+        # custom text encoders
+        if hasattr(self.text_encoder.config, "text_encoder_out_layers"):
+            self.text_encoder_out_layers = self.text_encoder.config.text_encoder_out_layers
+        else:
+            self.text_encoder_out_layers = (9, 18, 27)
+
         self._guidance_scale = None
         self._attention_kwargs = None
         self._num_timesteps = None
@@ -293,10 +302,10 @@ class Flux2KleinPipeline(
         text_encoder: Qwen3ForCausalLM,
         tokenizer: Qwen2TokenizerFast,
         prompt: str | list[str],
+        hidden_states_layers: tuple[int, ...],
         dtype: torch.dtype | None = None,
         device: torch.device | None = None,
         max_sequence_length: int = 512,
-        hidden_states_layers: list[int] = (9, 18, 27),
     ):
         dtype = text_encoder.dtype if dtype is None else dtype
         device = text_encoder.device if device is None else device
@@ -513,7 +522,6 @@ class Flux2KleinPipeline(
         num_images_per_prompt: int = 1,
         prompt_embeds: torch.Tensor | None = None,
         max_sequence_length: int = 512,
-        text_encoder_out_layers: tuple[int, ...] = (9, 18, 27),
     ):
         device = device or self._execution_device
 
@@ -527,9 +535,9 @@ class Flux2KleinPipeline(
                 text_encoder=self.text_encoder,
                 tokenizer=self.tokenizer,
                 prompt=prompt,
+                hidden_states_layers=self.text_encoder_out_layers,
                 device=device,
                 max_sequence_length=max_sequence_length,
-                hidden_states_layers=text_encoder_out_layers,
             )
 
         batch_size, seq_len, _ = prompt_embeds.shape
@@ -795,9 +803,6 @@ class Flux2KleinPipeline(
         latents = req.sampling_params.latents
         strength = req.sampling_params.strength if req.sampling_params.strength is not None else 1.0
 
-        text_encoder_out_layers: tuple[int, ...] = req.sampling_params.extra_args.get(
-            "text_encoder_out_layers", (9, 18, 27)
-        )
         padding_mask_crop: int | None = req.sampling_params.extra_args.get("padding_mask_crop")
         attention_kwargs: dict[str, Any] | None = None
         callback_on_step_end: Callable[[int, int, dict], None] | None = None
@@ -833,7 +838,6 @@ class Flux2KleinPipeline(
             device=device,
             num_images_per_prompt=num_images_per_prompt,
             max_sequence_length=max_sequence_length,
-            text_encoder_out_layers=text_encoder_out_layers,
         )
 
         if self.do_classifier_free_guidance:
@@ -846,7 +850,6 @@ class Flux2KleinPipeline(
                 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