use with cogview4 transformers forward twice of u and uc

Author: zRzRzRzRzRzRzR

docs/source/en/_toctree.yml

@@ -276,6 +276,8 @@
         title: ConsisIDTransformer3DModel
       - local: api/models/cogview3plus_transformer2d
         title: CogView3PlusTransformer2DModel
+      - local: api/models/cogview4_transformer2d
+        title: CogView4Transformer2DModel
       - local: api/models/dit_transformer2d
         title: DiTTransformer2DModel
       - local: api/models/flux_transformer

docs/source/en/api/models/cogview4_transformer2d.md

@@ -0,0 +1,30 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License. -->
+
+# CogView4Transformer2DModel
+
+A Diffusion Transformer model for 2D data from [CogView4]()
+
+The model can be loaded with the following code snippet.
+
+```python
+from diffusers import CogView3PlusTransformer2DModel
+
+transformer = CogView3PlusTransformer2DModel.from_pretrained("THUDM/CogView4-6B", subfolder="transformer", torch_dtype=torch.bfloat16).to("cuda")
+```
+
+## CogView4Transformer2DModel
+
+[[autodoc]] CogView4Transformer2DModel
+
+## Transformer2DModelOutput
+
+[[autodoc]] models.modeling_outputs.Transformer2DModelOutput

scripts/convert_cogview4_to_diffusers.py

@@ -31,7 +31,7 @@
 from accelerate import init_empty_weights
 from transformers import PreTrainedTokenizerFast, GlmForCausalLM
 
-from diffusers import AutoencoderKL, CogView4DDIMScheduler, CogView4Pipeline, CogView3PlusTransformer2DModel
+from diffusers import AutoencoderKL, CogView4DDIMScheduler, CogView4Pipeline, CogView4Transformer2DModel
 from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint
 from diffusers.utils.import_utils import is_accelerate_available
 
@@ -168,7 +168,7 @@ def main(args):
         converted_transformer_state_dict = convert_cogview4_transformer_checkpoint_to_diffusers(
             args.transformer_checkpoint_path
         )
