Update model_loading.py

cogvideox_fun/fun_pab_transformer_3d.py

@@ -0,0 +1,741 @@
+# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and 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.
+
+from typing import Any, Dict, Optional, Tuple, Union
+
+import os
+import json
+import torch
+import glob
+import torch.nn.functional as F
+from torch import nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import is_torch_version, logging
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention import Attention, FeedForward
+#from diffusers.models.attention_processor import AttentionProcessor
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps, get_3d_sincos_pos_embed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+#from diffusers.models.normalization import AdaLayerNorm, CogVideoXLayerNormZero
+
+from ..videosys.modules.normalization import AdaLayerNorm, CogVideoXLayerNormZero
+from ..videosys.modules.embeddings import apply_rotary_emb
+from ..videosys.core.pab_mgr import enable_pab, if_broadcast_spatial
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+try:
+    from sageattention import sageattn
+    SAGEATTN_IS_AVAVILABLE = True
+    logger.info("Using sageattn")
+except:
+    logger.info("sageattn not found, using sdpa")
+    SAGEATTN_IS_AVAVILABLE = False
+
+class CogVideoXAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
+    query and key vectors, but does not include spatial normalization.
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        # if attn.parallel_manager.sp_size > 1:
+        #     assert (
+        #         attn.heads % attn.parallel_manager.sp_size == 0
+        #     ), f"Number of heads {attn.heads} must be divisible by sequence parallel size {attn.parallel_manager.sp_size}"
+        #     attn_heads = attn.heads // attn.parallel_manager.sp_size
+        #     query, key, value = map(
+        #         lambda x: all_to_all_comm(x, attn.parallel_manager.sp_group, scatter_dim=2, gather_dim=1),
+        #         [query, key, value],
+        #     )
+        
+        attn_heads = attn.heads
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn_heads
+
+        query = query.view(batch_size, -1, attn_heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn_heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn_heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply RoPE if needed
+        if image_rotary_emb is not None:
+            emb_len = image_rotary_emb[0].shape[0]
+            query[:, :, text_seq_length : emb_len + text_seq_length] = apply_rotary_emb(
+                query[:, :, text_seq_length : emb_len + text_seq_length], image_rotary_emb
+            )
+            if not attn.is_cross_attention:
+                key[:, :, text_seq_length : emb_len + text_seq_length] = apply_rotary_emb(
+                    key[:, :, text_seq_length : emb_len + text_seq_length], image_rotary_emb
+                )
+
+        if SAGEATTN_IS_AVAVILABLE:
+            hidden_states = sageattn(query, key, value, is_causal=False)
+        else:
+            hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn_heads * head_dim)
+
+        #if attn.parallel_manager.sp_size > 1:
+        #    hidden_states = all_to_all_comm(hidden_states, attn.parallel_manager.sp_group, scatter_dim=1, gather_dim=2)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+class FusedCogVideoXAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
+    query and key vectors, but does not include spatial normalization.
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        qkv = attn.to_qkv(hidden_states)
+        split_size = qkv.shape[-1] // 3
+        query, key, value = torch.split(qkv, split_size, dim=-1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply RoPE if needed
+        if image_rotary_emb is not None:
+            query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
+            if not attn.is_cross_attention:
+                key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
+
+        if SAGEATTN_IS_AVAVILABLE:
+            hidden_states = sageattn(query, key, value, is_causal=False)
+        else:
+            hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+class CogVideoXPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 2,
+        in_channels: int = 16,
+        embed_dim: int = 1920,
+        text_embed_dim: int = 4096,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias
+        )
+        self.text_proj = nn.Linear(text_embed_dim, embed_dim)
+
+    def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):
+        r"""
+        Args:
+            text_embeds (`torch.Tensor`):
+                Input text embeddings. Expected shape: (batch_size, seq_length, embedding_dim).
+            image_embeds (`torch.Tensor`):
+                Input image embeddings. Expected shape: (batch_size, num_frames, channels, height, width).
+        """
+        text_embeds = self.text_proj(text_embeds)
+
+        batch, num_frames, channels, height, width = image_embeds.shape
+        image_embeds = image_embeds.reshape(-1, channels, height, width)
+        image_embeds = self.proj(image_embeds)
+        image_embeds = image_embeds.view(batch, num_frames, *image_embeds.shape[1:])
+        image_embeds = image_embeds.flatten(3).transpose(2, 3)  # [batch, num_frames, height x width, channels]
+        image_embeds = image_embeds.flatten(1, 2)  # [batch, num_frames x height x width, channels]
+
+        embeds = torch.cat(
+            [text_embeds, image_embeds], dim=1
+        ).contiguous()  # [batch, seq_length + num_frames x height x width, channels]
+        return embeds
+
+@maybe_allow_in_graph
+class CogVideoXBlock(nn.Module):
+    r"""
+    Transformer block used in [CogVideoX](https://github.com/THUDM/CogVideo) model.
+
+    Parameters:
+        dim (`int`):
+            The number of channels in the input and output.
+        num_attention_heads (`int`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`):
+            The number of channels in each head.
+        time_embed_dim (`int`):
+            The number of channels in timestep embedding.
+        dropout (`float`, defaults to `0.0`):
+            The dropout probability to use.
+        activation_fn (`str`, defaults to `"gelu-approximate"`):
+            Activation function to be used in feed-forward.
+        attention_bias (`bool`, defaults to `False`):
+            Whether or not to use bias in attention projection layers.
+        qk_norm (`bool`, defaults to `True`):
+            Whether or not to use normalization after query and key projections in Attention.
+        norm_elementwise_affine (`bool`, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_eps (`float`, defaults to `1e-5`):
+            Epsilon value for normalization layers.
+        final_dropout (`bool` defaults to `False`):
+            Whether to apply a final dropout after the last feed-forward layer.
+        ff_inner_dim (`int`, *optional*, defaults to `None`):
+            Custom hidden dimension of Feed-forward layer. If not provided, `4 * dim` is used.
+        ff_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in Feed-forward layer.
+        attention_out_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in Attention output projection layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        time_embed_dim: int,
+        dropout: float = 0.0,
+        activation_fn: str = "gelu-approximate",
+        attention_bias: bool = False,
+        qk_norm: bool = True,
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        final_dropout: bool = True,
+        ff_inner_dim: Optional[int] = None,
+        ff_bias: bool = True,
+        attention_out_bias: bool = True,
+        block_idx: int = 0,
+    ):
+        super().__init__()
+
+        # 1. Self Attention
+        self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
+
+        self.attn1 = Attention(
+            query_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            qk_norm="layer_norm" if qk_norm else None,
+            eps=1e-6,
+            bias=attention_bias,
+            out_bias=attention_out_bias,
+            processor=CogVideoXAttnProcessor2_0(),
+        )
+
+        # parallel
+        #self.attn1.parallel_manager = None
+
+        # 2. Feed Forward
+        self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
+
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+            inner_dim=ff_inner_dim,
+            bias=ff_bias,
+        )
+
+        # pab
+        self.attn_count = 0
+        self.last_attn = None
+        self.block_idx = block_idx
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        timestep=None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        # norm & modulate
+        norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
+            hidden_states, encoder_hidden_states, temb
+        )
+
+        # attention
+        if enable_pab():
+            broadcast_attn, self.attn_count = if_broadcast_spatial(int(timestep[0]), self.attn_count, self.block_idx)
+        if enable_pab() and broadcast_attn:
+            attn_hidden_states, attn_encoder_hidden_states = self.last_attn
+        else:
+            attn_hidden_states, attn_encoder_hidden_states = self.attn1(
+                hidden_states=norm_hidden_states,
+                encoder_hidden_states=norm_encoder_hidden_states,
+                image_rotary_emb=image_rotary_emb,
+            )
+            if enable_pab():
+                self.last_attn = (attn_hidden_states, attn_encoder_hidden_states)
+
+        hidden_states = hidden_states + gate_msa * attn_hidden_states
+        encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states
+
+        # norm & modulate
+        norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
+            hidden_states, encoder_hidden_states, temb
+        )
+
+        # feed-forward
+        norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
+        ff_output = self.ff(norm_hidden_states)
+
+        hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]
+        encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]
+
+        return hidden_states, encoder_hidden_states
+
+
+class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin):
+    """
+    A Transformer model for video-like data in [CogVideoX](https://github.com/THUDM/CogVideo).
+
+    Parameters:
+        num_attention_heads (`int`, defaults to `30`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, defaults to `64`):
+            The number of channels in each head.
+        in_channels (`int`, defaults to `16`):
+            The number of channels in the input.
+        out_channels (`int`, *optional*, defaults to `16`):
+            The number of channels in the output.
+        flip_sin_to_cos (`bool`, defaults to `True`):
+            Whether to flip the sin to cos in the time embedding.
+        time_embed_dim (`int`, defaults to `512`):
+            Output dimension of timestep embeddings.
+        text_embed_dim (`int`, defaults to `4096`):
+            Input dimension of text embeddings from the text encoder.
+        num_layers (`int`, defaults to `30`):
+            The number of layers of Transformer blocks to use.
+        dropout (`float`, defaults to `0.0`):
+            The dropout probability to use.
+        attention_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in the attention projection layers.
+        sample_width (`int`, defaults to `90`):
+            The width of the input latents.
+        sample_height (`int`, defaults to `60`):
+            The height of the input latents.
+        sample_frames (`int`, defaults to `49`):
+            The number of frames in the input latents. Note that this parameter was incorrectly initialized to 49
+            instead of 13 because CogVideoX processed 13 latent frames at once in its default and recommended settings,
+            but cannot be changed to the correct value to ensure backwards compatibility. To create a transformer with
+            K latent frames, the correct value to pass here would be: ((K - 1) * temporal_compression_ratio + 1).
+        patch_size (`int`, defaults to `2`):
+            The size of the patches to use in the patch embedding layer.
+        temporal_compression_ratio (`int`, defaults to `4`):
+            The compression ratio across the temporal dimension. See documentation for `sample_frames`.
+        max_text_seq_length (`int`, defaults to `226`):
+            The maximum sequence length of the input text embeddings.
+        activation_fn (`str`, defaults to `"gelu-approximate"`):
+            Activation function to use in feed-forward.
+        timestep_activation_fn (`str`, defaults to `"silu"`):
+            Activation function to use when generating the timestep embeddings.
+        norm_elementwise_affine (`bool`, defaults to `True`):
+            Whether or not to use elementwise affine in normalization layers.
+        norm_eps (`float`, defaults to `1e-5`):
+            The epsilon value to use in normalization layers.
+        spatial_interpolation_scale (`float`, defaults to `1.875`):
+            Scaling factor to apply in 3D positional embeddings across spatial dimensions.
+        temporal_interpolation_scale (`float`, defaults to `1.0`):
+            Scaling factor to apply in 3D positional embeddings across temporal dimensions.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 30,
+        attention_head_dim: int = 64,
+        in_channels: int = 16,
+        out_channels: Optional[int] = 16,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        time_embed_dim: int = 512,
+        text_embed_dim: int = 4096,
+        num_layers: int = 30,
+        dropout: float = 0.0,
+        attention_bias: bool = True,
+        sample_width: int = 90,
+        sample_height: int = 60,
+        sample_frames: int = 49,
+        patch_size: int = 2,
+        temporal_compression_ratio: int = 4,
+        max_text_seq_length: int = 226,
+        activation_fn: str = "gelu-approximate",
+        timestep_activation_fn: str = "silu",
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        spatial_interpolation_scale: float = 1.875,
+        temporal_interpolation_scale: float = 1.0,
+        use_rotary_positional_embeddings: bool = False,
+        add_noise_in_inpaint_model: bool = False,
+    ):
+        super().__init__()
+        inner_dim = num_attention_heads * attention_head_dim
+
+        post_patch_height = sample_height // patch_size
+        post_patch_width = sample_width // patch_size
+        post_time_compression_frames = (sample_frames - 1) // temporal_compression_ratio + 1
+        self.num_patches = post_patch_height * post_patch_width * post_time_compression_frames
+        self.post_patch_height = post_patch_height
+        self.post_patch_width = post_patch_width
+        self.post_time_compression_frames = post_time_compression_frames
+        self.patch_size = patch_size
+
+        # 1. Patch embedding
+        self.patch_embed = CogVideoXPatchEmbed(patch_size, in_channels, inner_dim, text_embed_dim, bias=True)
+        self.embedding_dropout = nn.Dropout(dropout)
+
+        # 2. 3D positional embeddings
+        spatial_pos_embedding = get_3d_sincos_pos_embed(
+            inner_dim,
+            (post_patch_width, post_patch_height),
+            post_time_compression_frames,
+            spatial_interpolation_scale,
+            temporal_interpolation_scale,
+        )
+        spatial_pos_embedding = torch.from_numpy(spatial_pos_embedding).flatten(0, 1)
+        pos_embedding = torch.zeros(1, max_text_seq_length + self.num_patches, inner_dim, requires_grad=False)
+        pos_embedding.data[:, max_text_seq_length:].copy_(spatial_pos_embedding)
+        self.register_buffer("pos_embedding", pos_embedding, persistent=False)
+
+        # 3. Time embeddings
+        self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift)
+        self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn)
+
+        # 4. Define spatio-temporal transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                CogVideoXBlock(
+                    dim=inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    time_embed_dim=time_embed_dim,
+                    dropout=dropout,
+                    activation_fn=activation_fn,
+                    attention_bias=attention_bias,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                    norm_eps=norm_eps,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm_final = nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine)
+
+        # 5. Output blocks
+        self.norm_out = AdaLayerNorm(
+            embedding_dim=time_embed_dim,
+            output_dim=2 * inner_dim,
+            norm_elementwise_affine=norm_elementwise_affine,
+            norm_eps=norm_eps,
+            chunk_dim=1,
+        )
+        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        self.gradient_checkpointing = value
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedCogVideoXAttnProcessor2_0
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedCogVideoXAttnProcessor2_0())
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: Union[int, float, torch.LongTensor],
+        timestep_cond: Optional[torch.Tensor] = None,
+        inpaint_latents: Optional[torch.Tensor] = None,
+        control_latents: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        return_dict: bool = True,
+    ):
+        batch_size, num_frames, channels, height, width = hidden_states.shape
+
+        # 1. Time embedding
+        timesteps = timestep
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=hidden_states.dtype)
+        emb = self.time_embedding(t_emb, timestep_cond)
+
+        # 2. Patch embedding
+        if inpaint_latents is not None:
+            hidden_states = torch.concat([hidden_states, inpaint_latents], 2)
+        if control_latents is not None:
+            hidden_states = torch.concat([hidden_states, control_latents], 2)
+        hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)
+
+        # 3. Position embedding
+        text_seq_length = encoder_hidden_states.shape[1]
+        if not self.config.use_rotary_positional_embeddings:
+            seq_length = height * width * num_frames // (self.config.patch_size**2)
+            # pos_embeds = self.pos_embedding[:, : text_seq_length + seq_length]
+            pos_embeds = self.pos_embedding
+            emb_size = hidden_states.size()[-1]
+            pos_embeds_without_text = pos_embeds[:, text_seq_length: ].view(1, self.post_time_compression_frames, self.post_patch_height, self.post_patch_width, emb_size)
+            pos_embeds_without_text = pos_embeds_without_text.permute([0, 4, 1, 2, 3])
+            pos_embeds_without_text = F.interpolate(pos_embeds_without_text,size=[self.post_time_compression_frames, height // self.config.patch_size, width // self.config.patch_size],mode='trilinear',align_corners=False)
+            pos_embeds_without_text = pos_embeds_without_text.permute([0, 2, 3, 4, 1]).view(1, -1, emb_size)
+            pos_embeds = torch.cat([pos_embeds[:, :text_seq_length], pos_embeds_without_text], dim = 1)
+            pos_embeds = pos_embeds[:, : text_seq_length + seq_length]
+            hidden_states = hidden_states + pos_embeds
+            hidden_states = self.embedding_dropout(hidden_states)
+
+        encoder_hidden_states = hidden_states[:, :text_seq_length]
+        hidden_states = hidden_states[:, text_seq_length:]
+
+        # 4. Transformer blocks
+        
+        for i, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                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,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+            else:
+                hidden_states, encoder_hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=emb,
+                    image_rotary_emb=image_rotary_emb,
+                    timestep=timestep,
+                )
+
+        if not self.config.use_rotary_positional_embeddings:
+            # CogVideoX-2B
+            hidden_states = self.norm_final(hidden_states)
+        else:
+            # CogVideoX-5B
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+            hidden_states = self.norm_final(hidden_states)
+            hidden_states = hidden_states[:, text_seq_length:]
+
+        # 5. Final block
+        hidden_states = self.norm_out(hidden_states, temb=emb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # 6. Unpatchify
+        p = self.config.patch_size
+        output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, channels, p, p)
+        output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    @classmethod
+    def from_pretrained_2d(cls, pretrained_model_path, subfolder=None, transformer_additional_kwargs={}):
+        if subfolder is not None:
+            pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
+        print(f"loaded 3D transformer's pretrained weights from {pretrained_model_path} ...")
+
+        config_file = os.path.join(pretrained_model_path, 'config.json')
+        if not os.path.isfile(config_file):
+            raise RuntimeError(f"{config_file} does not exist")
+        with open(config_file, "r") as f:
+            config = json.load(f)
+
+        from diffusers.utils import WEIGHTS_NAME
+        model = cls.from_config(config, **transformer_additional_kwargs)
+        model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
+        model_file_safetensors = model_file.replace(".bin", ".safetensors")
+        if os.path.exists(model_file):
+            state_dict = torch.load(model_file, map_location="cpu")
+        elif os.path.exists(model_file_safetensors):
+            from safetensors.torch import load_file, safe_open
+            state_dict = load_file(model_file_safetensors)
+        else:
+            from safetensors.torch import load_file, safe_open
+            model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+            state_dict = {}
+            for model_file_safetensors in model_files_safetensors:
+                _state_dict = load_file(model_file_safetensors)
+                for key in _state_dict:
+                    state_dict[key] = _state_dict[key]
+        
+        if model.state_dict()['patch_embed.proj.weight'].size() != state_dict['patch_embed.proj.weight'].size():
+            new_shape   = model.state_dict()['patch_embed.proj.weight'].size()
+            if len(new_shape) == 5:
+                state_dict['patch_embed.proj.weight'] = state_dict['patch_embed.proj.weight'].unsqueeze(2).expand(new_shape).clone()
+                state_dict['patch_embed.proj.weight'][:, :, :-1] = 0
+            else:
+                if model.state_dict()['patch_embed.proj.weight'].size()[1] > state_dict['patch_embed.proj.weight'].size()[1]:
+                    model.state_dict()['patch_embed.proj.weight'][:, :state_dict['patch_embed.proj.weight'].size()[1], :, :] = state_dict['patch_embed.proj.weight']
+                    model.state_dict()['patch_embed.proj.weight'][:, state_dict['patch_embed.proj.weight'].size()[1]:, :, :] = 0
+                    state_dict['patch_embed.proj.weight'] = model.state_dict()['patch_embed.proj.weight']
+                else:
+                    model.state_dict()['patch_embed.proj.weight'][:, :, :, :] = state_dict['patch_embed.proj.weight'][:, :model.state_dict()['patch_embed.proj.weight'].size()[1], :, :]
+                    state_dict['patch_embed.proj.weight'] = model.state_dict()['patch_embed.proj.weight']
+
+        tmp_state_dict = {} 
+        for key in state_dict:
+            if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size():
+                tmp_state_dict[key] = state_dict[key]
+            else:
+                print(key, "Size don't match, skip")
+        state_dict = tmp_state_dict
+
+        m, u = model.load_state_dict(state_dict, strict=False)
+        print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
+        print(m)
+        
+        params = [p.numel() if "mamba" in n else 0 for n, p in model.named_parameters()]
+        print(f"### Mamba Parameters: {sum(params) / 1e6} M")
+
+        params = [p.numel() if "attn1." in n else 0 for n, p in model.named_parameters()]
+        print(f"### attn1 Parameters: {sum(params) / 1e6} M")
+        
+        return model

cogvideox_fun/pipeline_cogvideox_control.py

@@ -33,6 +33,10 @@ from diffusers.video_processor import VideoProcessor
 from diffusers.image_processor import VaeImageProcessor
 from einops import rearrange
 
+from ..videosys.core.pipeline import VideoSysPipeline
+from ..videosys.cogvideox_transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelPAB
+from ..videosys.core.pab_mgr import set_pab_manager
+
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -154,7 +158,7 @@ class CogVideoX_Fun_PipelineOutput(BaseOutput):
     videos: torch.Tensor
 
 
-class CogVideoX_Fun_Pipeline_Control(DiffusionPipeline):
+class CogVideoX_Fun_Pipeline_Control(VideoSysPipeline):
     r"""
     Pipeline for text-to-video generation using CogVideoX.
 
@@ -184,6 +188,7 @@ class CogVideoX_Fun_Pipeline_Control(DiffusionPipeline):
         vae: AutoencoderKLCogVideoX,
         transformer: CogVideoXTransformer3DModel,
         scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+        pab_config = None
     ):
         super().__init__()
 
@@ -205,6 +210,9 @@ class CogVideoX_Fun_Pipeline_Control(DiffusionPipeline):
             vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True
         )
 
+        if pab_config is not None:
+            set_pab_manager(pab_config)
+
     def prepare_latents(
         self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, timesteps, denoise_strength, num_inference_steps,
          latents=None, freenoise=True, context_size=None, context_overlap=None
@@ -340,6 +348,16 @@ class CogVideoX_Fun_Pipeline_Control(DiffusionPipeline):
         if accepts_generator:
             extra_step_kwargs["generator"] = generator
         return extra_step_kwargs
+    
+    def _gaussian_weights(self, t_tile_length, t_batch_size):
+        from numpy import pi, exp, sqrt
+
+        var = 0.01
+        midpoint = (t_tile_length - 1) / 2  # -1 because index goes from 0 to latent_width - 1
+        t_probs = [exp(-(t-midpoint)*(t-midpoint)/(t_tile_length*t_tile_length)/(2*var)) / sqrt(2*pi*var) for t in range(t_tile_length)]
+        weights = torch.tensor(t_probs)
+        weights = weights.unsqueeze(0).unsqueeze(2).unsqueeze(3).unsqueeze(4).repeat(1, t_batch_size,1, 1, 1)
+        return weights
 
     # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
     def check_inputs(
@@ -679,15 +697,24 @@ class CogVideoX_Fun_Pipeline_Control(DiffusionPipeline):
         # 8. Denoising loop
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
 
-        if context_schedule is not None:
+        # 8.5. Temporal tiling prep
+        if context_schedule is not None and context_schedule == "temporal_tiling":
+            t_tile_length = context_frames
+            t_tile_overlap = context_overlap
+            t_tile_weights = self._gaussian_weights(t_tile_length=t_tile_length, t_batch_size=1).to(latents.device).to(self.vae.dtype)
+            use_temporal_tiling = True
+            print("Temporal tiling enabled")
+        elif context_schedule is not None:
             print(f"Context schedule enabled: {context_frames} frames, {context_stride} stride, {context_overlap} overlap")
+            use_temporal_tiling = False
             use_context_schedule = True
             from .context import get_context_scheduler
             context = get_context_scheduler(context_schedule)
 
         else:
+            use_temporal_tiling = False
             use_context_schedule = False
-            print(" context schedule disabled")
+            print("Temporal tiling and context schedule disabled")
             # 7. Create rotary embeds if required
             image_rotary_emb = (
                 self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device)
@@ -708,7 +735,88 @@ class CogVideoX_Fun_Pipeline_Control(DiffusionPipeline):
             for i, t in enumerate(timesteps):
                 if self.interrupt:
                     continue
-                if use_context_schedule:
+
+                if use_temporal_tiling and isinstance(self.scheduler, CogVideoXDDIMScheduler):
+                    #temporal tiling code based on https://github.com/mayuelala/FollowYourEmoji/blob/main/models/video_pipeline.py
+                    # =====================================================
+                    grid_ts = 0
+                    cur_t = 0
+                    while cur_t < latents.shape[1]:
+                        cur_t = max(grid_ts * t_tile_length - t_tile_overlap * grid_ts, 0) + t_tile_length
+                        grid_ts += 1
+
+                    all_t = latents.shape[1]
+                    latents_all_list = []
+                    # =====================================================
+
+                    image_rotary_emb = (
+                            self._prepare_rotary_positional_embeddings(height, width, context_frames, device)
+                            if self.transformer.config.use_rotary_positional_embeddings
+                            else None
+                        )
+
+                    for t_i in range(grid_ts):
+                        if t_i < grid_ts - 1:
+                            ofs_t = max(t_i * t_tile_length - t_tile_overlap * t_i, 0)
+                        if t_i == grid_ts - 1:
+                            ofs_t = all_t - t_tile_length
+
+                        input_start_t = ofs_t
+                        input_end_t = ofs_t + t_tile_length
+
+                        latents_tile = latents[:, input_start_t:input_end_t,:, :, :]
+                        control_latents_tile = control_latents[:, input_start_t:input_end_t, :, :, :]
+
+                        latent_model_input_tile = torch.cat([latents_tile] * 2) if do_classifier_free_guidance else latents_tile
+                        latent_model_input_tile = self.scheduler.scale_model_input(latent_model_input_tile, t)
+
+                        #t_input = t[None].to(device)
+                        t_input = t.expand(latent_model_input_tile.shape[0]) # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                
+                        # predict noise model_output
+                        noise_pred = self.transformer(
+                            hidden_states=latent_model_input_tile,
+                            encoder_hidden_states=prompt_embeds,
+                            timestep=t_input,
+                            image_rotary_emb=image_rotary_emb,
+                            return_dict=False,
+                            control_latents=control_latents_tile,
+                        )[0]
+                        noise_pred = noise_pred.float()                  
+
+                        if do_classifier_free_guidance:
+                            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                            noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                        # compute the previous noisy sample x_t -> x_t-1
+                        latents_tile = self.scheduler.step(noise_pred, t, latents_tile.to(self.vae.dtype), **extra_step_kwargs, return_dict=False)[0]
+                        latents_all_list.append(latents_tile)
+
+                    # ==========================================
+                    latents_all = torch.zeros(latents.shape, device=latents.device, dtype=self.vae.dtype)
+                    contributors = torch.zeros(latents.shape, device=latents.device, dtype=self.vae.dtype)
+                    # Add each tile contribution to overall latents
+                    for t_i in range(grid_ts):
+                        if t_i < grid_ts - 1:
+                            ofs_t = max(t_i * t_tile_length - t_tile_overlap * t_i, 0)
+                        if t_i == grid_ts - 1:
+                            ofs_t = all_t - t_tile_length
+
+                        input_start_t = ofs_t
+                        input_end_t = ofs_t + t_tile_length
+
+                        latents_all[:, input_start_t:input_end_t,:, :, :] += latents_all_list[t_i] * t_tile_weights
+                        contributors[:, input_start_t:input_end_t,:, :, :] += t_tile_weights
+                    
+                    latents_all /= contributors
+
+                    latents = latents_all                    
+                    
+                    if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                        progress_bar.update()
+                        pbar.update(1)
+                    # ==========================================
+                elif use_context_schedule:
                     
                     latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

cogvideox_fun/pipeline_cogvideox_inpaint.py

@@ -33,6 +33,11 @@ from diffusers.video_processor import VideoProcessor
 from diffusers.image_processor import VaeImageProcessor
 from einops import rearrange
 
+from ..videosys.core.pipeline import VideoSysPipeline
+from ..videosys.cogvideox_transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelPAB
+from ..videosys.core.pab_mgr import set_pab_manager
+
+
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
@@ -201,7 +206,7 @@ class CogVideoX_Fun_PipelineOutput(BaseOutput):
     videos: torch.Tensor
 
 
-class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
+class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline):
     r"""
     Pipeline for text-to-video generation using CogVideoX.
 
@@ -231,6 +236,7 @@ class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
         vae: AutoencoderKLCogVideoX,
         transformer: CogVideoXTransformer3DModel,
         scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+        pab_config = None
     ):
         super().__init__()
 
