From c8af080794b782619346e16e8e370a4a7065e9cd Mon Sep 17 00:00:00 2001 From: wutachiang Date: Wed, 15 Jul 2026 12:50:22 +0800 Subject: [PATCH] Fix Qwen Image multimodal token type forwarding --- .../qwenimage/pipeline_qwenimage_edit.py | 8 +++ .../pipeline_qwenimage_edit_inpaint.py | 8 +++ .../qwenimage/pipeline_qwenimage_edit_plus.py | 8 +++ .../qwenimage/test_qwenimage_edit_plus.py | 59 +++++++++++++++++++ 4 files changed, 83 insertions(+) diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py index 85abb815cf23..939dce9707a6 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py @@ -245,12 +245,20 @@ def _get_qwen_prompt_embeds( return_tensors="pt", ).to(device) + text_encoder_kwargs = {} + if ( + "mm_token_type_ids" in model_inputs + and "mm_token_type_ids" in inspect.signature(self.text_encoder.forward).parameters + ): + text_encoder_kwargs["mm_token_type_ids"] = model_inputs["mm_token_type_ids"] + outputs = self.text_encoder( input_ids=model_inputs["input_ids"], attention_mask=model_inputs["attention_mask"], pixel_values=model_inputs.get("pixel_values"), image_grid_thw=model_inputs.get("image_grid_thw"), output_hidden_states=True, + **text_encoder_kwargs, ) hidden_states = outputs.hidden_states[-1] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py index 57d1fdaaf99f..92d6d52cc18f 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py @@ -256,12 +256,20 @@ def _get_qwen_prompt_embeds( return_tensors="pt", ).to(device) + text_encoder_kwargs = {} + if ( + "mm_token_type_ids" in model_inputs + and "mm_token_type_ids" in inspect.signature(self.text_encoder.forward).parameters + ): + text_encoder_kwargs["mm_token_type_ids"] = model_inputs["mm_token_type_ids"] + outputs = self.text_encoder( input_ids=model_inputs["input_ids"], attention_mask=model_inputs["attention_mask"], pixel_values=model_inputs.get("pixel_values"), image_grid_thw=model_inputs.get("image_grid_thw"), output_hidden_states=True, + **text_encoder_kwargs, ) hidden_states = outputs.hidden_states[-1] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py index 84d1b60152b1..f02c5bacf51c 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py @@ -258,12 +258,20 @@ def _get_qwen_prompt_embeds( return_tensors="pt", ).to(device) + text_encoder_kwargs = {} + if ( + "mm_token_type_ids" in model_inputs + and "mm_token_type_ids" in inspect.signature(self.text_encoder.forward).parameters + ): + text_encoder_kwargs["mm_token_type_ids"] = model_inputs["mm_token_type_ids"] + outputs = self.text_encoder( input_ids=model_inputs["input_ids"], attention_mask=model_inputs["attention_mask"], pixel_values=model_inputs.get("pixel_values"), image_grid_thw=model_inputs.get("image_grid_thw"), output_hidden_states=True, + **text_encoder_kwargs, ) hidden_states = outputs.hidden_states[-1] diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py index f240c9e02fc1..017400a75f4d 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py +++ b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py @@ -171,6 +171,65 @@ def test_inference(self): generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) + def test_mm_token_type_ids_forwarding_is_backward_compatible(self): + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to("cpu") + + image = Image.new("RGB", (32, 32)) + original_processor = pipe.processor + + class ProcessorWithMMTokenTypeIds: + def __call__(self, *args, **kwargs): + model_inputs = original_processor(*args, **kwargs) + model_inputs["mm_token_type_ids"] = torch.zeros_like(model_inputs["input_ids"]) + return model_inputs + + pipe.processor = ProcessorWithMMTokenTypeIds() + original_forward = pipe.text_encoder.forward + received_mm_token_type_ids = [] + + def current_forward( + input_ids, + attention_mask, + pixel_values, + image_grid_thw, + output_hidden_states, + mm_token_type_ids=None, + ): + received_mm_token_type_ids.append(mm_token_type_ids) + return original_forward( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + output_hidden_states=output_hidden_states, + mm_token_type_ids=mm_token_type_ids, + ) + + pipe.text_encoder.forward = current_forward + pipe.encode_prompt(prompt="dance monkey", image=[image, image], device="cpu") + + self.assertEqual(len(received_mm_token_type_ids), 1) + self.assertIsNotNone(received_mm_token_type_ids[0]) + + legacy_forward_called = [] + + def legacy_forward(input_ids, attention_mask, pixel_values, image_grid_thw, output_hidden_states): + legacy_forward_called.append(True) + return original_forward( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + output_hidden_states=output_hidden_states, + ) + + pipe.text_encoder.forward = legacy_forward + pipe.encode_prompt(prompt="dance monkey", image=[image, image], device="cpu") + + self.assertEqual(legacy_forward_called, [True]) + def test_attention_slicing_forward_pass( self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 ):