[davrens101]

Fixes #15161.

What was wrong

DreamLiteUNetModel relies on DreamLiteAttnProcessor2_0, which repeats K/V heads up to the query head count for grouped-query attention. UNet2DConditionModel.set_attention_slice(None) walks every Attention module and calls Attention.set_attention_slice(None), which installs the stock AttnProcessor2_0 — so pipeline.disable_attention_slicing() breaks the next forward pass with a shape error even if slicing was never enabled:

RuntimeError: shape '[1, -1, 8, 0]' is invalid for input of size 2048

(query.view(batch_size, -1, attn.heads, head_dim) with head_dim = query.shape[-1] // attn.heads == 0.)

Sliced attention itself is not usable here either: SlicedAttnProcessor has no GQA path, so enable_attention_slicing() silently produced broken inference.

Fix

Override DreamLiteUNetModel.set_attention_slice:

• None (the path used by disable_attention_slicing()) restores DreamLiteAttnProcessor2_0 via the existing set_default_attn_processor() override, without touching anything else.
• enabling slicing raises a ValueError explaining that DreamLite's GQA layout is not supported by SlicedAttnProcessor, instead of failing later inside attention.

Output before/after set_attention_slice(None) on a tiny GQA (num_kv_heads=1) UNet is now bit-identical.

Tests

Added TestDreamLiteAttentionSlicing in tests/pipelines/dreamlite/test_pipeline_dreamlite.py: disable keeps the output identical, enable raises.

python -m pytest tests/pipelines/dreamlite/test_pipeline_dreamlite.py -k AttentionSlicing → 2 passed
ruff check / ruff format --check on both files → clean

AI assistance was used for this change; I reviewed and ran the tests above.

Coordination status

The issue author has already investigated this report and is awaiting maintainer feedback. Linking the existing report here does not claim issue assignment or maintainer acknowledgment.

[github-actions]

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