@@ -252,6 +258,9 @@ class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
             vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True
         )
 
+        if pab_config is not None:
+            set_pab_manager(pab_config)
+
     def prepare_latents(
         self, 
         batch_size,
@@ -424,6 +433,16 @@ class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
             extra_step_kwargs["generator"] = generator
         return extra_step_kwargs
     
+    def _gaussian_weights(self, t_tile_length, t_batch_size):
+        from numpy import pi, exp, sqrt
+
+        var = 0.01
+        midpoint = (t_tile_length - 1) / 2  # -1 because index goes from 0 to latent_width - 1
+        t_probs = [exp(-(t-midpoint)*(t-midpoint)/(t_tile_length*t_tile_length)/(2*var)) / sqrt(2*pi*var) for t in range(t_tile_length)]
+        weights = torch.tensor(t_probs)
+        weights = weights.unsqueeze(0).unsqueeze(2).unsqueeze(3).unsqueeze(4).repeat(1, t_batch_size,1, 1, 1)
+        return weights
+    
     # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
     def check_inputs(
         self,
@@ -739,8 +758,6 @@ class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
         num_channels_transformer = self.transformer.config.in_channels
         return_image_latents = num_channels_transformer == num_channels_latents
 
-        self.vae.to(device)
-
         latents_outputs = self.prepare_latents(
             batch_size * num_videos_per_prompt,
             num_channels_latents,
@@ -842,9 +859,6 @@ class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
                 mask = rearrange(mask, "b c f h w -> b f c h w")
 
                 inpaint_latents = None
-
-        self.vae.to(torch.device("cpu"))
-
         if comfyui_progressbar:
             pbar.update(1)
 
@@ -852,14 +866,22 @@ class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
         # 7. Create rotary embeds if required
-        if context_schedule is not None:
+        if context_schedule is not None and context_schedule == "temporal_tiling":
+            t_tile_length = context_frames
+            t_tile_overlap = context_overlap
+            t_tile_weights = self._gaussian_weights(t_tile_length=t_tile_length, t_batch_size=1).to(latents.device).to(self.vae.dtype)
+            use_temporal_tiling = True
+            print("Temporal tiling enabled")
+        elif context_schedule is not None:
             print(f"Context schedule enabled: {context_frames} frames, {context_stride} stride, {context_overlap} overlap")
+            use_temporal_tiling = False
             use_context_schedule = True
             from .context import get_context_scheduler
             context = get_context_scheduler(context_schedule)
         else:
+            use_temporal_tiling = False
             use_context_schedule = False
-            print("context schedule disabled")
+            print("Temporal tiling and context schedule disabled")
             # 7. Create rotary embeds if required
             image_rotary_emb = (
                 self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device)
@@ -893,7 +915,87 @@ class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
                 if self.interrupt:
                     continue
 
-                if use_context_schedule:
+                if use_temporal_tiling and isinstance(self.scheduler, CogVideoXDDIMScheduler):
+                    #temporal tiling code based on https://github.com/mayuelala/FollowYourEmoji/blob/main/models/video_pipeline.py
+                    # =====================================================
+                    grid_ts = 0
+                    cur_t = 0
+                    while cur_t < latents.shape[1]:
+                        cur_t = max(grid_ts * t_tile_length - t_tile_overlap * grid_ts, 0) + t_tile_length
+                        grid_ts += 1
+
+                    all_t = latents.shape[1]
+                    latents_all_list = []
+                    # =====================================================
+
+                    image_rotary_emb = (
+                            self._prepare_rotary_positional_embeddings(height, width, t_tile_length, device)
+                            if self.transformer.config.use_rotary_positional_embeddings
+                            else None
+                        )
+
+                    for t_i in range(grid_ts):
+                        if t_i < grid_ts - 1:
+                            ofs_t = max(t_i * t_tile_length - t_tile_overlap * t_i, 0)
+                        if t_i == grid_ts - 1:
+                            ofs_t = all_t - t_tile_length
+
+                        input_start_t = ofs_t
+                        input_end_t = ofs_t + t_tile_length
+
+                        latents_tile = latents[:, input_start_t:input_end_t,:, :, :]
+                        inpaint_latents_tile = inpaint_latents[:, input_start_t:input_end_t, :, :, :]
+
+                        latent_model_input_tile = torch.cat([latents_tile] * 2) if do_classifier_free_guidance else latents_tile
+                        latent_model_input_tile = self.scheduler.scale_model_input(latent_model_input_tile, t)
+
+                        #t_input = t[None].to(device)
+                        t_input = t.expand(latent_model_input_tile.shape[0]) # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                
+                        # predict noise model_output
+                        noise_pred = self.transformer(
+                            hidden_states=latent_model_input_tile,
+                            encoder_hidden_states=prompt_embeds,
+                            timestep=t_input,
+                            image_rotary_emb=image_rotary_emb,
+                            return_dict=False,
+                            inpaint_latents=inpaint_latents_tile,
+                        )[0]
+                        noise_pred = noise_pred.float()                  
+
+                        if do_classifier_free_guidance:
+                            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                            noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                        # compute the previous noisy sample x_t -> x_t-1
+                        latents_tile = self.scheduler.step(noise_pred, t, latents_tile.to(self.vae.dtype), **extra_step_kwargs, return_dict=False)[0]
+                        latents_all_list.append(latents_tile)
+
+                    # ==========================================
+                    latents_all = torch.zeros(latents.shape, device=latents.device, dtype=self.vae.dtype)
+                    contributors = torch.zeros(latents.shape, device=latents.device, dtype=self.vae.dtype)
+                    # Add each tile contribution to overall latents
+                    for t_i in range(grid_ts):
+                        if t_i < grid_ts - 1:
+                            ofs_t = max(t_i * t_tile_length - t_tile_overlap * t_i, 0)
+                        if t_i == grid_ts - 1:
+                            ofs_t = all_t - t_tile_length
+
+                        input_start_t = ofs_t
+                        input_end_t = ofs_t + t_tile_length
+
+                        latents_all[:, input_start_t:input_end_t,:, :, :] += latents_all_list[t_i] * t_tile_weights
+                        contributors[:, input_start_t:input_end_t,:, :, :] += t_tile_weights
+                    
+                    latents_all /= contributors
+
+                    latents = latents_all                    
+                    
+                    if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                        progress_bar.update()
+                        pbar.update(1)
+                    # ==========================================
+                elif use_context_schedule:
                     latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
 
@@ -1031,7 +1133,18 @@ class CogVideoX_Fun_Pipeline_Inpaint(DiffusionPipeline):
                         else:
                             pbar.update(1)
 
+        # if output_type == "numpy":
+        #     video = self.decode_latents(latents)
+        # elif not output_type == "latent":
+        #     video = self.decode_latents(latents)
+        #     video = self.video_processor.postprocess_video(video=video, output_type=output_type)
+        # else:
+        #     video = latents
+
         # Offload all models
         self.maybe_free_model_hooks()
 
+        # if not return_dict:
+        #     video = torch.from_numpy(video)
+
         return latents

cogvideox_fun/transformer_3d.py

@@ -37,9 +37,11 @@ logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 from einops import rearrange
 try:
     from sageattention import sageattn
-    SAGEATTN_IS_AVAILABLE = True
+    SAGEATTN_IS_AVAVILABLE = True
+    logger.info("Using sageattn")
 except:
-    SAGEATTN_IS_AVAILABLE = False
+    logger.info("sageattn not found, using sdpa")
+    SAGEATTN_IS_AVAVILABLE = False
 
 def fft(tensor):
     tensor_fft = torch.fft.fft2(tensor)
@@ -75,7 +77,6 @@ class CogVideoXAttnProcessor2_0:
         encoder_hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
-        attention_mode: Optional[str] = None,
     ) -> torch.Tensor:
         text_seq_length = encoder_hidden_states.size(1)
 
@@ -112,12 +113,83 @@ class CogVideoXAttnProcessor2_0:
             query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
             if not attn.is_cross_attention:
                 key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
-                 
-        if attention_mode == "sageattn":
-            if SAGEATTN_IS_AVAILABLE:
-               hidden_states = sageattn(query, key, value, attn_mask=attention_mask, dropout_p=0.0,is_causal=False)
-            else:
-                raise ImportError("sageattn not found")
+      
+        if SAGEATTN_IS_AVAVILABLE:
+            hidden_states = sageattn(query, key, value, is_causal=False)
+        else:
+            hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+class FusedCogVideoXAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
+    query and key vectors, but does not include spatial normalization.
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        qkv = attn.to_qkv(hidden_states)
+        split_size = qkv.shape[-1] // 3
+        query, key, value = torch.split(qkv, split_size, dim=-1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply RoPE if needed
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
+            if not attn.is_cross_attention:
+                key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
+
+        if SAGEATTN_IS_AVAVILABLE:
+            hidden_states = sageattn(query, key, value, is_causal=False)
         else:
             hidden_states = F.scaled_dot_product_attention(
             query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
@@ -226,7 +298,6 @@ class CogVideoXBlock(nn.Module):
         ff_inner_dim: Optional[int] = None,
         ff_bias: bool = True,
         attention_out_bias: bool = True,
-        attention_mode: Optional[str] = None,
     ):
         super().__init__()
 
@@ -255,10 +326,7 @@ class CogVideoXBlock(nn.Module):
             inner_dim=ff_inner_dim,
             bias=ff_bias,
         )
-        self.cached_hidden_states = []
-        self.cached_encoder_hidden_states = []
-        self.attention_mode = attention_mode
-        
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -267,7 +335,6 @@ class CogVideoXBlock(nn.Module):
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         video_flow_feature: Optional[torch.Tensor] = None,
         fuser=None,
-        block_use_fastercache=False,
         fastercache_counter=0,
         fastercache_start_step=15,
         fastercache_device="cuda:0",
@@ -285,43 +352,33 @@ class CogVideoXBlock(nn.Module):
             h = rearrange(norm_hidden_states, "B (T H W) C -> (B T) C H W", H=H, W=W)
             h = fuser(h, video_flow_feature.to(h), T=T)
             norm_hidden_states = rearrange(h, "(B T) C H W ->  B (T H W) C", T=T)
-            del h, fuser        
-        
-        #region fastercache
-        if block_use_fastercache:
-            B = norm_hidden_states.shape[0]
-            if fastercache_counter >= fastercache_start_step + 3 and fastercache_counter%3!=0 and self.cached_hidden_states[-1].shape[0] >= B:
-                attn_hidden_states = (
-                    self.cached_hidden_states[1][:B] + 
-                    (self.cached_hidden_states[1][:B] - self.cached_hidden_states[0][:B]) 
-                    * 0.3
-                    ).to(norm_hidden_states.device, non_blocking=True)
-                attn_encoder_hidden_states = (
-                    self.cached_encoder_hidden_states[1][:B] + 
-                    (self.cached_encoder_hidden_states[1][:B] - self.cached_encoder_hidden_states[0][:B])
-                    * 0.3
-                    ).to(norm_hidden_states.device, non_blocking=True)
-            else:
-                attn_hidden_states, attn_encoder_hidden_states = self.attn1(
-                    hidden_states=norm_hidden_states,
-                    encoder_hidden_states=norm_encoder_hidden_states,
-                    image_rotary_emb=image_rotary_emb,
-                    attention_mode=self.attention_mode,
-                )
-                if fastercache_counter == fastercache_start_step:
-                    self.cached_hidden_states = [attn_hidden_states.to(fastercache_device), attn_hidden_states.to(fastercache_device)]
-                    self.cached_encoder_hidden_states = [attn_encoder_hidden_states.to(fastercache_device), attn_encoder_hidden_states.to(fastercache_device)]
-                elif fastercache_counter > fastercache_start_step:
-                    self.cached_hidden_states[-1].copy_(attn_hidden_states.to(fastercache_device))
-                    self.cached_encoder_hidden_states[-1].copy_(attn_encoder_hidden_states.to(fastercache_device))
+            del h, fuser      
+        #fastercache
+        B = norm_hidden_states.shape[0]
+        if fastercache_counter >= fastercache_start_step + 3 and fastercache_counter%3!=0 and self.cached_hidden_states[-1].shape[0] >= B:
+            attn_hidden_states = (
+                self.cached_hidden_states[1][:B] + 
+                (self.cached_hidden_states[1][:B] - self.cached_hidden_states[0][:B]) 
+                * 0.3
+                ).to(norm_hidden_states.device, non_blocking=True)
+            attn_encoder_hidden_states = (
+                self.cached_encoder_hidden_states[1][:B] + 
+                (self.cached_encoder_hidden_states[1][:B] - self.cached_encoder_hidden_states[0][:B])
+                 * 0.3
+                ).to(norm_hidden_states.device, non_blocking=True)
         else:
             attn_hidden_states, attn_encoder_hidden_states = self.attn1(
                 hidden_states=norm_hidden_states,
                 encoder_hidden_states=norm_encoder_hidden_states,
                 image_rotary_emb=image_rotary_emb,
-                attention_mode=self.attention_mode,
             )
-        
+            if fastercache_counter == fastercache_start_step:
+                self.cached_hidden_states = [attn_hidden_states.to(fastercache_device), attn_hidden_states.to(fastercache_device)]
+                self.cached_encoder_hidden_states = [attn_encoder_hidden_states.to(fastercache_device), attn_encoder_hidden_states.to(fastercache_device)]
+            elif fastercache_counter > fastercache_start_step:
+                self.cached_hidden_states[-1].copy_(attn_hidden_states.to(fastercache_device))
+                self.cached_encoder_hidden_states[-1].copy_(attn_encoder_hidden_states.to(fastercache_device))
+
         hidden_states = hidden_states + gate_msa * attn_hidden_states
         encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states
 
@@ -424,7 +481,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin):
         temporal_interpolation_scale: float = 1.0,
         use_rotary_positional_embeddings: bool = False,
         add_noise_in_inpaint_model: bool = False,
-        attention_mode: Optional[str] = None,
     ):
         super().__init__()
         inner_dim = num_attention_heads * attention_head_dim
@@ -498,8 +554,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin):
         self.fastercache_lf_step = 40
         self.fastercache_hf_step = 30
         self.fastercache_device = "cuda"
-        self.fastercache_num_blocks_to_cache = len(self.transformer_blocks)
-        self.attention_mode = attention_mode
 
     def _set_gradient_checkpointing(self, module, value=False):
         self.gradient_checkpointing = value
@@ -666,8 +720,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin):
                     image_rotary_emb=image_rotary_emb,
                     video_flow_feature=video_flow_features[i][:1] if video_flow_features is not None else None,
                     fuser = self.fuser_list[i] if self.fuser_list is not None else None,
-                    block_use_fastercache = i <= self.fastercache_num_blocks_to_cache,
-                    fastercache_start_step = self.fastercache_start_step,
                     fastercache_counter = self.fastercache_counter,
                     fastercache_device = self.fastercache_device
                 )
@@ -718,9 +770,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin):
                     image_rotary_emb=image_rotary_emb,
                     video_flow_feature=video_flow_features[i] if video_flow_features is not None else None,
                     fuser = self.fuser_list[i] if self.fuser_list is not None else None,
-                    block_use_fastercache = i <= self.fastercache_num_blocks_to_cache,
                     fastercache_counter = self.fastercache_counter,
-                    fastercache_start_step = self.fastercache_start_step,
                     fastercache_device = self.fastercache_device
                 )
 

convert_weight_sat2hf.py

@@ -1,303 +0,0 @@
-"""
-
-The script demonstrates how to convert the weights of the CogVideoX model from SAT to Hugging Face format.
-This script supports the conversion of the following models:
-- CogVideoX-2B
-- CogVideoX-5B, CogVideoX-5B-I2V
-- CogVideoX1.1-5B, CogVideoX1.1-5B-I2V
-
-Original Script:
-https://github.com/huggingface/diffusers/blob/main/scripts/convert_cogvideox_to_diffusers.py
-
-"""
-import argparse
-from typing import Any, Dict
-
-import torch
-from transformers import T5EncoderModel, T5Tokenizer
-
-from diffusers import (
-    AutoencoderKLCogVideoX,
-    CogVideoXDDIMScheduler,
-    CogVideoXImageToVideoPipeline,
-    CogVideoXPipeline,
-    #CogVideoXTransformer3DModel,
-)
-from custom_cogvideox_transformer_3d import CogVideoXTransformer3DModel
-
-
-def reassign_query_key_value_inplace(key: str, state_dict: Dict[str, Any]):
-    to_q_key = key.replace("query_key_value", "to_q")
-    to_k_key = key.replace("query_key_value", "to_k")
-    to_v_key = key.replace("query_key_value", "to_v")
-    to_q, to_k, to_v = torch.chunk(state_dict[key], chunks=3, dim=0)
-    state_dict[to_q_key] = to_q
-    state_dict[to_k_key] = to_k
-    state_dict[to_v_key] = to_v
-    state_dict.pop(key)
-
-
-def reassign_query_key_layernorm_inplace(key: str, state_dict: Dict[str, Any]):
-    layer_id, weight_or_bias = key.split(".")[-2:]
-
-    if "query" in key:
-        new_key = f"transformer_blocks.{layer_id}.attn1.norm_q.{weight_or_bias}"
-    elif "key" in key:
-        new_key = f"transformer_blocks.{layer_id}.attn1.norm_k.{weight_or_bias}"
-
-    state_dict[new_key] = state_dict.pop(key)
-
-
-def reassign_adaln_norm_inplace(key: str, state_dict: Dict[str, Any]):
-    layer_id, _, weight_or_bias = key.split(".")[-3:]
-
-    weights_or_biases = state_dict[key].chunk(12, dim=0)
-    norm1_weights_or_biases = torch.cat(weights_or_biases[0:3] + weights_or_biases[6:9])
-    norm2_weights_or_biases = torch.cat(weights_or_biases[3:6] + weights_or_biases[9:12])
-
-    norm1_key = f"transformer_blocks.{layer_id}.norm1.linear.{weight_or_bias}"
-    state_dict[norm1_key] = norm1_weights_or_biases
-
-    norm2_key = f"transformer_blocks.{layer_id}.norm2.linear.{weight_or_bias}"
-    state_dict[norm2_key] = norm2_weights_or_biases
-
-    state_dict.pop(key)
-
-
-def remove_keys_inplace(key: str, state_dict: Dict[str, Any]):
-    state_dict.pop(key)
-
-
-def replace_up_keys_inplace(key: str, state_dict: Dict[str, Any]):
-    key_split = key.split(".")
-    layer_index = int(key_split[2])
-    replace_layer_index = 4 - 1 - layer_index
-
-    key_split[1] = "up_blocks"
-    key_split[2] = str(replace_layer_index)
-    new_key = ".".join(key_split)
-
-    state_dict[new_key] = state_dict.pop(key)
-
-
-TRANSFORMER_KEYS_RENAME_DICT = {
-    "transformer.final_layernorm": "norm_final",
-    "transformer": "transformer_blocks",
-    "attention": "attn1",
-    "mlp": "ff.net",
-    "dense_h_to_4h": "0.proj",
-    "dense_4h_to_h": "2",
-    ".layers": "",
-    "dense": "to_out.0",
-    "input_layernorm": "norm1.norm",
-    "post_attn1_layernorm": "norm2.norm",
-    "time_embed.0": "time_embedding.linear_1",
-    "time_embed.2": "time_embedding.linear_2",
-    "mixins.patch_embed": "patch_embed",
-    "mixins.final_layer.norm_final": "norm_out.norm",
-    "mixins.final_layer.linear": "proj_out",
-    "mixins.final_layer.adaLN_modulation.1": "norm_out.linear",
-    "mixins.pos_embed.pos_embedding": "patch_embed.pos_embedding",  # Specific to CogVideoX-5b-I2V
-}
-
-TRANSFORMER_SPECIAL_KEYS_REMAP = {
-    "query_key_value": reassign_query_key_value_inplace,
-    "query_layernorm_list": reassign_query_key_layernorm_inplace,
-    "key_layernorm_list": reassign_query_key_layernorm_inplace,
-    "adaln_layer.adaLN_modulations": reassign_adaln_norm_inplace,
-    "embed_tokens": remove_keys_inplace,
-    "freqs_sin": remove_keys_inplace,
-    "freqs_cos": remove_keys_inplace,
-    "position_embedding": remove_keys_inplace,
-}
-
-VAE_KEYS_RENAME_DICT = {
-    "block.": "resnets.",
-    "down.": "down_blocks.",
-    "downsample": "downsamplers.0",
-    "upsample": "upsamplers.0",
-    "nin_shortcut": "conv_shortcut",
-    "encoder.mid.block_1": "encoder.mid_block.resnets.0",
-    "encoder.mid.block_2": "encoder.mid_block.resnets.1",
-    "decoder.mid.block_1": "decoder.mid_block.resnets.0",
-    "decoder.mid.block_2": "decoder.mid_block.resnets.1",
-}
-
-VAE_SPECIAL_KEYS_REMAP = {
-    "loss": remove_keys_inplace,
-    "up.": replace_up_keys_inplace,
-}
-
-TOKENIZER_MAX_LENGTH = 226
-
-
-def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
-    state_dict = saved_dict
-    if "model" in saved_dict.keys():
-        state_dict = state_dict["model"]
-    if "module" in saved_dict.keys():
-        state_dict = state_dict["module"]
-    if "state_dict" in saved_dict.keys():
-        state_dict = state_dict["state_dict"]
-    return state_dict
-
-
-def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
-    state_dict[new_key] = state_dict.pop(old_key)
-
-
-def convert_transformer(
-        ckpt_path: str,
-        num_layers: int,
-        num_attention_heads: int,
-        use_rotary_positional_embeddings: bool,
-        i2v: bool,
-        dtype: torch.dtype,
-):
-    PREFIX_KEY = "model.diffusion_model."
-
-    original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True))
-    transformer = CogVideoXTransformer3DModel(
-        in_channels=32 if i2v else 16,
-        num_layers=num_layers,
-        num_attention_heads=num_attention_heads,
-        use_rotary_positional_embeddings=use_rotary_positional_embeddings,
-        use_learned_positional_embeddings=i2v,
-    ).to(dtype=dtype)
-
-    for key in list(original_state_dict.keys()):
-        new_key = key[len(PREFIX_KEY):]
-        for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():
-            new_key = new_key.replace(replace_key, rename_key)
-        update_state_dict_inplace(original_state_dict, key, new_key)
-
-    for key in list(original_state_dict.keys()):
-        for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items():
-            if special_key not in key:
-                continue
-            handler_fn_inplace(key, original_state_dict)
-    transformer.load_state_dict(original_state_dict, strict=True)
-    return transformer
-
-
-def convert_vae(ckpt_path: str, scaling_factor: float, dtype: torch.dtype):
-    original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True))
-    vae = AutoencoderKLCogVideoX(scaling_factor=scaling_factor).to(dtype=dtype)
-
-    for key in list(original_state_dict.keys()):
-        new_key = key[:]
-        for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items():
-            new_key = new_key.replace(replace_key, rename_key)
-        update_state_dict_inplace(original_state_dict, key, new_key)
-
-    for key in list(original_state_dict.keys()):
-        for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items():
-            if special_key not in key:
-                continue
-            handler_fn_inplace(key, original_state_dict)
-
-    vae.load_state_dict(original_state_dict, strict=True)
-    return vae
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint"
-    )
-    parser.add_argument("--vae_ckpt_path", type=str, default=None, help="Path to original vae checkpoint")
-    parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved")
-    parser.add_argument("--fp16", action="store_true", default=False, help="Whether to save the model weights in fp16")
-    parser.add_argument("--bf16", action="store_true", default=False, help="Whether to save the model weights in bf16")
-    parser.add_argument(
-        "--push_to_hub", action="store_true", default=False, help="Whether to push to HF Hub after saving"
-    )
-    parser.add_argument(
-        "--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory"
-    )
-    # For CogVideoX-2B, num_layers is 30. For 5B, it is 42
-    parser.add_argument("--num_layers", type=int, default=30, help="Number of transformer blocks")
-    # For CogVideoX-2B, num_attention_heads is 30. For 5B, it is 48
-    parser.add_argument("--num_attention_heads", type=int, default=30, help="Number of attention heads")
-    # For CogVideoX-2B, use_rotary_positional_embeddings is False. For 5B, it is True
-    parser.add_argument(
-        "--use_rotary_positional_embeddings", action="store_true", default=False, help="Whether to use RoPE or not"
-    )
-    # For CogVideoX-2B, scaling_factor is 1.15258426. For 5B, it is 0.7
-    parser.add_argument("--scaling_factor", type=float, default=1.15258426, help="Scaling factor in the VAE")
-    # For CogVideoX-2B, snr_shift_scale is 3.0. For 5B, it is 1.0
-    parser.add_argument("--snr_shift_scale", type=float, default=3.0, help="Scaling factor in the VAE")
-    parser.add_argument("--i2v", action="store_true", default=False, help="Whether to save the model weights in fp16")
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = get_args()
-
-    transformer = None
-    vae = None
-
-    if args.fp16 and args.bf16:
-        raise ValueError("You cannot pass both --fp16 and --bf16 at the same time.")
-
-    dtype = torch.float16 if args.fp16 else torch.bfloat16 if args.bf16 else torch.float32
-
-    if args.transformer_ckpt_path is not None:
-        transformer = convert_transformer(
-            args.transformer_ckpt_path,
-            args.num_layers,
-            args.num_attention_heads,
-            args.use_rotary_positional_embeddings,
-            args.i2v,
-            dtype,
-        )
-    if args.vae_ckpt_path is not None:
-        vae = convert_vae(args.vae_ckpt_path, args.scaling_factor, dtype)
-
-    #text_encoder_id = "/share/official_pretrains/hf_home/t5-v1_1-xxl"
-    #tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH)
-    #text_encoder = T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir)
-
-    # Apparently, the conversion does not work anymore without this :shrug:
-    #for param in text_encoder.parameters():
-    #    param.data = param.data.contiguous()
-
-    scheduler = CogVideoXDDIMScheduler.from_config(
-        {
-            "snr_shift_scale": args.snr_shift_scale,
-            "beta_end": 0.012,
-            "beta_schedule": "scaled_linear",
-            "beta_start": 0.00085,
-            "clip_sample": False,
-            "num_train_timesteps": 1000,
-            "prediction_type": "v_prediction",
-            "rescale_betas_zero_snr": True,
-            "set_alpha_to_one": True,
-            "timestep_spacing": "trailing",
-        }
-    )
-    if args.i2v:
-        pipeline_cls = CogVideoXImageToVideoPipeline
-    else:
-        pipeline_cls = CogVideoXPipeline
-
-    pipe = pipeline_cls(
-        tokenizer=None,
-        text_encoder=None,
-        vae=vae,
-        transformer=transformer,
-        scheduler=scheduler,
-    )
-
-    if args.fp16:
-        pipe = pipe.to(dtype=torch.float16)
-    if args.bf16:
-        pipe = pipe.to(dtype=torch.bfloat16)
-
-    # We don't use variant here because the model must be run in fp16 (2B) or bf16 (5B). It would be weird
-    # for users to specify variant when the default is not fp32 and they want to run with the correct default (which
-    # is either fp16/bf16 here).
-
-    # This is necessary This is necessary for users with insufficient memory,
-    # such as those using Colab and notebooks, as it can save some memory used for model loading.
-    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB", push_to_hub=args.push_to_hub)