-        transformer = CogView3PlusTransformer2DModel(
+        transformer = CogView4Transformer2DModel(
             patch_size=2,
             in_channels=16,
             num_layers=28,
@@ -209,23 +209,27 @@ def main(args):
         if dtype is not None:
             vae = vae.to(dtype=dtype)
 
-    # text_encoder_id = "THUDM/glm-4-9b-hf"
-    # tokenizer = PreTrainedTokenizerFast.from_pretrained(text_encoder_id)
-    # text_encoder = GlmForCausalLM.from_pretrained(
-    #     text_encoder_id,
-    #     cache_dir=args.text_encoder_cache_dir,
-    #     torch_dtype=torch.bfloat16 if args.dtype == "bf16" else torch.float32,
-    # )
-    from transformers import AutoTokenizer,AutoModel
-    text_encoder_id = "/share/home/zyx/Models/Megatron-VLM/examples/dit/ckpts/glm-4-9b"
-    tokenizer = AutoTokenizer.from_pretrained(text_encoder_id,trust_remote_code=True)
-    text_encoder = AutoModel.from_pretrained(
+    text_encoder_id = "/share/home/zyx/Models/glm-4-9b-hf"
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(text_encoder_id)
+    text_encoder = GlmForCausalLM.from_pretrained(
         text_encoder_id,
         cache_dir=args.text_encoder_cache_dir,
         torch_dtype=torch.bfloat16 if args.dtype == "bf16" else torch.float32,
-        trust_remote_code = True
     )
-    # Apparently, the conversion does not work anymore without this :shrug:
+
+    # TODO: This is for Older GLM-4 as https://huggingface.co/THUDM/glm-4-9b, will use https://huggingface.co/THUDM/glm-4-9b-hf for new transformers version format.
+    # TODO: Remove it later
+
+    # from transformers import AutoTokenizer,AutoModel
+    # text_encoder_id = "/share/home/zyx/Models/Megatron-VLM/examples/dit/ckpts/glm-4-9b"
+    # tokenizer = AutoTokenizer.from_pretrained(text_encoder_id,trust_remote_code=True)
+    # text_encoder = AutoModel.from_pretrained(
+    #     text_encoder_id,
+    #     cache_dir=args.text_encoder_cache_dir,
+    #     torch_dtype=torch.bfloat16 if args.dtype == "bf16" else torch.float32,
+    #     trust_remote_code = True
+    # )
+
     for param in text_encoder.parameters():
         param.data = param.data.contiguous()
 

src/diffusers/__init__.py

@@ -92,6 +92,7 @@
             "AutoencoderTiny",
             "CogVideoXTransformer3DModel",
             "CogView3PlusTransformer2DModel",
+            "CogView4Transformer2DModel",
             "ConsisIDTransformer3DModel",
             "ConsistencyDecoderVAE",
             "ControlNetModel",
@@ -606,6 +607,7 @@
             AutoencoderTiny,
             CogVideoXTransformer3DModel,
             CogView3PlusTransformer2DModel,
+            CogView4Transformer2DModel,
             ConsisIDTransformer3DModel,
             ConsistencyDecoderVAE,
             ControlNetModel,

src/diffusers/models/__init__.py

@@ -68,6 +68,7 @@
     _import_structure["transformers.transformer_2d"] = ["Transformer2DModel"]
     _import_structure["transformers.transformer_allegro"] = ["AllegroTransformer3DModel"]
     _import_structure["transformers.transformer_cogview3plus"] = ["CogView3PlusTransformer2DModel"]
+    _import_structure["transformers.transformer_cogview4"] = ["CogView4Transformer2DModel"]
     _import_structure["transformers.transformer_flux"] = ["FluxTransformer2DModel"]
     _import_structure["transformers.transformer_hunyuan_video"] = ["HunyuanVideoTransformer3DModel"]
     _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"]
@@ -130,6 +131,7 @@
             AuraFlowTransformer2DModel,
             CogVideoXTransformer3DModel,
             CogView3PlusTransformer2DModel,
+            CogView4Transformer2DModel,
             ConsisIDTransformer3DModel,
             DiTTransformer2DModel,
             DualTransformer2DModel,

src/diffusers/models/transformers/__init__.py

@@ -18,6 +18,7 @@
     from .transformer_2d import Transformer2DModel
     from .transformer_allegro import AllegroTransformer3DModel
     from .transformer_cogview3plus import CogView3PlusTransformer2DModel
+    from .transformer_cogview4 import CogView4Transformer2DModel
     from .transformer_flux import FluxTransformer2DModel
     from .transformer_hunyuan_video import HunyuanVideoTransformer3DModel
     from .transformer_ltx import LTXVideoTransformer3DModel

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, CogView4PatchEmbed
+from ..embeddings import CogView3CombinedTimestepSizeEmbeddings, CogView3PlusPatchEmbed
 from ..modeling_outputs import Transformer2DModelOutput
 from ..normalization import CogView3PlusAdaLayerNormZeroTextImage
 
@@ -84,7 +84,6 @@ 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)
 
@@ -104,7 +103,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, **kwargs
+            hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states
         )
 
         hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_hidden_states
@@ -167,7 +166,8 @@ class CogView3PlusTransformer2DModel(ModelMixin, ConfigMixin):
     """
 
     _supports_gradient_checkpointing = True
-    _no_split_modules = ["CogView3PlusTransformerBlock", "CogView3PlusPatchEmbed", "CogView4PlusPatchEmbed"]
+    _skip_layerwise_casting_patterns = ["patch_embed", "norm"]
+    _no_split_modules = ["CogView3PlusTransformerBlock", "CogView3PlusPatchEmbed"]
 
     @register_to_config
     def __init__(
@@ -192,16 +192,7 @@ def __init__(
         # Each of these are sincos embeddings of shape 2 * condition_dim
         self.pooled_projection_dim = 3 * 2 * condition_dim
 
-        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(
+        self.patch_embed = CogView3PlusPatchEmbed(
             in_channels=in_channels,
             hidden_size=self.inner_dim,
             patch_size=patch_size,
@@ -232,8 +223,7 @@ def __init__(
             embedding_dim=self.inner_dim,
             conditioning_embedding_dim=time_embed_dim,
             elementwise_affine=False,
-            # eps=1e-6,
-            eps=1e-5,
+            eps=1e-6,
         )
         self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
 
@@ -303,55 +293,10 @@ 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,
-        prompt_embeds: torch.Tensor,
-        negative_prompt_embeds: torch.Tensor | None,
+        encoder_hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         original_size: torch.Tensor,
         target_size: torch.Tensor,
@@ -386,103 +331,58 @@ def forward(
             `torch.Tensor` or [`~models.transformer_2d.Transformer2DModelOutput`]:
                 The denoised latents using provided inputs as conditioning.
         """
-        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
-
-        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"
+        height, width = hidden_states.shape[-2:]
+        text_seq_length = encoder_hidden_states.shape[1]
 
-        hidden_states, prompt_embeds, negative_prompt_embeds = self.patch_embed(
-            hidden_states, prompt_embeds, negative_prompt_embeds
-        )
+        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)
 
