From 93d7a06e7fc9a7e886e9dacd7127a99cd85dd0c1 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Fri, 18 Sep 2026 13:50:36 +0800 Subject: [PATCH] Use pre-norm text features for Qwen-Image 2.1 (#7765) Signed-off-by: Tianyu Guo --- .../models/qwen_image_21/pipeline_qwen_image_21.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py index 369f20b54..8f40c94c1 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py +++ b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py @@ -541,7 +541,12 @@ class QwenImage21Pipeline( if getattr(model_inputs, "mm_token_type_ids", None) is not None: forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids - outputs = self.text_encoder(**forward_kwargs) + text_model = getattr(self.text_encoder.model, "language_model", self.text_encoder.model) + handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) + try: + outputs = self.text_encoder(**forward_kwargs) + finally: + handle.remove() hidden_states = outputs.hidden_states[-1] split_hidden_states = list(self._extract_masked_hidden(hidden_states, model_inputs.attention_mask))