ConsisIDTransformer3DModel

A Diffusion Transformer model for 3D data from ConsisID was introduced in Identity-Preserving Text-to-Video Generation by Frequency Decomposition by Peking University & University of Rochester & etc.

The model can be loaded with the following code snippet.

from mindone.diffusers import ConsisIDTransformer3DModel

transformer = ConsisIDTransformer3DModel.from_pretrained("BestWishYsh/ConsisID-preview", subfolder="transformer", mindspore_dtype=mindspore.bfloat16)

mindone.diffusers.models.transformers.ConsisIDTransformer3DModel

Bases: ModelMixin, ConfigMixin, PeftAdapterMixin

A Transformer model for video-like data in ConsisID.

PARAMETER	DESCRIPTION
num_attention_heads	The number of heads to use for multi-head attention. TYPE: `int`, defaults to `30` DEFAULT: 30
attention_head_dim	The number of channels in each head. TYPE: `int`, defaults to `64` DEFAULT: 64
in_channels	The number of channels in the input. TYPE: `int`, defaults to `16` DEFAULT: 16
out_channels	The number of channels in the output. TYPE: `int`, *optional*, defaults to `16` DEFAULT: 16
flip_sin_to_cos	Whether to flip the sin to cos in the time embedding. TYPE: `bool`, defaults to `True` DEFAULT: True
time_embed_dim	Output dimension of timestep embeddings. TYPE: `int`, defaults to `512` DEFAULT: 512
text_embed_dim	Input dimension of text embeddings from the text encoder. TYPE: `int`, defaults to `4096` DEFAULT: 4096
num_layers	The number of layers of Transformer blocks to use. TYPE: `int`, defaults to `30` DEFAULT: 30
dropout	The dropout probability to use. TYPE: `float`, defaults to `0.0` DEFAULT: 0.0
attention_bias	Whether to use bias in the attention projection layers. TYPE: `bool`, defaults to `True` DEFAULT: True
sample_width	The width of the input latents. TYPE: `int`, defaults to `90` DEFAULT: 90
sample_height	The height of the input latents. TYPE: `int`, defaults to `60` DEFAULT: 60
sample_frames	The number of frames in the input latents. Note that this parameter was incorrectly initialized to 49 instead of 13 because ConsisID 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). TYPE: `int`, defaults to `49` DEFAULT: 49
patch_size	The size of the patches to use in the patch embedding layer. TYPE: `int`, defaults to `2` DEFAULT: 2
temporal_compression_ratio	The compression ratio across the temporal dimension. See documentation for sample_frames. TYPE: `int`, defaults to `4` DEFAULT: 4
max_text_seq_length	The maximum sequence length of the input text embeddings. TYPE: `int`, defaults to `226` DEFAULT: 226
activation_fn	Activation function to use in feed-forward. TYPE: `str`, defaults to `"gelu-approximate"` DEFAULT: 'gelu-approximate'
timestep_activation_fn	Activation function to use when generating the timestep embeddings. TYPE: `str`, defaults to `"silu"` DEFAULT: 'silu'
norm_elementwise_affine	Whether to use elementwise affine in normalization layers. TYPE: `bool`, defaults to `True` DEFAULT: True
norm_eps	The epsilon value to use in normalization layers. TYPE: `float`, defaults to `1e-5` DEFAULT: 1e-05
spatial_interpolation_scale	Scaling factor to apply in 3D positional embeddings across spatial dimensions. TYPE: `float`, defaults to `1.875` DEFAULT: 1.875
temporal_interpolation_scale	Scaling factor to apply in 3D positional embeddings across temporal dimensions. TYPE: `float`, defaults to `1.0` DEFAULT: 1.0
is_train_face	Whether to use enable the identity-preserving module during the training process. When set to True, the model will focus on identity-preserving tasks. TYPE: `bool`, defaults to `False` DEFAULT: False
is_kps	Whether to enable keypoint for global facial extractor. If True, keypoints will be in the model. TYPE: `bool`, defaults to `False` DEFAULT: False
cross_attn_interval	The interval between cross-attention layers in the Transformer architecture. A larger value may reduce the frequency of cross-attention computations, which can help reduce computational overhead. TYPE: `int`, defaults to `2` DEFAULT: 2
cross_attn_dim_head	The dimensionality of each attention head in the cross-attention layers of the Transformer architecture. A larger value increases the capacity to attend to more complex patterns, but also increases memory and computation costs. TYPE: `int`, optional, defaults to `128` DEFAULT: 128
cross_attn_num_heads	The number of attention heads in the cross-attention layers. More heads allow for more parallel attention mechanisms, capturing diverse relationships between different components of the input, but can also increase computational requirements. TYPE: `int`, optional, defaults to `16` DEFAULT: 16
LFE_id_dim	The dimensionality of the identity vector used in the Local Facial Extractor (LFE). This vector represents the identity features of a face, which are important for tasks like face recognition and identity preservation across different frames. TYPE: `int`, optional, defaults to `1280` DEFAULT: 1280
LFE_vit_dim	The dimension of the vision transformer (ViT) output used in the Local Facial Extractor (LFE). This value dictates the size of the transformer-generated feature vectors that will be processed for facial feature extraction. TYPE: `int`, optional, defaults to `1024` DEFAULT: 1024
LFE_depth	The number of layers in the Local Facial Extractor (LFE). Increasing the depth allows the model to capture more complex representations of facial features, but also increases the computational load. TYPE: `int`, optional, defaults to `10` DEFAULT: 10
LFE_dim_head	The dimensionality of each attention head in the Local Facial Extractor (LFE). This parameter affects how finely the model can process and focus on different parts of the facial features during the extraction process. TYPE: `int`, optional, defaults to `64` DEFAULT: 64
LFE_num_heads	The number of attention heads in the Local Facial Extractor (LFE). More heads can improve the model's ability to capture diverse facial features, but at the cost of increased computational complexity. TYPE: `int`, optional, defaults to `16` DEFAULT: 16
LFE_num_id_token	The number of identity tokens used in the Local Facial Extractor (LFE). This defines how many identity-related tokens the model will process to ensure face identity preservation during feature extraction. TYPE: `int`, optional, defaults to `5` DEFAULT: 5
LFE_num_querie	The number of query tokens used in the Local Facial Extractor (LFE). These tokens are used to capture high-frequency face-related information that aids in accurate facial feature extraction. TYPE: `int`, optional, defaults to `32` DEFAULT: 32
LFE_output_dim	The output dimension of the Local Facial Extractor (LFE). This dimension determines the size of the feature vectors produced by the LFE module, which will be used for subsequent tasks such as face recognition or tracking. TYPE: `int`, optional, defaults to `2048` DEFAULT: 2048
LFE_ff_mult	The multiplication factor applied to the feed-forward network's hidden layer size in the Local Facial Extractor (LFE). A higher value increases the model's capacity to learn more complex facial feature transformations, but also increases the computation and memory requirements. TYPE: `int`, optional, defaults to `4` DEFAULT: 4
LFE_num_scale	The number of different scales visual feature. A higher value increases the model's capacity to learn more complex facial feature transformations, but also increases the computation and memory requirements. TYPE: `int`, optional, defaults to `5` DEFAULT: 5
local_face_scale	A scaling factor used to adjust the importance of local facial features in the model. This can influence how strongly the model focuses on high frequency face-related content. TYPE: `float`, defaults to `1.0` DEFAULT: 1.0