custom_cogvideox_transformer_3d.py

@@ -27,28 +27,23 @@ from diffusers.utils import logging
 from diffusers.utils.torch_utils import maybe_allow_in_graph
 from diffusers.models.attention import Attention, FeedForward
 from diffusers.models.attention_processor import AttentionProcessor
-from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.embeddings import CogVideoXPatchEmbed, TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.modeling_utils import ModelMixin
 from diffusers.models.normalization import AdaLayerNorm, CogVideoXLayerNormZero
 from diffusers.loaders import PeftAdapterMixin
-from diffusers.models.embeddings import apply_rotary_emb
-from .embeddings import CogVideoXPatchEmbed
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 try:
     from sageattention import sageattn
-    
     SAGEATTN_IS_AVAILABLE = True
+    logger.info("Using sageattn")
 except:
+    logger.info("sageattn not found, using sdpa")
     SAGEATTN_IS_AVAILABLE = False
 
-@torch.compiler.disable()
-def sageattn_func(query, key, value, attn_mask=None, dropout_p=0.0,is_causal=False):
-    return sageattn(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p,is_causal=is_causal)
-
 def fft(tensor):
     tensor_fft = torch.fft.fft2(tensor)
     tensor_fft_shifted = torch.fft.fftshift(tensor_fft)
@@ -75,8 +70,7 @@ class CogVideoXAttnProcessor2_0:
     def __init__(self):
         if not hasattr(F, "scaled_dot_product_attention"):
             raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
-        
-    #@torch.compiler.disable()
+    @torch.compiler.disable()
     def __call__(
         self,
         attn: Attention,
@@ -84,7 +78,6 @@ class CogVideoXAttnProcessor2_0:
         encoder_hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
-        attention_mode: Optional[str] = None,
     ) -> torch.Tensor:
         text_seq_length = encoder_hidden_states.size(1)
 
@@ -98,14 +91,9 @@ class CogVideoXAttnProcessor2_0:
             attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
             attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
 
-        if attention_mode != "fused_sdpa" or attention_mode != "fused_sageattn":
-            query = attn.to_q(hidden_states)
-            key = attn.to_k(hidden_states)
-            value = attn.to_v(hidden_states)
-        else:
-            qkv = attn.to_qkv(hidden_states)
-            split_size = qkv.shape[-1] // 3
-            query, key, value = torch.split(qkv, split_size, dim=-1)
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
 
         inner_dim = key.shape[-1]
         head_dim = inner_dim // attn.heads
@@ -121,18 +109,18 @@ class CogVideoXAttnProcessor2_0:
 
         # Apply RoPE if needed
         if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
             query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
             if not attn.is_cross_attention:
                 key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
-                 
-        if attention_mode == "sageattn" or attention_mode == "fused_sageattn":
-            hidden_states = sageattn_func(query, key, value, attn_mask=attention_mask, dropout_p=0.0,is_causal=False)
+
+        if SAGEATTN_IS_AVAILABLE:
+            hidden_states = sageattn(query, key, value, is_causal=False)
         else:
             hidden_states = F.scaled_dot_product_attention(
-                query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-                )
-        #if torch.isinf(hidden_states).any():
-        #    raise ValueError(f"hidden_states after dot product has inf")
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
 
         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
 
@@ -146,7 +134,80 @@ class CogVideoXAttnProcessor2_0:
         )
         return hidden_states, encoder_hidden_states
 
-#region Blocks
+
+class FusedCogVideoXAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
+    query and key vectors, but does not include spatial normalization.
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+    @torch.compiler.disable()
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        qkv = attn.to_qkv(hidden_states)
+        split_size = qkv.shape[-1] // 3
+        query, key, value = torch.split(qkv, split_size, dim=-1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply RoPE if needed
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
+            if not attn.is_cross_attention:
+                key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
+
+        if SAGEATTN_IS_AVAILABLE:
+            hidden_states = sageattn(query, key, value, is_causal=False)
+        else:
+            hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+    
 @maybe_allow_in_graph
 class CogVideoXBlock(nn.Module):
     
@@ -205,7 +266,6 @@ class CogVideoXBlock(nn.Module):
 
         # 1. Self Attention
         self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
-       
 
         self.attn1 = Attention(
             query_dim=dim,
@@ -240,21 +300,15 @@ class CogVideoXBlock(nn.Module):
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         video_flow_feature: Optional[torch.Tensor] = None,
         fuser=None,
-        block_use_fastercache=False,
         fastercache_counter=0,
         fastercache_start_step=15,
         fastercache_device="cuda:0",
-        attention_mode="sdpa",
     ) -> torch.Tensor:
-        #print("hidden_states in block: ", hidden_states.shape) #1.5: torch.Size([2, 3200, 3072]) 10.: torch.Size([2, 6400, 3072])
         text_seq_length = encoder_hidden_states.size(1)
-
         # norm & modulate
         norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
             hidden_states, encoder_hidden_states, temb
         )
-        #print("norm_hidden_states in block: ", norm_hidden_states.shape) #torch.Size([2, 3200, 3072])
-      
         # Tora Motion-guidance Fuser
         if video_flow_feature is not None:
             H, W = video_flow_feature.shape[-2:]
@@ -264,46 +318,36 @@ class CogVideoXBlock(nn.Module):
             norm_hidden_states = rearrange(h, "(B T) C H W ->  B (T H W) C", T=T)
             del h, fuser        
         
-        #region fastercache
-        if block_use_fastercache:
-            B = norm_hidden_states.shape[0]
-            if fastercache_counter >= fastercache_start_step + 3 and fastercache_counter%3!=0 and self.cached_hidden_states[-1].shape[0] >= B:
-                attn_hidden_states = (
-                    self.cached_hidden_states[1][:B] + 
-                    (self.cached_hidden_states[1][:B] - self.cached_hidden_states[0][:B]) 
-                    * 0.3
-                    ).to(norm_hidden_states.device, non_blocking=True)
-                attn_encoder_hidden_states = (
-                    self.cached_encoder_hidden_states[1][:B] + 
-                    (self.cached_encoder_hidden_states[1][:B] - self.cached_encoder_hidden_states[0][:B])
-                    * 0.3
-                    ).to(norm_hidden_states.device, non_blocking=True)
-            else:
-                attn_hidden_states, attn_encoder_hidden_states = self.attn1(
-                    hidden_states=norm_hidden_states,
-                    encoder_hidden_states=norm_encoder_hidden_states,
-                    image_rotary_emb=image_rotary_emb,
-                    attention_mode=attention_mode,
-                )
-                if fastercache_counter == fastercache_start_step:
-                    self.cached_hidden_states = [attn_hidden_states.to(fastercache_device), attn_hidden_states.to(fastercache_device)]
-                    self.cached_encoder_hidden_states = [attn_encoder_hidden_states.to(fastercache_device), attn_encoder_hidden_states.to(fastercache_device)]
-                elif fastercache_counter > fastercache_start_step:
-                    self.cached_hidden_states[-1].copy_(attn_hidden_states.to(fastercache_device))
-                    self.cached_encoder_hidden_states[-1].copy_(attn_encoder_hidden_states.to(fastercache_device))
+        #fastercache
+        B = norm_hidden_states.shape[0]
+        if fastercache_counter >= fastercache_start_step + 3 and fastercache_counter%3!=0 and self.cached_hidden_states[-1].shape[0] >= B:
+            attn_hidden_states = (
+                self.cached_hidden_states[1][:B] + 
+                (self.cached_hidden_states[1][:B] - self.cached_hidden_states[0][:B]) 
+                * 0.3
+                ).to(norm_hidden_states.device, non_blocking=True)
+            attn_encoder_hidden_states = (
+                self.cached_encoder_hidden_states[1][:B] + 
+                (self.cached_encoder_hidden_states[1][:B] - self.cached_encoder_hidden_states[0][:B])
+                 * 0.3
+                ).to(norm_hidden_states.device, non_blocking=True)
         else:
             attn_hidden_states, attn_encoder_hidden_states = self.attn1(
                 hidden_states=norm_hidden_states,
                 encoder_hidden_states=norm_encoder_hidden_states,
                 image_rotary_emb=image_rotary_emb,
-                attention_mode=attention_mode,
             )
-        
+            if fastercache_counter == fastercache_start_step:
+                self.cached_hidden_states = [attn_hidden_states.to(fastercache_device), attn_hidden_states.to(fastercache_device)]
+                self.cached_encoder_hidden_states = [attn_encoder_hidden_states.to(fastercache_device), attn_encoder_hidden_states.to(fastercache_device)]
+            elif fastercache_counter > fastercache_start_step:
+                self.cached_hidden_states[-1].copy_(attn_hidden_states.to(fastercache_device))
+                self.cached_encoder_hidden_states[-1].copy_(attn_encoder_hidden_states.to(fastercache_device))
+
         hidden_states = hidden_states + gate_msa * attn_hidden_states
         encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states
 
         # norm & modulate
-       
         norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
             hidden_states, encoder_hidden_states, temb
         )
@@ -317,7 +361,7 @@ class CogVideoXBlock(nn.Module):
 
         return hidden_states, encoder_hidden_states
 
