Fix Qwen Image edit pipelines to forward multimodal token type IDs

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]

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]

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]

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
     ):