draft of cogview3plus

Author: zRzRzRzRzRzRzR

scripts/convert_cogview3_to_diffusers.py

@@ -0,0 +1,25 @@
+import argparse
+from typing import Any, Dict
+
+import torch
+from transformers import T5EncoderModel, T5Tokenizer
+
+from diffusers import (
+    CogView3PlusTransformer2DModel,
+    CogView3PlusPipeline,
+)
+
+TRANSFORMER_KEYS_RENAME_DICT = {
+    "transformer": "transformer_blocks",
+    "attention": "attn1",
+    "mlp": "ff.net",
+    "dense_h_to_4h": "0.proj",
+    "dense_4h_to_h": "2",
+    ".layers": "",
+    "dense": "to_out.0",
+    "patch_embed": "norm1.norm",
+    "post_attn1_layernorm": "norm2.norm",
+    "mixins.patch_embed": "patch_embed",
+    "mixins.final_layer.adaln": "norm_out",
+    "mixins.final_layer.linear": "proj_out",
+}

src/diffusers/__init__.py

@@ -84,6 +84,7 @@
             "AutoencoderOobleck",
             "AutoencoderTiny",
             "CogVideoXTransformer3DModel",
+            "CogView3PlusTransformer2DModel",
             "ConsistencyDecoderVAE",
             "ControlNetModel",
             "ControlNetXSAdapter",
@@ -258,6 +259,7 @@
             "CogVideoXImageToVideoPipeline",
             "CogVideoXPipeline",
             "CogVideoXVideoToVideoPipeline",
+            "CogView3PlusPipeline",
             "CycleDiffusionPipeline",
             "FluxControlNetImg2ImgPipeline",
             "FluxControlNetInpaintPipeline",
@@ -558,6 +560,7 @@
             AutoencoderOobleck,
             AutoencoderTiny,
             CogVideoXTransformer3DModel,
+            CogView3PlusTransformer2DModel,
             ConsistencyDecoderVAE,
             ControlNetModel,
             ControlNetXSAdapter,
@@ -710,6 +713,7 @@
             CogVideoXImageToVideoPipeline,
             CogVideoXPipeline,
             CogVideoXVideoToVideoPipeline,
+            CogView3PlusPipeline,
             CycleDiffusionPipeline,
             FluxControlNetImg2ImgPipeline,
             FluxControlNetInpaintPipeline,

src/diffusers/models/__init__.py

@@ -44,6 +44,7 @@
     _import_structure["modeling_utils"] = ["ModelMixin"]
     _import_structure["transformers.auraflow_transformer_2d"] = ["AuraFlowTransformer2DModel"]
     _import_structure["transformers.cogvideox_transformer_3d"] = ["CogVideoXTransformer3DModel"]
+    _import_structure["transformers.transformer_cogview3dplus"] = ["CogView3PlusTransformer2DModel"]
     _import_structure["transformers.dit_transformer_2d"] = ["DiTTransformer2DModel"]
     _import_structure["transformers.dual_transformer_2d"] = ["DualTransformer2DModel"]
     _import_structure["transformers.hunyuan_transformer_2d"] = ["HunyuanDiT2DModel"]
