(x, heads)
| 4 | from diffusers.models.modeling_utils import ModelMixin |
| 5 | |
| 6 | def reshape_tensor(x, heads): |
| 7 | bs, length, width = x.shape |
| 8 | x = x.view(bs, length, heads, -1) |
| 9 | x = x.transpose(1, 2) |
| 10 | x = x.reshape(bs, heads, length, -1) |
| 11 | return x |
| 12 | |
| 13 | class PerceiverAttention(nn.Module): |
| 14 | def __init__(self, *, dim, dim_head=64, heads=8): |