-#region Transformer
+
 class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
     """
     A Transformer model for video-like data in [CogVideoX](https://github.com/THUDM/CogVideo).
@@ -384,7 +428,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         flip_sin_to_cos: bool = True,
         freq_shift: int = 0,
         time_embed_dim: int = 512,
-        ofs_embed_dim: Optional[int] = None,
         text_embed_dim: int = 4096,
         num_layers: int = 30,
         dropout: float = 0.0,
@@ -393,7 +436,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         sample_height: int = 60,
         sample_frames: int = 49,
         patch_size: int = 2,
-        patch_size_t: int = None,
         temporal_compression_ratio: int = 4,
         max_text_seq_length: int = 226,
         activation_fn: str = "gelu-approximate",
@@ -404,7 +446,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         temporal_interpolation_scale: float = 1.0,
         use_rotary_positional_embeddings: bool = False,
         use_learned_positional_embeddings: bool = False,
-        patch_bias: bool = True,
     ):
         super().__init__()
         inner_dim = num_attention_heads * attention_head_dim
@@ -419,11 +460,10 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         # 1. Patch embedding
         self.patch_embed = CogVideoXPatchEmbed(
             patch_size=patch_size,
-            patch_size_t=patch_size_t,
             in_channels=in_channels,
             embed_dim=inner_dim,
             text_embed_dim=text_embed_dim,
-            bias=patch_bias,
+            bias=True,
             sample_width=sample_width,
             sample_height=sample_height,
             sample_frames=sample_frames,
@@ -440,13 +480,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift)
         self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn)
 
-        self.ofs_proj = None
-        self.ofs_embedding = None
-
-        if ofs_embed_dim:
-            self.ofs_proj = Timesteps(ofs_embed_dim, flip_sin_to_cos, freq_shift)
-            self.ofs_embedding = TimestepEmbedding(ofs_embed_dim, ofs_embed_dim, timestep_activation_fn) # same as time embeddings, for ofs
-
         # 3. Define spatio-temporal transformers blocks
         self.transformer_blocks = nn.ModuleList(
             [
@@ -474,14 +507,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
             norm_eps=norm_eps,
             chunk_dim=1,
         )
-        if patch_size_t is None:
-            # For CogVideox 1.0
-            output_dim = patch_size * patch_size * out_channels
-        else:
-            # For CogVideoX 1.5
-            output_dim = patch_size * patch_size * patch_size_t * out_channels
-
-        self.proj_out = nn.Linear(inner_dim, output_dim)
+        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
 
         self.gradient_checkpointing = False
 
@@ -492,9 +518,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         self.fastercache_lf_step = 40
         self.fastercache_hf_step = 30
         self.fastercache_device = "cuda"
-        self.fastercache_num_blocks_to_cache = len(self.transformer_blocks)
-        self.attention_mode = "sdpa"
-        
 
     def _set_gradient_checkpointing(self, module, value=False):
         self.gradient_checkpointing = value
@@ -559,6 +582,45 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         for name, module in self.named_children():
             fn_recursive_attn_processor(name, module, processor)
 
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedCogVideoXAttnProcessor2_0
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedCogVideoXAttnProcessor2_0())
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
 
     def forward(
         self,
@@ -566,7 +628,6 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         encoder_hidden_states: torch.Tensor,
         timestep: Union[int, float, torch.LongTensor],
         timestep_cond: Optional[torch.Tensor] = None,
-        ofs: Optional[Union[int, float, torch.LongTensor]] = None,
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         controlnet_states: torch.Tensor = None,
         controlnet_weights: Optional[Union[float, int, list, np.ndarray, torch.FloatTensor]] = 1.0,
@@ -574,7 +635,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         return_dict: bool = True,
     ):
         batch_size, num_frames, channels, height, width = hidden_states.shape
-   
+
         # 1. Time embedding
         timesteps = timestep
         t_emb = self.time_proj(timesteps)
@@ -583,29 +644,15 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         # but time_embedding might actually be running in fp16. so we need to cast here.
         # there might be better ways to encapsulate this.
         t_emb = t_emb.to(dtype=hidden_states.dtype)
-
         emb = self.time_embedding(t_emb, timestep_cond)
-        if self.ofs_embedding is not None: #1.5 I2V
-            ofs_emb = self.ofs_proj(ofs)
-            ofs_emb = ofs_emb.to(dtype=hidden_states.dtype)
-            ofs_emb = self.ofs_embedding(ofs_emb)
-            emb = emb + ofs_emb
 
         # 2. Patch embedding
-        p = self.config.patch_size
-        p_t = self.config.patch_size_t
-
-        #print("hidden_states before patch_embedding", hidden_states.shape) #torch.Size([2, 4, 16, 60, 90])
-        
         hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)
-        #print("hidden_states after patch_embedding", hidden_states.shape) #1.5: torch.Size([2, 2926, 3072]) #1.0: torch.Size([2, 5626, 3072])
         hidden_states = self.embedding_dropout(hidden_states)
 
         text_seq_length = encoder_hidden_states.shape[1]
         encoder_hidden_states = hidden_states[:, :text_seq_length]
         hidden_states = hidden_states[:, text_seq_length:]
-        #print("hidden_states after split", hidden_states.shape) #1.5: torch.Size([2, 2700, 3072]) #1.0: torch.Size([2, 5400, 3072])
-      
         if self.use_fastercache:
             self.fastercache_counter+=1
         if self.fastercache_counter >= self.fastercache_start_step + 3 and self.fastercache_counter % 5 !=0:
@@ -618,11 +665,8 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
                         image_rotary_emb=image_rotary_emb,
                         video_flow_feature=video_flow_features[i][:1] if video_flow_features is not None else None,
                         fuser = self.fuser_list[i] if self.fuser_list is not None else None,
-                        block_use_fastercache = i <= self.fastercache_num_blocks_to_cache,
                         fastercache_counter = self.fastercache_counter,
-                        fastercache_start_step = self.fastercache_start_step,
-                        fastercache_device = self.fastercache_device,
-                        attention_mode = self.attention_mode
+                        fastercache_device = self.fastercache_device
                     )
 
                     if (controlnet_states is not None) and (i < len(controlnet_states)):
@@ -652,15 +696,9 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
             # Note: we use `-1` instead of `channels`:
             #   - It is okay to `channels` use for CogVideoX-2b and CogVideoX-5b (number of input channels is equal to output channels)
             #   - However, for CogVideoX-5b-I2V also takes concatenated input image latents (number of input channels is twice the output channels)
-            
-            if p_t is None:
-                output = hidden_states.reshape(1, num_frames, height // p, width // p, -1, p, p)
-                output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
-            else:
-                output = hidden_states.reshape(
-                    1, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p
-                )
-                output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2)
+            p = self.config.patch_size
+            output = hidden_states.reshape(1, num_frames, height // p, width // p, -1, p, p)
+            output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
             
             (bb, tt, cc, hh, ww) = output.shape
             cond = rearrange(output, "B T C H W -> (B T) C H W", B=bb, C=cc, T=tt, H=hh, W=ww)
@@ -689,15 +727,9 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
                     image_rotary_emb=image_rotary_emb,
                     video_flow_feature=video_flow_features[i] if video_flow_features is not None else None,
                     fuser = self.fuser_list[i] if self.fuser_list is not None else None,
-                    block_use_fastercache = i <= self.fastercache_num_blocks_to_cache,
                     fastercache_counter = self.fastercache_counter,
-                    fastercache_start_step = self.fastercache_start_step,
-                    fastercache_device = self.fastercache_device,
-                    attention_mode = self.attention_mode
+                    fastercache_device = self.fastercache_device
                 )
-                #has_nan = torch.isnan(hidden_states).any()
-                #if has_nan:
-                #    raise ValueError(f"block output hidden_states has nan: {has_nan}")
 
             if (controlnet_states is not None) and (i < len(controlnet_states)):
                 controlnet_states_block = controlnet_states[i]
@@ -726,15 +758,9 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
             # Note: we use `-1` instead of `channels`:
             #   - It is okay to `channels` use for CogVideoX-2b and CogVideoX-5b (number of input channels is equal to output channels)
             #   - However, for CogVideoX-5b-I2V also takes concatenated input image latents (number of input channels is twice the output channels)
-           
-            if p_t is None:
-                output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)
-                output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
-            else:
-                output = hidden_states.reshape(
-                    batch_size, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p
-                )
-                output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2)
+            p = self.config.patch_size
+            output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)
+            output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
 
             if self.fastercache_counter >= self.fastercache_start_step + 1: 
                 (bb, tt, cc, hh, ww) = output.shape

embeddings.py

@@ -1,226 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-from typing import Tuple, Union, Optional
-from diffusers.models.embeddings import get_3d_sincos_pos_embed, get_1d_rotary_pos_embed
-
-
-class CogVideoXPatchEmbed(nn.Module):
-    def __init__(
-        self,
-        patch_size: int = 2,
-        patch_size_t: Optional[int] = None,
-        in_channels: int = 16,
-        embed_dim: int = 1920,
-        text_embed_dim: int = 4096,
-        bias: bool = True,
-        sample_width: int = 90,
-        sample_height: int = 60,
-        sample_frames: int = 49,
-        temporal_compression_ratio: int = 4,
-        max_text_seq_length: int = 226,
-        spatial_interpolation_scale: float = 1.875,
-        temporal_interpolation_scale: float = 1.0,
-        use_positional_embeddings: bool = True,
-        use_learned_positional_embeddings: bool = True,
-    ) -> None:
-        super().__init__()
-
-        self.patch_size = patch_size
-        self.patch_size_t = patch_size_t
-        self.embed_dim = embed_dim
-        self.sample_height = sample_height
-        self.sample_width = sample_width
-        self.sample_frames = sample_frames
-        self.temporal_compression_ratio = temporal_compression_ratio
-        self.max_text_seq_length = max_text_seq_length
-        self.spatial_interpolation_scale = spatial_interpolation_scale
-        self.temporal_interpolation_scale = temporal_interpolation_scale
-        self.use_positional_embeddings = use_positional_embeddings
-        self.use_learned_positional_embeddings = use_learned_positional_embeddings
-
-        if patch_size_t is None:
-            # CogVideoX 1.0 checkpoints
-            self.proj = nn.Conv2d(
-                in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias
-            )
-        else:
-            # CogVideoX 1.5 checkpoints
-            self.proj = nn.Linear(in_channels * patch_size * patch_size * patch_size_t, embed_dim)
-
-        self.text_proj = nn.Linear(text_embed_dim, embed_dim)
-
-        if use_positional_embeddings or use_learned_positional_embeddings:
-            persistent = use_learned_positional_embeddings
-            pos_embedding = self._get_positional_embeddings(sample_height, sample_width, sample_frames)
-            self.register_buffer("pos_embedding", pos_embedding, persistent=persistent)
-
-    def _get_positional_embeddings(self, sample_height: int, sample_width: int, sample_frames: int) -> torch.Tensor:
-        post_patch_height = sample_height // self.patch_size
-        post_patch_width = sample_width // self.patch_size
-        post_time_compression_frames = (sample_frames - 1) // self.temporal_compression_ratio + 1
-        num_patches = post_patch_height * post_patch_width * post_time_compression_frames
-
-        pos_embedding = get_3d_sincos_pos_embed(
-            self.embed_dim,
-            (post_patch_width, post_patch_height),
-            post_time_compression_frames,
-            self.spatial_interpolation_scale,
-            self.temporal_interpolation_scale,
-        )
-        pos_embedding = torch.from_numpy(pos_embedding).flatten(0, 1)
-        joint_pos_embedding = torch.zeros(
-            1, self.max_text_seq_length + num_patches, self.embed_dim, requires_grad=False
-        )
-        joint_pos_embedding.data[:, self.max_text_seq_length :].copy_(pos_embedding)
-
-        return joint_pos_embedding
-
-    def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):
-        r"""
-        Args:
-            text_embeds (`torch.Tensor`):
-                Input text embeddings. Expected shape: (batch_size, seq_length, embedding_dim).
-            image_embeds (`torch.Tensor`):
-                Input image embeddings. Expected shape: (batch_size, num_frames, channels, height, width).
-        """
-        text_embeds = self.text_proj(text_embeds)
-
-        batch_size, num_frames, channels, height, width = image_embeds.shape
-
-        if self.patch_size_t is None:
-            image_embeds = image_embeds.reshape(-1, channels, height, width)
-            image_embeds = self.proj(image_embeds)
-            image_embeds = image_embeds.view(batch_size, num_frames, *image_embeds.shape[1:])
-            image_embeds = image_embeds.flatten(3).transpose(2, 3)  # [batch, num_frames, height x width, channels]
-            image_embeds = image_embeds.flatten(1, 2)  # [batch, num_frames x height x width, channels]
-        else:
-            p = self.patch_size
-            p_t = self.patch_size_t
-
-            image_embeds = image_embeds.permute(0, 1, 3, 4, 2)
-            image_embeds = image_embeds.reshape(
-                batch_size, num_frames // p_t, p_t, height // p, p, width // p, p, channels
-            )
-            image_embeds = image_embeds.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(4, 7).flatten(1, 3)
-            image_embeds = self.proj(image_embeds)
-
-        embeds = torch.cat(
-            [text_embeds, image_embeds], dim=1
-        ).contiguous()  # [batch, seq_length + num_frames x height x width, channels]
-
-        if self.use_positional_embeddings or self.use_learned_positional_embeddings:
-            if self.use_learned_positional_embeddings and (self.sample_width != width or self.sample_height != height):
-                raise ValueError(
-                    "It is currently not possible to generate videos at a different resolution that the defaults. This should only be the case with 'THUDM/CogVideoX-5b-I2V'."
-                    "If you think this is incorrect, please open an issue at https://github.com/huggingface/diffusers/issues."
-                )
-
-            pre_time_compression_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
-
-            if (
-                self.sample_height != height
-                or self.sample_width != width
-                or self.sample_frames != pre_time_compression_frames
-            ):
-                pos_embedding = self._get_positional_embeddings(height, width, pre_time_compression_frames)
-                pos_embedding = pos_embedding.to(embeds.device, dtype=embeds.dtype)
-            else:
-                pos_embedding = self.pos_embedding
-
-            embeds = embeds + pos_embedding
-
-        return embeds
-
-def get_3d_rotary_pos_embed(
-    embed_dim,
-    crops_coords,
-    grid_size,
-    temporal_size,
-    theta: int = 10000,
-    use_real: bool = True,
-    grid_type: str = "linspace",
-    max_size: Optional[Tuple[int, int]] = None,
-) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-    """
-    RoPE for video tokens with 3D structure.
-
-    Args:
-    embed_dim: (`int`):
-        The embedding dimension size, corresponding to hidden_size_head.
-    crops_coords (`Tuple[int]`):
-        The top-left and bottom-right coordinates of the crop.
-    grid_size (`Tuple[int]`):
-        The grid size of the spatial positional embedding (height, width).
-    temporal_size (`int`):
-        The size of the temporal dimension.
-    theta (`float`):
-        Scaling factor for frequency computation.
-    grid_type (`str`):
-        Whether to use "linspace" or "slice" to compute grids.
-
-    Returns:
-        `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.
-    """
-    if use_real is not True:
-        raise ValueError(" `use_real = False` is not currently supported for get_3d_rotary_pos_embed")
-
-    if grid_type == "linspace":
-        start, stop = crops_coords
-        grid_size_h, grid_size_w = grid_size
-        grid_h = np.linspace(start[0], stop[0], grid_size_h, endpoint=False, dtype=np.float32)
-        grid_w = np.linspace(start[1], stop[1], grid_size_w, endpoint=False, dtype=np.float32)
-        grid_t = np.arange(temporal_size, dtype=np.float32)
-        grid_t = np.linspace(0, temporal_size, temporal_size, endpoint=False, dtype=np.float32)
-    elif grid_type == "slice":
-        max_h, max_w = max_size
-        grid_size_h, grid_size_w = grid_size
-        grid_h = np.arange(max_h, dtype=np.float32)
-        grid_w = np.arange(max_w, dtype=np.float32)
-        grid_t = np.arange(temporal_size, dtype=np.float32)
-    else:
-        raise ValueError("Invalid value passed for `grid_type`.")
-
-    # Compute dimensions for each axis
-    dim_t = embed_dim // 4
-    dim_h = embed_dim // 8 * 3
-    dim_w = embed_dim // 8 * 3
-
-    # Temporal frequencies
-    freqs_t = get_1d_rotary_pos_embed(dim_t, grid_t, use_real=True)
-    # Spatial frequencies for height and width
-    freqs_h = get_1d_rotary_pos_embed(dim_h, grid_h, use_real=True)
-    freqs_w = get_1d_rotary_pos_embed(dim_w, grid_w, use_real=True)
-
-    # BroadCast and concatenate temporal and spaial frequencie (height and width) into a 3d tensor
-    def combine_time_height_width(freqs_t, freqs_h, freqs_w):
-        freqs_t = freqs_t[:, None, None, :].expand(
-            -1, grid_size_h, grid_size_w, -1
-        )  # temporal_size, grid_size_h, grid_size_w, dim_t
-        freqs_h = freqs_h[None, :, None, :].expand(
-            temporal_size, -1, grid_size_w, -1
-        )  # temporal_size, grid_size_h, grid_size_2, dim_h
-        freqs_w = freqs_w[None, None, :, :].expand(
-            temporal_size, grid_size_h, -1, -1
-        )  # temporal_size, grid_size_h, grid_size_2, dim_w
-
-        freqs = torch.cat(
-            [freqs_t, freqs_h, freqs_w], dim=-1
-        )  # temporal_size, grid_size_h, grid_size_w, (dim_t + dim_h + dim_w)
-        freqs = freqs.view(
-            temporal_size * grid_size_h * grid_size_w, -1
-        )  # (temporal_size * grid_size_h * grid_size_w), (dim_t + dim_h + dim_w)
-        return freqs
-
-    t_cos, t_sin = freqs_t  # both t_cos and t_sin has shape: temporal_size, dim_t
-    h_cos, h_sin = freqs_h  # both h_cos and h_sin has shape: grid_size_h, dim_h
-    w_cos, w_sin = freqs_w  # both w_cos and w_sin has shape: grid_size_w, dim_w
-
-    if grid_type == "slice":
-        t_cos, t_sin = t_cos[:temporal_size], t_sin[:temporal_size]
-        h_cos, h_sin = h_cos[:grid_size_h], h_sin[:grid_size_h]
-        w_cos, w_sin = w_cos[:grid_size_w], w_sin[:grid_size_w]
-
-    cos = combine_time_height_width(t_cos, h_cos, w_cos)
-    sin = combine_time_height_width(t_sin, h_sin, w_sin)
-    return cos, sin

fp8_optimization.py

@@ -36,12 +36,10 @@ def fp8_linear_forward(cls, original_dtype, input):
     else:
         return cls.original_forward(input)
 
-def convert_fp8_linear(module, original_dtype, params_to_keep={}):
+def convert_fp8_linear(module, original_dtype):
     setattr(module, "fp8_matmul_enabled", True)
-   
     for name, module in module.named_modules():
-        if not any(keyword in name for keyword in params_to_keep):
-            if isinstance(module, nn.Linear):
-                original_forward = module.forward
-                setattr(module, "original_forward", original_forward)
-                setattr(module, "forward", lambda input, m=module: fp8_linear_forward(m, original_dtype, input))
+        if isinstance(module, nn.Linear):
+            original_forward = module.forward
+            setattr(module, "original_forward", original_forward)
+            setattr(module, "forward", lambda input, m=module: fp8_linear_forward(m, original_dtype, input))

model_loading.py

@@ -1,41 +1,11 @@
 import os
+import torch
+import torch.nn as nn
 import json
 import folder_paths
 import comfy.model_management as mm
-from typing import Union
-
-def patched_write_atomic(
-    path_: str,
-    content: Union[str, bytes],
-    make_dirs: bool = False,
-    encode_utf_8: bool = False,
-) -> None:
-    # Write into temporary file first to avoid conflicts between threads
-    # Avoid using a named temporary file, as those have restricted permissions
-    from pathlib import Path
-    import os
-    import shutil
-    import threading
-    assert isinstance(
-        content, (str, bytes)
-    ), "Only strings and byte arrays can be saved in the cache"
-    path = Path(path_)
-    if make_dirs:
-        path.parent.mkdir(parents=True, exist_ok=True)
-    tmp_path = path.parent / f".{os.getpid()}.{threading.get_ident()}.tmp"
-    write_mode = "w" if isinstance(content, str) else "wb"
-    with tmp_path.open(write_mode, encoding="utf-8" if encode_utf_8 else None) as f:
-        f.write(content)
-    shutil.copy2(src=tmp_path, dst=path) #changed to allow overwriting cache files
-    os.remove(tmp_path)
-try:
-    import torch._inductor.codecache
-    torch._inductor.codecache.write_atomic = patched_write_atomic
-except:
-    pass
-
-import torch
-import torch.nn as nn
+from .utils import check_diffusers_version, remove_specific_blocks, log
+check_diffusers_version()
 
 from diffusers.models import AutoencoderKLCogVideoX
 from diffusers.schedulers import CogVideoXDDIMScheduler
@@ -44,12 +14,14 @@ from .pipeline_cogvideox import CogVideoXPipeline
 from contextlib import nullcontext
 
 from .cogvideox_fun.transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelFun
+from .cogvideox_fun.fun_pab_transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelFunPAB
 from .cogvideox_fun.autoencoder_magvit import AutoencoderKLCogVideoX as AutoencoderKLCogVideoXFun
 
 from .cogvideox_fun.pipeline_cogvideox_inpaint import CogVideoX_Fun_Pipeline_Inpaint
 from .cogvideox_fun.pipeline_cogvideox_control import CogVideoX_Fun_Pipeline_Control
 
-from .utils import remove_specific_blocks, log
+from .videosys.cogvideox_transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelPAB
+
 from comfy.utils import load_torch_file
 
 script_directory = os.path.dirname(os.path.abspath(__file__))
@@ -89,8 +61,7 @@ class CogVideoLoraSelect:
         cog_loras_list.append(cog_lora)
         print(cog_loras_list)
         return (cog_loras_list,)
-    
-#region DownloadAndLoadCogVideoModel
+
 class DownloadAndLoadCogVideoModel:
     @classmethod
     def INPUT_TYPES(s):
@@ -101,8 +72,6 @@ class DownloadAndLoadCogVideoModel:
                         "THUDM/CogVideoX-2b",
                         "THUDM/CogVideoX-5b",
                         "THUDM/CogVideoX-5b-I2V",
-                        "kijai/CogVideoX-5b-1.5-T2V",
-                        "kijai/CogVideoX-5b-1.5-I2V",
                         "bertjiazheng/KoolCogVideoX-5b",
                         "kijai/CogVideoX-Fun-2b",
                         "kijai/CogVideoX-Fun-5b",
@@ -121,14 +90,13 @@ class DownloadAndLoadCogVideoModel:
                 "precision": (["fp16", "fp32", "bf16"],
                     {"default": "bf16", "tooltip": "official recommendation is that 2b model should be fp16, 5b model should be bf16"}
                 ),
-                "fp8_transformer": (['disabled', 'enabled', 'fastmode', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6"], {"default": 'disabled', "tooltip": "enabled casts the transformer to torch.float8_e4m3fn, fastmode is only for latest nvidia GPUs and requires torch 2.4.0 and cu124 minimum"}),
+                "fp8_transformer": (['disabled', 'enabled', 'fastmode'], {"default": 'disabled', "tooltip": "enabled casts the transformer to torch.float8_e4m3fn, fastmode is only for latest nvidia GPUs and requires torch 2.4.0 and cu124 minimum"}),
                 "compile": (["disabled","onediff","torch"], {"tooltip": "compile the model for faster inference, these are advanced options only available on Linux, see readme for more info"}),
                 "enable_sequential_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "significantly reducing memory usage and slows down the inference"}),
+                "pab_config": ("PAB_CONFIG", {"default": None}),
                 "block_edit": ("TRANSFORMERBLOCKS", {"default": None}),
                 "lora": ("COGLORA", {"default": None}),
                 "compile_args":("COMPILEARGS", ),
-                "attention_mode": (["sdpa", "sageattn", "fused_sdpa", "fused_sageattn"], {"default": "sdpa"}),
-                "load_device": (["main_device", "offload_device"], {"default": "main_device"}),
             }
         }
 
@@ -138,17 +106,12 @@ class DownloadAndLoadCogVideoModel:
     CATEGORY = "CogVideoWrapper"
     DESCRIPTION = "Downloads and loads the selected CogVideo model from Huggingface to 'ComfyUI/models/CogVideo'"
 
-    def loadmodel(self, model, precision, fp8_transformer="disabled", compile="disabled", 
-                  enable_sequential_cpu_offload=False, block_edit=None, lora=None, compile_args=None, 
-                  attention_mode="sdpa", load_device="main_device"):
+    def loadmodel(self, model, precision, fp8_transformer="disabled", compile="disabled", enable_sequential_cpu_offload=False, pab_config=None, block_edit=None, lora=None, compile_args=None):
         
-        if precision == "fp16" and "1.5" in model:
-            raise ValueError("1.5 models do not currently work in fp16")
+        check_diffusers_version()
 
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
-        manual_offloading = True
-        transformer_load_device = device if load_device == "main_device" else offload_device
         mm.soft_empty_cache()
 
         dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
@@ -171,8 +134,6 @@ class DownloadAndLoadCogVideoModel:
                 if not os.path.exists(base_path):
                     base_path = os.path.join(download_path, (model.split("/")[-1]))
                 download_path = base_path
-            subfolder = "transformer"
-            allow_patterns = ["*transformer*", "*scheduler*", "*vae*"]
 
         elif "2b" in model:
             if 'img2vid' in model:
@@ -183,52 +144,46 @@ class DownloadAndLoadCogVideoModel:
                 base_path = os.path.join(download_path, "CogVideo2B")
                 download_path = base_path
                 repo_id = model
-            subfolder = "transformer"
-            allow_patterns = ["*transformer*", "*scheduler*", "*vae*"]
-        elif "1.5-T2V" in model or "1.5-I2V" in model:
-            base_path = os.path.join(download_path, "CogVideoX-5b-1.5")
-            download_path = base_path
-            subfolder = "transformer_T2V" if "1.5-T2V" in model else "transformer_I2V"
-            allow_patterns = [f"*{subfolder}*", "*vae*", "*scheduler*"]
-            repo_id = "kijai/CogVideoX-5b-1.5"
         else:
             base_path = os.path.join(download_path, (model.split("/")[-1]))
             download_path = base_path
             repo_id = model
-            subfolder = "transformer"
-            allow_patterns = ["*transformer*", "*scheduler*", "*vae*"]
+            
 
         if "2b" in model:
             scheduler_path = os.path.join(script_directory, 'configs', 'scheduler_config_2b.json')
         else:
             scheduler_path = os.path.join(script_directory, 'configs', 'scheduler_config_5b.json')
         
-        if not os.path.exists(base_path) or not os.path.exists(os.path.join(base_path, subfolder)):
+        if not os.path.exists(base_path) or not os.path.exists(os.path.join(base_path, "transformer")):
             log.info(f"Downloading model to: {base_path}")
             from huggingface_hub import snapshot_download
 
             snapshot_download(
                 repo_id=repo_id,
-                allow_patterns=allow_patterns,
                 ignore_patterns=["*text_encoder*", "*tokenizer*"],
                 local_dir=download_path,
                 local_dir_use_symlinks=False,
             )
 
-        #transformer
+        # transformer
         if "Fun" in model:
-            transformer = CogVideoXTransformer3DModelFun.from_pretrained(base_path, subfolder=subfolder)
+            if pab_config is not None:
+                transformer = CogVideoXTransformer3DModelFunPAB.from_pretrained(base_path, subfolder="transformer")
+            else:
+                transformer = CogVideoXTransformer3DModelFun.from_pretrained(base_path, subfolder="transformer")
         else:
-            transformer = CogVideoXTransformer3DModel.from_pretrained(base_path, subfolder=subfolder)
+            if pab_config is not None:
+                transformer = CogVideoXTransformer3DModelPAB.from_pretrained(base_path, subfolder="transformer")
+            else:
+                transformer = CogVideoXTransformer3DModel.from_pretrained(base_path, subfolder="transformer")
         
-        transformer = transformer.to(dtype).to(transformer_load_device)
-
-        if "1.5" in model:
-            transformer.config.sample_height = 300
-            transformer.config.sample_width = 300
+        transformer = transformer.to(dtype).to(offload_device)
 
         if block_edit is not None:
             transformer = remove_specific_blocks(transformer, block_edit)
+        
+        
 
         with open(scheduler_path) as f:
             scheduler_config = json.load(f)
@@ -238,12 +193,12 @@ class DownloadAndLoadCogVideoModel:
         if "Fun" in model:
             vae = AutoencoderKLCogVideoXFun.from_pretrained(base_path, subfolder="vae").to(dtype).to(offload_device)
             if "Pose" in model:
-                pipe = CogVideoX_Fun_Pipeline_Control(vae, transformer, scheduler)
+                pipe = CogVideoX_Fun_Pipeline_Control(vae, transformer, scheduler, pab_config=pab_config)
             else:
-                pipe = CogVideoX_Fun_Pipeline_Inpaint(vae, transformer, scheduler)
+                pipe = CogVideoX_Fun_Pipeline_Inpaint(vae, transformer, scheduler, pab_config=pab_config)
         else:
             vae = AutoencoderKLCogVideoX.from_pretrained(base_path, subfolder="vae").to(dtype).to(offload_device)
-            pipe = CogVideoXPipeline(vae, transformer, scheduler)
+            pipe = CogVideoXPipeline(vae, transformer, scheduler, pab_config=pab_config)
             if "cogvideox-2b-img2vid" in model:
                 pipe.input_with_padding = False
         
@@ -275,101 +230,26 @@ class DownloadAndLoadCogVideoModel:
                 for l in lora:
                     pipe.set_adapters(adapter_list, adapter_weights=adapter_weights)
                 if fuse:
-                    lora_scale = 1
-                    dimension_loras = ["orbit", "dimensionx"] # for now dimensionx loras need scaling
-                    if any(item in lora[-1]["path"].lower() for item in dimension_loras):
-                        lora_scale = lora_scale / lora_rank
-                    pipe.fuse_lora(lora_scale=lora_scale, components=["transformer"])
+                    pipe.fuse_lora(lora_scale=lora[-1]["strength"] / lora_rank, components=["transformer"])
 
         #fp8
         if fp8_transformer == "enabled" or fp8_transformer == "fastmode":
-            params_to_keep = {"patch_embed", "lora", "pos_embedding", "time_embedding", "norm_k", "norm_q", "to_k.bias", "to_q.bias", "to_v.bias"}
-            if "1.5" in model:
-                    params_to_keep.update({"norm1.linear.weight", "ofs_embedding", "norm_final", "norm_out", "proj_out"}) 
             for name, param in pipe.transformer.named_parameters():
+                params_to_keep = {"patch_embed", "lora", "pos_embedding"}
                 if not any(keyword in name for keyword in params_to_keep):
                     param.data = param.data.to(torch.float8_e4m3fn)
         
             if fp8_transformer == "fastmode":
                 from .fp8_optimization import convert_fp8_linear
-                if "1.5" in model:
-                    params_to_keep.update({"ff"}) #otherwise NaNs
-                convert_fp8_linear(pipe.transformer, dtype, params_to_keep=params_to_keep)
+                convert_fp8_linear(pipe.transformer, dtype)
 
-        elif "torchao" in fp8_transformer:
-            try:
-                from torchao.quantization import (
-                quantize_,
-                fpx_weight_only,
-                float8_dynamic_activation_float8_weight,
-                int8_dynamic_activation_int8_weight
-            )
-            except:
-                raise ImportError("torchao is not installed, please install torchao to use fp8dq")
-
-            def filter_fn(module: nn.Module, fqn: str) -> bool:
-                target_submodules = {'attn1', 'ff'} # avoid norm layers, 1.5 at least won't work with quantized norm1 #todo: test other models
-                if any(sub in fqn for sub in target_submodules):
-                    return isinstance(module, nn.Linear)
-                return False
-            
-            if "fp6" in fp8_transformer: #slower for some reason on 4090
-                quant_func = fpx_weight_only(3, 2)
-            elif "fp8dq" in fp8_transformer: #very fast on 4090 when compiled
-                quant_func = float8_dynamic_activation_float8_weight()
-            elif 'fp8dqrow' in fp8_transformer:
-                from torchao.quantization.quant_api import PerRow
-                quant_func = float8_dynamic_activation_float8_weight(granularity=PerRow())
-            elif 'int8dq' in fp8_transformer:
-                quant_func = int8_dynamic_activation_int8_weight()
-        
-            for i, block in enumerate(pipe.transformer.transformer_blocks):
-                if "CogVideoXBlock" in str(block):
-                    quantize_(block, quant_func, filter_fn=filter_fn)
-                        
-            manual_offloading = False # to disable manual .to(device) calls
-        
         if enable_sequential_cpu_offload:
             pipe.enable_sequential_cpu_offload()
-            manual_offloading = False
-
-        # CogVideoXBlock(
-        #     (norm1): CogVideoXLayerNormZero(
-        #         (silu): SiLU()
-        #         (linear): Linear(in_features=512, out_features=18432, bias=True)
-        #         (norm): LayerNorm((3072,), eps=1e-05, elementwise_affine=True)
-        #     )
-        #     (attn1): Attention(
-        #         (norm_q): LayerNorm((64,), eps=1e-06, elementwise_affine=True)
-        #         (norm_k): LayerNorm((64,), eps=1e-06, elementwise_affine=True)
-        #         (to_q): Linear(in_features=3072, out_features=3072, bias=True)
-        #         (to_k): Linear(in_features=3072, out_features=3072, bias=True)
-        #         (to_v): Linear(in_features=3072, out_features=3072, bias=True)
-        #         (to_out): ModuleList(
-        #         (0): Linear(in_features=3072, out_features=3072, bias=True)
-        #         (1): Dropout(p=0.0, inplace=False)
-        #         )
-        #     )
-        #     (norm2): CogVideoXLayerNormZero(
-        #         (silu): SiLU()
-        #         (linear): Linear(in_features=512, out_features=18432, bias=True)
-        #         (norm): LayerNorm((3072,), eps=1e-05, elementwise_affine=True)
-        #     )
-        #     (ff): FeedForward(
-        #         (net): ModuleList(
-        #         (0): GELU(
-        #             (proj): Linear(in_features=3072, out_features=12288, bias=True)
-        #         )
-        #         (1): Dropout(p=0.0, inplace=False)
-        #         (2): Linear(in_features=12288, out_features=3072, bias=True)
-        #         (3): Dropout(p=0.0, inplace=False)
-        #         )
-        #     )
-        #     )       
             
+        
         # compilation
         if compile == "torch":
-            #pipe.transformer.to(memory_format=torch.channels_last)
+            pipe.transformer.to(memory_format=torch.channels_last)
             if compile_args is not None:
                 torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"]
                 for i, block in enumerate(pipe.transformer.transformer_blocks):
@@ -379,16 +259,9 @@ class DownloadAndLoadCogVideoModel:
                 for i, block in enumerate(pipe.transformer.transformer_blocks):
                     if "CogVideoXBlock" in str(block):
                         pipe.transformer.transformer_blocks[i] = torch.compile(block, fullgraph=False, dynamic=False, backend="inductor")
+  
         
-        transformer.attention_mode = attention_mode
-
-        if "fused" in attention_mode:
-            from diffusers.models.attention import Attention
-            transformer.fuse_qkv_projections = True
-            for module in transformer.modules():
-                if isinstance(module, Attention):
-                    module.fuse_projections(fuse=True)
-
+            
         elif compile == "onediff":
             from onediffx import compile_pipe
             os.environ['NEXFORT_FX_FORCE_TRITON_SDPA'] = '1'
@@ -398,22 +271,21 @@ class DownloadAndLoadCogVideoModel:
             backend="nexfort",
             options= {"mode": "max-optimize:max-autotune:max-autotune", "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": False, "triton.fuse_attention_allow_fp16_reduction": False}},
             ignores=["vae"],
-            fuse_qkv_projections= False,
+            fuse_qkv_projections=True if pab_config is None else False,
             )          
-        
+
         pipeline = {
             "pipe": pipe,
             "dtype": dtype,
             "base_path": base_path,
             "onediff": True if compile == "onediff" else False,
             "cpu_offloading": enable_sequential_cpu_offload,
-            "manual_offloading": manual_offloading,
             "scheduler_config": scheduler_config,
-            "model_name": model,
+            "model_name": model
         }
 
         return (pipeline,)
-#region GGUF
+
 class DownloadAndLoadCogVideoGGUFModel:
     @classmethod
     def INPUT_TYPES(s):
@@ -423,12 +295,12 @@ class DownloadAndLoadCogVideoGGUFModel:
                     [
                         "CogVideoX_5b_GGUF_Q4_0.safetensors",
                         "CogVideoX_5b_I2V_GGUF_Q4_0.safetensors",
-                        "CogVideoX_5b_1_5_I2V_GGUF_Q4_0.safetensors",
                         "CogVideoX_5b_fun_GGUF_Q4_0.safetensors",
                         "CogVideoX_5b_fun_1_1_GGUF_Q4_0.safetensors",
                         "CogVideoX_5b_fun_1_1_Pose_GGUF_Q4_0.safetensors",
                         "CogVideoX_5b_Interpolation_GGUF_Q4_0.safetensors",
                         "CogVideoX_5b_Tora_GGUF_Q4_0.safetensors",
+
                     ],
                 ),
             "vae_precision": (["fp16", "fp32", "bf16"], {"default": "bf16", "tooltip": "VAE dtype"}),
@@ -437,10 +309,9 @@ class DownloadAndLoadCogVideoGGUFModel:
             "enable_sequential_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "significantly reducing memory usage and slows down the inference"}),
             },
             "optional": {
+                "pab_config": ("PAB_CONFIG", {"default": None}),
                 "block_edit": ("TRANSFORMERBLOCKS", {"default": None}),
-                #"lora": ("COGLORA", {"default": None}),
                 "compile": (["disabled","torch"], {"tooltip": "compile the model for faster inference, these are advanced options only available on Linux, see readme for more info"}),
-                "attention_mode": (["sdpa", "sageattn"], {"default": "sdpa"}),
             }
         }
 
@@ -449,8 +320,9 @@ class DownloadAndLoadCogVideoGGUFModel:
     FUNCTION = "loadmodel"
     CATEGORY = "CogVideoWrapper"
 
-    def loadmodel(self, model, vae_precision, fp8_fastmode, load_device, enable_sequential_cpu_offload, 
-                  block_edit=None, compile="disabled", attention_mode="sdpa"):
+    def loadmodel(self, model, vae_precision, fp8_fastmode, load_device, enable_sequential_cpu_offload, pab_config=None, block_edit=None, compile="disabled"):
+
+        check_diffusers_version()
 
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
@@ -486,7 +358,7 @@ class DownloadAndLoadCogVideoGGUFModel:
         with open(transformer_path) as f:
             transformer_config = json.load(f)
 
-        
+        sd = load_torch_file(gguf_path)
 
         from . import mz_gguf_loader
         import importlib
@@ -498,46 +370,45 @@ class DownloadAndLoadCogVideoGGUFModel:
                     transformer_config["in_channels"] = 32
                 else:
                     transformer_config["in_channels"] = 33
-                transformer = CogVideoXTransformer3DModelFun.from_config(transformer_config)
+                if pab_config is not None:
+                    transformer = CogVideoXTransformer3DModelFunPAB.from_config(transformer_config)
+                else:
+                    transformer = CogVideoXTransformer3DModelFun.from_config(transformer_config)
             elif "I2V" in model or "Interpolation" in model:
                 transformer_config["in_channels"] = 32
-                if "1_5" in model:
-                    transformer_config["ofs_embed_dim"] = 512
-                    transformer_config["use_learned_positional_embeddings"] = False
-                    transformer_config["patch_size_t"] = 2
-                    transformer_config["patch_bias"] = False
-                    transformer_config["sample_height"] = 300
-                    transformer_config["sample_width"] = 300
-                transformer = CogVideoXTransformer3DModel.from_config(transformer_config)
+                if pab_config is not None:
+                    transformer = CogVideoXTransformer3DModelPAB.from_config(transformer_config)
+                else:
+                    transformer = CogVideoXTransformer3DModel.from_config(transformer_config)
             else:
                 transformer_config["in_channels"] = 16
-                transformer = CogVideoXTransformer3DModel.from_config(transformer_config)
-
-            params_to_keep = {"patch_embed", "pos_embedding", "time_embedding"}
-            if "2b" in model:
-                cast_dtype = torch.float16
-            elif "1_5" in model:
-                params_to_keep = {"norm1.linear.weight", "patch_embed", "time_embedding", "ofs_embedding", "norm_final", "norm_out", "proj_out"}
-                cast_dtype = torch.bfloat16
-            for name, param in transformer.named_parameters():
-                if not any(keyword in name for keyword in params_to_keep):
-                    param.data = param.data.to(torch.float8_e4m3fn)
+                if pab_config is not None:
+                    transformer = CogVideoXTransformer3DModelPAB.from_config(transformer_config)
                 else:
-                    param.data = param.data.to(cast_dtype)
-            #for name, param in transformer.named_parameters():
-             #       print(name, param.data.dtype)
-           
+                    transformer = CogVideoXTransformer3DModel.from_config(transformer_config)
+
+            if "2b" in model:
+                for name, param in transformer.named_parameters():
+                    if name != "pos_embedding":
+                        param.data = param.data.to(torch.float8_e4m3fn)
+                    else:
+                        param.data = param.data.to(torch.float16)
+            else:
+                transformer.to(torch.float8_e4m3fn)
+
             if block_edit is not None:
                 transformer = remove_specific_blocks(transformer, block_edit)
 
-        transformer.attention_mode = attention_mode
+            transformer = mz_gguf_loader.quantize_load_state_dict(transformer, sd, device="cpu")
+            if load_device == "offload_device":
+                transformer.to(offload_device)
+            else:
+                transformer.to(device)
+        
         
         if fp8_fastmode:
-           params_to_keep = {"patch_embed", "lora", "pos_embedding", "time_embedding"}
-           if "1.5" in model:
-                params_to_keep.update({"ff","norm1.linear.weight", "norm_k", "norm_q","ofs_embedding", "norm_final", "norm_out", "proj_out"})   
            from .fp8_optimization import convert_fp8_linear
-           convert_fp8_linear(transformer, vae_dtype, params_to_keep=params_to_keep)
+           convert_fp8_linear(transformer, vae_dtype)
 
         if compile == "torch":
             # compilation
@@ -569,55 +440,17 @@ class DownloadAndLoadCogVideoGGUFModel:
             vae = AutoencoderKLCogVideoXFun.from_config(vae_config).to(vae_dtype).to(offload_device)
             vae.load_state_dict(vae_sd)
             if "Pose" in model:
-                pipe = CogVideoX_Fun_Pipeline_Control(vae, transformer, scheduler)
+                pipe = CogVideoX_Fun_Pipeline_Control(vae, transformer, scheduler, pab_config=pab_config)
             else:
-                pipe = CogVideoX_Fun_Pipeline_Inpaint(vae, transformer, scheduler)
+                pipe = CogVideoX_Fun_Pipeline_Inpaint(vae, transformer, scheduler, pab_config=pab_config)
         else:
             vae = AutoencoderKLCogVideoX.from_config(vae_config).to(vae_dtype).to(offload_device)
             vae.load_state_dict(vae_sd)
-            pipe = CogVideoXPipeline(vae, transformer, scheduler)
+            pipe = CogVideoXPipeline(vae, transformer, scheduler, pab_config=pab_config)
 
         if enable_sequential_cpu_offload:
             pipe.enable_sequential_cpu_offload()
 
-        sd = load_torch_file(gguf_path)
-
-        # #LoRAs
-        # if lora is not None:        
-        #     if "fun" in model.lower():
-        #         raise NotImplementedError("LoRA with GGUF is not supported for Fun models")
-        #         from .lora_utils import merge_lora#, load_lora_into_transformer
-        #         #for l in lora:
-        #         #    log.info(f"Merging LoRA weights from {l['path']} with strength {l['strength']}")
-        #         #    pipe.transformer = merge_lora(pipe.transformer, l["path"], l["strength"])
-        #     else:
-        #         adapter_list = []
-        #         adapter_weights = []
-        #         for l in lora:
-        #             lora_sd = load_torch_file(l["path"])             
-        #             for key, val in lora_sd.items():
-        #                 if "lora_B" in key:
-        #                     lora_rank = val.shape[1]
-        #                     break
-        #             log.info(f"Loading rank {lora_rank} LoRA weights from {l['path']} with strength {l['strength']}")
-        #             adapter_name = l['path'].split("/")[-1].split(".")[0]
-        #             adapter_weight = l['strength']
-        #             pipe.load_lora_weights(l['path'], weight_name=l['path'].split("/")[-1], lora_rank=lora_rank, adapter_name=adapter_name)
-                    
-        #             #transformer = load_lora_into_transformer(lora, transformer)
-        #             adapter_list.append(adapter_name)
-        #             adapter_weights.append(adapter_weight)
-        #         for l in lora:
-        #             pipe.set_adapters(adapter_list, adapter_weights=adapter_weights)
-        #         #pipe.fuse_lora(lora_scale=1 / lora_rank, components=["transformer"])
-        
-        pipe.transformer = mz_gguf_loader.quantize_load_state_dict(pipe.transformer, sd, device="cpu")
-        if load_device == "offload_device":
-            pipe.transformer.to(offload_device)
-        else:
-            pipe.transformer.to(device)
-
-
         pipeline = {
             "pipe": pipe,
             "dtype": vae_dtype,
@@ -625,50 +458,11 @@ class DownloadAndLoadCogVideoGGUFModel:
             "onediff": False,
             "cpu_offloading": enable_sequential_cpu_offload,
             "scheduler_config": scheduler_config,
-            "model_name": model,
-            "manual_offloading": True,
+            "model_name": model
         }
 
         return (pipeline,)
-    
-#revion VAE
 
-class CogVideoXVAELoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model_name": (folder_paths.get_filename_list("vae"), {"tooltip": "The name of the checkpoint (vae) to load."}),
-            },
-            "optional": {
-                "precision": (["fp16", "fp32", "bf16"],
-                    {"default": "bf16"}
-                ),
-            }
-        }
-
-    RETURN_TYPES = ("VAE",)
-    RETURN_NAMES = ("vae", )
-    FUNCTION = "loadmodel"
-    CATEGORY = "CogVideoWrapper"
-    DESCRIPTION = "Loads CogVideoX VAE model from 'ComfyUI/models/vae'"
-
-    def loadmodel(self, model_name, precision):
-        device = mm.get_torch_device()
-        offload_device = mm.unet_offload_device()
-
-        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
-        with open(os.path.join(script_directory, 'configs', 'vae_config.json')) as f:
-            vae_config = json.load(f)
-        model_path = folder_paths.get_full_path("vae", model_name)
-        vae_sd = load_torch_file(model_path)
-
-        vae = AutoencoderKLCogVideoX.from_config(vae_config).to(dtype).to(offload_device)
-        vae.load_state_dict(vae_sd)
-
-        return (vae,)
-    
-#region Tora
 class DownloadAndLoadToraModel:
     @classmethod
     def INPUT_TYPES(s):
@@ -689,6 +483,9 @@ class DownloadAndLoadToraModel:
     DESCRIPTION = "Downloads and loads the the Tora model from Huggingface to 'ComfyUI/models/CogVideo/CogVideoX-5b-Tora'"
 
     def loadmodel(self, model):
+        
+        check_diffusers_version()
+
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
         mm.soft_empty_cache()
@@ -773,7 +570,7 @@ class DownloadAndLoadToraModel:
         }
 
         return (toramodel,)
-#region controlnet
+
 class DownloadAndLoadCogVideoControlNet:
     @classmethod
     def INPUT_TYPES(s):
@@ -828,7 +625,6 @@ NODE_CLASS_MAPPINGS = {
     "DownloadAndLoadCogVideoControlNet": DownloadAndLoadCogVideoControlNet,
     "DownloadAndLoadToraModel": DownloadAndLoadToraModel,
     "CogVideoLoraSelect": CogVideoLoraSelect,
-    "CogVideoXVAELoader": CogVideoXVAELoader,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "DownloadAndLoadCogVideoModel": "(Down)load CogVideo Model",
@@ -836,5 +632,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "DownloadAndLoadCogVideoControlNet": "(Down)load CogVideo ControlNet",
     "DownloadAndLoadToraModel": "(Down)load Tora Model",
     "CogVideoLoraSelect": "CogVideo LoraSelect",
-    "CogVideoXVAELoader": "CogVideoX VAE Loader",
-    }
+    }

mz_gguf_loader.py

@@ -19,21 +19,17 @@ class quantize_lazy_load():
 
 
 def quantize_load_state_dict(model, state_dict, device="cpu"):
-    quant_keys = []
+    Q4_0_qkey = []
     for key in state_dict.keys():
         if key.endswith(".Q4_0_qweight"):
-            quant_keys.append(key.replace(".Q4_0_qweight", ""))
-            qtype = "Q4_0"
-        elif key.endswith(".Q8_0_qweight"):
-            quant_keys.append(key.replace(".Q8_0_qweight", ""))
-            qtype = "Q8_0"
+            Q4_0_qkey.append(key.replace(".Q4_0_qweight", ""))
 
     for name, module in model.named_modules():
-        if name in quant_keys:
+        if name in Q4_0_qkey:
             q_linear = WQLinear_GGUF.from_linear(
                 linear=module,
                 device=device,
-                qtype=qtype,
+                qtype="Q4_0",
             )
             set_op_by_name(model, name, q_linear)
 
@@ -121,14 +117,14 @@ class WQLinear_GGUF(nn.Module):
 
     @torch.no_grad()
     def forward(self, x):
+        # x = torch.matmul(x, dequantize_blocks_Q4_0(self.qweight))
         if self.qtype == "Q4_0":
-            dequant = dequantize_blocks_Q4_0(self.Q4_0_qweight, x.dtype)
-        elif self.qtype == "Q8_0":
-            dequant = dequantize_blocks_Q8_0(self.Q8_0_qweight, x.dtype)
+            x = F.linear(x, dequantize_blocks_Q4_0(
+                self.Q4_0_qweight, x.dtype), self.bias.to(x.dtype) if self.bias is not None else None)
         else:
             raise ValueError(f"Unknown qtype: {self.qtype}")
-        
-        return F.linear(x, dequant, bias=self.bias.to(x.dtype) if self.bias is not None else None)
+
+        return x
 
 
 def split_block_dims(blocks, *args):
@@ -157,7 +153,6 @@ def quant_shape_from_byte_shape(shape, qtype) -> tuple[int, ...]:
 
 GGML_QUANT_SIZES = {
     "Q4_0": (32, 2 + 16),
-    "Q8_0": (32, 2 + 32),
 }
 
 
@@ -191,31 +186,3 @@ def dequantize_blocks_Q4_0(data, dtype=torch.float16):
     )).to(dtype)
     return out
 
-def dequantize_blocks_Q8_0(data, dtype=torch.float16):
-    block_size, type_size = GGML_QUANT_SIZES["Q8_0"]
-
-    data = data.to(torch.uint8)
-    shape = data.shape
-
-    rows = data.reshape(
-        (-1, data.shape[-1])
-    ).view(torch.uint8)
-
-    n_blocks = rows.numel() // type_size
-    blocks = data.reshape((n_blocks, type_size))
-
-    n_blocks = blocks.shape[0]
-
-    d, qs = split_block_dims(blocks, 2)
-    d = d.view(torch.float16).to(torch.float32)
-
-    qs = qs.view(torch.int8).to(torch.float32)
-
-    out = (d * qs)
-
-    out = out.reshape(quant_shape_from_byte_shape(
-        shape,
-        qtype="Q8_0",
-    )).to(dtype)
-    return out
-

nodes.py

@@ -5,8 +5,6 @@ import comfy.model_management as mm
 from einops import rearrange
 from contextlib import nullcontext
 
-from .utils import log, check_diffusers_version, print_memory
-check_diffusers_version()
 from diffusers.schedulers import (
     CogVideoXDDIMScheduler, 
     CogVideoXDPMScheduler, 
@@ -44,6 +42,8 @@ from PIL import Image
 import numpy as np
 import json
 
+from .utils import log, check_diffusers_version
+
 script_directory = os.path.dirname(os.path.abspath(__file__))
 
 if not "CogVideo" in folder_paths.folder_names_and_paths:
@@ -51,11 +51,61 @@ if not "CogVideo" in folder_paths.folder_names_and_paths:
 if not "cogvideox_loras" in folder_paths.folder_names_and_paths:
     folder_paths.add_model_folder_path("cogvideox_loras", os.path.join(folder_paths.models_dir, "CogVideo", "loras"))
 
+#PAB
+from .videosys.pab import CogVideoXPABConfig
+
+class CogVideoPABConfig:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "spatial_broadcast": ("BOOLEAN", {"default": True, "tooltip": "Enable Spatial PAB, highest impact"}),
+            "spatial_threshold_start": ("INT", {"default": 850, "min": 0, "max": 1000, "tooltip": "PAB Start Timestep"} ),
+            "spatial_threshold_end": ("INT", {"default": 100, "min": 0, "max": 1000, "tooltip": "PAB End Timestep"} ),
+            "spatial_range": ("INT", {"default": 2, "min": 0, "max": 10, "tooltip": "Broadcast timesteps range, higher values are faster but quality may suffer"} ),
+            "temporal_broadcast": ("BOOLEAN", {"default": False, "tooltip": "Enable Temporal PAB, medium impact"}),
+            "temporal_threshold_start": ("INT", {"default": 850, "min": 0, "max": 1000, "tooltip": "PAB Start Timestep"} ),
+            "temporal_threshold_end": ("INT", {"default": 100, "min": 0, "max": 1000, "tooltip": "PAB End Timestep"} ),
+            "temporal_range": ("INT", {"default": 4, "min": 0, "max": 10, "tooltip": "Broadcast timesteps range, higher values are faster but quality may suffer"} ),
+            "cross_broadcast": ("BOOLEAN", {"default": False, "tooltip": "Enable Cross Attention PAB, low impact"}),
+            "cross_threshold_start": ("INT", {"default": 850, "min": 0, "max": 1000, "tooltip": "PAB Start Timestep"} ),
+            "cross_threshold_end": ("INT", {"default": 100, "min": 0, "max": 1000, "tooltip": "PAB End Timestep"} ),
+            "cross_range": ("INT", {"default": 6, "min": 0, "max": 10, "tooltip": "Broadcast timesteps range, higher values are faster but quality may suffer"} ),
+
+            "steps": ("INT", {"default": 50, "min": 0, "max": 1000, "tooltip": "Should match the sampling steps"} ),
+            }
+        }
+
+    RETURN_TYPES = ("PAB_CONFIG",)
+    RETURN_NAMES = ("pab_config", )
+    FUNCTION = "config"
+    CATEGORY = "CogVideoWrapper"
+    DESCRIPTION = "EXPERIMENTAL:Pyramid Attention Broadcast (PAB) speeds up inference by mitigating redundant attention computation. Increases memory use"
+
+    def config(self, spatial_broadcast, spatial_threshold_start, spatial_threshold_end, spatial_range, 
+               temporal_broadcast, temporal_threshold_start, temporal_threshold_end, temporal_range, 
+               cross_broadcast, cross_threshold_start, cross_threshold_end, cross_range, steps):
+        
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+        pab_config = CogVideoXPABConfig(
+            steps=steps, 
+            spatial_broadcast=spatial_broadcast, 
+            spatial_threshold=[spatial_threshold_end, spatial_threshold_start], 
+            spatial_range=spatial_range,
+            temporal_broadcast=temporal_broadcast,
+            temporal_threshold=[temporal_threshold_end, temporal_threshold_start],
+            temporal_range=temporal_range,
+            cross_broadcast=cross_broadcast,
+            cross_threshold=[cross_threshold_end, cross_threshold_start],
+            cross_range=cross_range
+            )
+
+        return (pab_config, )
+
 class CogVideoContextOptions:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
-            "context_schedule": (["uniform_standard", "uniform_looped", "static_standard"],),
+            "context_schedule": (["uniform_standard", "uniform_looped", "static_standard", "temporal_tiling"],),
             "context_frames": ("INT", {"default": 48, "min": 2, "max": 100, "step": 1, "tooltip": "Number of pixel frames in the context, NOTE: the latent space has 4 frames in 1"} ),
             "context_stride": ("INT", {"default": 4, "min": 4, "max": 100, "step": 1, "tooltip": "Context stride as pixel frames, NOTE: the latent space has 4 frames in 1"} ),
             "context_overlap": ("INT", {"default": 4, "min": 4, "max": 100, "step": 1, "tooltip": "Context overlap as pixel frames, NOTE: the latent space has 4 frames in 1"} ),
@@ -246,7 +296,7 @@ class CogVideoTextEncode:
         
         embeds = clip.encode_from_tokens(tokens, return_pooled=False, return_dict=False)
 
-        if embeds.shape[1] > max_tokens:
+        if embeds.shape[1] > 226:
             raise ValueError(f"Prompt is too long, max tokens supported is {max_tokens} or less, got {embeds.shape[1]}")
         embeds *= strength
         if force_offload:
@@ -294,11 +344,9 @@ class CogVideoImageEncode:
             "image": ("IMAGE", ),
             },
             "optional": {
-                "chunk_size": ("INT", {"default": 16, "min": 4, "tooltip": "How many images to encode at once, lower values use less memory"}),
+                "chunk_size": ("INT", {"default": 16, "min": 4}),
                 "enable_tiling": ("BOOLEAN", {"default": False, "tooltip": "Enable tiling for the VAE to reduce memory usage"}),
                 "mask": ("MASK", ),
-                "noise_aug_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001, "tooltip": "Augment image with noise"}),
-                "vae_override" : ("VAE", {"default": None, "tooltip": "Override the VAE model in the pipeline"}),
             },
         }
 
@@ -307,21 +355,15 @@ class CogVideoImageEncode:
     FUNCTION = "encode"
     CATEGORY = "CogVideoWrapper"
 
-    def encode(self, pipeline, image, chunk_size=8, enable_tiling=False, mask=None, noise_aug_strength=0.0, vae_override=None):
+    def encode(self, pipeline, image, chunk_size=8, enable_tiling=False, mask=None):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
         generator = torch.Generator(device=device).manual_seed(0)
 
         B, H, W, C = image.shape
 
-        vae = pipeline["pipe"].vae if vae_override is None else vae_override
+        vae = pipeline["pipe"].vae
         vae.enable_slicing()
-        model_name = pipeline.get("model_name", "")
-
-        if ("1.5" in model_name or "1_5" in model_name) and image.shape[0] == 1:
-            vae_scaling_factor = 1 #/ vae.config.scaling_factor
-        else:
-            vae_scaling_factor = vae.config.scaling_factor
         
         if enable_tiling:
             from .mz_enable_vae_encode_tiling import enable_vae_encode_tiling
@@ -346,17 +388,11 @@ class CogVideoImageEncode:
             # input_image = input_image * (1 -mask)
         else:
             pipeline["pipe"].original_mask = None
-        #input_image = input_image.permute(0, 3, 1, 2)  # B, C, H, W
-        #input_image = pipeline["pipe"].video_processor.preprocess(input_image).to(device, dtype=vae.dtype)
-        #input_image = input_image.unsqueeze(2)
-
+            
         input_image = input_image * 2.0 - 1.0
         input_image = input_image.to(vae.dtype).to(device)
         input_image = input_image.unsqueeze(0).permute(0, 4, 1, 2, 3) # B, C, T, H, W
-      
         B, C, T, H, W = input_image.shape
-        if noise_aug_strength > 0:
-            input_image = add_noise_to_reference_video(input_image, ratio=noise_aug_strength)
 
         latents_list = []
         # Loop through the temporal dimension in chunks of 16
@@ -376,13 +412,12 @@ class CogVideoImageEncode:
             elif hasattr(latents, "latents"):
                 latents = latents.latents
 
+            latents = vae.config.scaling_factor * latents
             latents = latents.permute(0, 2, 1, 3, 4)  # B, T_chunk, C, H, W
             latents_list.append(latents)
 
         # Concatenate all the chunks along the temporal dimension
         final_latents = torch.cat(latents_list, dim=1)
-        final_latents = final_latents * vae_scaling_factor
-        
         log.info(f"Encoded latents shape: {final_latents.shape}")
         if not pipeline["cpu_offloading"]:
             vae.to(offload_device)
@@ -400,8 +435,6 @@ class CogVideoImageInterpolationEncode:
             "optional": {
                 "enable_tiling": ("BOOLEAN", {"default": False, "tooltip": "Enable tiling for the VAE to reduce memory usage"}),
                 "mask": ("MASK", ),
-                "vae_override" : ("VAE", {"default": None, "tooltip": "Override the VAE model in the pipeline"}),
-
             },
         }
 
@@ -410,21 +443,14 @@ class CogVideoImageInterpolationEncode:
     FUNCTION = "encode"
     CATEGORY = "CogVideoWrapper"
 
-    def encode(self, pipeline, start_image, end_image, enable_tiling=False, mask=None, vae_override=None):
+    def encode(self, pipeline, start_image, end_image, chunk_size=8, enable_tiling=False, mask=None):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
         generator = torch.Generator(device=device).manual_seed(0)
 
         B, H, W, C = start_image.shape
 
-        vae = pipeline["pipe"].vae if vae_override is None else vae_override
-        vae.enable_slicing()
-        model_name = pipeline.get("model_name", "")
-
-        if ("1.5" in model_name or "1_5" in model_name):
-            vae_scaling_factor = 1 / vae.config.scaling_factor
-        else:
-            vae_scaling_factor = vae.config.scaling_factor
+        vae = pipeline["pipe"].vae
         vae.enable_slicing()
         
         if enable_tiling:
@@ -457,8 +483,8 @@ class CogVideoImageInterpolationEncode:
         latents_list = []           
 
         # Encode the chunk of images
-        start_latents = vae.encode(start_image).latent_dist.sample(generator) * vae_scaling_factor
-        end_latents = vae.encode(end_image).latent_dist.sample(generator) * vae_scaling_factor
+        start_latents = vae.encode(start_image).latent_dist.sample(generator) * vae.config.scaling_factor
+        end_latents = vae.encode(end_image).latent_dist.sample(generator) * vae.config.scaling_factor
 
         start_latents = start_latents.permute(0, 2, 1, 3, 4)  # B, T, C, H, W
         end_latents = end_latents.permute(0, 2, 1, 3, 4)  # B, T, C, H, W
@@ -547,18 +573,16 @@ class ToraEncodeTrajectory:
             vae.to(device)
 
         video_flow = vae.encode(video_flow).latent_dist.sample(generator) * vae.config.scaling_factor
-        log.info(f"video_flow shape after encoding: {video_flow.shape}") #torch.Size([1, 16, 4, 80, 80])
 
         if not pipeline["cpu_offloading"]:
             vae.to(offload_device)
-        #print("video_flow shape before traj_extractor: ", video_flow.shape) #torch.Size([1, 16, 4, 80, 80])
+
         video_flow_features = tora_model["traj_extractor"](video_flow.to(torch.float32))
         video_flow_features = torch.stack(video_flow_features)
-        #print("video_flow_features after traj_extractor: ", video_flow_features.shape) #torch.Size([42, 4, 128, 40, 40])
 
         video_flow_features = video_flow_features * strength
 
-        
+        log.info(f"video_flow shape: {video_flow.shape}")
 
         tora = {
             "video_flow_features" : video_flow_features,
@@ -728,7 +752,6 @@ class CogVideoXFasterCache:
                 "hf_step": ("INT", {"default": 30, "min": 0, "max": 1024, "step": 1}),
                 "lf_step": ("INT", {"default": 40, "min": 0, "max": 1024, "step": 1}),
                 "cache_device": (["main_device", "offload_device", "cuda:1"], {"default": "main_device", "tooltip": "The device to use for the cache, main_device is on GPU and uses a lot of VRAM"}),
-                "num_blocks_to_cache": ("INT", {"default": 42, "min": 0, "max": 1024, "step": 1, "tooltip": "Number of transformer blocks to cache, 5b model has 42 blocks, tradeoff between speed and memory"}),
             },
         }
 
@@ -737,7 +760,7 @@ class CogVideoXFasterCache:
     FUNCTION = "args"
     CATEGORY = "CogVideoWrapper"
 
-    def args(self, start_step, hf_step, lf_step, cache_device, num_blocks_to_cache):
+    def args(self, start_step, hf_step, lf_step, cache_device):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
         if cache_device == "cuda:1":
@@ -746,8 +769,7 @@ class CogVideoXFasterCache:
             "start_step" : start_step,
             "hf_step" : hf_step,
             "lf_step" : lf_step,
-            "cache_device" : device if cache_device != "offload_device" else offload_device,
-            "num_blocks_to_cache" : num_blocks_to_cache,
+            "cache_device" : device if cache_device != "offload_device" else offload_device
         }
         return (fastercache,)
 
@@ -760,9 +782,9 @@ class CogVideoSampler:
                 "pipeline": ("COGVIDEOPIPE",),
                 "positive": ("CONDITIONING", ),
                 "negative": ("CONDITIONING", ),
-                "width": ("INT", {"default": 720, "min": 128, "max": 2048, "step": 16}),
-                "height": ("INT", {"default": 480, "min": 128, "max": 2048, "step": 16}),
-                "num_frames": ("INT", {"default": 49, "min": 17, "max": 1024, "step": 4}),
+                "height": ("INT", {"default": 480, "min": 128, "max": 2048, "step": 8}),
+                "width": ("INT", {"default": 720, "min": 128, "max": 2048, "step": 8}),
+                "num_frames": ("INT", {"default": 49, "min": 16, "max": 1024, "step": 1}),
                 "steps": ("INT", {"default": 50, "min": 1}),
                 "cfg": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 30.0, "step": 0.01}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
@@ -772,14 +794,13 @@ class CogVideoSampler:
                     }),
             },
             "optional": {
-                "samples": ("LATENT", {"tooltip": "init Latents to use for video2video process"} ),
+                "samples": ("LATENT", ),
                 "denoise_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "image_cond_latents": ("LATENT",{"tooltip": "Latent to use for image2video conditioning"} ),
+                "image_cond_latents": ("LATENT", ),
                 "context_options": ("COGCONTEXT", ),
                 "controlnet": ("COGVIDECONTROLNET",),
                 "tora_trajectory": ("TORAFEATURES", ),
                 "fastercache": ("FASTERCACHEARGS", ),
-                #"sigmas": ("SIGMAS", ),
             }
         }
 
@@ -793,25 +814,9 @@ class CogVideoSampler:
         mm.soft_empty_cache()
 
         base_path = pipeline["base_path"]
-        model_name = pipeline.get("model_name", "")
-        supports_image_conds = True if "I2V" in model_name or "interpolation" in model_name.lower() else False
 
         assert "fun" not in base_path.lower(), "'Fun' models not supported in 'CogVideoSampler', use the 'CogVideoXFunSampler'"
-        assert (
-            "I2V" not in model_name or
-            "1.5" in model_name or
-            "1_5" in model_name or
-            num_frames == 49 or
-            context_options is not None
-        ), "1.0 I2V model can only do 49 frames"
-        if image_cond_latents is not None:
-            assert supports_image_conds, "Image condition latents only supported for I2V and Interpolation models"
-            # if "I2V" in model_name:
-            #     assert image_cond_latents["samples"].shape[1] == 1, "I2V model only supports single image condition latent"
-            # elif "interpolation" in model_name.lower():
-            #     assert image_cond_latents["samples"].shape[1] == 2, "Interpolation model needs two image condition latents"
-        else:
-            assert not supports_image_conds, "Image condition latents required for I2V models"
+        assert ("I2V" not in pipeline.get("model_name","") or num_frames == 49 or context_options is not None), "I2V model can only do 49 frames"
 
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
@@ -819,7 +824,7 @@ class CogVideoSampler:
         dtype = pipeline["dtype"]
         scheduler_config = pipeline["scheduler_config"]
         
-        if not pipeline["cpu_offloading"] and pipeline["manual_offloading"]:
+        if not pipeline["cpu_offloading"]:
             pipe.transformer.to(device)
         generator = torch.Generator(device=torch.device("cpu")).manual_seed(seed)
 
@@ -851,26 +856,12 @@ class CogVideoSampler:
             pipe.transformer.fastercache_lf_step = fastercache["lf_step"]
             pipe.transformer.fastercache_hf_step = fastercache["hf_step"]
             pipe.transformer.fastercache_device = fastercache["cache_device"]
-            pipe.transformer.fastercache_num_blocks_to_cache = fastercache["num_blocks_to_cache"]
-            log.info(f"FasterCache enabled for {pipe.transformer.fastercache_num_blocks_to_cache} blocks out of {len(pipe.transformer.transformer_blocks)}")
         else:
             pipe.transformer.use_fastercache = False
             pipe.transformer.fastercache_counter = 0
 
-        if not isinstance(cfg, list):
-            cfg = [cfg for _ in range(steps)]
-        else:
-            assert len(cfg) == steps, "Length of cfg list must match number of steps"
-
-        # if sigmas is not None:
-        #     sigma_list = sigmas.tolist()
-        try:
-            torch.cuda.reset_peak_memory_stats(device)
-        except:
-            pass
-  
         autocastcondition = not pipeline["onediff"] or not dtype == torch.float32
-        autocast_context = torch.autocast(mm.get_autocast_device(device), dtype=dtype) if autocastcondition else nullcontext()
+        autocast_context = torch.autocast(mm.get_autocast_device(device)) if autocastcondition else nullcontext()
         with autocast_context:
             latents = pipeline["pipe"](
                 num_inference_steps=steps,
@@ -878,7 +869,6 @@ class CogVideoSampler:
                 width = width,
                 num_frames = num_frames,
                 guidance_scale=cfg,
-                #sigmas=sigma_list if sigmas is not None else None,
                 latents=samples["samples"] if samples is not None else None,
                 image_cond_latents=image_cond_latents["samples"] if image_cond_latents is not None else None,
                 denoise_strength=denoise_strength,
@@ -894,7 +884,7 @@ class CogVideoSampler:
                 controlnet=controlnet,
                 tora=tora_trajectory if tora_trajectory is not None else None,
             )
-        if not pipeline["cpu_offloading"] and pipeline["manual_offloading"]:
+        if not pipeline["cpu_offloading"]:
             pipe.transformer.to(offload_device)
 
         if fastercache is not None:
@@ -902,13 +892,8 @@ class CogVideoSampler:
                 if (hasattr, block, "cached_hidden_states") and block.cached_hidden_states is not None:
                     block.cached_hidden_states = None
                     block.cached_encoder_hidden_states = None
-
-        print_memory(device)
+                    
         mm.soft_empty_cache()
-        try:
-            torch.cuda.reset_peak_memory_stats(device)
-        except:
-            pass
 
         return (pipeline, {"samples": latents})
 
@@ -954,7 +939,7 @@ class CogVideoDecode:
         return {"required": {
             "pipeline": ("COGVIDEOPIPE",),
             "samples": ("LATENT", ),
-            "enable_vae_tiling": ("BOOLEAN", {"default": True, "tooltip": "Drastically reduces memory use but may introduce seams"}),
+            "enable_vae_tiling": ("BOOLEAN", {"default": False, "tooltip": "Drastically reduces memory use but may introduce seams"}),
             },
             "optional": {
             "tile_sample_min_height": ("INT", {"default": 240, "min": 16, "max": 2048, "step": 8, "tooltip": "Minimum tile height, default is half the height"}),
@@ -962,7 +947,6 @@ class CogVideoDecode:
             "tile_overlap_factor_height": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.001}),
             "tile_overlap_factor_width": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.001}),
             "auto_tile_size": ("BOOLEAN", {"default": True, "tooltip": "Auto size based on height and width, default is half the size"}),
-            "vae_override": ("VAE", {"default": None}),
             }
         }
 
@@ -971,14 +955,11 @@ class CogVideoDecode:
     FUNCTION = "decode"
     CATEGORY = "CogVideoWrapper"
 
-    def decode(self, pipeline, samples, enable_vae_tiling, tile_sample_min_height, tile_sample_min_width, tile_overlap_factor_height, tile_overlap_factor_width, 
-               auto_tile_size=True, vae_override=None):
+    def decode(self, pipeline, samples, enable_vae_tiling, tile_sample_min_height, tile_sample_min_width, tile_overlap_factor_height, tile_overlap_factor_width, auto_tile_size=True):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
         latents = samples["samples"]
-        vae = pipeline["pipe"].vae if vae_override is None else vae_override
-
-        additional_frames = getattr(pipeline["pipe"], "additional_frames", 0)
+        vae = pipeline["pipe"].vae
 
         vae.enable_slicing()
 
@@ -1003,15 +984,7 @@ class CogVideoDecode:
             vae._clear_fake_context_parallel_cache()
         except:
             pass
-        
-        try:
-            frames = vae.decode(latents[:, :, additional_frames:]).sample
-        except:
-            mm.soft_empty_cache()
-            log.warning("Failed to decode, retrying with tiling")
-            vae.enable_tiling()
-            frames = vae.decode(latents[:, :, additional_frames:]).sample
-
+        frames = vae.decode(latents).sample
         vae.disable_tiling()
         if not pipeline["cpu_offloading"]:
             vae.to(offload_device)
@@ -1100,6 +1073,10 @@ class CogVideoXFunSampler:
         base_path = pipeline["base_path"]
         assert "fun" in base_path.lower(), "'Unfun' models not supported in 'CogVideoXFunSampler', use the 'CogVideoSampler'"
         assert "pose" not in base_path.lower(), "'Pose' models not supported in 'CogVideoXFunSampler', use the 'CogVideoXFunControlSampler'"
+        
+
+        if not pipeline["cpu_offloading"]:
+            pipe.enable_model_cpu_offload(device=device)
 
         mm.soft_empty_cache()
 
@@ -1112,6 +1089,9 @@ class CogVideoXFunSampler:
             end_img = [to_pil(_end_img) for _end_img in end_img] if end_img is not None else None       
         
         # Load Sampler
+        if context_options is not None and context_options["context_schedule"] == "temporal_tiling":
+            log.info("Temporal tiling enabled, changing scheduler to CogVideoXDDIM")
+            scheduler="CogVideoXDDIM"
         scheduler_config = pipeline["scheduler_config"]
         if scheduler in scheduler_mapping:
             noise_scheduler = scheduler_mapping[scheduler].from_config(scheduler_config)
@@ -1119,8 +1099,8 @@ class CogVideoXFunSampler:
         else:
             raise ValueError(f"Unknown scheduler: {scheduler}")
 
-        if not pipeline["cpu_offloading"] and pipeline["manual_offloading"]:
-            pipe.transformer.to(device)
+        #if not pipeline["cpu_offloading"]:
+        #    pipe.transformer.to(device)
 
         if context_options is not None:
             context_frames = context_options["context_frames"] // 4
@@ -1139,8 +1119,6 @@ class CogVideoXFunSampler:
             pipe.transformer.fastercache_lf_step = fastercache["lf_step"]
             pipe.transformer.fastercache_hf_step = fastercache["hf_step"]
             pipe.transformer.fastercache_device = fastercache["cache_device"]
-            pipe.transformer.fastercache_num_blocks_to_cache = fastercache["num_blocks_to_cache"]
-            log.info(f"FasterCache enabled for {pipe.transformer.fastercache_num_blocks_to_cache} blocks out of {len(pipe.transformer.transformer_blocks)}")
         else:
             pipe.transformer.use_fastercache = False
             pipe.transformer.fastercache_counter = 0
@@ -1148,7 +1126,7 @@ class CogVideoXFunSampler:
         generator = torch.Generator(device=torch.device("cpu")).manual_seed(seed)
 
         autocastcondition = not pipeline["onediff"] or not dtype == torch.float32
-        autocast_context = torch.autocast(mm.get_autocast_device(device), dtype=dtype) if autocastcondition else nullcontext()
+        autocast_context = torch.autocast(mm.get_autocast_device(device)) if autocastcondition else nullcontext()
         with autocast_context:
             video_length = int((video_length - 1) // pipe.vae.config.temporal_compression_ratio * pipe.vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
             if vid2vid_images is not None:
@@ -1180,8 +1158,8 @@ class CogVideoXFunSampler:
                 noise_aug_strength = noise_aug_strength,
                 strength = vid2vid_denoise,
             )
-        if not pipeline["cpu_offloading"] and pipeline["manual_offloading"]:
-            pipe.transformer.to(offload_device)
+        #if not pipeline["cpu_offloading"]:
+        #     pipe.transformer.to(offload_device)
         #clear FasterCache
         if fastercache is not None:
             for block in pipe.transformer.transformer_blocks:
@@ -1239,7 +1217,7 @@ class CogVideoXFunControlSampler:
     CATEGORY = "CogVideoWrapper"
 
     def process(self, pipeline, positive, negative, seed, steps, cfg, scheduler, control_latents, 
-                control_strength=1.0, control_start_percent=0.0, control_end_percent=1.0, 
+                control_strength=1.0, control_start_percent=0.0, control_end_percent=1.0, t_tile_length=16, t_tile_overlap=8, 
                 samples=None, denoise_strength=1.0, context_options=None):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
@@ -1263,6 +1241,9 @@ class CogVideoXFunControlSampler:
 
         # Load Sampler
         scheduler_config = pipeline["scheduler_config"]
+        if context_options is not None and context_options["context_schedule"] == "temporal_tiling":
+            log.info("Temporal tiling enabled, changing scheduler to CogVideoXDDIM")
+            scheduler="CogVideoXDDIM"
         if scheduler in scheduler_mapping:
             noise_scheduler = scheduler_mapping[scheduler].from_config(scheduler_config)
             pipe.scheduler = noise_scheduler
@@ -1381,6 +1362,7 @@ NODE_CLASS_MAPPINGS = {
     "CogVideoXFunVid2VidSampler": CogVideoXFunVid2VidSampler,
     "CogVideoXFunControlSampler": CogVideoXFunControlSampler,
     "CogVideoTextEncodeCombine": CogVideoTextEncodeCombine,
+    "CogVideoPABConfig": CogVideoPABConfig,
     "CogVideoTransformerEdit": CogVideoTransformerEdit,
     "CogVideoControlImageEncode": CogVideoControlImageEncode,
     "CogVideoContextOptions": CogVideoContextOptions,
@@ -1403,6 +1385,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CogVideoXFunVid2VidSampler": "CogVideoXFun Vid2Vid Sampler",
     "CogVideoXFunControlSampler": "CogVideoXFun Control Sampler",
     "CogVideoTextEncodeCombine": "CogVideo TextEncode Combine",
+    "CogVideoPABConfig": "CogVideo PABConfig",
     "CogVideoTransformerEdit": "CogVideo TransformerEdit",
     "CogVideoControlImageEncode": "CogVideo Control ImageEncode",
     "CogVideoContextOptions": "CogVideo Context Options",

pipeline_cogvideox.py

@@ -20,22 +20,24 @@ import torch
 import torch.nn.functional as F
 import math
 
-from diffusers.models import AutoencoderKLCogVideoX
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.models import AutoencoderKLCogVideoX#, CogVideoXTransformer3DModel
 from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
 from diffusers.utils import logging
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
-#from diffusers.models.embeddings import get_3d_rotary_pos_embed
+from diffusers.models.embeddings import get_3d_rotary_pos_embed
 from diffusers.loaders import CogVideoXLoraLoaderMixin
 
-from .embeddings import get_3d_rotary_pos_embed
 from .custom_cogvideox_transformer_3d import CogVideoXTransformer3DModel
 
 from comfy.utils import ProgressBar
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
+from .videosys.core.pipeline import VideoSysPipeline
+from .videosys.cogvideox_transformer_3d import CogVideoXTransformer3DModel as CogVideoXTransformer3DModelPAB
+from .videosys.core.pab_mgr import set_pab_manager
+
 def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
     tw = tgt_width
     th = tgt_height
@@ -112,7 +114,7 @@ def retrieve_timesteps(
         timesteps = scheduler.timesteps
     return timesteps, num_inference_steps
 
-class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
+class CogVideoXPipeline(VideoSysPipeline, CogVideoXLoraLoaderMixin):
     r"""
     Pipeline for text-to-video generation using CogVideoX.
 
@@ -141,9 +143,10 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
     def __init__(
         self,
         vae: AutoencoderKLCogVideoX,
-        transformer: CogVideoXTransformer3DModel,
+        transformer: Union[CogVideoXTransformer3DModel, CogVideoXTransformer3DModelPAB],
         scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
         original_mask = None,
+        pab_config = None
     ):
         super().__init__()
 
@@ -155,14 +158,15 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
         )
         self.vae_scale_factor_temporal = (
             self.vae.config.temporal_compression_ratio if hasattr(self, "vae") and self.vae is not None else 4
-        )        
+        )
         self.original_mask = original_mask
         self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
-        self.video_processor.config.do_resize = False
+
+        if pab_config is not None:
+            set_pab_manager(pab_config)
 
         self.input_with_padding = True
 
-
     def prepare_latents(
         self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, timesteps, denoise_strength,
          num_inference_steps, latents=None, freenoise=True, context_size=None, context_overlap=None
@@ -282,6 +286,29 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
             self.scheduler.set_begin_index(t_start * self.scheduler.order)
 
         return timesteps.to(device), num_inference_steps - t_start
+    
+    def _gaussian_weights(self, t_tile_length, t_batch_size):
+        from numpy import pi, exp, sqrt
+
+        var = 0.01
+        midpoint = (t_tile_length - 1) / 2  # -1 because index goes from 0 to latent_width - 1
+        t_probs = [exp(-(t-midpoint)*(t-midpoint)/(t_tile_length*t_tile_length)/(2*var)) / sqrt(2*pi*var) for t in range(t_tile_length)]
+        weights = torch.tensor(t_probs)
+        weights = weights.unsqueeze(0).unsqueeze(2).unsqueeze(3).unsqueeze(4).repeat(1, t_batch_size,1, 1, 1)
+        return weights
+    
+    def fuse_qkv_projections(self) -> None:
+        r"""Enables fused QKV projections."""
+        self.fusing_transformer = True
+        self.transformer.fuse_qkv_projections()
+
+    def unfuse_qkv_projections(self) -> None:
+        r"""Disable QKV projection fusion if enabled."""
+        if not self.fusing_transformer:
+            logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.")
+        else:
+            self.transformer.unfuse_qkv_projections()
+            self.fusing_transformer = False
 
     def _prepare_rotary_positional_embeddings(
         self,
@@ -289,41 +316,34 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
         width: int,
         num_frames: int,
         device: torch.device,
+        start_frame: int = None,
+        end_frame: int = None,
         ) -> Tuple[torch.Tensor, torch.Tensor]:
         grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
         grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
+        base_size_width = 720 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
+        base_size_height = 480 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
 
-        p = self.transformer.config.patch_size
-        p_t = self.transformer.config.patch_size_t
+        grid_crops_coords = get_resize_crop_region_for_grid(
+            (grid_height, grid_width), base_size_width, base_size_height
+        )
+        freqs_cos, freqs_sin = get_3d_rotary_pos_embed(
+            embed_dim=self.transformer.config.attention_head_dim,
+            crops_coords=grid_crops_coords,
+            grid_size=(grid_height, grid_width),
+            temporal_size=num_frames,
+            use_real=True,
+        )
+        
+        if start_frame is not None:
+            freqs_cos = freqs_cos.view(num_frames, grid_height * grid_width, -1)
+            freqs_sin = freqs_sin.view(num_frames, grid_height * grid_width, -1)
 
-        if p_t is None:
-            # CogVideoX 1.0 I2V
-            base_size_width = self.transformer.config.sample_width // p
-            base_size_height = self.transformer.config.sample_height // p
+            freqs_cos = freqs_cos[start_frame:end_frame]
+            freqs_sin = freqs_sin[start_frame:end_frame]
 
-            grid_crops_coords = get_resize_crop_region_for_grid(
-                (grid_height, grid_width), base_size_width, base_size_height
-            )
-            freqs_cos, freqs_sin = get_3d_rotary_pos_embed(
-                embed_dim=self.transformer.config.attention_head_dim,
-                crops_coords=grid_crops_coords,
-                grid_size=(grid_height, grid_width),
-                temporal_size=num_frames,
-            )
-        else:
-            # CogVideoX 1.5 I2V
-            base_size_width = self.transformer.config.sample_width // p
-            base_size_height = self.transformer.config.sample_height // p
-            base_num_frames = (num_frames + p_t - 1) // p_t
-
-            freqs_cos, freqs_sin = get_3d_rotary_pos_embed(
-                embed_dim=self.transformer.config.attention_head_dim,
-                crops_coords=None,
-                grid_size=(grid_height, grid_width),
-                temporal_size=base_num_frames,
-                grid_type="slice",
-                max_size=(base_size_height, base_size_width),
-            )
+            freqs_cos = freqs_cos.view(-1, freqs_cos.shape[-1])
+            freqs_sin = freqs_sin.view(-1, freqs_sin.shape[-1])
 
         freqs_cos = freqs_cos.to(device=device)
         freqs_sin = freqs_sin.to(device=device)
@@ -350,11 +370,12 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
         height: int = 480,
         width: int = 720,
         num_frames: int = 48,
+        t_tile_length: int = 12,
+        t_tile_overlap: int = 4,
         num_inference_steps: int = 50,
         timesteps: Optional[List[int]] = None,
         guidance_scale: float = 6,
         denoise_strength: float = 1.0,
-        sigmas: Optional[List[float]] = None,
         num_videos_per_prompt: int = 1,
         eta: float = 0.0,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
@@ -415,13 +436,15 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
                 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                 argument.
         """
-        
+
+        #assert (
+        #    num_frames <= 48 and num_frames % fps == 0 and fps == 8
+        #), f"The number of frames must be divisible by {fps=} and less than 48 frames (for now). Other values are not supported in CogVideoX."
+
         height = height or self.transformer.config.sample_size * self.vae_scale_factor_spatial
         width = width or self.transformer.config.sample_size * self.vae_scale_factor_spatial
         num_videos_per_prompt = 1
 
-        self.num_frames = num_frames
-
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
             height,
@@ -439,36 +462,21 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
         # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
         # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
         # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale[0] > 1.0
+        do_classifier_free_guidance = guidance_scale > 1.0
 
         if do_classifier_free_guidance:
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
         prompt_embeds = prompt_embeds.to(self.vae.dtype)
 
         # 4. Prepare timesteps
-        if sigmas is None:
-            timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
-        else:
-            timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, sigmas=sigmas, device=device)
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
         self._num_timesteps = len(timesteps)
 
         # 5. Prepare latents.
         latent_channels = self.vae.config.latent_channels
-        latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-
-        # For CogVideoX 1.5, the latent frames should be padded to make it divisible by patch_size_t
-        patch_size_t = getattr(self.transformer.config, "patch_size_t", None)
-        if patch_size_t is None:
-            self.transformer.config.patch_size_t = None
-        ofs_embed_dim = getattr(self.transformer.config, "ofs_embed_dim", None)
-        if ofs_embed_dim is None:
-            self.transformer.config.ofs_embed_dim = None
-
-        self.additional_frames = 0
-        if patch_size_t is not None and latent_frames % patch_size_t != 0:
-            self.additional_frames = patch_size_t - latent_frames % patch_size_t
-            num_frames += self.additional_frames * self.vae_scale_factor_temporal
 
+        if latents is None and num_frames == t_tile_length:
+            num_frames += 1
 
         if self.original_mask is not None:
             image_latents = latents
@@ -492,6 +500,7 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
             freenoise=freenoise,
         )
         latents = latents.to(self.vae.dtype)
+        #print("latents", latents.shape)
 
         # 5.5.
         if image_cond_latents is not None:
@@ -505,11 +514,8 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
                 width // self.vae_scale_factor_spatial,
                 )
                 latent_padding = torch.zeros(padding_shape, device=device, dtype=self.vae.dtype)
-                image_cond_latents = torch.cat([image_cond_latents[:, 0, :, :, :].unsqueeze(1), latent_padding, image_cond_latents[:, -1, :, :, :].unsqueeze(1)], dim=1)
-                if self.transformer.config.patch_size_t is not None:
-                        first_frame = image_cond_latents[:, : image_cond_latents.size(1) % self.transformer.config.patch_size_t, ...]
-                        image_cond_latents = torch.cat([first_frame, image_cond_latents], dim=1)
 
+                image_cond_latents = torch.cat([image_cond_latents[:, 0, :, :, :].unsqueeze(1), latent_padding, image_cond_latents[:, -1, :, :, :].unsqueeze(1)], dim=1)
                 logger.info(f"image cond latents shape: {image_cond_latents.shape}")
             else:
                 logger.info("Only one image conditioning frame received, img2vid")
@@ -523,15 +529,18 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
                     )
                     latent_padding = torch.zeros(padding_shape, device=device, dtype=self.vae.dtype)
                     image_cond_latents = torch.cat([image_cond_latents, latent_padding], dim=1)
-                    # Select the first frame along the second dimension
-                    if self.transformer.config.patch_size_t is not None:
-                        first_frame = image_cond_latents[:, : image_cond_latents.size(1) % self.transformer.config.patch_size_t, ...]
-                        image_cond_latents = torch.cat([first_frame, image_cond_latents], dim=1)
                 else:
                     image_cond_latents = image_cond_latents.repeat(1, latents.shape[1], 1, 1, 1)
         # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
+        # 6.5. Create rotary embeds if required
+        image_rotary_emb = (
+            self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device)
+            if self.transformer.config.use_rotary_positional_embeddings
+            else None
+        )
+
         # masks
         if self.original_mask is not None:
             mask = self.original_mask.to(device)
@@ -545,35 +554,40 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
             logger.info(f"latents: {latents.shape}")
             logger.info(f"mask: {mask.shape}")
 
-      
+        # 7. Denoising loop
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        comfy_pbar = ProgressBar(num_inference_steps)
 
-        # 7. context schedule
-        if context_schedule is not None:
+        # 8. context schedule and temporal tiling
+        if context_schedule is not None and context_schedule == "temporal_tiling":
+            t_tile_length = context_frames
+            t_tile_overlap = context_overlap
+            t_tile_weights = self._gaussian_weights(t_tile_length=t_tile_length, t_batch_size=1).to(latents.device).to(self.vae.dtype)
+            use_temporal_tiling = True
+            logger.info("Temporal tiling enabled")
+        elif context_schedule is not None:
             if image_cond_latents is not None:
                 raise NotImplementedError("Context schedule not currently supported with image conditioning")
             logger.info(f"Context schedule enabled: {context_frames} frames, {context_stride} stride, {context_overlap} overlap")
+            use_temporal_tiling = False
             use_context_schedule = True
             from .cogvideox_fun.context import get_context_scheduler
             context = get_context_scheduler(context_schedule)
-            #todo ofs embeds?
 
         else:
+            use_temporal_tiling = False
             use_context_schedule = False
-            logger.info("Context schedule disabled")
-            # 7.5. Create rotary embeds if required
+            logger.info("Temporal tiling and context schedule disabled")
+            # 7. Create rotary embeds if required
             image_rotary_emb = (
                 self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device)
                 if self.transformer.config.use_rotary_positional_embeddings
                 else None
             )
-            # 7.6. Create ofs embeds if required
-            ofs_emb = None if self.transformer.config.ofs_embed_dim is None else latents.new_full((1,), fill_value=2.0)
-
             if tora is not None and do_classifier_free_guidance:
                 video_flow_features = tora["video_flow_features"].repeat(1, 2, 1, 1, 1).contiguous()
 
-        #8. Controlnet
+        # 9. Controlnet
         if controlnet is not None:
             self.controlnet = controlnet["control_model"].to(device)
             if self.transformer.dtype == torch.float8_e4m3fn:
@@ -602,26 +616,112 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
         if tora is not None:
             trajectory_length = tora["video_flow_features"].shape[1]
             logger.info(f"Tora trajectory length: {trajectory_length}")
-            #if trajectory_length != latents.shape[1]:
-            #    raise ValueError(f"Tora trajectory length {trajectory_length} does not match inpaint_latents count {latents.shape[2]}")
+            if trajectory_length != latents.shape[1]:
+                raise ValueError(f"Tora trajectory length {trajectory_length} does not match inpaint_latents count {latents.shape[2]}")
             for module in self.transformer.fuser_list:
                 for param in module.parameters():
                     param.data = param.data.to(device)
 
-        logger.info(f"Sampling {num_frames} frames in {latent_frames} latent frames at {width}x{height} with {num_inference_steps} inference steps")
-
-        from .latent_preview import prepare_callback
-        callback = prepare_callback(self.transformer, num_inference_steps)
-
-        # 9. Denoising loop
-        comfy_pbar = ProgressBar(len(timesteps))
-        with self.progress_bar(total=len(timesteps)) as progress_bar:    
+        # 10. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:    
             old_pred_original_sample = None # for DPM-solver++
             for i, t in enumerate(timesteps):
                 if self.interrupt:
                     continue
-                # region context schedule sampling
-                if use_context_schedule:
+                if use_temporal_tiling and isinstance(self.scheduler, CogVideoXDDIMScheduler):
+                    #temporal tiling code based on https://github.com/mayuelala/FollowYourEmoji/blob/main/models/video_pipeline.py
+                    # =====================================================
+                    grid_ts = 0
+                    cur_t = 0
+                    while cur_t < latents.shape[1]:
+                        cur_t = max(grid_ts * t_tile_length - t_tile_overlap * grid_ts, 0) + t_tile_length
+                        grid_ts += 1
+
+                    all_t = latents.shape[1]
+                    latents_all_list = []
+                    # =====================================================
+
+                    for t_i in range(grid_ts):
+                        if t_i < grid_ts - 1:
+                            ofs_t = max(t_i * t_tile_length - t_tile_overlap * t_i, 0)
+                        if t_i == grid_ts - 1:
+                            ofs_t = all_t - t_tile_length
+
+                        input_start_t = ofs_t
+                        input_end_t = ofs_t + t_tile_length
+
+                        #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)
+
+                        image_rotary_emb = (
+                            self._prepare_rotary_positional_embeddings(height, width, t_tile_length, device)
+                            if self.transformer.config.use_rotary_positional_embeddings
+                            else None
+                        )
+
+                        latents_tile = latents[:, input_start_t:input_end_t,:, :, :]
+                        latent_model_input_tile = torch.cat([latents_tile] * 2) if do_classifier_free_guidance else latents_tile
+                        latent_model_input_tile = self.scheduler.scale_model_input(latent_model_input_tile, t)
+
+                        #t_input = t[None].to(device)
+                        t_input = t.expand(latent_model_input_tile.shape[0]) # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                
+                        # predict noise model_output
+                        noise_pred = self.transformer(
+                            hidden_states=latent_model_input_tile,
+                            encoder_hidden_states=prompt_embeds,
+                            timestep=t_input,
+                            image_rotary_emb=image_rotary_emb,
+                            return_dict=False,
+                        )[0]
+                        noise_pred = noise_pred.float()                  
+
+                        if self.do_classifier_free_guidance:
+                            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                            noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                        # compute the previous noisy sample x_t -> x_t-1
+                        latents_tile = self.scheduler.step(noise_pred, t, latents_tile.to(self.vae.dtype), **extra_step_kwargs, return_dict=False)[0]
+                        latents_all_list.append(latents_tile)
+
+                    # ==========================================
+                    latents_all = torch.zeros(latents.shape, device=latents.device, dtype=self.vae.dtype)
+                    contributors = torch.zeros(latents.shape, device=latents.device, dtype=self.vae.dtype)
+                    # Add each tile contribution to overall latents
+                    for t_i in range(grid_ts):
+                        if t_i < grid_ts - 1:
+                            ofs_t = max(t_i * t_tile_length - t_tile_overlap * t_i, 0)
+                        if t_i == grid_ts - 1:
+                            ofs_t = all_t - t_tile_length
+
+                        input_start_t = ofs_t
+                        input_end_t = ofs_t + t_tile_length
+
+                        latents_all[:, input_start_t:input_end_t,:, :, :] += latents_all_list[t_i] * t_tile_weights
+                        contributors[:, input_start_t:input_end_t,:, :, :] += t_tile_weights
+                    
+                    latents_all /= contributors
+
+                    latents = latents_all
+                    #print("latents",latents.shape)
+                    # start diff diff
+                    if i < len(timesteps) - 1 and self.original_mask is not None:
+                        noise_timestep = timesteps[i + 1]
+                        image_latent = self.scheduler.add_noise(original_image_latents, noise, torch.tensor([noise_timestep])
+                        )
+                        mask = mask.to(latents)
+                        ts_from = timesteps[0]
+                        ts_to = timesteps[-1]
+                        threshold = (t - ts_to) / (ts_from - ts_to)
+                        mask = torch.where(mask >= threshold, mask, torch.zeros_like(mask))
+                        latents = image_latent * mask + latents * (1 - mask)
+                        # end diff diff
+                    
+                    if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                        progress_bar.update()
+                        comfy_pbar.update(1)
+                    # ==========================================
+                elif use_context_schedule:
                     latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                     counter = torch.zeros_like(latent_model_input)
@@ -719,7 +819,7 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
                     noise_pred /= counter
                     if do_classifier_free_guidance:
                         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                        noise_pred = noise_pred_uncond + self._guidance_scale[i] * (noise_pred_text - noise_pred_uncond)
+                        noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)
                        
                     # compute the previous noisy sample x_t -> x_t-1
                     if not isinstance(self.scheduler, CogVideoXDPMScheduler):
@@ -739,8 +839,7 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
                     if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                         progress_bar.update()
                         comfy_pbar.update(1)
-
-                # region sampling
+    
                 else:
                     latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
@@ -778,21 +877,22 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
                         encoder_hidden_states=prompt_embeds,
                         timestep=timestep,
                         image_rotary_emb=image_rotary_emb,
-                        ofs=ofs_emb,
                         return_dict=False,
                         controlnet_states=controlnet_states,
                         controlnet_weights=control_weights,
                         video_flow_features=video_flow_features if (tora is not None and tora["start_percent"] <= current_step_percentage <= tora["end_percent"]) else None,
+
                     )[0]
                     noise_pred = noise_pred.float()
+
                     if isinstance(self.scheduler, CogVideoXDPMScheduler):
-                        self._guidance_scale[i] = 1 + guidance_scale[i] * (
+                        self._guidance_scale = 1 + guidance_scale * (
                             (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
                         )
                     
                     if do_classifier_free_guidance:
                         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                        noise_pred = noise_pred_uncond + self._guidance_scale[i] * (noise_pred_text - noise_pred_uncond)
+                        noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                     # compute the previous noisy sample x_t -> x_t-1
                     if not isinstance(self.scheduler, CogVideoXDPMScheduler):
@@ -823,10 +923,7 @@ class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
 
                     if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                         progress_bar.update()
-                        if callback is not None:
-                            callback(i, latents.detach()[-1], None, num_inference_steps)
-                        else:
-                            comfy_pbar.update(1)
+                        comfy_pbar.update(1)
             
 
         # Offload all models

requirements.txt

@@ -3,4 +3,3 @@ diffusers>=0.31.0
 accelerate>=0.33.0
 einops
 peft
-opencv-python

utils.py

@@ -1,5 +1,5 @@
 import importlib.metadata
-import torch
+
 import logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 log = logging.getLogger(__name__)
@@ -19,12 +19,4 @@ def remove_specific_blocks(model, block_indices_to_remove):
     new_blocks = [block for i, block in enumerate(transformer_blocks) if i not in block_indices_to_remove]
     model.transformer_blocks = nn.ModuleList(new_blocks)
     
-    return model
-
-def print_memory(device):
-    memory = torch.cuda.memory_allocated(device) / 1024**3
-    max_memory = torch.cuda.max_memory_allocated(device) / 1024**3
-    max_reserved = torch.cuda.max_memory_reserved(device) / 1024**3
-    log.info(f"Allocated memory: {memory=:.3f} GB")
-    log.info(f"Max allocated memory: {max_memory=:.3f} GB")
-    log.info(f"Max reserved memory: {max_reserved=:.3f} GB")
+    return model

videosys/cogvideox_transformer_3d.py

@@ -0,0 +1,621 @@
+# Adapted from CogVideo
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# CogVideo: https://github.com/THUDM/CogVideo
+# diffusers: https://github.com/huggingface/diffusers
+# --------------------------------------------------------
+
+from typing import Any, Dict, Optional, Tuple, Union
+from einops import rearrange
+import torch
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention import Attention, FeedForward
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps, get_3d_sincos_pos_embed, CogVideoXPatchEmbed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import is_torch_version
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from torch import nn
+
+from .core.pab_mgr import enable_pab, if_broadcast_spatial
+from .modules.embeddings import apply_rotary_emb
+
+#from .modules.embeddings import CogVideoXPatchEmbed
+
+from .modules.normalization import AdaLayerNorm, CogVideoXLayerNormZero
+try:
+    from sageattention import sageattn
+    SAGEATTN_IS_AVAVILABLE = True
+except:
+    SAGEATTN_IS_AVAVILABLE = False
+
+class CogVideoXAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
+    query and key vectors, but does not include spatial normalization.
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+    @torch.compiler.disable()
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+        
+        attn_heads = attn.heads
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn_heads
+
+        query = query.view(batch_size, -1, attn_heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn_heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn_heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply RoPE if needed
+        if image_rotary_emb is not None:
+            emb_len = image_rotary_emb[0].shape[0]
+            query[:, :, text_seq_length : emb_len + text_seq_length] = apply_rotary_emb(
+                query[:, :, text_seq_length : emb_len + text_seq_length], image_rotary_emb
+            )
+            if not attn.is_cross_attention:
+                key[:, :, text_seq_length : emb_len + text_seq_length] = apply_rotary_emb(
+                    key[:, :, text_seq_length : emb_len + text_seq_length], image_rotary_emb
+                )
+
+        if SAGEATTN_IS_AVAVILABLE:
+            hidden_states = sageattn(query, key, value, is_causal=False)
+        else:
+            hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn_heads * head_dim)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+class FusedCogVideoXAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
+    query and key vectors, but does not include spatial normalization.
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+    @torch.compiler.disable()
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        qkv = attn.to_qkv(hidden_states)
+        split_size = qkv.shape[-1] // 3
+        query, key, value = torch.split(qkv, split_size, dim=-1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply RoPE if needed
+        if image_rotary_emb is not None:
+            query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
+            if not attn.is_cross_attention:
+                key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
+
+        if SAGEATTN_IS_AVAVILABLE:
+            hidden_states = sageattn(query, key, value, is_causal=False)
+        else:
+            hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+@maybe_allow_in_graph
+class CogVideoXBlock(nn.Module):
+    r"""
+    Transformer block used in [CogVideoX](https://github.com/THUDM/CogVideo) model.
+
+    Parameters:
+        dim (`int`):
+            The number of channels in the input and output.
+        num_attention_heads (`int`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`):
+            The number of channels in each head.
+        time_embed_dim (`int`):
+            The number of channels in timestep embedding.
+        dropout (`float`, defaults to `0.0`):
+            The dropout probability to use.
+        activation_fn (`str`, defaults to `"gelu-approximate"`):
+            Activation function to be used in feed-forward.
+        attention_bias (`bool`, defaults to `False`):
+            Whether or not to use bias in attention projection layers.
+        qk_norm (`bool`, defaults to `True`):
+            Whether or not to use normalization after query and key projections in Attention.
+        norm_elementwise_affine (`bool`, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_eps (`float`, defaults to `1e-5`):
+            Epsilon value for normalization layers.
+        final_dropout (`bool` defaults to `False`):
+            Whether to apply a final dropout after the last feed-forward layer.
+        ff_inner_dim (`int`, *optional*, defaults to `None`):
+            Custom hidden dimension of Feed-forward layer. If not provided, `4 * dim` is used.
+        ff_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in Feed-forward layer.
+        attention_out_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in Attention output projection layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        time_embed_dim: int,
+        dropout: float = 0.0,
+        activation_fn: str = "gelu-approximate",
+        attention_bias: bool = False,
+        qk_norm: bool = True,
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        final_dropout: bool = True,
+        ff_inner_dim: Optional[int] = None,
+        ff_bias: bool = True,
+        attention_out_bias: bool = True,
+        block_idx: int = 0,
+    ):
+        super().__init__()
+
+        # 1. Self Attention
+        self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
+
+        self.attn1 = Attention(
+            query_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            qk_norm="layer_norm" if qk_norm else None,
+            eps=1e-6,
+            bias=attention_bias,
+            out_bias=attention_out_bias,
+            processor=CogVideoXAttnProcessor2_0(),
+        )
+
+        # parallel
+        #self.attn1.parallel_manager = None
+
+        # 2. Feed Forward
+        self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
+
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+            inner_dim=ff_inner_dim,
+            bias=ff_bias,
+        )
+
+        # pab
+        self.attn_count = 0
+        self.last_attn = None
+        self.block_idx = block_idx
+    #@torch.compiler.disable()
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        timestep=None,
+        video_flow_feature: Optional[torch.Tensor] = None,
+        fuser=None,
+    ) -> torch.Tensor:
+        text_seq_length = encoder_hidden_states.size(1)
+
+        # norm & modulate
+        norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
+            hidden_states, encoder_hidden_states, temb
+        )
+         # Tora Motion-guidance Fuser
+        if video_flow_feature is not None:
+            H, W = video_flow_feature.shape[-2:]
+            T = norm_hidden_states.shape[1] // H // W
+            h = rearrange(norm_hidden_states, "B (T H W) C -> (B T) C H W", H=H, W=W)
+            h = fuser(h, video_flow_feature.to(h), T=T)
+            norm_hidden_states = rearrange(h, "(B T) C H W ->  B (T H W) C", T=T)
+            del h, fuser
+        # attention
+        if enable_pab():
+            broadcast_attn, self.attn_count = if_broadcast_spatial(int(timestep[0]), self.attn_count, self.block_idx)
+        if enable_pab() and broadcast_attn:
+            attn_hidden_states, attn_encoder_hidden_states = self.last_attn
+        else:
+            attn_hidden_states, attn_encoder_hidden_states = self.attn1(
+                hidden_states=norm_hidden_states,
+                encoder_hidden_states=norm_encoder_hidden_states,
+                image_rotary_emb=image_rotary_emb,
+            )
+            if enable_pab():
+                self.last_attn = (attn_hidden_states, attn_encoder_hidden_states)
+
+        hidden_states = hidden_states + gate_msa * attn_hidden_states
+        encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states
+
+        # norm & modulate
+        norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
+            hidden_states, encoder_hidden_states, temb
+        )
+
+        # feed-forward
+        norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
+        ff_output = self.ff(norm_hidden_states)
+
+        hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]
+        encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]
+
+        return hidden_states, encoder_hidden_states
+
+
+class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin):
+    """
+    A Transformer model for video-like data in [CogVideoX](https://github.com/THUDM/CogVideo).
+
+    Parameters:
+        num_attention_heads (`int`, defaults to `30`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, defaults to `64`):
+            The number of channels in each head.
+        in_channels (`int`, defaults to `16`):
+            The number of channels in the input.
+        out_channels (`int`, *optional*, defaults to `16`):
+            The number of channels in the output.
+        flip_sin_to_cos (`bool`, defaults to `True`):
+            Whether to flip the sin to cos in the time embedding.
+        time_embed_dim (`int`, defaults to `512`):
+            Output dimension of timestep embeddings.
+        text_embed_dim (`int`, defaults to `4096`):
+            Input dimension of text embeddings from the text encoder.
+        num_layers (`int`, defaults to `30`):
+            The number of layers of Transformer blocks to use.
+        dropout (`float`, defaults to `0.0`):
+            The dropout probability to use.
+        attention_bias (`bool`, defaults to `True`):
+            Whether or not to use bias in the attention projection layers.
+        sample_width (`int`, defaults to `90`):
+            The width of the input latents.
+        sample_height (`int`, defaults to `60`):
+            The height of the input latents.
+        sample_frames (`int`, defaults to `49`):
+            The number of frames in the input latents. Note that this parameter was incorrectly initialized to 49
+            instead of 13 because CogVideoX processed 13 latent frames at once in its default and recommended settings,
+            but cannot be changed to the correct value to ensure backwards compatibility. To create a transformer with
+            K latent frames, the correct value to pass here would be: ((K - 1) * temporal_compression_ratio + 1).
+        patch_size (`int`, defaults to `2`):
+            The size of the patches to use in the patch embedding layer.
+        temporal_compression_ratio (`int`, defaults to `4`):
+            The compression ratio across the temporal dimension. See documentation for `sample_frames`.
+        max_text_seq_length (`int`, defaults to `226`):
+            The maximum sequence length of the input text embeddings.
+        activation_fn (`str`, defaults to `"gelu-approximate"`):
+            Activation function to use in feed-forward.
+        timestep_activation_fn (`str`, defaults to `"silu"`):
+            Activation function to use when generating the timestep embeddings.
+        norm_elementwise_affine (`bool`, defaults to `True`):
+            Whether or not to use elementwise affine in normalization layers.
+        norm_eps (`float`, defaults to `1e-5`):
+            The epsilon value to use in normalization layers.
+        spatial_interpolation_scale (`float`, defaults to `1.875`):
+            Scaling factor to apply in 3D positional embeddings across spatial dimensions.
+        temporal_interpolation_scale (`float`, defaults to `1.0`):
+            Scaling factor to apply in 3D positional embeddings across temporal dimensions.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 30,
+        attention_head_dim: int = 64,
+        in_channels: int = 16,
+        out_channels: Optional[int] = 16,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        time_embed_dim: int = 512,
+        text_embed_dim: int = 4096,
+        num_layers: int = 30,
+        dropout: float = 0.0,
+        attention_bias: bool = True,
+        sample_width: int = 90,
+        sample_height: int = 60,
+        sample_frames: int = 49,
+        patch_size: int = 2,
+        temporal_compression_ratio: int = 4,
+        max_text_seq_length: int = 226,
+        activation_fn: str = "gelu-approximate",
+        timestep_activation_fn: str = "silu",
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        spatial_interpolation_scale: float = 1.875,
+        temporal_interpolation_scale: float = 1.0,
+        use_rotary_positional_embeddings: bool = False,
+        use_learned_positional_embeddings: bool = False,
+    ):
+        super().__init__()
+        inner_dim = num_attention_heads * attention_head_dim
+
+        post_patch_height = sample_height // patch_size
+        post_patch_width = sample_width // patch_size
+        post_time_compression_frames = (sample_frames - 1) // temporal_compression_ratio + 1
+        self.num_patches = post_patch_height * post_patch_width * post_time_compression_frames
+
+        # 1. Patch embedding
+        self.patch_embed = CogVideoXPatchEmbed(
+            patch_size=patch_size,
+            in_channels=in_channels,
+            embed_dim=inner_dim,
+            text_embed_dim=text_embed_dim,
+            bias=True,
+            sample_width=sample_width,
+            sample_height=sample_height,
+            sample_frames=sample_frames,
+            temporal_compression_ratio=temporal_compression_ratio,
+            max_text_seq_length=max_text_seq_length,
+            spatial_interpolation_scale=spatial_interpolation_scale,
+            temporal_interpolation_scale=temporal_interpolation_scale,
+            use_positional_embeddings=not use_rotary_positional_embeddings,
+            use_learned_positional_embeddings=use_learned_positional_embeddings,
+        )
+        self.embedding_dropout = nn.Dropout(dropout)
+
+        # 2. 3D positional embeddings
+        spatial_pos_embedding = get_3d_sincos_pos_embed(
+            inner_dim,
+            (post_patch_width, post_patch_height),
+            post_time_compression_frames,
+            spatial_interpolation_scale,
+            temporal_interpolation_scale,
+        )
+        spatial_pos_embedding = torch.from_numpy(spatial_pos_embedding).flatten(0, 1)
+        pos_embedding = torch.zeros(1, max_text_seq_length + self.num_patches, inner_dim, requires_grad=False)
+        pos_embedding.data[:, max_text_seq_length:].copy_(spatial_pos_embedding)
+        self.register_buffer("pos_embedding", pos_embedding, persistent=False)
+
+        # 3. Time embeddings
+        self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift)
+        self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn)
+
+        # 4. Define spatio-temporal transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                CogVideoXBlock(
+                    dim=inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    time_embed_dim=time_embed_dim,
+                    dropout=dropout,
+                    activation_fn=activation_fn,
+                    attention_bias=attention_bias,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                    norm_eps=norm_eps,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm_final = nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine)
+
+        # 5. Output blocks
+        self.norm_out = AdaLayerNorm(
+            embedding_dim=time_embed_dim,
+            output_dim=2 * inner_dim,
+            norm_elementwise_affine=norm_elementwise_affine,
+            norm_eps=norm_eps,
+            chunk_dim=1,
+        )
+        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
+
+        self.gradient_checkpointing = False
+
+        self.fuser_list = None
+
+        # parallel
+        #self.parallel_manager = None
+
+    # def enable_parallel(self, dp_size, sp_size, enable_cp):
+    #     # update cfg parallel
+    #     if enable_cp and sp_size % 2 == 0:
+    #         sp_size = sp_size // 2
+    #         cp_size = 2
+    #     else:
+    #         cp_size = 1
+
+    #     self.parallel_manager: ParallelManager = ParallelManager(dp_size, cp_size, sp_size)
+
+    #     for _, module in self.named_modules():
+    #         if hasattr(module, "parallel_manager"):
+    #             module.parallel_manager = self.parallel_manager
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        self.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: Union[int, float, torch.LongTensor],
+        timestep_cond: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        return_dict: bool = True,
+        controlnet_states: torch.Tensor = None,
+        controlnet_weights: Optional[Union[float, int, list, torch.FloatTensor]] = 1.0,
+        video_flow_features: Optional[torch.Tensor] = None,
+    ):
+        # if self.parallel_manager.cp_size > 1:
+        #     (
+        #         hidden_states,
+        #         encoder_hidden_states,
+        #         timestep,
+        #         timestep_cond,
+        #         image_rotary_emb,
+        #     ) = batch_func(
+        #         partial(split_sequence, process_group=self.parallel_manager.cp_group, dim=0),
+        #         hidden_states,
+        #         encoder_hidden_states,
+        #         timestep,
+        #         timestep_cond,
+        #         image_rotary_emb,
+        #     )
+        
+        batch_size, num_frames, channels, height, width = hidden_states.shape
+
+        # 1. Time embedding
+        timesteps = timestep
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=hidden_states.dtype)
+        emb = self.time_embedding(t_emb, timestep_cond)
+
+        # 2. Patch embedding
+        hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)
+
+        # 3. Position embedding
+        text_seq_length = encoder_hidden_states.shape[1]
+        if not self.config.use_rotary_positional_embeddings:
+            seq_length = height * width * num_frames // (self.config.patch_size**2)
+
+            pos_embeds = self.pos_embedding[:, : text_seq_length + seq_length]
+            hidden_states = hidden_states + pos_embeds
+            hidden_states = self.embedding_dropout(hidden_states)
+
+        encoder_hidden_states = hidden_states[:, :text_seq_length]
+        hidden_states = hidden_states[:, text_seq_length:]
+
+        # if self.parallel_manager.sp_size > 1:
+        #     set_pad("pad", hidden_states.shape[1], self.parallel_manager.sp_group)
+        #     hidden_states = split_sequence(hidden_states, self.parallel_manager.sp_group, dim=1, pad=get_pad("pad"))
+
+        # 4. Transformer blocks
+        for i, block in enumerate(self.transformer_blocks):
+            hidden_states, encoder_hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=emb,
+                image_rotary_emb=image_rotary_emb,
+                timestep=timesteps if enable_pab() else None,
+                video_flow_feature=video_flow_features[i] if video_flow_features is not None else None,
+                fuser = self.fuser_list[i] if self.fuser_list is not None else None,
+            )
+            if (controlnet_states is not None) and (i < len(controlnet_states)):
+                controlnet_states_block = controlnet_states[i]
+                controlnet_block_weight = 1.0
+                if isinstance(controlnet_weights, (list)) or torch.is_tensor(controlnet_weights):
+                    controlnet_block_weight = controlnet_weights[i]
+                elif isinstance(controlnet_weights, (float, int)):
+                    controlnet_block_weight = controlnet_weights
+                
+                hidden_states = hidden_states + controlnet_states_block * controlnet_block_weight
+
+        #if self.parallel_manager.sp_size > 1:
+        #    hidden_states = gather_sequence(hidden_states, self.parallel_manager.sp_group, dim=1, pad=get_pad("pad"))
+
+        if not self.config.use_rotary_positional_embeddings:
+            # CogVideoX-2B
+            hidden_states = self.norm_final(hidden_states)
+        else:
+            # CogVideoX-5B
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+            hidden_states = self.norm_final(hidden_states)
+            hidden_states = hidden_states[:, text_seq_length:]
+
+        # 5. Final block
+        hidden_states = self.norm_out(hidden_states, temb=emb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # 6. Unpatchify
+        p = self.config.patch_size
+        output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)
+        output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
+
+        #if self.parallel_manager.cp_size > 1:
+        #    output = gather_sequence(output, self.parallel_manager.cp_group, dim=0)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)

videosys/core/pab_mgr.py

@@ -0,0 +1,232 @@
+
+PAB_MANAGER = None
+
+
+class PABConfig:
+    def __init__(
+        self,
+        steps: int,
+        cross_broadcast: bool = False,
+        cross_threshold: list = None,
+        cross_range: int = None,
+        spatial_broadcast: bool = False,
+        spatial_threshold: list = None,
+        spatial_range: int = None,
+        temporal_broadcast: bool = False,
+        temporal_threshold: list = None,
+        temporal_range: int = None,
+        mlp_broadcast: bool = False,
+        mlp_spatial_broadcast_config: dict = None,
+        mlp_temporal_broadcast_config: dict = None,
+    ):
+        self.steps = steps
+
+        self.cross_broadcast = cross_broadcast
+        self.cross_threshold = cross_threshold
+        self.cross_range = cross_range
+
+        self.spatial_broadcast = spatial_broadcast
+        self.spatial_threshold = spatial_threshold
+        self.spatial_range = spatial_range
+
+        self.temporal_broadcast = temporal_broadcast
+        self.temporal_threshold = temporal_threshold
+        self.temporal_range = temporal_range
+
+        self.mlp_broadcast = mlp_broadcast
+        self.mlp_spatial_broadcast_config = mlp_spatial_broadcast_config
+        self.mlp_temporal_broadcast_config = mlp_temporal_broadcast_config
+        self.mlp_temporal_outputs = {}
+        self.mlp_spatial_outputs = {}
+
+
+class PABManager:
+    def __init__(self, config: PABConfig):
+        self.config: PABConfig = config
+
+        init_prompt = f"Init Pyramid Attention Broadcast. steps: {config.steps}."
+        init_prompt += f" spatial broadcast: {config.spatial_broadcast}, spatial range: {config.spatial_range}, spatial threshold: {config.spatial_threshold}."
+        init_prompt += f" temporal broadcast: {config.temporal_broadcast}, temporal range: {config.temporal_range}, temporal_threshold: {config.temporal_threshold}."
+        init_prompt += f" cross broadcast: {config.cross_broadcast}, cross range: {config.cross_range}, cross threshold: {config.cross_threshold}."
+        init_prompt += f" mlp broadcast: {config.mlp_broadcast}."
+        print(init_prompt)
+
+    def if_broadcast_cross(self, timestep: int, count: int):
+        if (
+            self.config.cross_broadcast
+            and (timestep is not None)
+            and (count % self.config.cross_range != 0)
+            and (self.config.cross_threshold[0] < timestep < self.config.cross_threshold[1])
+        ):
+            flag = True
+        else:
+            flag = False
+        count = (count + 1) % self.config.steps
+        return flag, count
+
+    def if_broadcast_temporal(self, timestep: int, count: int):
+        if (
+            self.config.temporal_broadcast
+            and (timestep is not None)
+            and (count % self.config.temporal_range != 0)
+            and (self.config.temporal_threshold[0] < timestep < self.config.temporal_threshold[1])
+        ):
+            flag = True
+        else:
+            flag = False
+        count = (count + 1) % self.config.steps
+        return flag, count
+
+    def if_broadcast_spatial(self, timestep: int, count: int, block_idx: int):
+        if (
+            self.config.spatial_broadcast
+            and (timestep is not None)
+            and (count % self.config.spatial_range != 0)
+            and (self.config.spatial_threshold[0] < timestep < self.config.spatial_threshold[1])
+        ):
+            flag = True
+        else:
+            flag = False
+        count = (count + 1) % self.config.steps
+        return flag, count
+
+    @staticmethod
+    def _is_t_in_skip_config(all_timesteps, timestep, config):
+        is_t_in_skip_config = False
+        skip_range = None
+        for key in config:
+            if key not in all_timesteps:
+                continue
+            index = all_timesteps.index(key)
+            skip_range = all_timesteps[index : index + 1 + int(config[key]["skip_count"])]
+            if timestep in skip_range:
+                is_t_in_skip_config = True
+                skip_range = [all_timesteps[index], all_timesteps[index + int(config[key]["skip_count"])]]
+                break
+        return is_t_in_skip_config, skip_range
+
+    def if_skip_mlp(self, timestep: int, count: int, block_idx: int, all_timesteps, is_temporal=False):
+        if not self.config.mlp_broadcast:
+            return False, None, False, None
+
+        if is_temporal:
+            cur_config = self.config.mlp_temporal_broadcast_config
+        else:
+            cur_config = self.config.mlp_spatial_broadcast_config
+
+        is_t_in_skip_config, skip_range = self._is_t_in_skip_config(all_timesteps, timestep, cur_config)
+        next_flag = False
+        if (
+            self.config.mlp_broadcast
+            and (timestep is not None)
+            and (timestep in cur_config)
+            and (block_idx in cur_config[timestep]["block"])
+        ):
+            flag = False
+            next_flag = True
+            count = count + 1
+        elif (
+            self.config.mlp_broadcast
+            and (timestep is not None)
+            and (is_t_in_skip_config)
+            and (block_idx in cur_config[skip_range[0]]["block"])
+        ):
+            flag = True
+            count = 0
+        else:
+            flag = False
+
+        return flag, count, next_flag, skip_range
+
+    def save_skip_output(self, timestep, block_idx, ff_output, is_temporal=False):
+        if is_temporal:
+            self.config.mlp_temporal_outputs[(timestep, block_idx)] = ff_output
+        else:
+            self.config.mlp_spatial_outputs[(timestep, block_idx)] = ff_output
+
+    def get_mlp_output(self, skip_range, timestep, block_idx, is_temporal=False):
+        skip_start_t = skip_range[0]
+        if is_temporal:
+            skip_output = (
+                self.config.mlp_temporal_outputs.get((skip_start_t, block_idx), None)
+                if self.config.mlp_temporal_outputs is not None
+                else None
+            )
+        else:
+            skip_output = (
+                self.config.mlp_spatial_outputs.get((skip_start_t, block_idx), None)
+                if self.config.mlp_spatial_outputs is not None
+                else None
+            )
+
+        if skip_output is not None:
+            if timestep == skip_range[-1]:
+                # TODO: save memory
+                if is_temporal:
+                    del self.config.mlp_temporal_outputs[(skip_start_t, block_idx)]
+                else:
+                    del self.config.mlp_spatial_outputs[(skip_start_t, block_idx)]
+        else:
+            raise ValueError(
+                f"No stored MLP output found | t {timestep} |[{skip_range[0]}, {skip_range[-1]}] | block {block_idx}"
+            )
+
+        return skip_output
+
+    def get_spatial_mlp_outputs(self):
+        return self.config.mlp_spatial_outputs
+
+    def get_temporal_mlp_outputs(self):
+        return self.config.mlp_temporal_outputs
+
+
+def set_pab_manager(config: PABConfig):
+    global PAB_MANAGER
+    PAB_MANAGER = PABManager(config)
+
+
+def enable_pab():
+    if PAB_MANAGER is None:
+        return False
+    return (
+        PAB_MANAGER.config.cross_broadcast
+        or PAB_MANAGER.config.spatial_broadcast
+        or PAB_MANAGER.config.temporal_broadcast
+    )
+
+
+def update_steps(steps: int):
+    if PAB_MANAGER is not None:
+        PAB_MANAGER.config.steps = steps
+
+
+def if_broadcast_cross(timestep: int, count: int):
+    if not enable_pab():
+        return False, count
+    return PAB_MANAGER.if_broadcast_cross(timestep, count)
+
+
+def if_broadcast_temporal(timestep: int, count: int):
+    if not enable_pab():
+        return False, count
+    return PAB_MANAGER.if_broadcast_temporal(timestep, count)
+
+
+def if_broadcast_spatial(timestep: int, count: int, block_idx: int):
+    if not enable_pab():
+        return False, count
+    return PAB_MANAGER.if_broadcast_spatial(timestep, count, block_idx)
+
+
+def if_broadcast_mlp(timestep: int, count: int, block_idx: int, all_timesteps, is_temporal=False):
+    if not enable_pab():
+        return False, count
+    return PAB_MANAGER.if_skip_mlp(timestep, count, block_idx, all_timesteps, is_temporal)
+
+
+def save_mlp_output(timestep: int, block_idx: int, ff_output, is_temporal=False):
+    return PAB_MANAGER.save_skip_output(timestep, block_idx, ff_output, is_temporal)
+
+
+def get_mlp_output(skip_range, timestep, block_idx: int, is_temporal=False):
+    return PAB_MANAGER.get_mlp_output(skip_range, timestep, block_idx, is_temporal)

videosys/core/pipeline.py

@@ -0,0 +1,44 @@
+import inspect
+from abc import abstractmethod
+
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+class VideoSysPipeline(DiffusionPipeline):
+    def __init__(self):
+        super().__init__()
+
+    @staticmethod
+    def set_eval_and_device(device: torch.device, *modules):
+        for module in modules:
+            module.eval()
+            module.to(device)
+
+    @abstractmethod
+    def generate(self, *args, **kwargs):
+        pass
+
+    def __call__(self, *args, **kwargs):
+        """
+        In diffusers, it is a convention to call the pipeline object.
+        But in VideoSys, we will use the generate method for better prompt.
+        This is a wrapper for the generate method to support the diffusers usage.
+        """
+        return self.generate(*args, **kwargs)
+
+    @classmethod
+    def _get_signature_keys(cls, obj):
+        parameters = inspect.signature(obj.__init__).parameters
+        required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}
+        optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})
+        expected_modules = set(required_parameters.keys()) - {"self"}
+        # modify: remove the config module from the expected modules
+        expected_modules = expected_modules - {"config"}
+
+        optional_names = list(optional_parameters)
+        for name in optional_names:
+            if name in cls._optional_components:
+                expected_modules.add(name)
+                optional_parameters.remove(name)
+
+        return expected_modules, optional_parameters

videosys/modules/activations.py

@@ -0,0 +1,3 @@
+import torch.nn as nn
+
+approx_gelu = lambda: nn.GELU(approximate="tanh")

videosys/modules/downsampling.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CogVideoXDownsample3D(nn.Module):
+    # Todo: Wait for paper relase.
+    r"""
+    A 3D Downsampling layer using in [CogVideoX]() by Tsinghua University & ZhipuAI
+
+    Args:
+        in_channels (`int`):
+            Number of channels in the input image.
+        out_channels (`int`):
+            Number of channels produced by the convolution.
+        kernel_size (`int`, defaults to `3`):
+            Size of the convolving kernel.
+        stride (`int`, defaults to `2`):
+            Stride of the convolution.
+        padding (`int`, defaults to `0`):
+            Padding added to all four sides of the input.
+        compress_time (`bool`, defaults to `False`):
+            Whether or not to compress the time dimension.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 2,
+        padding: int = 0,
+        compress_time: bool = False,
+    ):
+        super().__init__()
+
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
+        self.compress_time = compress_time
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.compress_time:
+            batch_size, channels, frames, height, width = x.shape
+
+            # (batch_size, channels, frames, height, width) -> (batch_size, height, width, channels, frames) -> (batch_size * height * width, channels, frames)
+            x = x.permute(0, 3, 4, 1, 2).reshape(batch_size * height * width, channels, frames)
+
+            if x.shape[-1] % 2 == 1:
+                x_first, x_rest = x[..., 0], x[..., 1:]
+                if x_rest.shape[-1] > 0:
+                    # (batch_size * height * width, channels, frames - 1) -> (batch_size * height * width, channels, (frames - 1) // 2)
+                    x_rest = F.avg_pool1d(x_rest, kernel_size=2, stride=2)
+
+                x = torch.cat([x_first[..., None], x_rest], dim=-1)
+                # (batch_size * height * width, channels, (frames // 2) + 1) -> (batch_size, height, width, channels, (frames // 2) + 1) -> (batch_size, channels, (frames // 2) + 1, height, width)
+                x = x.reshape(batch_size, height, width, channels, x.shape[-1]).permute(0, 3, 4, 1, 2)
+            else:
+                # (batch_size * height * width, channels, frames) -> (batch_size * height * width, channels, frames // 2)
+                x = F.avg_pool1d(x, kernel_size=2, stride=2)
+                # (batch_size * height * width, channels, frames // 2) -> (batch_size, height, width, channels, frames // 2) -> (batch_size, channels, frames // 2, height, width)
+                x = x.reshape(batch_size, height, width, channels, x.shape[-1]).permute(0, 3, 4, 1, 2)
+
+        # Pad the tensor
+        pad = (0, 1, 0, 1)
+        x = F.pad(x, pad, mode="constant", value=0)
+        batch_size, channels, frames, height, width = x.shape
+        # (batch_size, channels, frames, height, width) -> (batch_size, frames, channels, height, width) -> (batch_size * frames, channels, height, width)
+        x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * frames, channels, height, width)
+        x = self.conv(x)
+        # (batch_size * frames, channels, height, width) -> (batch_size, frames, channels, height, width) -> (batch_size, channels, frames, height, width)
+        x = x.reshape(batch_size, frames, x.shape[1], x.shape[2], x.shape[3]).permute(0, 2, 1, 3, 4)
+        return x

