Skip to content

SDXL: pooled prompt embeds need batch dim for batched writer#2846

Merged
bghira merged 2 commits into
mainfrom
bugfix/sdxl-pooled-embeds-batch-dim
Jul 17, 2026
Merged

SDXL: pooled prompt embeds need batch dim for batched writer#2846
bghira merged 2 commits into
mainfrom
bugfix/sdxl-pooled-embeds-batch-dim

Conversation

@bghira

@bghira bghira commented Jul 17, 2026

Copy link
Copy Markdown
Owner

This pull request updates the _slice_batch_output_for_cache function in simpletuner/helpers/caching/text_embeds.py to better handle tensor shapes and ensure compatibility with batch processing. The main changes focus on adjusting tensor dimensions when dealing with single-item batches and specific output keys.

Tensor shape handling improvements:

  • Added logic to unsqueeze tensors with mismatched batch dimensions when batch_size == 1, ensuring consistent tensor shapes for downstream processing.
  • For the keys pooled_prompt_embeds and negative_pooled_prompt_embeds, added a squeeze operation to remove unnecessary batch dimensions when the tensor is 2D, improving compatibility with expected output formats.

This comment was marked as resolved.

@bghira
bghira merged commit e66ac5d into main Jul 17, 2026
2 checks passed
@bghira
bghira deleted the bugfix/sdxl-pooled-embeds-batch-dim branch July 17, 2026 15:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants