Fix PriorTransformer Group Offloading Bug

[dg845]

What does this PR do?

This PR fixes a group offloading bug in pipelines that use a PriorTransformer component: because these pipelines call PriorTransformer.post_process_latents after PriorTransformer.forward, and post_process_latents uses the clip_mean and clip_std parameters directly, if group offloading is active, these parameters will be offloaded (since the offload hook is on forward and has fired), which causes a device mismatch error. This affects the following pipelines, whose group offloading tests should now pass:

• Stable Unclip (group offloading tests fail on main on GPU)
  • pytest tests/pipelines/stable_unclip/test_stable_unclip.py -k "group_offloading"
• Kandinsky prior pipelines (tests are currently skipped on main)
  • pytest tests/pipelines/kandinsky/test_kandinsky_prior.py -k "group_offloading"
  • pytest tests/pipelines/kandinsky2_2/test_kandinsky_prior.py -k "group_offloading"
  • pytest tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py -k "group_offloading"

See #14635 (comment) for more info.

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?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul
@DN6

[dg845]

Self-Review Report

Self-review

Diff reviewed: 5359ec41e, 4 files, +4/−120. Rubric: .ai/references/review-rules.md, plus code_style.md, models.md, testing.md, pitfalls.md.

Blocking issues

None.

Non-blocking issues

1. The src/ gap is now documented only for the two parameters that were fixed

src/diffusers/models/transformers/prior_transformer.py:321-322

The deleted xfail reason strings were the only place recording that PriorTransformer keeps model-level parameters group offloading onloads solely for the duration of forward. Deleting them is correct — testing.md prescribes exactly that ("prefer strict=True so the marker reports XPASS — and gets deleted — the day the pipeline is fixed") — and the new comment carries the reason forward for clip_mean / clip_std. What's gone is the general statement of the hazard. Low impact for this file, since the parameters are now safe; see the doc proposal below.

2. Same bug class, same pipeline, left unfixed

src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py:40-54

StableUnCLIPImageNormalizer.scale / unscale read root-level mean / std outside forward — the identical pattern this PR fixes. It's masked by pipeline_stable_unclip.py:625 calling self.image_normalizer.to(image_embeds.device), and that override swaps in new nn.Parameter objects, so after one run the module's params are stranded on the accelerator while the offload group manages orphaned CPU tensors (verified: differing id()s, group still holding cpu). The same non-standard signature also fails test_layerwise_casting_inference in both Stable unCLIP test files (TypeError: ... got an unexpected keyword argument 'dtype'), which fails identically on main. Out of scope for a group-offloading fix.

Dead code (advisory)

path:line	Verdict	Reason
MemoryTesterMixin._USE_STREAM (tests/pipelines/testing_utils/memory.py:355)	Used	Still consumed by motif_video, audioldm2 tests and the mixin's own tests
MemoryTesterMixin imports, 3 Kandinsky files	Used	Still the base class of each *PipelineMemory class
*_GROUP_OFFLOAD_XFAIL_REASON	Removed	0 remaining references repo-wide
pytest import, kandinsky/test_kandinsky_prior.py + kandinsky2_2/..._emb2emb.py	Removed	Became unused with the overrides; ruff --fix dropped both. The 2.2 non-emb2emb file still uses pytest elsewhere and keeps its import
post_process_latents	Used	5 in-tree pipelines + 2 community examples

Suggestions / additional info

Agent-doc proposal (per review-rules.md → "Agent docs"). This gotcha isn't written down and had to be reasoned out from group_offloading.py:785: a model-level nn.Parameter (one not owned by a submodule) is gathered into the "unmatched group" whose hook wraps the top-level forward, so any public non-forward method that reads it sees the offload device. testing.md documents the neighbouring leaf-level MultiheadAttention hazard but not this one. Worth a gotcha in models.md — it generalizes past priors, since VAE encode / decode are also non-forward entry points.

Checked and clean — so the reviewer knows these were considered:

• No dtype counterpart. Checked whether the fix should also pin dtype. It should not: _apply_layerwise_casting (hooks/layerwise_casting.py:184-189) only hooks _GO_LC_SUPPORTED_PYTORCH_LAYERS instances found via named_children(), so root parameters are never converted to storage dtype (verified — embedding_proj.weight → float8_e4m3fn, clip_std → float32). Adding .to(dtype=...) would be the defensive code code_style.md forbids.
• models.md gotcha (4) compliance. The fix derives device from the input tensor, matching how the same file already derives dtype (prior_transformer.py:244,266: .to(hidden_states.dtype)).
• Fix is complete within the model. positional_embedding / prd_embedding are read only between lines 244-308, inside forward, so they are correctly onloaded; clip_mean / clip_std are read only in post_process_latents.
• Comment is not ephemeral context. It states the mechanism, names no PR or reviewer, and stands alone.
• No # Copied from link on post_process_latents; utils/check_copies.py clean.
• No usage-doc impact — no public API, argument, or default changed.
• Test coverage — the group offloading tests that were failing are the guard; no new test needed. Note pipeline_unclip.py:354 also calls this method and has no test file at all, a pre-existing gap this PR doesn't widen.

Verification

# all four affected pipelines, offload tests
55 passed, 5 skipped          # no XPASS

# three Kandinsky files in full + model-level prior tests
131 passed, 26 skipped

ruff check / ruff format --check    clean (4 files)
utils/check_copies.py              clean

The two test_layerwise_casting_inference failures in tests/pipelines/stable_unclip/ are pre-existing (issue 2 above) and confirmed identical on main.

Summary

A one-line device fix in shared model code, plus removal of the strict xfail scaffolding it obsoletes in three Kandinsky prior test files — which is mandatory, not optional, since strict XPASS reports as a CI failure. The change follows the file's existing input-derived-cast idiom and the testing.md lifecycle for xfail markers. Both non-blocking items are context and scope questions, not defects in the diff.

Verdict: READY

• Fix before submitting — nothing.
• Leave for the actual review — issue 2 (StableUnCLIPImageNormalizer, worth its own PR: the layerwise-casting failure is user-visible and pre-existing) and the models.md gotcha proposal, which a maintainer should weigh in on before it's written.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

neato!

[dg845]

/diffusers-bot pytest tests/pipelines/stable_unclip/test_stable_unclip.py tests/pipelines/kandinsky/test_kandinsky_prior.py tests/pipelines/kandinsky2_2/test_kandinsky_prior.py tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py -k "group_offloading"

[github-actions]

✅ pytest tests/pipelines/stable_unclip/test_stable_unclip.py tests/pipelines/kandinsky/test_kandinsky_prior.py tests/pipelines/kandinsky2_2/test_kandinsky_prior.py tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py -k "group_offloading" passed on GPU — view logs.

[dg845]

Merging as the GPU test run from #14695 (comment) is green and the CI failures are unrelated.