Fix callback tensor inputs that are never bound in the denoising loop

src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py

@@ -219,7 +219,6 @@ class StableDiffusionXLControlNetUnionInpaintPipeline(
         "add_text_embeds",
         "add_time_ids",
         "mask",
-        "masked_image_latents",
         "control_image",
     ]
 

src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py

@@ -300,7 +300,7 @@ class LEditsPPPipelineStableDiffusion(
 
     model_cpu_offload_seq = "text_encoder->unet->vae"
     _exclude_from_cpu_offload = ["safety_checker"]
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+    _callback_tensor_inputs = ["latents", "negative_prompt_embeds"]
     _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
 
     def __init__(
@@ -1186,7 +1186,6 @@ def __call__(
                     callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
 
                     latents = callback_outputs.pop("latents", latents)
-                    # prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
                     negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
 
                 # call the callback, if provided

src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py

@@ -340,7 +340,6 @@ class LEditsPPPipelineStableDiffusionXL(
         "add_text_embeds",
         "add_time_ids",
         "negative_pooled_prompt_embeds",
-        "negative_add_time_ids",
     ]
 
     def __init__(
@@ -1337,7 +1336,6 @@ def __call__(
                         "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
                     )
                     add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
-                    # negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids)
 
                 # call the callback, if provided
                 if i == len(timesteps) - 1 or ((i + 1) > 0 and (i + 1) % self.scheduler.order == 0):

tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py

@@ -203,6 +203,28 @@ def test_ledits_pp_warmup_steps(self):
         inputs["edit_warmup_steps"] = [10, 5]
         pipe(**inputs).images
 
+    def test_callback_inputs(self):
+        device = "cpu"  # ensure determinism for the device-dependent torch.Generator
+        components = self.get_dummy_components()
+        pipe = LEditsPPPipelineStableDiffusion(**components)
+        pipe = pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        inversion_inputs = self.get_dummy_inversion_inputs(device)
+        inversion_inputs["image"] = inversion_inputs["image"][0]
+        pipe.invert(**inversion_inputs)
+
+        def callback_inputs_all(pipe, i, t, callback_kwargs):
+            for tensor_name in pipe._callback_tensor_inputs:
+                assert tensor_name in callback_kwargs
+
+            return callback_kwargs
+
+        inputs = self.get_dummy_inputs(device)
+        inputs["callback_on_step_end"] = callback_inputs_all
+        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
+        pipe(**inputs)
+
 
 @slow
 @require_torch_accelerator

tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion_xl.py

@@ -251,6 +251,28 @@ def test_ledits_pp_warmup_steps(self):
         inputs["edit_warmup_steps"] = [10, 5]
         pipe(**inputs).images
 
+    def test_callback_inputs(self):
+        device = "cpu"  # ensure determinism for the device-dependent torch.Generator
+        components = self.get_dummy_components()
+        pipe = LEditsPPPipelineStableDiffusionXL(**components)
+        pipe = pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        inversion_inputs = self.get_dummy_inversion_inputs(device)
+        inversion_inputs["image"] = inversion_inputs["image"][0]
+        pipe.invert(**inversion_inputs)
+
+        def callback_inputs_all(pipe, i, t, callback_kwargs):
+            for tensor_name in pipe._callback_tensor_inputs:
+                assert tensor_name in callback_kwargs
+
+            return callback_kwargs
+
+        inputs = self.get_dummy_inputs(device)
+        inputs["callback_on_step_end"] = callback_inputs_all
+        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
+        pipe(**inputs)
+
 
 @slow
 @require_torch_accelerator