Skip to content

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

Merged
dg845 merged 3 commits into
huggingface:mainfrom
LuShadowX:fix-callback-tensor-inputs
Aug 28, 2026
Merged

Fix callback tensor inputs that are never bound in the denoising loop#14416
dg845 merged 3 commits into
huggingface:mainfrom
LuShadowX:fix-callback-tensor-inputs

Conversation

@LuShadowX

Copy link
Copy Markdown
Contributor

What does this PR do?

Three pipelines list names in _callback_tensor_inputs that are never bound in their __call__, so callback_kwargs[k] = locals()[k] raises KeyError as soon as a user requests one:

pipeline name why it isn't a local
LEditsPPPipelineStableDiffusion prompt_embeds never assigned; the matching pop was already commented out
LEditsPPPipelineStableDiffusionXL negative_add_time_ids this pipeline has no negative time ids; pop also commented out
StableDiffusionXLControlNetUnionInpaintPipeline masked_image_latents mask, _ = self.prepare_mask_latents(...) discards it

On main:

pipe(
    ...,
    callback_on_step_end=lambda p, i, t, kw: kw,
    callback_on_step_end_tensor_inputs=["prompt_embeds"],
)
# KeyError: 'prompt_embeds'

Each allowlist now matches the actual denoising-loop locals. Both LEdits++ test classes are plain unittest.TestCase, so the PipelineTesterMixin callback tests never ran against them — this adds test_callback_inputs to each. StableDiffusionXLControlNetUnionInpaintPipeline has no test file, so it is covered by the fix alone; I found it with an AST sweep over src/diffusers/pipelines for the same defect, which now reports zero remaining instances.

Addresses Issue 4 of #13635 (coordination comment: #13635 (comment)). The commented-out check_inputs() call that issue also mentions is left alone.

Tests

DIFFUSERS_TEST_DEVICE=cpu pytest tests/pipelines/ledits_pp -q
# main:        2 failed, 6 passed, 2 skipped
#              E KeyError: 'prompt_embeds'
#              E KeyError: 'negative_add_time_ids'
# this branch: 8 passed, 2 skipped

DIFFUSERS_TEST_DEVICE=cpu pytest tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py \
                                tests/pipelines/controlnet/test_controlnet_sdxl.py -q
# 147 passed, 64 skipped

make quality
# All checks passed! / 2020 files already formatted

python utils/check_copies.py && python utils/check_dummies.py && python utils/check_forward_call_docstrings.py
# clean

utils/check_repo.py and utils/check_inits.py fail the same way on a pristine checkout of main on my machine (missing diffusers.models.auto, and a lookup for src/transformers), so they are unrelated to this diff.

Self-review notes

  • Fixed the allowlists to match real locals rather than inventing bindings for the missing names — per .ai/AGENTS.md, don't guess intent and silently correct behavior.
  • Deleted the two commented-out callback_outputs.pop(...) lines instead of leaving them next to the change, per the no-dead-code rule.
  • The new tests reuse the callback_inputs_all check from PipelineTesterMixin.test_callback_inputs rather than inventing a pattern, per .ai/testing.md.
  • Left deliberately: in both LEdits++ pipelines negative_prompt_embeds resolves to the __call__ argument, which is None unless the caller passes it — the real uncond tensor is uncond_embeddings in the SD pipeline. It doesn't crash, so I'd rather have your call on it than change it here.
  • Left deliberately: SDXL's commented-out check_inputs() call. Re-enabling it changes validation behavior well beyond this bug.

Before submitting

Who can review?

@yiyixuxu @asomoza

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.
@github-actions github-actions Bot added size/M PR with diff < 200 LOC tests pipelines and removed size/M PR with diff < 200 LOC labels Aug 7, 2026
@github-actions

github-actions Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Hi @LuShadowX, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

@github-actions

github-actions Bot commented Aug 22, 2026

Copy link
Copy Markdown
Contributor

Hi @LuShadowX, a follow-up on the reminder above: this PR still does not link an issue it fixes.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice. Adding a closing keyword (e.g. Fixes #1234) to the PR description, or a maintainer adding the no-issue-needed label, will prevent that.

@sayakpaul
sayakpaul requested a review from dg845 August 27, 2026 10:58
@github-actions github-actions Bot added the size/M PR with diff < 200 LOC label Aug 28, 2026

@dg845 dg845 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR!

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"]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I think it might make sense to replace "prompt_embeds" here with "text_embeddings" in a follow-up PR, since my understanding is that text_embeddings plays an analogous role in this pipeline with a non-standard name.

@dg845

dg845 commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator

Merging as the CI failures are unrelated.

@dg845
dg845 merged commit 265b3f4 into huggingface:main Aug 28, 2026
13 of 15 checks passed
@LuShadowX

Copy link
Copy Markdown
Contributor Author

Thanks for reviewing and merging this. Good catch on text_embeddings — that name is the odd one out in the pipeline.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pipelines size/M PR with diff < 200 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants