Run cleanup and honor return_dict for SDXL latent output

[Devam0311]

Issue link: this addresses Issue 3 of #13610.
I have deliberately not used a Fixes keyword: #13610 tracks 7 separate findings, so a closing keyword would close it as soon as this one merges while the others are still open. Happy for the no-issue-needed label to be applied, or I can add a closing keyword if you would rather the review issue be closed and reopened.

What does this PR do?

Fixes Issue 3 of #13610 (stable_diffusion_xl model/pipeline review).

For output_type="latent", the SDXL inpaint and instruct-pix2pix pipelines return early:

else:
    return StableDiffusionXLPipelineOutput(images=latents)

That return jumps over everything after it, so two things break:

• maybe_free_model_hooks() never runs — model offload cleanup is skipped entirely for latent output
• return_dict=False is ignored — the caller gets a StableDiffusionXLPipelineOutput where a tuple was requested

Reproduced on main:

StableDiffusionXLPipelineOutput {'cleanup': False}

Solution

Assign image = latents rather than returning, and guard watermarking, postprocessing, and the padding-mask overlay behind output_type != "latent" — so the shared cleanup and return_dict handling at the end of __call__ runs for every output type. This is the shape the text-to-image pipeline already uses.

The issue names two files. I swept for the same early return and found three more, all fixed together:

named in the issue	pipeline_stable_diffusion_xl_inpaint.py, pipeline_stable_diffusion_xl_instruct_pix2pix.py
found by sweep	pipeline_controlnet_inpaint_sd_xl.py, pipeline_controlnet_union_inpaint_sd_xl.py, pipeline_pag_sd_xl_inpaint.py

The inpaint pipelines also apply padding_mask_crop overlays, which operate on decoded PIL images, so that call moved inside the same guard.

Testing

test_latent_output_runs_cleanup_and_honors_return_dict stubs maybe_free_model_hooks to record whether it fired, requests output_type="latent" with return_dict=False, and asserts cleanup ran, a tuple came back, and the payload is still a latent tensor (by channel count — postprocess would have collapsed it to 3).

Verified it catches the bug: with the fix reverted it fails.

pytest tests/pipelines/stable_diffusion_xl/ -k "not ip_adapter"   -> 195 passed, 129 skipped
ruff check / ruff format --check                                  -> clean
git diff --check                                                  -> clean

The TestStableDiffusionXL*PipelineIPAdapter classes are deselected: they fail on a clean checkout of main in my environment with a ModuleNotFoundError from an optional dependency I do not have installed, unrelated to this change.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Was this discussed/approved via a GitHub issue? — stable_diffusion_xl model/pipeline review #13610
• Did you write any new necessary tests?

Self-review notes

Reviewed against .ai/references/review-rules.md and pipelines.md. No blocking issues.

For the reviewer:

1. return_dict=False with output_type="latent" now returns a tuple instead of an output object. That is the fix, but it is a return-type change for anyone who was relying on the old behaviour.
2. Three files beyond the issue's list, found by sweeping for the identical early return. Happy to split them out if you'd rather review them separately.
3. This overlaps in spirit with Do not watermark or postprocess SDXL latent output #14678 (Issue 4, watermarking latents in img2img): both add the output_type != "latent" guard, but to disjoint files, so the two do not conflict.
4. The padding_mask_crop overlay moving inside the guard is a deliberate part of the fix, not incidental reindentation — it operates on decoded images and would fail on latents.

Who can review?

@yiyixuxu @asomoza @hlky

[github-actions]

Hi @Devam0311, 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.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[adi-IL]

Tested locally on Linux x86_64.

Bypassing maybe_free_model_hooks() on output_type='latent' leaked device allocations during CPU/model offloading, and returning StableDiffusionXLPipelineOutput early broke return_dict=False. Setting image = latents and guarding watermarking/postprocessing lets execution reach hook cleanup and return_dict handling across the affected SDXL pipelines. The regression test in test_stable_diffusion_xl_inpaint.py verifies both behaviors cleanly.