-        hidden_states_cond, hidden_states_uncond = hidden_states.chunk(2)
-        emb_cond, emb_uncond = emb.chunk(2)
-
-        # 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
-        )
-
-        ######################
-        # prompt_embeds = torch.load("/home/lhy/code/cogview/c_condition_embedding.pt")
-        # negative_prompt_embeds = torch.load("/home/lhy/code/cogview/uc_condition_embedding.pt")
-        prompt_embeds = torch.load("/home/lhy/code/cogview/cp_condition_0_16.pt")[None, ::]
-        negative_prompt_embeds = torch.load("/home/lhy/code/cogview/cp_uncondition_16_32.pt")[None, ::]
+        encoder_hidden_states = hidden_states[:, :text_seq_length]
+        hidden_states = hidden_states[:, text_seq_length:]
 
-        hidden_states_cond = torch.load("/home/lhy/code/cogview/cp_vision_input_0_4096.pt")
-        hidden_states_uncond = torch.load("/home/lhy/code/cogview/cp_vision_input_4096:8192.pt")
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
 
-        emb_cond = torch.load("/home/lhy/code/cogview/time_embedding_0_1.pt")
-        emb_uncond = torch.load("/home/lhy/code/cogview/time_embedding_1_2.pt")
-        ######################
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
 
-        encoder_hidden_states_cond = prompt_embeds
-        encoder_hidden_states_uncond = negative_prompt_embeds
+                    return custom_forward
 
-        for index_block, block in enumerate(self.transformer_blocks):
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
-                ...
-            else:
-                hidden_states_cond, encoder_hidden_states_cond = block(
-                    hidden_states=hidden_states_cond,
-                    encoder_hidden_states=encoder_hidden_states_cond,
-                    emb=emb_cond,  # refactor later
-                    image_rotary_emb=image_rotary_emb,
-                    # image_rotary_emb=None,
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    emb,
+                    **ckpt_kwargs,
                 )
-                ###########################
-                # hidden_states_cond, encoder_hidden_states_cond = (
-                #     self.norm_out.norm(hidden_states_cond),
-                #     self.norm_out.norm(encoder_hidden_states_cond),
-                # )
-                ###########################
-
-                hidden_states_uncond, encoder_hidden_states_uncond = block(
-                    hidden_states=hidden_states_uncond,
-                    encoder_hidden_states=encoder_hidden_states_uncond,
-                    emb=emb_uncond,  # refactor later
-                    image_rotary_emb=image_rotary_emb,
-                    # image_rotary_emb=None,
+            else:
+                hidden_states, encoder_hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    emb=emb,
                 )
-                ###########################
-                # hidden_states_uncond, encoder_hidden_states_uncond = (
-                #     self.norm_out.norm(hidden_states_uncond),
-                #     self.norm_out.norm(encoder_hidden_states_uncond),
-                # )
-                ###########################
-
-        hidden_states_cond = self.norm_out(hidden_states_cond, emb_cond)  # 结果对应于megatron里的final_layer_input
-        hidden_states_uncond = self.norm_out(hidden_states_uncond, emb_uncond)  # 结果对应于megatron里的final_layer_input
-        hidden_states_cond = self.proj_out(hidden_states_cond)  # (batch_size, height*width, patch_size*patch_size*out_channels)
-        hidden_states_uncond = self.proj_out(hidden_states_uncond)  # (batch_size, height*width, patch_size*patch_size*out_channels)
+
+        hidden_states = self.norm_out(hidden_states, emb)
+        hidden_states = self.proj_out(hidden_states)  # (batch_size, height*width, patch_size*patch_size*out_channels)
 
         # unpatchify
         patch_size = self.config.patch_size
         height = height // patch_size
         width = width // patch_size
 
-        hidden_states_cond = hidden_states_cond.reshape(
-            shape=(hidden_states_cond.shape[0], height, width, self.out_channels, patch_size, patch_size)
-        )
-        hidden_states_cond = torch.einsum("nhwcpq->nchpwq", hidden_states_cond)
-        output_cond = hidden_states_cond.reshape(
-            shape=(hidden_states_cond.shape[0], self.out_channels, height * patch_size, width * patch_size)
-        )
-
-        hidden_states_uncond = hidden_states_uncond.reshape(
-            shape=(hidden_states_uncond.shape[0], height, width, self.out_channels, patch_size, patch_size)
+        hidden_states = hidden_states.reshape(
+            shape=(hidden_states.shape[0], height, width, self.out_channels, patch_size, patch_size)
         )
-        hidden_states_uncond = torch.einsum("nhwcpq->nchpwq", hidden_states_uncond)
-        output_uncond = hidden_states_uncond.reshape(
-            shape=(hidden_states_uncond.shape[0], self.out_channels, height * patch_size, width * patch_size)
+        hidden_states = torch.einsum("nhwcpq->nchpwq", hidden_states)
+        output = hidden_states.reshape(
+            shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
         )
 
         if not return_dict:
-            return (output_cond, output_uncond)
+            return (output,)
 
-        return Transformer2DModelOutput(sample=output_cond), Transformer2DModelOutput(sample=output_uncond)
+        return Transformer2DModelOutput(sample=output)