[WIP][cogview4]: implement initial CogView4 pipeline

Implement the basic CogView4 pipeline structure with the following changes:
- Add CogView4 pipeline implementation
- Implement DDIM scheduler for CogView4
- Add CogView3Plus transformer architecture
- Update embedding models

Current limitations:
- CFG implementation uses padding for sequence length alignment
- Need to verify transformer inference alignment with Megatron

TODO:
- Consider separate forward passes for condition/uncondition
  instead of padding approach

Author: OleehyO

src/diffusers/models/embeddings.py

@@ -812,6 +812,7 @@ def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tens
 
         return (hidden_states + pos_embed).to(hidden_states.dtype)
 
+
 class CogView4PatchEmbed(nn.Module):
     def __init__(
         self,
@@ -832,55 +833,35 @@ def __init__(
 
         # Linear projection for text embeddings
         self.text_proj = nn.Linear(text_hidden_size, hidden_size)
-        #TODO：这里需要改成RotaryEmbed
-        pos_embed = get_2d_sincos_pos_embed(
-            hidden_size, pos_embed_max_size, base_size=pos_embed_max_size, output_type="pt"
-        )
-        pos_embed = pos_embed.reshape(pos_embed_max_size, pos_embed_max_size, hidden_size)
-        self.register_buffer("pos_embed", pos_embed.float(), persistent=False)
 
-    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, hidden_states: torch.Tensor, prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor | None
+    ) -> torch.Tensor:
         batch_size, channel, height, width = hidden_states.shape
 
         if height % self.patch_size != 0 or width % self.patch_size != 0:
             raise ValueError("Height and width must be divisible by patch size")
 
-        height = height // self.patch_size
-        width = width // self.patch_size
-        hidden_states = hidden_states.view(batch_size, channel, height, self.patch_size, width, self.patch_size)
-        hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).contiguous()
-        hidden_states = hidden_states.view(batch_size, height * width, channel * self.patch_size * self.patch_size)
+        patch_height = height // self.patch_size
+        patch_width = width // self.patch_size
 
