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)