draft patch(not work)

Author: zRzRzRzRzRzRzR

src/diffusers/models/embeddings.py

@@ -812,6 +812,75 @@ 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,
+        in_channels: int = 16,
+        hidden_size: int = 2560,
+        patch_size: int = 2,
+        text_hidden_size: int = 4096,
+        pos_embed_max_size: int = 128,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+        self.patch_size = patch_size
+        self.text_hidden_size = text_hidden_size
+        self.pos_embed_max_size = pos_embed_max_size
+        # Linear projection for image patches
+        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
+
+        # 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:
+        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)
+
+        # 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
+        )
+        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
 
 def get_3d_rotary_pos_embed(
     embed_dim,

src/diffusers/models/transformers/transformer_cogview3plus.py

@@ -28,7 +28,7 @@
 from ...models.modeling_utils import ModelMixin
 from ...models.normalization import AdaLayerNormContinuous
 from ...utils import is_torch_version, logging
-from ..embeddings import CogView3CombinedTimestepSizeEmbeddings, CogView3PlusPatchEmbed
+from ..embeddings import CogView3CombinedTimestepSizeEmbeddings, CogView3PlusPatchEmbed, CogView4PatchEmbed
 from ..modeling_outputs import Transformer2DModelOutput
 from ..normalization import CogView3PlusAdaLayerNormZeroTextImage
 
@@ -166,7 +166,7 @@ class CogView3PlusTransformer2DModel(ModelMixin, ConfigMixin):
     """
 
     _supports_gradient_checkpointing = True
-    _no_split_modules = ["CogView3PlusTransformerBlock", "CogView3PlusPatchEmbed"]
+    _no_split_modules = ["CogView3PlusTransformerBlock", "CogView3PlusPatchEmbed", "CogView4PlusPatchEmbed"]
 
     @register_to_config
     def __init__(
@@ -191,7 +191,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(
+        # 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.patch_embed = CogView4PatchEmbed(
             in_channels=in_channels,
             hidden_size=self.inner_dim,
             patch_size=patch_size,

src/diffusers/pipelines/cogview4/pipeline_cogview4.py

@@ -311,6 +311,24 @@ def encode_prompt(
                 device=device,
                 dtype=dtype,
             )
+
+        #TODO: 先pad 0 ，后续再处理不同长度的问题
+        seq_len_prompt = prompt_embeds.shape[1]
+        seq_len_neg = negative_prompt_embeds.shape[1]
+        if seq_len_neg < seq_len_prompt:
+                # 创建一个新的张量，大小为 [batch_size, seq_len_prompt, hidden_size]
+                batch_size = negative_prompt_embeds.shape[0]
+                hidden_size = negative_prompt_embeds.shape[2]
+                # 填充后的张量
+                padded_negative_prompt_embeds = torch.zeros(
+                    batch_size,
+                    seq_len_prompt,
+                    hidden_size,
+                    dtype=negative_prompt_embeds.dtype,
+                    device=negative_prompt_embeds.device
+                )
+                padded_negative_prompt_embeds[:, :seq_len_neg, :] = negative_prompt_embeds
+                negative_prompt_embeds = padded_negative_prompt_embeds
         return prompt_embeds, negative_prompt_embeds
 
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
@@ -582,7 +600,7 @@ def __call__(
             device=device,
         )
         if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=1)
+            prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds], dim=0)
 
         # 4. Prepare timesteps
         timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
@@ -594,7 +612,6 @@ def __call__(
         sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])  # Append zero at the end
 
         self.sigmas = time_shift(mu, 1.0, sigmas).to(torch.long).to("cpu")  # This is for noisy control of cogview4
-
         self._num_timesteps = len(timesteps)
 
         # 5. Prepare latents.
@@ -635,11 +652,8 @@ def __call__(
             for i, t in enumerate(timesteps):
                 if self.interrupt:
                     continue
-
-                # latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
-                latent_model_input = latents  # For CogView4 concat the text embed and only use prompt
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
                 latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
                 # Use sigma instead of timestep directly
                 sigma = self.sigmas[i]  # Get the corresponding sigma value
                 timestep = sigma.expand(latent_model_input.shape[0]).to(device)  # Use sigma to scale the timestep