Fix swapped indices and names in 2D sincos pos_embed

[Seas0]

What does this PR do?

CURRENTLY STILL A DRAFT, I'M NOT SURE ALL NAMES ARE FIXED!

In get_2d_sincos_pos_embed and get_2d_sincos_pos_embed_from_grid the conventional embedding scheme originated from the mae models seems incorrectly swapped the order of h and w when calling meshgrid, and causing emb_h to actually encode width information, while emb_w to actually encode height information, which, given the nature of permutational equivariance found in Attention mechanism, should not cause any practical differences when training or inferencing, but still confuse people reading the code for a comprehensive understanding of model's internal structure.

This commit should fix the confusion and the swapped names in these functions, while keep being compatibles with pre-trained models using the old code.
Also this commit removed unnecessary stacking and reshaping of the grid tensors.

For reference of the behavior of the meshgrid function, see also: https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html https://pytorch.org/docs/stable/generated/torch.meshgrid.html

Also referenced to the original PR to reimplement sincos pos_embed in Pytorch #10156 written kindly by @hlky

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed.

@yiyixuxu

[Seas0]

Sample code for visualizing 2d positional embeddings

# Visualization of the 2D positional embeddings

import matplotlib.pyplot as plt
import seaborn as sns

from matplotlib_inline.backend_inline import set_matplotlib_formats
set_matplotlib_formats("svg")

from diffusers.models.embeddings import get_2d_sincos_pos_embed

# Create a grid for visualization
h = 16
w = 32
grid_size = (h, w)
embed_dim = 128  # Example embedding dimension
base_size = 16
interpolation_scale = 1.0

# Get positional embeddings
# First half of the hidden dimension is for the width, second half is for the height
pos_embed = get_2d_sincos_pos_embed(
    embed_dim, (h, w), base_size=base_size, interpolation_scale=interpolation_scale
)
pos_embed = pos_embed.reshape(h, w, embed_dim)

# Visualize first few dimensions
num_dims_to_show = 2
assert num_dims_to_show <= embed_dim
assert num_dims_to_show % 2 == 0
fig, axes = plt.subplots(1, 2, figsize=(15, 5))
axes = axes.flatten()

for i in range(num_dims_to_show):
    # dim = embed_dim//2 - 1 - i if i < num_dims_to_show//2 else embed_dim - 1 - i
    dim = i if i < num_dims_to_show // 2 else embed_dim // 2 + i - num_dims_to_show // 2
    sns.heatmap(
        pos_embed[:, :, dim],
        ax=axes[i],
        cmap="viridis",
        square=True,
        cbar=False,
        vmin=-1,
        vmax=1,
    )
    print(pos_embed[:, :, dim])
    axes[i].set_title(f"Dimension {dim}")
    axes[i].set_xlabel("Width")
    axes[i].set_ylabel("Height")

plt.tight_layout()
plt.show()