[Diffusion] Make Flux2Pipeline text_encoder_out_layers configurable (#6390)

Co-authored-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

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),
         },
     )

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