Fix prior preservation shape mismatch in TimestepEmbedding and PixArtAlphaTextProjection

[Liauuu]

Summary

This PR fixes a RuntimeError: mat1 and mat2 shapes cannot be multiplied that occurs during FLUX DreamBooth training when --with_prior_preservation is enabled (reported in #12494).

The failure happens inside PixArtAlphaTextProjection.linear_1 (via CombinedTimestepGuidanceTextProjEmbeddings / CombinedTimestepTextProjEmbeddings in the FLUX transformer), with shapes such as (2×1536) passed to a layer expecting (N×768).

Root cause

Prior preservation training concatenates instance and class samples so the transformer can process both in a single forward pass. In affected code paths, pooled text embeddings are sometimes concatenated along the feature dimension (dim=-1) instead of the batch dimension (dim=0).

For example:

• Expected: [batch_size * 2, 768] (instance + class stacked on batch)
• Observed: [batch_size, 1536] (instance + class concatenated horizontally)

Because PixArtAlphaTextProjection and TimestepEmbedding define linear_1 with in_features=768 (or the configured channel size), a last dimension of 1536 (768 × 2) triggers the matmul failure:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (2×1536 and 768×3072)

This matches the stack trace in #12494 from train_dreambooth_lora_flux_advanced.py with --with_prior_preservation.

Solution

Introduce a small shared helper, _unstack_doubled_features, in src/diffusers/models/embeddings.py:

def _unstack_doubled_features(tensor, expected_features):
    if tensor.shape[-1] == expected_features * 2:
        first, second = tensor.chunk(2, dim=-1)
        return torch.cat([first, second], dim=0)
    return tensor

Apply it at the start of:

• TimestepEmbedding.forward — on sample, and on condition when cond_proj is used
• PixArtAlphaTextProjection.forward — on caption (pooled projections)

When the last dimension is exactly 2 × in_features, the helper splits on dim=-1 and re-stacks on dim=0, converting [B, 2F] → [2B, F] before the linear layers run. Downstream modules then receive batch-doubled embeddings as intended for prior preservation.

When inputs already have the correct shape ([B, F]), the helper is a no-op, so normal inference and training without prior preservation are unchanged.

Changes

• src/diffusers/models/embeddings.py
  • Add _unstack_doubled_features
  • Call it from TimestepEmbedding.forward and PixArtAlphaTextProjection.forward

Test plan

• Reproduce i think there is something wrong with new/latest scripts. RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x1536 and 768x3072) #12494: run train_dreambooth_lora_flux_advanced.py (or train_dreambooth_lora_flux.py) with --with_prior_preservation, --class_data_dir, and --class_prompt; confirm training starts without the (2×1536) matmul error
• Run the same script without --with_prior_preservation; confirm behavior is unchanged
• Run a short FLUX inference pass to confirm no regression in standard (non-training) usage

Fixes #12494