fix: correct cross-frame attention repeat factor in MemoryEfficientCrossAttention

[Mr-Neutr0n]

Summary

MemoryEfficientCrossAttention has a bug in its cross-frame attention implementation (based on Text2Video-Zero). The repeat factor n used when expanding k and v tensors is incorrect.

The Bug

The n_cp calculation was commented out, and n_times_crossframe_attn_in_self was used as the repeat count instead:

# n_cp = x.shape[0]//n_times_crossframe_attn_in_self   # commented out!
k = repeat(k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_times_crossframe_attn_in_self)
v = repeat(v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_times_crossframe_attn_in_self)

The slicing k[::n_times_crossframe_attn_in_self] selects every N-th frame, producing batch_size / n_times_crossframe_attn_in_self entries. To restore the original batch dimension, the repeat factor must be n_cp = batch_size // n_times_crossframe_attn_in_self, not n_times_crossframe_attn_in_self itself.

Using n_times_crossframe_attn_in_self as the repeat factor only produces the correct result when batch_size == n_times_crossframe_attn_in_self^2, which is not generally the case. For all other batch sizes, the output tensor has the wrong batch dimension, leading to shape mismatches or silently incorrect attention.

The Fix

Uncomment n_cp and use it as the repeat factor, consistent with the existing correct implementation in CrossAttention:

n_cp = x.shape[0] // n_times_crossframe_attn_in_self
k = repeat(k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp)
v = repeat(v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp)

[Mr-Neutr0n]

Closing this one to tidy up my open pull requests.

It's been open around five months with no review activity, and I'd rather withdraw it than leave it sitting in your queue. Nothing needed from you, and no hard feelings at all.

If the fix is still wanted, this can be reopened, or I'm happy to redo it properly against current main. Apologies for the noise.