videosys/modules/embeddings.py

@@ -0,0 +1,308 @@
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+
+class CogVideoXPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 2,
+        in_channels: int = 16,
+        embed_dim: int = 1920,
+        text_embed_dim: int = 4096,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias
+        )
+        self.text_proj = nn.Linear(text_embed_dim, embed_dim)
+
+    def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):
+        r"""
+        Args:
+            text_embeds (`torch.Tensor`):
+                Input text embeddings. Expected shape: (batch_size, seq_length, embedding_dim).
+            image_embeds (`torch.Tensor`):
+                Input image embeddings. Expected shape: (batch_size, num_frames, channels, height, width).
+        """
+        text_embeds = self.text_proj(text_embeds)
+
+        batch, num_frames, channels, height, width = image_embeds.shape
+        image_embeds = image_embeds.reshape(-1, channels, height, width)
+        image_embeds = self.proj(image_embeds)
+        image_embeds = image_embeds.view(batch, num_frames, *image_embeds.shape[1:])
+        image_embeds = image_embeds.flatten(3).transpose(2, 3)  # [batch, num_frames, height x width, channels]
+        image_embeds = image_embeds.flatten(1, 2)  # [batch, num_frames x height x width, channels]
+
+        embeds = torch.cat(
+            [text_embeds, image_embeds], dim=1
+        ).contiguous()  # [batch, seq_length + num_frames x height x width, channels]
+        return embeds
+
+
+class OpenSoraPatchEmbed3D(nn.Module):
+    """Video to Patch Embedding.
+
+    Args:
+        patch_size (int): Patch token size. Default: (2,4,4).
+        in_chans (int): Number of input video channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(
+        self,
+        patch_size=(2, 4, 4),
+        in_chans=3,
+        embed_dim=96,
+        norm_layer=None,
+        flatten=True,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.flatten = flatten
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        """Forward function."""
+        # padding
+        _, _, D, H, W = x.size()
+        if W % self.patch_size[2] != 0:
+            x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
+        if H % self.patch_size[1] != 0:
+            x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
+        if D % self.patch_size[0] != 0:
+            x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))
+
+        x = self.proj(x)  # (B C T H W)
+        if self.norm is not None:
+            D, Wh, Ww = x.size(2), x.size(3), x.size(4)
+            x = x.flatten(2).transpose(1, 2)
+            x = self.norm(x)
+            x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCTHW -> BNC
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
+        freqs = freqs.to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t, dtype):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        if t_freq.dtype != dtype:
+            t_freq = t_freq.to(dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class SizeEmbedder(TimestepEmbedder):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__(hidden_size=hidden_size, frequency_embedding_size=frequency_embedding_size)
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+        self.outdim = hidden_size
+
+    def forward(self, s, bs):
+        if s.ndim == 1:
+            s = s[:, None]
+        assert s.ndim == 2
+        if s.shape[0] != bs:
+            s = s.repeat(bs // s.shape[0], 1)
+            assert s.shape[0] == bs
+        b, dims = s.shape[0], s.shape[1]
+        s = rearrange(s, "b d -> (b d)")
+        s_freq = self.timestep_embedding(s, self.frequency_embedding_size).to(self.dtype)
+        s_emb = self.mlp(s_freq)
+        s_emb = rearrange(s_emb, "(b d) d2 -> b (d d2)", b=b, d=dims, d2=self.outdim)
+        return s_emb
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+def get_3d_rotary_pos_embed(
+    embed_dim, crops_coords, grid_size, temporal_size, theta: int = 10000, use_real: bool = True
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    RoPE for video tokens with 3D structure.
+
+    Args:
+    embed_dim: (`int`):
+        The embedding dimension size, corresponding to hidden_size_head.
+    crops_coords (`Tuple[int]`):
+        The top-left and bottom-right coordinates of the crop.
+    grid_size (`Tuple[int]`):
+        The grid size of the spatial positional embedding (height, width).
+    temporal_size (`int`):
+        The size of the temporal dimension.
+    theta (`float`):
+        Scaling factor for frequency computation.
+    use_real (`bool`):
+        If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+
+    Returns:
+        `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.
+    """
+    start, stop = crops_coords
+    grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32)
+    grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32)
+    grid_t = np.linspace(0, temporal_size, temporal_size, endpoint=False, dtype=np.float32)
+
+    # Compute dimensions for each axis
+    dim_t = embed_dim // 4
+    dim_h = embed_dim // 8 * 3
+    dim_w = embed_dim // 8 * 3
+
+    # Temporal frequencies
+    freqs_t = 1.0 / (theta ** (torch.arange(0, dim_t, 2).float() / dim_t))
+    grid_t = torch.from_numpy(grid_t).float()
+    freqs_t = torch.einsum("n , f -> n f", grid_t, freqs_t)
+    freqs_t = freqs_t.repeat_interleave(2, dim=-1)
+
+    # Spatial frequencies for height and width
+    freqs_h = 1.0 / (theta ** (torch.arange(0, dim_h, 2).float() / dim_h))
+    freqs_w = 1.0 / (theta ** (torch.arange(0, dim_w, 2).float() / dim_w))
+    grid_h = torch.from_numpy(grid_h).float()
+    grid_w = torch.from_numpy(grid_w).float()
+    freqs_h = torch.einsum("n , f -> n f", grid_h, freqs_h)
+    freqs_w = torch.einsum("n , f -> n f", grid_w, freqs_w)
+    freqs_h = freqs_h.repeat_interleave(2, dim=-1)
+    freqs_w = freqs_w.repeat_interleave(2, dim=-1)
+
+    # Broadcast and concatenate tensors along specified dimension
+    def broadcast(tensors, dim=-1):
+        num_tensors = len(tensors)
+        shape_lens = {len(t.shape) for t in tensors}
+        assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+        shape_len = list(shape_lens)[0]
+        dim = (dim + shape_len) if dim < 0 else dim
+        dims = list(zip(*(list(t.shape) for t in tensors)))
+        expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+        assert all(
+            [*(len(set(t[1])) <= 2 for t in expandable_dims)]
+        ), "invalid dimensions for broadcastable concatenation"
+        max_dims = [(t[0], max(t[1])) for t in expandable_dims]
+        expanded_dims = [(t[0], (t[1],) * num_tensors) for t in max_dims]
+        expanded_dims.insert(dim, (dim, dims[dim]))
+        expandable_shapes = list(zip(*(t[1] for t in expanded_dims)))
+        tensors = [t[0].expand(*t[1]) for t in zip(tensors, expandable_shapes)]
+        return torch.cat(tensors, dim=dim)
+
+    freqs = broadcast((freqs_t[:, None, None, :], freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim=-1)
+
+    t, h, w, d = freqs.shape
+    freqs = freqs.view(t * h * w, d)
+
+    # Generate sine and cosine components
+    sin = freqs.sin()
+    cos = freqs.cos()
+
+    if use_real:
+        return cos, sin
+    else:
+        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+        return freqs_cis
+
+
+def apply_rotary_emb(
+    x: torch.Tensor,
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+    use_real: bool = True,
+    use_real_unbind_dim: int = -1,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+    to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+    reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+    tensors contain rotary embeddings and are returned as real tensors.
+
+    Args:
+        x (`torch.Tensor`):
+            Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
+        freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+    """
+    if use_real:
+        cos, sin = freqs_cis  # [S, D]
+        cos = cos[None, None]
+        sin = sin[None, None]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+
+        if use_real_unbind_dim == -1:
+            # Use for example in Lumina
+            x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+        elif use_real_unbind_dim == -2:
+            # Use for example in Stable Audio
+            x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
+            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+        else:
+            raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+        return out
+    else:
+        x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        freqs_cis = freqs_cis.unsqueeze(2)
+        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+        return x_out.type_as(x)

videosys/modules/normalization.py

@@ -0,0 +1,85 @@
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+
+class CogVideoXLayerNormZero(nn.Module):
+    def __init__(
+        self,
+        conditioning_dim: int,
+        embedding_dim: int,
+        elementwise_affine: bool = True,
+        eps: float = 1e-5,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(conditioning_dim, 6 * embedding_dim, bias=bias)
+        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
+
+    def forward(
+        self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        shift, scale, gate, enc_shift, enc_scale, enc_gate = self.linear(self.silu(temb)).chunk(6, dim=1)
+        hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
+        encoder_hidden_states = self.norm(encoder_hidden_states) * (1 + enc_scale)[:, None, :] + enc_shift[:, None, :]
+        return hidden_states, encoder_hidden_states, gate[:, None, :], enc_gate[:, None, :]
+
+
+class AdaLayerNorm(nn.Module):
+    r"""
+    Norm layer modified to incorporate timestep embeddings.
+
+    Parameters:
+        embedding_dim (`int`): The size of each embedding vector.
+        num_embeddings (`int`, *optional*): The size of the embeddings dictionary.
+        output_dim (`int`, *optional*):
+        norm_elementwise_affine (`bool`, defaults to `False):
+        norm_eps (`bool`, defaults to `False`):
+        chunk_dim (`int`, defaults to `0`):
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_embeddings: Optional[int] = None,
+        output_dim: Optional[int] = None,
+        norm_elementwise_affine: bool = False,
+        norm_eps: float = 1e-5,
+        chunk_dim: int = 0,
+    ):
+        super().__init__()
+
+        self.chunk_dim = chunk_dim
+        output_dim = output_dim or embedding_dim * 2
+
+        if num_embeddings is not None:
+            self.emb = nn.Embedding(num_embeddings, embedding_dim)
+        else:
+            self.emb = None
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, output_dim)
+        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
+
+    def forward(
+        self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        if self.emb is not None:
+            temb = self.emb(timestep)
+
+        temb = self.linear(self.silu(temb))
+
+        if self.chunk_dim == 1:
+            # This is a bit weird why we have the order of "shift, scale" here and "scale, shift" in the
+            # other if-branch. This branch is specific to CogVideoX for now.
+            shift, scale = temb.chunk(2, dim=1)
+            shift = shift[:, None, :]
+            scale = scale[:, None, :]
+        else:
+            scale, shift = temb.chunk(2, dim=0)
+
+        x = self.norm(x) * (1 + scale) + shift
+        return x

videosys/modules/upsampling.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CogVideoXUpsample3D(nn.Module):
+    r"""
+    A 3D Upsample layer using in CogVideoX by Tsinghua University & ZhipuAI # Todo: Wait for paper relase.
+
+    Args:
+        in_channels (`int`):
+            Number of channels in the input image.
+        out_channels (`int`):
+            Number of channels produced by the convolution.
+        kernel_size (`int`, defaults to `3`):
+            Size of the convolving kernel.
+        stride (`int`, defaults to `1`):
+            Stride of the convolution.
+        padding (`int`, defaults to `1`):
+            Padding added to all four sides of the input.
+        compress_time (`bool`, defaults to `False`):
+            Whether or not to compress the time dimension.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: int = 1,
+        compress_time: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
+        self.compress_time = compress_time
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        if self.compress_time:
+            if inputs.shape[2] > 1 and inputs.shape[2] % 2 == 1:
+                # split first frame
+                x_first, x_rest = inputs[:, :, 0], inputs[:, :, 1:]
+
+                x_first = F.interpolate(x_first, scale_factor=2.0)
+                x_rest = F.interpolate(x_rest, scale_factor=2.0)
+                x_first = x_first[:, :, None, :, :]
+                inputs = torch.cat([x_first, x_rest], dim=2)
+            elif inputs.shape[2] > 1:
+                inputs = F.interpolate(inputs, scale_factor=2.0)
+            else:
+                inputs = inputs.squeeze(2)
+                inputs = F.interpolate(inputs, scale_factor=2.0)
+                inputs = inputs[:, :, None, :, :]
+        else:
+            # only interpolate 2D
+            b, c, t, h, w = inputs.shape
+            inputs = inputs.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+            inputs = F.interpolate(inputs, scale_factor=2.0)
+            inputs = inputs.reshape(b, t, c, *inputs.shape[2:]).permute(0, 2, 1, 3, 4)
+
+        b, c, t, h, w = inputs.shape
+        inputs = inputs.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+        inputs = self.conv(inputs)
+        inputs = inputs.reshape(b, t, *inputs.shape[1:]).permute(0, 2, 1, 3, 4)
+
+        return inputs

videosys/pab.py

@@ -0,0 +1,64 @@
+class PABConfig:
+    def __init__(
+        self,
+        steps: int,
+        cross_broadcast: bool = False,
+        cross_threshold: list = None,
+        cross_range: int = None,
+        spatial_broadcast: bool = False,
+        spatial_threshold: list = None,
+        spatial_range: int = None,
+        temporal_broadcast: bool = False,
+        temporal_threshold: list = None,
+        temporal_range: int = None,
+        mlp_broadcast: bool = False,
+        mlp_spatial_broadcast_config: dict = None,
+        mlp_temporal_broadcast_config: dict = None,
+    ):
+        self.steps = steps
+
+        self.cross_broadcast = cross_broadcast
+        self.cross_threshold = cross_threshold
+        self.cross_range = cross_range
+
+        self.spatial_broadcast = spatial_broadcast
+        self.spatial_threshold = spatial_threshold
+        self.spatial_range = spatial_range
+
+        self.temporal_broadcast = temporal_broadcast
+        self.temporal_threshold = temporal_threshold
+        self.temporal_range = temporal_range
+
+        self.mlp_broadcast = mlp_broadcast
+        self.mlp_spatial_broadcast_config = mlp_spatial_broadcast_config
+        self.mlp_temporal_broadcast_config = mlp_temporal_broadcast_config
+        self.mlp_temporal_outputs = {}
+        self.mlp_spatial_outputs = {}
+
+class CogVideoXPABConfig(PABConfig):
+    def __init__(
+        self,
+        steps: int = 50,
+        spatial_broadcast: bool = True,
+        spatial_threshold: list = [100, 850],
+        spatial_range: int = 2,
+        temporal_broadcast: bool = False,
+        temporal_threshold: list = [100, 850],
+        temporal_range: int = 4,
+        cross_broadcast: bool = False,
+        cross_threshold: list = [100, 850],
+        cross_range: int = 6,
+    ):
+        super().__init__(
+            steps=steps,
+            spatial_broadcast=spatial_broadcast,
+            spatial_threshold=spatial_threshold,
+            spatial_range=spatial_range,
+            temporal_broadcast=temporal_broadcast,
+            temporal_threshold=temporal_threshold,
+            temporal_range=temporal_range,
+            cross_broadcast=cross_broadcast,
+            cross_threshold=cross_threshold,
+            cross_range=cross_range
+
+        )