Source code in mindone/diffusers/models/transformers/consisid_transformer_3d.py

class ConsisIDTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
    """
    A Transformer model for video-like data in [ConsisID](https://github.com/PKU-YuanGroup/ConsisID).

    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 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 ConsisID 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 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.
        is_train_face (`bool`, defaults to `False`):
            Whether to use enable the identity-preserving module during the training process. When set to `True`, the
            model will focus on identity-preserving tasks.
        is_kps (`bool`, defaults to `False`):
            Whether to enable keypoint for global facial extractor. If `True`, keypoints will be in the model.
        cross_attn_interval (`int`, defaults to `2`):
            The interval between cross-attention layers in the Transformer architecture. A larger value may reduce the
            frequency of cross-attention computations, which can help reduce computational overhead.
        cross_attn_dim_head (`int`, optional, defaults to `128`):
            The dimensionality of each attention head in the cross-attention layers of the Transformer architecture. A
            larger value increases the capacity to attend to more complex patterns, but also increases memory and
            computation costs.
        cross_attn_num_heads (`int`, optional, defaults to `16`):
            The number of attention heads in the cross-attention layers. More heads allow for more parallel attention
            mechanisms, capturing diverse relationships between different components of the input, but can also
            increase computational requirements.
        LFE_id_dim (`int`, optional, defaults to `1280`):
            The dimensionality of the identity vector used in the Local Facial Extractor (LFE). This vector represents
            the identity features of a face, which are important for tasks like face recognition and identity
            preservation across different frames.
        LFE_vit_dim (`int`, optional, defaults to `1024`):
            The dimension of the vision transformer (ViT) output used in the Local Facial Extractor (LFE). This value
            dictates the size of the transformer-generated feature vectors that will be processed for facial feature
            extraction.
        LFE_depth (`int`, optional, defaults to `10`):
            The number of layers in the Local Facial Extractor (LFE). Increasing the depth allows the model to capture
            more complex representations of facial features, but also increases the computational load.
        LFE_dim_head (`int`, optional, defaults to `64`):
            The dimensionality of each attention head in the Local Facial Extractor (LFE). This parameter affects how
            finely the model can process and focus on different parts of the facial features during the extraction
            process.
        LFE_num_heads (`int`, optional, defaults to `16`):
            The number of attention heads in the Local Facial Extractor (LFE). More heads can improve the model's
            ability to capture diverse facial features, but at the cost of increased computational complexity.
        LFE_num_id_token (`int`, optional, defaults to `5`):
            The number of identity tokens used in the Local Facial Extractor (LFE). This defines how many
            identity-related tokens the model will process to ensure face identity preservation during feature
            extraction.
        LFE_num_querie (`int`, optional, defaults to `32`):
            The number of query tokens used in the Local Facial Extractor (LFE). These tokens are used to capture
            high-frequency face-related information that aids in accurate facial feature extraction.
        LFE_output_dim (`int`, optional, defaults to `2048`):
            The output dimension of the Local Facial Extractor (LFE). This dimension determines the size of the feature
            vectors produced by the LFE module, which will be used for subsequent tasks such as face recognition or
            tracking.
        LFE_ff_mult (`int`, optional, defaults to `4`):
            The multiplication factor applied to the feed-forward network's hidden layer size in the Local Facial
            Extractor (LFE). A higher value increases the model's capacity to learn more complex facial feature
            transformations, but also increases the computation and memory requirements.
        LFE_num_scale (`int`, optional, defaults to `5`):
            The number of different scales visual feature. A higher value increases the model's capacity to learn more
            complex facial feature transformations, but also increases the computation and memory requirements.
        local_face_scale (`float`, defaults to `1.0`):
            A scaling factor used to adjust the importance of local facial features in the model. This can influence
            how strongly the model focuses on high frequency face-related content.
    """

    _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,
        is_train_face: bool = False,
        is_kps: bool = False,
        cross_attn_interval: int = 2,
        cross_attn_dim_head: int = 128,
        cross_attn_num_heads: int = 16,
        LFE_id_dim: int = 1280,
        LFE_vit_dim: int = 1024,
        LFE_depth: int = 10,
        LFE_dim_head: int = 64,
        LFE_num_heads: int = 16,
        LFE_num_id_token: int = 5,
        LFE_num_querie: int = 32,
        LFE_output_dim: int = 2048,
        LFE_ff_mult: int = 4,
        LFE_num_scale: int = 5,
        local_face_scale: float = 1.0,
    ):
        super().__init__()
        inner_dim = num_attention_heads * attention_head_dim

        if not use_rotary_positional_embeddings and use_learned_positional_embeddings:
            raise ValueError(
                "There are no ConsisID checkpoints available with disable rotary embeddings and learned positional "
                "embeddings. If you're using a custom model and/or believe this should be supported, please open an "
                "issue at https://github.com/huggingface/diffusers/issues."
            )

        # 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 = mint.nn.Dropout(dropout)

        # 2. 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)

        # 3. Define spatio-temporal transformers blocks
        self.transformer_blocks = nn.CellList(
            [
                ConsisIDBlock(
                    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 = mint.nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine)

        # 4. 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 = mint.nn.Linear(inner_dim, patch_size * patch_size * out_channels)

        self.is_train_face = is_train_face
        self.is_kps = is_kps

        # 5. Define identity-preserving config
        if is_train_face:
            # LFE configs
            self.LFE_id_dim = LFE_id_dim
            self.LFE_vit_dim = LFE_vit_dim
            self.LFE_depth = LFE_depth
            self.LFE_dim_head = LFE_dim_head
            self.LFE_num_heads = LFE_num_heads
            self.LFE_num_id_token = LFE_num_id_token
            self.LFE_num_querie = LFE_num_querie
            self.LFE_output_dim = LFE_output_dim
            self.LFE_ff_mult = LFE_ff_mult
            self.LFE_num_scale = LFE_num_scale
            # cross configs
            self.inner_dim = inner_dim
            self.cross_attn_interval = cross_attn_interval
            self.num_cross_attn = num_layers // cross_attn_interval
            self.cross_attn_dim_head = cross_attn_dim_head
            self.cross_attn_num_heads = cross_attn_num_heads
            self.cross_attn_kv_dim = int(self.inner_dim / 3 * 2)
            self.local_face_scale = local_face_scale
            # face modules
            self._init_face_inputs()

        self.gradient_checkpointing = False

        self.p = self.config.patch_size

    def _init_face_inputs(self):
        self.local_facial_extractor = LocalFacialExtractor(
            id_dim=self.LFE_id_dim,
            vit_dim=self.LFE_vit_dim,
            depth=self.LFE_depth,
            dim_head=self.LFE_dim_head,
            heads=self.LFE_num_heads,
            num_id_token=self.LFE_num_id_token,
            num_queries=self.LFE_num_querie,
            output_dim=self.LFE_output_dim,
            ff_mult=self.LFE_ff_mult,
            num_scale=self.LFE_num_scale,
        )
        self.perceiver_cross_attention = nn.CellList(
            [
                PerceiverCrossAttention(
                    dim=self.inner_dim,
                    dim_head=self.cross_attn_dim_head,
                    heads=self.cross_attn_num_heads,
                    kv_dim=self.cross_attn_kv_dim,
                )
                for _ in range(self.num_cross_attn)
            ]
        )

    @property
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
    def attn_processors(self) -> Dict[str, AttentionProcessor]:
        r"""
        Returns:
            `dict` of attention processors: A dictionary containing all attention processors used in the model with
            indexed by its weight name.
        """
        # set recursively
        processors = {}

        def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str, AttentionProcessor]):
            if hasattr(module, "get_processor"):
                processors[f"{name}.processor"] = module.get_processor()

            for sub_name, child in module.name_cells().items():
                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)

            return processors

        for name, module in self.name_cells().items():
            fn_recursive_add_processors(name, module, processors)

        return processors

    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
        r"""
        Sets the attention processor to use to compute attention.

        Parameters:
            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
                The instantiated processor class or a dictionary of processor classes that will be set as the processor
                for **all** `Attention` layers.

                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
                processor. This is strongly recommended when setting trainable attention processors.

        """
        count = len(self.attn_processors.keys())

        if isinstance(processor, dict) and len(processor) != count:
            raise ValueError(
                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
            )

        def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
            if hasattr(module, "set_processor"):
                if not isinstance(processor, dict):
                    module.set_processor(processor)
                else:
                    module.set_processor(processor.pop(f"{name}.processor"))

            for sub_name, child in module.name_cells().items():
                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

        for name, module in self.name_cells().items():
            fn_recursive_attn_processor(name, module, processor)

    def construct(
        self,
        hidden_states: ms.Tensor,
        encoder_hidden_states: ms.Tensor,
        timestep: Union[int, float, ms.Tensor],
        timestep_cond: Optional[ms.Tensor] = None,
        image_rotary_emb: Optional[Tuple[ms.Tensor, ms.Tensor]] = None,
        attention_kwargs: Optional[Dict[str, Any]] = None,
        id_cond: Optional[ms.Tensor] = None,
        id_vit_hidden: Optional[ms.Tensor] = None,
        return_dict: bool = True,
    ):
        if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
            # weight the lora layers by setting `lora_scale` for each PEFT layer here
            # and remove `lora_scale` from each PEFT layer at the end.
            # scale_lora_layers & unscale_lora_layers maybe contains some operation forbidden in graph mode
            raise RuntimeError(
                f"You are trying to set scaling of lora layer by passing {attention_kwargs['scale']=}. "
                f"However it's not allowed in on-the-fly model forwarding. "
                f"Please manually call `scale_lora_layers(model, lora_scale)` before model forwarding and "
                f"`unscale_lora_layers(model, lora_scale)` after model forwarding. "
                f"For example, it can be done in a pipeline call like `StableDiffusionPipeline.__call__`."
            )

        # fuse clip and insightface
        valid_face_emb = None
        if self.is_train_face:
            id_cond = id_cond.to(dtype=hidden_states.dtype)
            id_vit_hidden = [tensor.to(dtype=hidden_states.dtype) for tensor in id_vit_hidden]
            valid_face_emb = self.local_facial_extractor(
                id_cond, id_vit_hidden
            )  # torch.Size([1, 1280]), list[5](torch.Size([1, 577, 1024]))  ->  torch.Size([1, 32, 2048])

        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
        # torch.Size([1, 226, 4096])   torch.Size([1, 13, 32, 60, 90])
        hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)  # torch.Size([1, 17776, 3072])
        hidden_states = self.embedding_dropout(hidden_states)  # torch.Size([1, 17776, 3072])

        text_seq_length = encoder_hidden_states.shape[1]
        encoder_hidden_states = hidden_states[:, :text_seq_length]  # torch.Size([1, 226, 3072])
        hidden_states = hidden_states[:, text_seq_length:]  # torch.Size([1, 17550, 3072])

        # 3. Transformer blocks
        ca_idx = 0
        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,
            )

            if self.is_train_face:
                if i % self.cross_attn_interval == 0 and valid_face_emb is not None:
                    hidden_states = hidden_states + self.local_face_scale * self.perceiver_cross_attention[ca_idx](
                        valid_face_emb, hidden_states
                    )  # torch.Size([2, 32, 2048])  torch.Size([2, 17550, 3072])
                    ca_idx += 1

        hidden_states = mint.cat([encoder_hidden_states, hidden_states], dim=1)
        hidden_states = self.norm_final(hidden_states)
        hidden_states = hidden_states[:, text_seq_length:]

        # 4. Final block
        hidden_states = self.norm_out(hidden_states, temb=emb)
        hidden_states = self.proj_out(hidden_states)

        # 5. Unpatchify
        # Note: we use `-1` instead of `channels`:
        #   - It is okay to `channels` use for ConsisID (number of input channels is equal to output channels)
        output = hidden_states.reshape(batch_size, num_frames, height // self.p, width // self.p, -1, self.p, self.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)

mindone.diffusers.models.transformers.ConsisIDTransformer3DModel.attn_processors property

RETURNS	DESCRIPTION
Dict[str, AttentionProcessor]	dict of attention processors: A dictionary containing all attention processors used in the model with
Dict[str, AttentionProcessor]	indexed by its weight name.

mindone.diffusers.models.transformers.ConsisIDTransformer3DModel.set_attn_processor(processor)

Sets the attention processor to use to compute attention.

PARAMETER	DESCRIPTION
processor	The instantiated processor class or a dictionary of processor classes that will be set as the processor for all Attention layers. If processor is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors. TYPE: `dict` of `AttentionProcessor` or only `AttentionProcessor`

Source code in mindone/diffusers/models/transformers/consisid_transformer_3d.py

def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
    r"""
    Sets the attention processor to use to compute attention.

    Parameters:
        processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
            The instantiated processor class or a dictionary of processor classes that will be set as the processor
            for **all** `Attention` layers.

            If `processor` is a dict, the key needs to define the path to the corresponding cross attention
            processor. This is strongly recommended when setting trainable attention processors.

    """
    count = len(self.attn_processors.keys())

    if isinstance(processor, dict) and len(processor) != count:
        raise ValueError(
            f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
            f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
        )

    def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
        if hasattr(module, "set_processor"):
            if not isinstance(processor, dict):
                module.set_processor(processor)
            else:
                module.set_processor(processor.pop(f"{name}.processor"))

        for sub_name, child in module.name_cells().items():
            fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

    for name, module in self.name_cells().items():
        fn_recursive_attn_processor(name, module, processor)

mindone.diffusers.models.modeling_outputs.Transformer2DModelOutput dataclass

Bases: BaseOutput

The output of [Transformer2DModel].

PARAMETER	DESCRIPTION
`(batch	The hidden states output conditioned on the encoder_hidden_states input. If discrete, returns probability distributions for the unnoised latent pixels. TYPE: size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete

Source code in mindone/diffusers/models/modeling_outputs.py

@dataclass
class Transformer2DModelOutput(BaseOutput):
    """
    The output of [`Transformer2DModel`].

    Args:
        sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` or
        `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
            The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
            distributions for the unnoised latent pixels.
    """

    sample: "ms.Tensor"  # noqa: F821