feat(embeddings): add CogView 2D rotary positional embedding

Author: OleehyO

src/diffusers/models/embeddings.py

@@ -2611,3 +2611,78 @@ def forward(self, image_embeds: List[torch.Tensor]):
             projected_image_embeds.append(image_embed)
 
         return projected_image_embeds
+
+
+class CogViewRotary2DEmbedding(nn.Module):
+    def __init__(
+        self,
+        kv_channels: int,
+        rotary_percent: float,
+        max_h: int = 128,
+        max_w: int = 128,
+        rotary_interleaved: bool = False,
+        seq_len_interpolation_factor: float = None,
+        inner_interp: bool = False,
+        rotary_base: int = 10000,
+    ) -> None:
+        super().__init__()
+
+        dim = kv_channels
+        if rotary_percent < 1.0:
+            dim = int(dim * rotary_percent)
+        self.rotary_interleaved = rotary_interleaved
+
+        self.seq_len_interpolation_factor = seq_len_interpolation_factor
+        self.inner_interp = inner_interp
+
+        dim_h = kv_channels // 2
+        dim_w = kv_channels // 2
+
+        device = torch.cuda.current_device()
+        h_inv_freq = 1.0 / (
+            rotary_base
+            ** (torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[: (dim_h // 2)].float() / dim_h)
+        )
+        w_inv_freq = 1.0 / (
+            rotary_base
+            ** (torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[: (dim_w // 2)].float() / dim_w)
+        )
+
+        h_seq = torch.arange(max_h, device=device, dtype=h_inv_freq.dtype)
+        w_seq = torch.arange(max_w, device=device, dtype=w_inv_freq.dtype)
+
+        self.freqs_h = torch.outer(h_seq, h_inv_freq)
+        self.freqs_w = torch.outer(w_seq, w_inv_freq)
+        self.max_h = max_h
+        self.max_w = max_w
+
+    def forward(
+        self,
+        h_idx: torch.Tensor,
+        w_idx: torch.Tensor,
+        target_h: torch.Tensor = None,
+        target_w: torch.Tensor = None,
+        mask: torch.Tensor = None,
+    ) -> torch.Tensor:
+        if self.inner_interp:
+            inner_h_idx = (h_idx * self.max_h) // target_h
+            inner_w_idx = (w_idx * self.max_w) // target_w
+
+            h_emb = self.freqs_h[inner_h_idx]
+            w_emb = self.freqs_w[inner_w_idx]
+
+        else:
+            h_emb = self.freqs_h[h_idx]
+            w_emb = self.freqs_w[w_idx]
+
+        mask = (mask == 1).unsqueeze(-1)
+
+        emb = torch.cat([h_emb, w_emb], dim=-1) * mask
+
+        assert emb.ndim == 2, f"expected emb to have 2 dimensions, got {emb.ndim}"
+        if not self.rotary_interleaved:
+            emb = torch.repeat_interleave(emb, 2, dim=0)
+        else:
+            emb = torch.repeat_interleave(emb, 2, dim=1)
+
+        return emb