udit_models.py

# Copyright 2024 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
# --------------------------------------------------------
# References:
# GLIDE: https://github.com/openai/glide-text2im
# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
# --------------------------------------------------------

import torch
import torch.nn as nn
import numpy as np
import math
import torch.nn.functional as F
# from timm.models.vision_transformer import PatchEmbed
# from timm.models.layers import LayerNorm2d

from pdb import set_trace as stx
import numbers
# from torchprofile import profile_macs
from einops import rearrange
import einops

class LayerNorm2d(nn.LayerNorm):
    def __init__(self, num_channels, eps=1e-6, affine=True):
        super().__init__(num_channels, eps=eps, elementwise_affine=affine)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.permute(0, 2, 3, 1)
        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        x = x.permute(0, 3, 1, 2)
        return x

def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/weight_init.py
    # Cut & paste from Pytorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn(
            'mean is more than 2 std from [a, b] in nn.init.trunc_normal_. '
            'The distribution of values may be incorrect.',
            stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        low = norm_cdf((a - mean) / std)
        up = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [low, up], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * low - 1, 2 * up - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution.

    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/weight_init.py

    The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.

    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value

    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.trunc_normal_(w)
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)

#######################################3###################################
def precompute_freqs_cis_2d(dim: int, end: int, theta: float = 10000.0, scale=1.0, use_cls=False):
    H = int( end**0.5 )
    # assert  H * H == end
    flat_patch_pos = torch.arange(0 if not use_cls else -1, end) # N = end
    x_pos = flat_patch_pos % H # N
    y_pos = flat_patch_pos // H # N
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) # Hc/4
    x_freqs = torch.outer(x_pos, freqs).float() # N Hc/4
    y_freqs = torch.outer(y_pos, freqs).float() # N Hc/4
    x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
    y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
    freqs_cis = torch.cat([x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1) # N,Hc/4,2
    freqs_cis = freqs_cis.reshape(end if not use_cls else end + 1, -1)
    # we need to think how to implement this for multi heads.
    # freqs_cis = torch.cat([x_cis, y_cis], dim=-1) # N, Hc/2
    return freqs_cis


def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    # x: B N H Hc/2
    # freqs_cis:  N, H*Hc/2 or  N Hc/2
    ndim = x.ndim
    assert 0 <= 1 < ndim

    if freqs_cis.shape[-1] == x.shape[-1]:
        shape = [1 if i == 2 or i == 0 else d for i, d in enumerate(x.shape)]  # 1, N, 1, Hc/2
    else:
        shape = [d if i != 0 else 1 for i, d in enumerate(x.shape)] # 1, N, H, Hc/2
        # B, N, Hc/2
    return freqs_cis.view(*shape)

def apply_rotary_emb(
        xq: torch.Tensor,
        xk: torch.Tensor,
        freqs_cis: torch.Tensor,
):
    # xq : B N Head Ch_per_Head
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # B N H Hc/2
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) # B, N, H, Hc
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)


class DownSampler(nn.Module):
    def __init__(self, *args, **kwargs):
        ''' Required kwargs: down_factor, downsampler'''
        super().__init__()
        self.down_factor = kwargs.pop('down_factor')
        self.downsampler = kwargs.pop('downsampler')
        self.down_shortcut = kwargs.pop('down_shortcut')
        self.layer = nn.Conv2d(*args, **kwargs)

    def forward(self, x):
        x = self.layer(x) + (x if self.down_shortcut else 0)
        return rearrange(x, 'b d (h dh) (w dw) -> b (dh dw) (h w) d', dh=self.down_factor, dw=self.down_factor)


