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