Update src/diffusers/models/embeddings.py

Co-authored-by: YiYi Xu <[email protected]>

Author: a-r-r-o-w

src/diffusers/models/embeddings.py

@@ -2611,78 +2611,3 @@ 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