From b0d55ec64c9cd77247419f93d5f3dfb94a516589 Mon Sep 17 00:00:00 2001 From: LuShadowX Date: Fri, 7 Aug 2026 21:53:54 +0530 Subject: [PATCH] Drop callback tensor inputs that are never bound in the denoising loop LEditsPPPipelineStableDiffusion allows prompt_embeds, LEditsPPPipelineStableDiffusionXL allows negative_add_time_ids and StableDiffusionXLControlNetUnionInpaintPipeline allows masked_image_latents, but none of those names exist as locals in the respective __call__, so callback_kwargs[k] = locals()[k] raises KeyError as soon as a user requests them. Add a test_callback_inputs fast test to both LEdits++ test classes, which are plain unittest.TestCase and therefore never ran the PipelineTesterMixin callback tests. --- ...pipeline_controlnet_union_inpaint_sd_xl.py | 1 - .../pipeline_leditspp_stable_diffusion.py | 3 +-- .../pipeline_leditspp_stable_diffusion_xl.py | 2 -- .../test_ledits_pp_stable_diffusion.py | 22 +++++++++++++++++++ .../test_ledits_pp_stable_diffusion_xl.py | 22 +++++++++++++++++++ 5 files changed, 45 insertions(+), 5 deletions(-) diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py index 511611f036b4..6d7e1f62beb3 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py +++ b/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", ] diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py index 9048b56b01f4..c6cdb127309b 100644 --- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +++ b/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 diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py index c2e6f3caaf47..6e97b1ea2ed4 100644 --- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +++ b/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): diff --git a/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py b/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py index 6db20a464f19..de3ea6cc5152 100644 --- a/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py +++ b/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 diff --git a/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion_xl.py b/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion_xl.py index 06c1ceb0cf5a..f5d20d58dee8 100644 --- a/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion_xl.py +++ b/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