-        # Project the patches
-        hidden_states = self.proj(hidden_states)
-        prompt_encoder_hidden_states = []
-        negative_prompt_encoder_hidden_states = []
-
-        for i in range(0, batch_size, 2):
-            prompt_embeds = encoder_hidden_states[i, :, :]  # [seq_len, hidden_size]
-            negative_embeds = encoder_hidden_states[i + 1, :, :]  # [seq_len, hidden_size]
-            mask = negative_embeds.abs().sum(dim=-1) > 0
-            seq_len_neg = mask.sum().item()  # 非零部分的数量
-            negative_embeds_valid = negative_embeds[:seq_len_neg, :]  # [seq_len_neg, hidden_size]
-            prompt_encoder_hidden_states.append(prompt_embeds)
-            negative_prompt_encoder_hidden_states.append(negative_embeds_valid)
-        prompt_encoder_hidden_states = torch.stack(prompt_encoder_hidden_states, dim=0)
-        negative_prompt_encoder_hidden_states = torch.stack(negative_prompt_encoder_hidden_states, dim=0)
-        prompt_text_length = prompt_encoder_hidden_states.shape[1]
-        negative_prompt_text_length =  negative_prompt_encoder_hidden_states.shape[1]
-        image_pos_embed = self.pos_embed[:height, :width].reshape(height * width, -1)
-        prompt_text_pos_embed = torch.zeros(
-            (prompt_text_length, self.hidden_size), dtype=image_pos_embed.dtype, device=image_pos_embed.device
+        # b, c, h, w -> b, c, patch_height, patch_size, patch_width, patch_size
+        #            -> b, patch_height, patch_width, c, patch_size, patch_size
+        #            -> b, patch_height * patch_width, c * patch_size * patch_size
+        hidden_states = (
+            hidden_states.reshape(batch_size, channel, patch_height, self.patch_size, patch_width, self.patch_size)
+            .permute(0, 2, 4, 1, 3, 5)
+            .reshape(batch_size, patch_height * patch_width, channel * self.patch_size * self.patch_size)
         )
-        negative_prompt_text_pos_embed = torch.zeros(
-            (negative_prompt_text_length, self.hidden_size), dtype=image_pos_embed.dtype, device=image_pos_embed.device
-        )
-        prompt_pos_embed = torch.cat([prompt_text_pos_embed, image_pos_embed], dim=0)[None, ...]
-        negative_prompt_pos_embed = torch.cat([negative_prompt_text_pos_embed, image_pos_embed], dim=0)[None, ...]
-        # TODO: 拼接哼一个完整的 pos_embed 以及拼接 Rope Embed
-        pos_embed = torch.cat([prompt_pos_embed, negative_prompt_pos_embed], dim=0)
-        hidden_states = hidden_states + pos_embed.to(hidden_states.dtype)
-        return hidden_states
+
+        # project
+        hidden_states = self.proj(hidden_states)  # embed_dim: 64 -> 4096
+        prompt_embeds = self.text_proj(prompt_embeds)  # embed_dim: 4096 -> 4096
+        if negative_prompt_embeds is not None:
+            negative_prompt_embeds = self.text_proj(negative_prompt_embeds)  # embed_dim: 4096 -> 4096
+
+        return hidden_states, prompt_embeds, negative_prompt_embeds
+
 
 def get_3d_rotary_pos_embed(
     embed_dim,

src/diffusers/models/transformers/transformer_cogview3plus.py

@@ -84,6 +84,7 @@ def forward(
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         emb: torch.Tensor,
+        **kwargs,
     ) -> torch.Tensor:
         text_seq_length = encoder_hidden_states.size(1)
 
@@ -103,7 +104,7 @@ def forward(
 
         # attention
         attn_hidden_states, attn_encoder_hidden_states = self.attn1(
-            hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states
+            hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, **kwargs
         )
 
         hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_hidden_states
@@ -191,14 +192,15 @@ def __init__(
         # Each of these are sincos embeddings of shape 2 * condition_dim
         self.pooled_projection_dim = 3 * 2 * condition_dim
 
-        # self.patch_embed = CogView3PlusPatchEmbed(
-        #     in_channels=in_channels,
-        #     hidden_size=self.inner_dim,
-        #     patch_size=patch_size,
-        #     text_hidden_size=text_embed_dim,
-        #     pos_embed_max_size=pos_embed_max_size,
-        # )
-        # TODO: 兼容性适配
+        self.max_h = 256
+        self.max_w = 256
+        self.rope = self.prepare_rope(
+            embed_dim=self.config.attention_head_dim,
+            max_h=self.max_h,
+            max_w=self.max_w,
+            rotary_base=10000
+        )
+
         self.patch_embed = CogView4PatchEmbed(
             in_channels=in_channels,
             hidden_size=self.inner_dim,
@@ -300,10 +302,55 @@ def _set_gradient_checkpointing(self, module, value=False):
         if hasattr(module, "gradient_checkpointing"):
             module.gradient_checkpointing = value
 
+    @staticmethod
+    def prepare_rope(embed_dim, max_h, max_w, rotary_base):
+        dim_h = embed_dim // 2
+        dim_w = embed_dim // 2
+        h_inv_freq = 1.0 / (
+            rotary_base ** (torch.arange(0, dim_h, 2, dtype=torch.float32)[: (dim_h // 2)].float() / dim_h)
+        )
+        w_inv_freq = 1.0 / (
+            rotary_base ** (torch.arange(0, dim_w, 2, dtype=torch.float32)[: (dim_w // 2)].float() / dim_w)
+        )
+        h_seq = torch.arange(max_h, dtype=h_inv_freq.dtype)
+        w_seq = torch.arange(max_w, dtype=w_inv_freq.dtype)
+        freqs_h = torch.outer(h_seq, h_inv_freq)
+        freqs_w = torch.outer(w_seq, w_inv_freq)
+        return (freqs_h, freqs_w)
+
+    def get_rope_embedding(self, height, width, target_h, target_w, device):
+        # Get pre-computed frequencies
+        freqs_h, freqs_w = self.rope
+
+        h_idx = torch.arange(height)
+        w_idx = torch.arange(width)
+        inner_h_idx = (h_idx * self.max_h) // target_h
+        inner_w_idx = (w_idx * self.max_w) // target_w
+
+        freqs_h = freqs_h[inner_h_idx].to(device)
+        freqs_w = freqs_w[inner_w_idx].to(device)
+
+        # Create position matrices for height and width
+        # [height, 1, dim//4] and [1, width, dim//4]
+        freqs_h = freqs_h.unsqueeze(1)
+        freqs_w = freqs_w.unsqueeze(0)
+        # Broadcast freqs_h and freqs_w to [height, width, dim//4]
+        freqs_h = freqs_h.expand(height, width, -1)
+        freqs_w = freqs_w.expand(height, width, -1)
+
+        # Concatenate along last dimension to get [height, width, dim//2]
+        freqs = torch.cat([freqs_h, freqs_w], dim=-1)
+
+        freqs = torch.cat([freqs, freqs], dim=-1)  # [height, width, dim]
+        freqs = freqs.reshape(height*width, -1)
+
+        return freqs.cos(), freqs.sin()
+
     def forward(
         self,
         hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
+        prompt_embeds: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
         timestep: torch.LongTensor,
         original_size: torch.Tensor,
         target_size: torch.Tensor,
@@ -338,16 +385,27 @@ def forward(
             `torch.Tensor` or [`~models.transformer_2d.Transformer2DModelOutput`]:
                 The denoised latents using provided inputs as conditioning.
         """
-        height, width = hidden_states.shape[-2:]
-        text_seq_length = encoder_hidden_states.shape[1]
+        batch_size, channel, height, width = hidden_states.shape
+        patch_height, patch_width = height // self.config.patch_size, width // self.config.patch_size
+        do_cfg = negative_prompt_embeds is not None
 
-        hidden_states = self.patch_embed(
-            hidden_states, encoder_hidden_states
-        )  # takes care of adding positional embeddings too.
-        emb = self.time_condition_embed(timestep, original_size, target_size, crop_coords, hidden_states.dtype)
+        if do_cfg:
+            assert batch_size == prompt_embeds.shape[0] + negative_prompt_embeds.shape[0], "batch size mismatch in CFG mode"
+        else:
+            assert batch_size == prompt_embeds.shape[0], "batch size mismatch in non-CFG mode"
+
+        hidden_states, prompt_embeds, negative_prompt_embeds = self.patch_embed(
+            hidden_states, prompt_embeds, negative_prompt_embeds
+        )
 
-        encoder_hidden_states = hidden_states[:, :text_seq_length]
-        hidden_states = hidden_states[:, text_seq_length:]
+        encoder_hidden_states = torch.cat([prompt_embeds, negative_prompt_embeds], dim=0)
+
+        # prepare image_rotary__emb
+        image_rotary_emb = self.get_rope_embedding(
+            patch_height, patch_width, target_h=patch_height, target_w=patch_width, device=hidden_states.device
+        )
+
+        emb = self.time_condition_embed(timestep, original_size, target_size, crop_coords, hidden_states.dtype)
 
         for index_block, block in enumerate(self.transformer_blocks):
             if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -363,17 +421,19 @@ def custom_forward(*inputs):
                     create_custom_forward(block),
                     hidden_states,
                     encoder_hidden_states,
-                    emb,
+                    emb=emb,
+                    image_rotary_emb=image_rotary_emb,
                     **ckpt_kwargs,
                 )
             else:
                 hidden_states, encoder_hidden_states = block(
                     hidden_states=hidden_states,
                     encoder_hidden_states=encoder_hidden_states,
                     emb=emb,
+                    image_rotary_emb=image_rotary_emb,
                 )
 
-        hidden_states = self.norm_out(hidden_states, emb)
+        hidden_states = self.norm_out(hidden_states, emb)  # 结果对应于megatron里的final_layer_input
         hidden_states = self.proj_out(hidden_states)  # (batch_size, height*width, patch_size*patch_size*out_channels)
 
         # unpatchify