@@ -98,6 +99,7 @@
         from .transformers import (
             AuraFlowTransformer2DModel,
             CogVideoXTransformer3DModel,
+            CogView3PlusTransformer2DModel,
             DiTTransformer2DModel,
             DualTransformer2DModel,
             FluxTransformer2DModel,

src/diffusers/models/embeddings.py

@@ -714,6 +714,114 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor:
         return freqs_cos, freqs_sin
 
 
+class CogView3PlusPosEmbed(nn.Module):
+    def __init__(
+        self,
+        max_height: int = 128,
+        max_width: int = 128,
+        hidden_size: int = 2560,
+        text_length: int = 0,
+        block_size: int = 16,
+    ):
+        super().__init__()
+        self.max_height = max_height
+        self.max_width = max_width
+        self.hidden_size = hidden_size
+        self.text_length = text_length
+        self.block_size = block_size
+
+        # Initialize the positional embedding as a non-trainable parameter
+        self.image_pos_embedding = nn.Parameter(
+            torch.zeros(self.max_height, self.max_width, hidden_size), requires_grad=False
+        )
+        # Reinitialize the positional embedding using a sin-cos function
+        self.reinit()
+
+    def forward(self, target_size: List[int]) -> torch.Tensor:
+        ret = []
+        for h, w in target_size:
+            # Scale height and width according to the block size
+            h, w = h // self.block_size, w // self.block_size
+
+            # Reshape the image positional embedding for the target size
+            image_pos_embed = self.image_pos_embedding[:h, :w].reshape(h * w, -1)
+
+            # Combine the text positional embedding and image positional embedding
+            pos_embed = torch.cat(
+                [
+                    torch.zeros(
+                        (self.text_length, self.hidden_size),
+                        dtype=image_pos_embed.dtype,
+                        device=image_pos_embed.device,
+                    ),
+                    image_pos_embed,
+                ],
+                dim=0,
+            )
+
+            ret.append(pos_embed[None, ...])  # Add a batch dimension
+
+        return torch.cat(ret, dim=0)  # Concatenate along the batch dimension
+
+    def reinit(self):
+        # Initialize the positional embedding using a 2D sin-cos function
+        pos_embed_np = self.get_2d_sincos_pos_embed(self.hidden_size, self.max_height, self.max_width)
+        self.image_pos_embedding.data.copy_(torch.from_numpy(pos_embed_np).float())
+
+
+class CogView3PlusImagePatchEmbedding(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 128,
+        hidden_size: int = 128,
+        patch_size: int = 2,
+        text_hidden_size: int = 4096,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+        self.patch_size = patch_size
+        self.text_hidden_size = text_hidden_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)
+
+    def forward(self, images: torch.Tensor, encoder_outputs: torch.Tensor = None) -> torch.Tensor:
+        # Rearrange the images
+        # patches_images = rearrange(images, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=self.patch_size, p2=self.patch_size)
+
+        b, c, h, w = images.shape
+        p1, p2 = self.patch_size, self.patch_size
+        assert h % p1 == 0 and w % p2 == 0, "Height and width must be divisible by patch size"
+
+        images = images.view(b, c, h // p1, p1, w // p2, p2)
+        patches_images = images.permute(0, 2, 4, 1, 3, 5).contiguous()
+        patches_images = patches_images.view(b, (h // p1) * (w // p2), c * p1 * p2)
+
+        # Project the patches
+        image_emb = self.proj(patches_images)
+
+        # If text embeddings are provided, project and concatenate them
+        if self.text_hidden_size is not None and encoder_outputs is not None:
+            text_emb = self.text_proj(encoder_outputs)
+            emb = torch.cat([text_emb, image_emb], dim=1)
+        else:
+            emb = image_emb
+
+        return emb
+
+    def reinit(self, parent_model=None):
+        # Reinitialize the projection weights
+        nn.init.xavier_uniform_(self.proj.weight)
+        nn.init.constant_(self.proj.bias, 0)
+        if self.text_hidden_size is not None:
+            nn.init.xavier_uniform_(self.text_proj.weight)
+            nn.init.constant_(self.text_proj.bias, 0)
+
+
 class TimestepEmbedding(nn.Module):
     def __init__(
         self,

src/diffusers/models/transformers/__init__.py

@@ -16,4 +16,5 @@
     from .transformer_2d import Transformer2DModel
     from .transformer_flux import FluxTransformer2DModel
     from .transformer_sd3 import SD3Transformer2DModel
+    from .transformer_cogview3plus import CogView3PlusTransformer2DModel
     from .transformer_temporal import TransformerTemporalModel