class DownSample_Attn(nn.Module):
    def __init__(self, dim, input_size, down_factor, num_heads,
                bias=False, posemb_type=None, attn_type=None, downsampler=None, down_shortcut=False, **kwargs):
        super(DownSample_Attn, self).__init__()
        if kwargs != dict(): # is not empty
            print(f'Kwargs: {kwargs}')

        self.input_resolution = (input_size, input_size)
        self.dim = dim
        self.heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5
        self.dh, self.dw = down_factor, down_factor
        self.to_qkv = nn.Linear(dim, dim * 3, bias=bias)
        
        self.to_out = nn.Conv2d(dim, dim, 1)

        # v2
        self.attn_type = attn_type
        if attn_type == 'v2':
            self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)

        # posemb
        self.posemb_type = posemb_type
        # if posemb_type is not None:
        #     print(f'Pos Emb {posemb_type} activated...')
        
        # posemb type
        if self.posemb_type == 'rope2d':
            self.freqs_cis = None

        # downsampler
        downsampler_ker_size = int(downsampler[-1])
        downsampler_padding = (int(downsampler[-1])-1)//2
        self.downsampler = DownSampler(dim, dim, kernel_size=downsampler_ker_size, stride=1, padding=downsampler_padding, groups=dim, down_factor=down_factor, downsampler=downsampler, down_shortcut=down_shortcut)
        


    def forward(self, x):
        b, _, h, w = x.size()
        x = self.downsampler(x)
        qkv = self.to_qkv(x).chunk(3, dim=-1)

        if self.posemb_type == 'rope2d':
            N = h * w // self.dh // self.dw
            if self.freqs_cis is None or self.freqs_cis.shape[0] != N:
                self.freqs_cis = precompute_freqs_cis_2d(self.dim // self.heads, N).to(x.device)
            # q, k input shape: B N H Hc
            q, k = map(lambda t: rearrange(t, 'b p n (h d) -> (b p) n h d', h=self.heads), qkv[:-1])
            v = rearrange(qkv[2], 'b p n (h d) -> b p h n d', h=self.heads)
            q, k = apply_rotary_emb(q, k, freqs_cis=self.freqs_cis)
            # reshape back
            q = rearrange(q, '(b p) n h d -> b p h n d', b=b)
            k = rearrange(k, '(b p) n h d -> b p h n d', b=b)
        else:
            q, k, v = map(lambda t: rearrange(t, 'b p n (h d) -> b p h n d', h=self.heads), qkv)

        if self.attn_type is None: # v1 attention
            attn = (q @ k.transpose(-2, -1))
            attn = attn * self.scale
        elif self.attn_type == 'v2': # v2 attention
            attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
            logit_scale = torch.clamp(self.logit_scale, max=4.6052).exp()
            attn = attn * logit_scale

        attn = attn.softmax(dim=-1)
        x = (attn @ v)
        # upsample
        x = rearrange(x, 'b (dh dw) he (h w) d -> b (he d) (h dh) (w dw)', h=h // self.dh, w=w // self.dw, dh=self.dh, dw=self.dw)
        
        x = self.to_out(x)
        return x


class FeedForward(nn.Module):
    def __init__(self, dim, ffn_expansion_factor, bias=True, rep=False):
        super(FeedForward, self).__init__()

        self.rep = rep
        self.bias = bias

        self.hidden_features = hidden_features = int(dim*ffn_expansion_factor)

        # buffer
        self.buffer = False # False: empty
        self.project_in_weight = None
        self.project_in_bias = None
        self.dwconv_weight = None
        self.dwconv_bias = None
        self.project_out_weight = None
        self.project_out_bias = None

        if rep:
            self.project_in = nn.Sequential(nn.Conv2d(dim, hidden_features//2, kernel_size=1, bias=bias),
                                            nn.Conv2d(hidden_features//2, hidden_features, kernel_size=1, bias=bias))

            self.dwconv = nn.ModuleList([
                nn.Conv2d(hidden_features, hidden_features, kernel_size=5, stride=1, padding=2, dilation=1, groups=hidden_features, bias=bias),
                nn.Conv2d(hidden_features, hidden_features, kernel_size=3, stride=1, padding=1, dilation=1, groups=hidden_features, bias=bias),
                nn.Conv2d(hidden_features, hidden_features, kernel_size=1, stride=1, padding=0, dilation=1, groups=hidden_features, bias=bias)
            ])

            self.project_out = nn.Sequential(nn.Conv2d(hidden_features, hidden_features//2, kernel_size=1, bias=bias),
                                             nn.Conv2d(hidden_features//2, dim, kernel_size=1, bias=bias))

        else:
            self.project_in = nn.Conv2d(dim, hidden_features, kernel_size=1, bias=bias)
            self.dwconv = nn.Conv2d(hidden_features, hidden_features, kernel_size=5, stride=1, padding=2, groups=hidden_features, bias=bias)
            self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)

    def clear_buffer(self):
        self.buffer = False # False: empty
        self.project_in_weight = None
        self.project_in_bias = None
        self.dwconv_weight = None
        self.dwconv_bias = None
        self.project_out_weight = None
        self.project_out_bias = None

    def add_buffer(self):
        self.buffer = True

        w1 = self.project_in[0].weight.squeeze().detach()
        w2 = self.project_in[1].weight.squeeze().detach()
        self.project_in_weight = (w2 @ w1).unsqueeze(-1).unsqueeze(-1)
        self.project_in_bias = None

        if self.bias:
            b1 = self.project_in[0].bias.detach()
            b2 = self.project_in[1].bias.detach()
            self.project_in_bias = (w2 @ b1 + b2)

        self.dwconv_weight = self.dwconv[0].weight.detach()
        self.dwconv_weight[:, :, 1:4, 1:4] += self.dwconv[1].weight.detach()
        self.dwconv_weight[:, :, 2:3, 2:3] += (self.dwconv[2].weight.detach() + 1.) # skip connection
        
        self.dwconv_bias = None
        if self.bias:
            self.dwconv_bias = self.dwconv[0].bias.detach() + self.dwconv[1].bias.detach() + self.dwconv[2].bias.detach()

        w1 = self.project_out[0].weight.squeeze().detach()
        w2 = self.project_out[1].weight.squeeze().detach()
        self.project_out_weight = (w2 @ w1).unsqueeze(-1).unsqueeze(-1)
        self.project_out_bias = None

        if self.bias:
            b1 = self.project_out[0].bias.detach()
            b2 = self.project_out[1].bias.detach()
            self.project_out_bias = (w2 @ b1 + b2)


    def train(self, mode: bool = True):
        if not isinstance(mode, bool):
            raise ValueError("training mode is expected to be boolean")
        self.training = mode
        if mode:
            self.clear_buffer() # added: clear rep buffer
        else:
            self.add_buffer() # added: add rep buffer
        for module in self.children():
            module.train(mode)
        return self
    

    def forward(self, x):
        if (not self.rep) or (self.rep and self.training): # debugged: 2024/6/4
            x = self.project_in(x)
            x = F.gelu(x)
            if self.rep:
                out = x
                for module in self.dwconv:
                    out = out + module(x)
            else:
                out = self.dwconv(x)
            x = self.project_out(out)
        else: # eval & rep
            x = F.conv2d(x, self.project_in_weight, self.project_in_bias) # project_in
            x = F.gelu(x) # debugged: 2024/6/4
            x = F.conv2d(x, self.dwconv_weight, self.dwconv_bias, padding=2, groups=self.hidden_features)
            x = F.conv2d(x, self.project_out_weight, self.project_out_bias) # project_out

        return x




class Mlp(nn.Module):
    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
    """
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        bias = (bias, bias)
        drop_probs = (drop, drop)

        self.fc1 = nn.Conv2d(in_features, hidden_features, 1, bias=bias[0])
        self.act = act_layer()
        self.drop1 = nn.Dropout(drop_probs[0])
        self.fc2 = nn.Conv2d(hidden_features, out_features, 1, bias=bias[1])
        self.drop2 = nn.Dropout(drop_probs[1])

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop1(x)
        x = self.fc2(x)
        x = self.drop2(x)
        return x


def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(-1).unsqueeze(-1)) + shift.unsqueeze(-1).unsqueeze(-1)

#################################################################################
#               Embedding Layers for Timesteps and Class Labels                 #
#################################################################################

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
        ).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):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        t_emb = self.mlp(t_freq)
        return t_emb


class LabelEmbedder(nn.Module):
    """
    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
    """
    def __init__(self, num_classes, hidden_size, dropout_prob):
        super().__init__()
        use_cfg_embedding = dropout_prob > 0
        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
        self.num_classes = num_classes
        self.dropout_prob = dropout_prob

    def token_drop(self, labels, force_drop_ids=None):
        """
        Drops labels to enable classifier-free guidance.
        """
        if force_drop_ids is None:
            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
        else:
            drop_ids = force_drop_ids == 1
        labels = torch.where(drop_ids, self.num_classes, labels)
        return labels

    def forward(self, labels, train, force_drop_ids=None):
        use_dropout = self.dropout_prob > 0
        if (train and use_dropout) or (force_drop_ids is not None):
            labels = self.token_drop(labels, force_drop_ids)
        embeddings = self.embedding_table(labels)
        return embeddings


#################################################################################
#                                 Core IPT Model                                #
#################################################################################

class U_DiTBlock(nn.Module):
    """
    A IPT block with adaptive layer norm zero (adaLN-Zero) conIPTioning.
    """
    def __init__(self, input_size, hidden_size, num_heads, mlp_ratio=4.0, down_factor=2, rep=1, ffn_type='rep', **kwargs):
        '''
        Args:
        dw_expand: For NAFNet
        dw_type: NAFNet dwconv type
        posemb_type: pos emb type of IPT related
        '''
        super().__init__()
        
        
        self.norm11 = LayerNorm2d(hidden_size, affine=False, eps=1e-6)
        
        self.attn1 = DownSample_Attn(hidden_size, input_size, down_factor=down_factor, num_heads=num_heads, bias=True, **kwargs)

        self.norm21 = LayerNorm2d(hidden_size, affine=False, eps=1e-6)
        if ffn_type == 'rep':
            self.mlp1 = FeedForward(hidden_size, mlp_ratio, rep=rep)
        elif ffn_type == 'basic':
            approx_gelu = lambda: nn.GELU(approximate="tanh")
            mlp_hidden_dim = int(hidden_size * mlp_ratio)
            self.mlp1 = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
        else:
            raise NotImplementedError(f'FFN type not implemented: {ffn_type}!')
        self.adaLN_modulation1 = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

        
        self.norm12 = LayerNorm2d(hidden_size, affine=False, eps=1e-6)
        
        self.attn2 = DownSample_Attn(hidden_size, input_size, down_factor=down_factor, num_heads=num_heads, bias=True, **kwargs)

        self.norm22 = LayerNorm2d(hidden_size, affine=False, eps=1e-6)
        if ffn_type == 'rep':
            self.mlp2 = FeedForward(hidden_size, mlp_ratio, rep=rep)
        elif ffn_type == 'basic':
            approx_gelu = lambda: nn.GELU(approximate="tanh")
            mlp_hidden_dim = int(hidden_size * mlp_ratio)
            self.mlp2 = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
        else:
            raise NotImplementedError(f'FFN type not implemented: {ffn_type}!')
        self.adaLN_modulation2 = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(self, x, c):

        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation1(c).chunk(6, dim=1)
        x = x + gate_msa.unsqueeze(-1).unsqueeze(-1) * self.attn1(modulate(self.norm11(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(-1).unsqueeze(-1) * self.mlp1(modulate(self.norm21(x), shift_mlp, scale_mlp))

        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation2(c).chunk(6, dim=1)
        x = x + gate_msa.unsqueeze(-1).unsqueeze(-1) * self.attn2(modulate(self.norm12(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(-1).unsqueeze(-1) * self.mlp2(modulate(self.norm22(x), shift_mlp, scale_mlp))
        return x

class FinalLayer(nn.Module):
    """
    The final layer of IPT.
    """
    def __init__(self, hidden_size, out_channels):
        super().__init__()
        self.norm_final = LayerNorm2d(hidden_size, affine=False, eps=1e-6)
        self.out_proj = nn.Conv2d(hidden_size, out_channels, kernel_size=3, stride=1, padding=1, bias=True)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
        x = modulate(self.norm_final(x), shift, scale)
        x = self.out_proj(x)
        return x

class OverlapPatchEmbed(nn.Module):
    def __init__(self, in_c=3, embed_dim=48, bias=False):
        super(OverlapPatchEmbed, self).__init__()

        self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=3, stride=1, padding=1, bias=bias)

    def forward(self, x):
        x = self.proj(x)

        return x

class Downsample(nn.Module):
    def __init__(self, n_feat):
        super(Downsample, self).__init__()

        self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat // 2, kernel_size=3, stride=1, padding=1, bias=False),
                                  nn.PixelUnshuffle(2))

    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.Conv2d(n_feat, n_feat * 2, kernel_size=3, stride=1, padding=1, bias=False),
                                  nn.PixelShuffle(2))

    def forward(self, x):
        return self.body(x)

class U_DiT(nn.Module):
    """
    Diffusion UNet model with a Transformer backbone.
    """
    def __init__(
        self,
        input_size=32,
        down_factor=2,
        in_channels=4,
        hidden_size=1152,
        depth=[2,5,8,5,2],
        num_heads=16,
        mlp_ratio=4,
        class_dropout_prob=0.1,
        num_classes=1000,
        learn_sigma=True,
        rep=1,
        ffn_type='rep',
        **kwargs
    ):
        super().__init__()
        self.learn_sigma = learn_sigma
        self.in_channels = in_channels
        self.out_channels = in_channels * 2 if learn_sigma else in_channels
        self.num_heads = num_heads

        down_factor = down_factor if isinstance(down_factor, list) else [down_factor] * 5
        

        self.x_embedder = OverlapPatchEmbed(in_channels, hidden_size, bias=True)
        self.t_embedder_1 = TimestepEmbedder(hidden_size)
        self.y_embedder_1 = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)

        self.t_embedder_2 = TimestepEmbedder(hidden_size*2)
        self.y_embedder_2 = LabelEmbedder(num_classes, hidden_size*2, class_dropout_prob)

        self.t_embedder_3 = TimestepEmbedder(hidden_size*4)
        self.y_embedder_3 = LabelEmbedder(num_classes, hidden_size*4, class_dropout_prob)

        # encoder-1
        self.encoder_level_1 = nn.ModuleList([
            U_DiTBlock(input_size, hidden_size, num_heads, mlp_ratio=mlp_ratio, down_factor=down_factor[0], rep=rep, ffn_type=ffn_type, **kwargs) for _ in range(depth[0])
        ])
        self.down1_2 = Downsample(hidden_size) 

        # encoder-2
        self.encoder_level_2 = nn.ModuleList([
            U_DiTBlock(input_size//2, hidden_size*2, num_heads, mlp_ratio=mlp_ratio, down_factor=down_factor[1], rep=rep, ffn_type=ffn_type, **kwargs) for _ in range(depth[1])
        ])
        self.down2_3 = Downsample(hidden_size*2) 

        # latent
        self.latent = nn.ModuleList([
            U_DiTBlock(input_size//4, hidden_size*4, num_heads, mlp_ratio=mlp_ratio, down_factor=down_factor[2], rep=rep, ffn_type=ffn_type, **kwargs) for _ in range(depth[2])
        ])

        # decoder-2
        self.up3_2 = Upsample(int(hidden_size*4))  ## From Level 4 to Level 3
        self.reduce_chan_level2 = nn.Conv2d(int(hidden_size*4), int(hidden_size*2), kernel_size=1, bias=True)
        self.decoder_level_2 = nn.ModuleList([
            U_DiTBlock(input_size//2, hidden_size*2, num_heads, mlp_ratio=mlp_ratio, down_factor=down_factor[3], rep=rep, ffn_type=ffn_type, **kwargs) for _ in range(depth[3])
        ])

        # decoder-1
        self.up2_1 = Upsample(int(hidden_size*2))  ## From Level 4 to Level 3
        self.reduce_chan_level1 = nn.Conv2d(int(hidden_size*2), int(hidden_size*2), kernel_size=1, bias=True)
        self.decoder_level_1 = nn.ModuleList([
            U_DiTBlock(input_size, hidden_size*2, num_heads, mlp_ratio=mlp_ratio, down_factor=down_factor[4], rep=rep, ffn_type=ffn_type, **kwargs) for _ in range(depth[4])
        ])

        self.output = nn.Conv2d(int(hidden_size*2), int(hidden_size*2), kernel_size=3, stride=1, padding=1, bias=True)

        self.final_layer = FinalLayer(hidden_size*2, self.out_channels)
        self.initialize_weights()

    def initialize_weights(self):
        # Initialize transformer layers:
        def _basic_init(module):
            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
        self.apply(_basic_init)

        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
        w = self.x_embedder.proj.weight.data
        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
        nn.init.constant_(self.x_embedder.proj.bias, 0)

        # Initialize label embedding table:
        nn.init.normal_(self.y_embedder_1.embedding_table.weight, std=0.02)
        nn.init.normal_(self.y_embedder_2.embedding_table.weight, std=0.02)
        nn.init.normal_(self.y_embedder_3.embedding_table.weight, std=0.02)

        # Initialize timestep embedding MLP:
        nn.init.normal_(self.t_embedder_1.mlp[0].weight, std=0.02)
        nn.init.normal_(self.t_embedder_1.mlp[2].weight, std=0.02)

        nn.init.normal_(self.t_embedder_2.mlp[0].weight, std=0.02)
        nn.init.normal_(self.t_embedder_2.mlp[2].weight, std=0.02)

        nn.init.normal_(self.t_embedder_3.mlp[0].weight, std=0.02)
        nn.init.normal_(self.t_embedder_3.mlp[2].weight, std=0.02)

        # Zero-out adaLN modulation layers in IPT blocks:
        blocks = self.encoder_level_1 + self.encoder_level_2 + self.latent + self.decoder_level_2 + self.decoder_level_1
        for block in blocks:
            nn.init.constant_(block.adaLN_modulation1[-1].weight, 0)
            nn.init.constant_(block.adaLN_modulation1[-1].bias, 0)
            nn.init.constant_(block.adaLN_modulation2[-1].weight, 0)
            nn.init.constant_(block.adaLN_modulation2[-1].bias, 0)

        # Zero-out output layers:
        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
        nn.init.constant_(self.final_layer.out_proj.weight, 0)
        nn.init.constant_(self.final_layer.out_proj.bias, 0)

    def forward(self, x, t, y):
        """
        Forward pass of U-DiT.
        x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
        t: (N,) tensor of diffusion timesteps
        y: (N,) tensor of class labels
        """
        x = self.x_embedder(x)                   # (N, C, H, W)
        t1 = self.t_embedder_1(t)    # (N, C, 1, 1)
        y1 = self.y_embedder_1(y, self.training)    # (N, C, 1, 1)
        c1 = t1 + y1                                # (N, D, 1, 1)

        t2 = self.t_embedder_2(t)    # (N, C, 1, 1)
        y2 = self.y_embedder_2(y, self.training)    # (N, C, 1, 1)
        c2 = t2 + y2                                # (N, D, 1, 1)

        t3 = self.t_embedder_3(t)    # (N, C, 1, 1)
        y3 = self.y_embedder_3(y, self.training)    # (N, C, 1, 1)
        c3 = t3 + y3                                # (N, D, 1, 1)

        # encoder_1
        out_enc_level1 = x
        for block in self.encoder_level_1:
            out_enc_level1 = block(out_enc_level1, c1)
        inp_enc_level2 = self.down1_2(out_enc_level1)

        # encoder_2
        out_enc_level2 = inp_enc_level2
        for block in self.encoder_level_2:
            out_enc_level2 = block(out_enc_level2, c2)
        inp_enc_level3 = self.down2_3(out_enc_level2)

        # latent
        latent = inp_enc_level3
        for block in self.latent:
            latent = block(latent, c3)

        # decoder_2
        inp_dec_level2 = self.up3_2(latent)
        inp_dec_level2 = torch.cat([inp_dec_level2, out_enc_level2], 1)
        inp_dec_level2 = self.reduce_chan_level2(inp_dec_level2)
        out_dec_level2 = inp_dec_level2
        for block in self.decoder_level_2:
            out_dec_level2 = block(out_dec_level2, c2)

        # decoder_1
        inp_dec_level1 = self.up2_1(out_dec_level2)
        inp_dec_level1 = torch.cat([inp_dec_level1, out_enc_level1], 1)
        inp_dec_level1 = self.reduce_chan_level1(inp_dec_level1)
        out_dec_level1 = inp_dec_level1
        for block in self.decoder_level_1:
            out_dec_level1 = block(out_dec_level1, c2)

        # output
        x = self.output(out_dec_level1)

        x = self.final_layer(x, c2)                # (N, T, patch_size ** 2 * out_channels)

        return x

    def forward_with_cfg(self, x, t, y, cfg_scale):
        """
        Forward pass of IPT, but also batches the unconIPTional forward pass for classifier-free guidance.
        """
        # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
        half = x[: len(x) // 2]
        combined = torch.cat([half, half], dim=0)
        model_out = self.forward(combined, t, y)
        # For exact reproducibility reasons, we apply classifier-free guidance on only
        # three channels by default. The standard approach to cfg applies it to all channels.
        # This can be done by uncommenting the following line and commenting-out the line following that.
        # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
        eps, rest = model_out[:, :3], model_out[:, 3:]
        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
        half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
        eps = torch.cat([half_eps, half_eps], dim=0)
        return torch.cat([eps, rest], dim=1)

#################################################################################
#                                   U-DITs Configs                                  #
#################################################################################
def U_DiT_custom(**kwargs):
    return U_DiT(**kwargs)

def U_DiT_S(**kwargs):
    return U_DiT(down_factor=2, hidden_size=96, num_heads=4, depth=[2,5,8,5,2], ffn_type='rep', rep=1, mlp_ratio=2, attn_type='v2', posemb_type='rope2d', downsampler='dwconv5', down_shortcut=1)

def U_DiT_B(**kwargs):
    return U_DiT(down_factor=2, hidden_size=192, num_heads=8, depth=[2,5,8,5,2], ffn_type='rep', rep=1, mlp_ratio=2, attn_type='v2', posemb_type='rope2d', downsampler='dwconv5', down_shortcut=1)

def U_DiT_L(**kwargs):
    return U_DiT(down_factor=2, hidden_size=384, num_heads=16, depth=[2,5,8,5,2], ffn_type='rep', rep=1, mlp_ratio=2, attn_type='v2', posemb_type='rope2d', downsampler='dwconv5', down_shortcut=1)

DiT_models = {
    'U-DiT-custom': U_DiT_custom,
    'U-DiT-S': U_DiT_S, # U-DiT-S
    'U-DiT-B': U_DiT_B, # U-DiT-B
    'U-DiT-L': U_DiT_L, # U-DiT-L
}


if __name__=="__main__":
    from torchprofile import profile_macs
    import warnings

    # Ablations on downsampler
    model = U_DiT_S()
    model.load_state_dict(torch.load('pt/newpt/U-DiT/U-DiT-S-400k.pt'))

    model.cuda()
    model.eval()

    inputs = torch.rand(1, 4, 32, 32).cuda()
    t = torch.ones(1).int().cuda()
    y = torch.ones(1).int().cuda()
    
    model(inputs, t, y)
    out = model(inputs, t, y)

    flops = profile_macs(model, (inputs, t, y))
    print(f'FLOPS: {flops/1e6:.2f}')


    model.train()
    print(f"output : {out.size()}")

    # backward test
    out = model(inputs, t, y)
    gt = torch.rand(1, 8, 32, 32).cuda()

    loss = torch.mean(out-gt)
    loss.backward()