Upload folder using huggingface_hub

ummdit_ds1_small_singlenorm_v5.py

@@ -0,0 +1,681 @@
+# apply pos emb to downsampled and upsampled feats
+# add bias and scale to blockwise AdaIN params
+# subattn to subsampled feat
+# block list [4, 16, 4]
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models.transformers import SD3Transformer2DModel
+from diffusers.configuration_utils import register_to_config
+# from diffusers.models.attention import JointTransformerBlock
+from diffusers.utils import is_torch_version, logging
+from diffusers.models.embeddings import PatchEmbed, get_2d_sincos_pos_embed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.normalization import AdaLayerNormSingle
+
+from diffusers.models.attention_processor import Attention, JointAttnProcessor2_0
+from diffusers.models.normalization import SD35AdaLayerNormZeroX
+from diffusers.models.attention import FeedForward, _chunked_feed_forward
+
+
+from einops import rearrange
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+def cropped_pos_embed(pos_embed, height, width, patch_size=1, pos_embed_max_size=96):
+    """Crops positional embeddings for SD3 compatibility."""
+    if pos_embed_max_size is None:
+        raise ValueError("`pos_embed_max_size` must be set for cropping.")
+
+    height = height // patch_size
+    width = width // patch_size
+    if height > pos_embed_max_size:
+        raise ValueError(
+            f"Height ({height}) cannot be greater than `pos_embed_max_size`: {pos_embed_max_size}."
+        )
+    if width > pos_embed_max_size:
+        raise ValueError(
+            f"Width ({width}) cannot be greater than `pos_embed_max_size`: {pos_embed_max_size}."
+        )
+
+    top = (pos_embed_max_size - height) // 2
+    left = (pos_embed_max_size - width) // 2
+    spatial_pos_embed = pos_embed.reshape(1, pos_embed_max_size, pos_embed_max_size, -1)
+    spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]
+    # spatial_pos_embed = torch.permute(spatial_pos_embed, [0, 3, 1, 2])
+    # spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])
+    return spatial_pos_embed
+
+
+class JointTransformerBlockSingleNorm(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://huggingface.co/papers/2403.03206
+
+    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.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        context_pre_only: bool = False,
+        qk_norm: Optional[str] = None,
+        use_dual_attention: bool = False,
+        subsample_ratio = 1,
+        subsample_seq_len = 1,
+    ):
+        super().__init__()
+
+        self.use_dual_attention = use_dual_attention
+        self.context_pre_only = context_pre_only
+        context_norm_type = "ada_norm_continous" if context_pre_only else "ada_norm_single"
+
+        if use_dual_attention:
+            self.norm1 = SD35AdaLayerNormZeroX(dim)
+        else:
+            # self.norm1 = AdaLayerNormZero(dim)
+            self.norm1 = nn.LayerNorm(dim)
+        
+        assert subsample_ratio >= 1 and subsample_seq_len >= 1
+        self.subsample_ratio = subsample_ratio
+        self.subsample_seq_len = subsample_seq_len
+
+        print(self.subsample_ratio, self.subsample_seq_len)
+
+        # if context_norm_type == "ada_norm_continous":
+        #     # self.norm1_context = AdaLayerNormContinuous(
+        #     #     dim, dim, elementwise_affine=False, eps=1e-6, bias=True, norm_type="layer_norm"
+        #     # )
+        # elif context_norm_type == "ada_norm_single":
+        #     # self.norm1_context = AdaLayerNormZero(dim)
+        #     self.norm1_context = nn.LayerNorm(dim)
+        # else:
+        #     raise ValueError(
+        #         f"Unknown context_norm_type: {context_norm_type}, currently only support `ada_norm_continous`, `ada_norm_zero`"
+        #     )
+        self.norm1_context = nn.LayerNorm(dim)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            processor = JointAttnProcessor2_0()
+        else:
+            raise ValueError(
+                "The current PyTorch version does not support the `scaled_dot_product_attention` function."
+            )
+
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=context_pre_only,
+            bias=True,
+            processor=processor,
+            qk_norm=qk_norm,
+            eps=1e-6,
+        )
+
+        if use_dual_attention:
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=None,
+                dim_head=attention_head_dim,
+                heads=num_attention_heads,
+                out_dim=dim,
+                bias=True,
+                processor=processor,
+                qk_norm=qk_norm,
+                eps=1e-6,
+            )
+        else:
+            self.attn2 = None
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        if not context_pre_only:
+            self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+            self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        else:
+            self.norm2_context = None
+            self.ff_context = None
+
+
+        self.scale_shift_bias = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+        self.scale_shift_scale = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+        
+
+        if not context_pre_only:
+            self.scale_shift_bias_c = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+            self.scale_shift_scale_c = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    # Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        embedded_timestep: torch.FloatTensor = None,
+    ):
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        if self.use_dual_attention:
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1(
+                hidden_states, emb=temb
+            )
+        else:
+            # norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+            batch_size = hidden_states.shape[0]
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_bias[None] + temb.reshape(batch_size, 6, -1)*(1+self.scale_shift_scale[None])
+            ).chunk(6, dim=1)
+            norm_hidden_states = self.norm1(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+
+        if self.context_pre_only:
+            norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
+            # norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, embedded_timestep)
+        else:
+            # norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+            #     encoder_hidden_states, emb=temb
+            # )
+            batch_size = hidden_states.shape[0]
+            c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = (
+                self.scale_shift_bias_c[None] + temb.reshape(batch_size, 6, -1)*(1+self.scale_shift_scale_c)
+            ).chunk(6, dim=1)
+            norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
+            norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_msa) + c_shift_msa
+
+        if self.subsample_ratio > 1:
+            norm_hidden_states = rearrange(norm_hidden_states, 
+                                           'b (l s n) c -> (b s) (l n) c', 
+                                           n=self.subsample_seq_len, s=self.subsample_ratio)
+            norm_encoder_hidden_states = rearrange(norm_encoder_hidden_states, 
+                                           'b (l s n) c -> (b s) (l n) c', 
+                                           n=self.subsample_seq_len, s=self.subsample_ratio)
+
+        # Attention.
+        
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            **joint_attention_kwargs,
+        )
+        if self.subsample_ratio > 1:
+            attn_output = rearrange(attn_output, 
+                                           '(b s) (l n) c -> b (l s n) c', 
+                                           n=self.subsample_seq_len, s=self.subsample_ratio)
+            context_attn_output = rearrange(context_attn_output, 
+                                           '(b s) (l n) c -> b (l s n) c', 
+                                           n=self.subsample_seq_len, s=self.subsample_ratio)
+        # attn_output = norm_hidden_states
+        # context_attn_output = norm_encoder_hidden_states
+
+
+        # Process attention outputs for the `hidden_states`.
+        attn_output = gate_msa * attn_output
+        hidden_states = hidden_states + attn_output
+
+        if self.use_dual_attention:
+            attn_output2 = self.attn2(hidden_states=norm_hidden_states2, **joint_attention_kwargs)
+            attn_output2 = gate_msa2 * attn_output2
+            hidden_states = hidden_states + attn_output2
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
+        else:
+            ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp * ff_output
+
+        hidden_states = hidden_states + ff_output
+
+        # Process attention outputs for the `encoder_hidden_states`.
+        if self.context_pre_only:
+            encoder_hidden_states = None
+        else:
+            context_attn_output = c_gate_msa * context_attn_output
+            # print(context_attn_output.shape, encoder_hidden_states.shape)
+            encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+            norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+            norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp
+            if self._chunk_size is not None:
+                # "feed_forward_chunk_size" can be used to save memory
+                context_ff_output = _chunked_feed_forward(
+                    self.ff_context, norm_encoder_hidden_states, self._chunk_dim, self._chunk_size
+                )
+            else:
+                context_ff_output = self.ff_context(norm_encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
+
+        return encoder_hidden_states, hidden_states
+
+# class TimestepEmbeddings(nn.Module):
+#     def __init__(self, embedding_dim):
+#         super().__init__()
+
+#         self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+#         self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+#     def forward(self, timestep, dtype):
+#         timesteps_proj = self.time_proj(timestep)
+#         timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=dtype))  # (N, D)
+
+#         return timesteps_emb
+
+class Downsample(nn.Module):
+    def __init__(self, n_feat):
+        super(Downsample, self).__init__()
+
+        self.body = nn.Sequential(
+                                  nn.PixelUnshuffle(2),
+                                  nn.Conv2d(n_feat*4, n_feat, kernel_size=1, stride=1, padding=0, bias=True),
+                                  torch.nn.GELU('tanh'),
+                                  nn.Conv2d(n_feat, n_feat, kernel_size=1, stride=1, padding=0, bias=True))
+
+    def forward(self, x):
+        return self.body(x)
+    
+class Upsample(nn.Module):
+    def __init__(self, n_feat):
+        super(Upsample, self).__init__()
+
+        self.body = nn.Sequential(nn.PixelShuffle(2),
+                                  nn.Conv2d(n_feat//4, n_feat, kernel_size=1, stride=1, padding=0, bias=True),
+                                  torch.nn.GELU('tanh'),
+                                  nn.Conv2d(n_feat, n_feat, kernel_size=1, stride=1, padding=0, bias=True))
+
+    def forward(self, x):
+        return self.body(x)
+    
+class MMDiTTransformer2DModel(SD3Transformer2DModel):
+    """
+    The Transformer model introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        sample_size (`int`): The width of the latent images. This is fixed during training since
+            it is used to learn a number of position embeddings.
+        patch_size (`int`): Patch size to turn the input data into small patches.
+        in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
+        num_layers (`int`, *optional*, defaults to 18): The number of layers of Transformer blocks to use.
+        attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
+        num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
+        cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        caption_projection_dim (`int`): Number of dimensions to use when projecting the `encoder_hidden_states`.
+        out_channels (`int`, defaults to 16): Number of output channels.
+
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 128,
+        patch_size: int = 2,
+        in_channels: int = 16,
+        num_layers: int = 24,
+        attention_head_dim: int = 32,
+        num_attention_heads: int = 24,
+        caption_channels: int = 4096,
+        caption_projection_dim: int = 768,
+        out_channels: int = 16,
+        interpolation_scale: int = None,
+        pos_embed_max_size: int = 96,
+        dual_attention_layers: Tuple[
+            int, ...
+        ] = (),  # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5
+        qk_norm: Optional[str] = None,
+        repa_depth = -1,
+        projector_dim=2048,
+        z_dims=[768]
+    ):
+        super().__init__(
+            sample_size=sample_size,
+            patch_size=patch_size,
+            in_channels=in_channels,
+            num_layers=num_layers,
+            attention_head_dim=attention_head_dim,
+            num_attention_heads=num_attention_heads,
+            caption_projection_dim=caption_projection_dim,
+            out_channels=out_channels,
+            pos_embed_max_size=pos_embed_max_size,
+            dual_attention_layers=dual_attention_layers,
+            qk_norm=qk_norm,
+        )
+
+        self.time_text_embed = None
+
+        self.patch_mixer_depth = None # initially no masking applied
+        self.mask_ratio = 0
+
+        # self.block_split_stage = [2, 20, 2]
+        self.block_split_stage = [4, 16, 4]
+        # self.block_split_stage = [12, 1, 12]
+
+        default_out_channels = in_channels
+        self.out_channels = out_channels if out_channels is not None else default_out_channels
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+
+        if repa_depth != -1:
+            from core.models.projector import build_projector
+            self.projectors = nn.ModuleList([
+                build_projector(self.inner_dim, projector_dim, z_dim) for z_dim in z_dims
+                ])
+            
+            assert repa_depth >= 0 and repa_depth < num_layers
+            self.repa_depth = repa_depth
+
+
+        interpolation_scale = (
+            self.config.interpolation_scale
+            if self.config.interpolation_scale is not None
+            else max(self.config.sample_size // 16, 1)
+        )
+
+        self.pos_embed = PatchEmbed(
+            height=self.config.sample_size,
+            width=self.config.sample_size,
+            patch_size=self.config.patch_size,
+            in_channels=self.config.in_channels,
+            embed_dim=self.inner_dim,
+            interpolation_scale=interpolation_scale,
+            pos_embed_max_size=pos_embed_max_size,  # hard-code for now.
+        )
+
+        pos_embed_lv0 = get_2d_sincos_pos_embed(
+                self.inner_dim, pos_embed_max_size, base_size=self.config.sample_size // self.config.patch_size, 
+                interpolation_scale=interpolation_scale, output_type='pt'
+            ) # [grid_size**2, embed_dim]
+        
+        pos_embed_lv0 = cropped_pos_embed(pos_embed_lv0,
+                                                self.config.sample_size, 
+                                                self.config.sample_size, 
+                                                patch_size=1, pos_embed_max_size=pos_embed_max_size)
+
+
+        pos_embed_lv1 = pos_embed_lv0.clone()[:, ::2, ::2, :]
+
+        pos_embed_lv0 = pos_embed_lv0.reshape(1, -1, pos_embed_lv0.shape[-1])
+        pos_embed_lv1 = pos_embed_lv1.reshape(1, -1, pos_embed_lv1.shape[-1])
+        
+
+
+        self.register_buffer("pos_embed_lv0", pos_embed_lv0.float(), persistent=False)
+        self.register_buffer("pos_embed_lv1", pos_embed_lv1.float(), persistent=False)
+        
+        # self.time_text_embed = TimestepEmbeddings(embedding_dim=self.inner_dim)
+        self.context_embedder = nn.Linear(self.config.caption_channels, self.config.caption_projection_dim)
+
+        self.adaln_single = AdaLayerNormSingle(
+            self.inner_dim, use_additional_conditions=False
+        )
+
+        self.transformer_blocks = None
+
+        subample_ratio_list = [1, 4, 4]
+        seq_len_list = [1, 1, 4]
+        cur_ind = 0
+
+        self.block_groups = nn.ModuleList()
+        for grp_ids, cur_bks in enumerate(self.block_split_stage):
+            # cur_subample_ratio = 1
+            # seq_len_list = [1]
+            # if grp_ids == 1:
+            #     cur_subample_ratio = 4
+            #     seq_len_list = [1, 4]
+            cur_group = []
+            for i in range(cur_bks):
+                cur_group.append(JointTransformerBlockSingleNorm(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                    context_pre_only=(grp_ids==len(self.block_split_stage)-1) \
+                                        and (i == cur_bks - 1),
+                    qk_norm=qk_norm,
+                    use_dual_attention=False,
+                    subsample_ratio=subample_ratio_list[cur_ind%len(subample_ratio_list)],
+                    subsample_seq_len=seq_len_list[cur_ind%len(seq_len_list)],
+                ))
+                cur_ind += 1
+
+            cur_group = nn.ModuleList(cur_group)
+
+
+            # cur_group = nn.ModuleList(
+            # [
+            #     JointTransformerBlockSingleNorm(
+            #         dim=self.inner_dim,
+            #         num_attention_heads=self.config.num_attention_heads,
+            #         attention_head_dim=self.config.attention_head_dim,
+            #         context_pre_only=(grp_ids==len(self.block_split_stage)-1) \
+            #                             and (i == cur_bks - 1),
+            #         qk_norm=qk_norm,
+            #         use_dual_attention=False,
+            #         subsample_ratio=cur_subample_ratio,
+            #         subsample_seq_len=seq_len_list[i%len(seq_len_list)],
+            #     )
+            #     for i in range(cur_bks)
+            # ])
+            self.block_groups.append(cur_group)
+
+        ds_num = int(len(self.block_split_stage) // 2)
+        self.downsamplers = nn.ModuleList()
+        for _ in range(ds_num):
+            self.downsamplers.append(Downsample(self.inner_dim))
+        self.upsamplers = nn.ModuleList()
+        for _ in range(ds_num):
+            self.upsamplers.append(Upsample(self.inner_dim))
+        self.mergers = nn.ModuleList()
+        for _ in range(ds_num):
+            # self.mergers.append(nn.Linear(self.inner_dim*2, self.inner_dim))
+            self.mergers.append(nn.Sequential(
+                                  nn.Linear(self.inner_dim*2, self.inner_dim),
+                                  torch.nn.GELU('tanh'),
+                                  nn.Linear(self.inner_dim, self.inner_dim)))
+
+
+        self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        timestep: torch.LongTensor = None,
+        block_controlnet_hidden_states: List = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+        skip_layers: Optional[List[int]] = None,
+        **kwargs,
+    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+        """
+        The [`SD3Transformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            timestep (`torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+            skip_layers (`list` of `int`, *optional*):
+                A list of layer indices to skip during the forward pass.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+
+        height, width = hidden_states.shape[-2:]
+
+        cur_height = height // self.config.patch_size
+        cur_width = width // self.config.patch_size
+
+        hidden_states = self.pos_embed(hidden_states)  # takes care of adding positional embeddings too.
+        # temb = self.time_text_embed(timestep, dtype=encoder_hidden_states.dtype)
+        temb, embedded_timestep = self.adaln_single(
+            timestep, None, batch_size=hidden_states.shape[0], hidden_dtype=hidden_states.dtype
+        )
+
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        ids_keep = None  
+        len_keep = hidden_states.shape[1]
+        zs = None
+
+        ds_num = int(len(self.block_split_stage) // 2)
+        encoder_feats = []
+        for grp_ids, blocks in enumerate(self.block_groups):
+            # for encoders
+            for index_block, block in enumerate(blocks):
+                # Skip specified layers
+                is_skip = True if skip_layers is not None and index_block in skip_layers else False
+
+                if torch.is_grad_enabled() and self.gradient_checkpointing and not is_skip:
+
+                    def create_custom_forward(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        hidden_states,
+                        encoder_hidden_states,
+                        temb,
+                        joint_attention_kwargs,
+                        **ckpt_kwargs,
+                    )
+                elif not is_skip:
+                    encoder_hidden_states, hidden_states = block(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        temb=temb,
+                        joint_attention_kwargs=joint_attention_kwargs,
+                    )
+                
+                if grp_ids == 1 and index_block==self.repa_depth-self.block_split_stage[0]-1:
+                    if self.training and (self.repa_depth != -1):
+                        reshaped_out = rearrange(hidden_states, "n (h w) c -> n c h w", h=cur_height, w=cur_width)
+                        upsampled_out = torch.nn.functional.interpolate(reshaped_out, size=(cur_height*2, cur_width*2))
+                        out_1d = rearrange(upsampled_out, "n c h w  -> n (h w) c", h=cur_height*2, w=cur_width*2)
+                        zs = [projector(out_1d) for projector in self.projectors]
+            if grp_ids < ds_num:
+                encoder_feats.append(hidden_states)
+
+                hidden_states = self.downsamplers[grp_ids](rearrange(hidden_states, "n (h w) c -> n c h w", h=cur_height, w=cur_width))
+                cur_height = int(cur_height / 2)
+                cur_width = int(cur_width / 2)
+                hidden_states = rearrange(hidden_states, "n c h w  -> n (h w) c", h=cur_height, w=cur_width)
+                hidden_states = hidden_states + self.pos_embed_lv1
+            elif grp_ids < len(self.block_split_stage)-1:
+                hidden_states  = self.upsamplers[grp_ids-ds_num](rearrange(hidden_states, "n (h w) c -> n c h w", h=cur_height, w=cur_width))
+                cur_height = int(cur_height * 2)
+                cur_width = int(cur_width * 2)
+                hidden_states = rearrange(hidden_states, "n c h w  -> n (h w) c", h=cur_height, w=cur_width)
+                
+                hidden_states = torch.cat([hidden_states, encoder_feats[len(encoder_feats)-1-(grp_ids-ds_num)]], dim=2)
+                hidden_states = self.mergers[grp_ids-ds_num](hidden_states)
+                hidden_states = hidden_states + self.pos_embed_lv0
+
+        # print(hidden_states.shape, temb.shape)
+        hidden_states = self.norm_out(hidden_states)
+        hidden_states = self.proj_out(hidden_states)
+
+        if not self.training:
+            # unpatchify
+            patch_size = self.config.patch_size
+            height = height // patch_size
+            width = width // patch_size
+
+            hidden_states = hidden_states.reshape(
+                shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels)
+            )
+            hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+            output = hidden_states.reshape(
+                shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
+            )
+
+            if not return_dict:
+                return (output,)
+
+            return Transformer2DModelOutput(sample=output)
+        
+        else:
+            return hidden_states, ids_keep, zs
+        
+
+    def enable_masking(self, depth, mask_ratio):
+        # depth: apply masking after block_[depth]. should be [0, nblks-1]
+        assert depth >= 0 and depth < len(self.transformer_blocks)
+        self.patch_mixer_depth = depth
+        assert mask_ratio >= 0 and mask_ratio <= 1
+        self.mask_ratio = mask_ratio
+
+    def disable_masking(self):
+        self.patch_mixer_depth = None
+
+    def enable_gradient_checkpointing(self, nblocks_to_apply_grad_checkpointing):
+        N = len(self.transformer_blocks)
+
+        if nblocks_to_apply_grad_checkpointing == -1:
+            nblocks_to_apply_grad_checkpointing = N
+        nblocks_to_apply_grad_checkpointing = min(N, nblocks_to_apply_grad_checkpointing)
+
+        # Apply to blocks evenly spaced out
+        step = N / nblocks_to_apply_grad_checkpointing if nblocks_to_apply_grad_checkpointing > 0 else 0
+        indices = [int((i+0.5)*step) for i in range(nblocks_to_apply_grad_checkpointing)]
+
+        self.gradient_checkpointing = True
+        for blk_ind, block in enumerate(self.transformer_blocks):
+            block.gradient_checkpointing = (blk_ind in indices)
+            print(f"Block {blk_ind} grad checkpointing set to {block.gradient_checkpointing}")

ummdit_small_ds1_singlenorm_v5_newdata_ft512_ema_checkpoint-48000_model_ema.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:366656475662eca6791d6945cf428f01406e57c7387d77f9879e3f92c4594786
+size 1415138120

ummdit_small_ds1_singlenorm_v5_newdata_newsetting_ft512.yaml

@@ -0,0 +1,24 @@
+exp_name: 'ummdit_small_ds1_singlenorm_v5_newdata_ft512_ema'
+model:
+  flashSA: 'Joint_SA' # choose from ['SACA', 'SA', None]
+  arch: 'ummdit_small_ds1_singlenorm_v5'
+  caption_max_seq_length: 128
+
+training:
+  save_freq: 10000
+  max_iters: 50000
+  transformer_ckpt: '/root/tongs/projects/efficient_diffusion_training/ummdit_small_ds1_singlenorm_v5_newdata_newsetting/checkpoints/checkpoint-100000/model.safetensors'
+  use_ema: True
+
+dataset:
+  datasets: [
+            '/root/tongs/data/jDB/mds_latents_dcvae_flant5large/train',
+            '/root/tongs/data/jDB/mds_latents_dcvae_flant5large/valid',
+            '/root/tongs/data/flux_gen_data/imglatent_mds_UCSC_part0/',
+            '/root/tongs/data/flux_gen_data/imglatent_mds_UCSC_part1/',
+            '/root/tongs/data/flux_gen_data/imglatent_mds_UCSC_part2/',
+            '/root/tongs/data/flux_gen_data/imglatent_mds_UCSC_part3/',
+            '/root/tongs/data/flux_gen_data/imglatent_mds_UCSC_part4/',
+            '/root/tongs/data/flux_gen_data/imglatent_mds_UCSC_part5/',
+            '/root/tongs/data/flux_gen_data/imglatent_mds_diffdb/'
+            ]