import torch
import torch.nn as nn
import torch.utils.checkpoint
import math
from .utils.modules import PatchEmbed, TimestepEmbedder
from .utils.modules import PE_wrapper, RMSNorm
from .blocks import DiTBlock, JointDiTBlock, FinalBlock
class UDiT(nn.Module):
def __init__(self,
img_size=224, patch_size=16, in_chans=3,
input_type='2d', out_chans=None,
embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.,
qkv_bias=False, qk_scale=None, qk_norm=None,
act_layer='gelu', norm_layer='layernorm',
context_norm=False,
use_checkpoint=False,
# time fusion ada or token
time_fusion='token',
ada_lora_rank=None, ada_lora_alpha=None,
cls_dim=None,
# max length is only used for concat
context_dim=768, context_fusion='concat',
context_max_length=128, context_pe_method='sinu',
pe_method='abs', rope_mode='none',
use_conv=True,
skip=True, skip_norm=True):
super().__init__()
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
# input
self.in_chans = in_chans
self.input_type = input_type
if self.input_type == '2d':
num_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size)
elif self.input_type == '1d':
num_patches = img_size // patch_size
self.patch_embed = PatchEmbed(patch_size=patch_size, in_chans=in_chans,
embed_dim=embed_dim, input_type=input_type)
out_chans = in_chans if out_chans is None else out_chans
self.out_chans = out_chans
# position embedding
self.rope = rope_mode
self.x_pe = PE_wrapper(dim=embed_dim, method=pe_method,
length=num_patches)
print(f'x position embedding: {pe_method}')
print(f'rope mode: {self.rope}')
# time embed
self.time_embed = TimestepEmbedder(embed_dim)
self.time_fusion = time_fusion
self.use_adanorm = False
# cls embed
if cls_dim is not None:
self.cls_embed = nn.Sequential(
nn.Linear(cls_dim, embed_dim, bias=True),
nn.SiLU(),
nn.Linear(embed_dim, embed_dim, bias=True),)
else:
self.cls_embed = None
# time fusion
if time_fusion == 'token':
# put token at the beginning of sequence
self.extras = 2 if self.cls_embed else 1
self.time_pe = PE_wrapper(dim=embed_dim, method='abs', length=self.extras)
elif time_fusion in ['ada', 'ada_single', 'ada_lora', 'ada_lora_bias']:
self.use_adanorm = True
# aviod repetitive silu for each adaln block
self.time_act = nn.SiLU()
self.extras = 0
self.time_ada_final = nn.Linear(embed_dim, 2 * embed_dim, bias=True)
if time_fusion in ['ada_single', 'ada_lora', 'ada_lora_bias']:
# shared adaln
self.time_ada = nn.Linear(embed_dim, 6 * embed_dim, bias=True)
else:
self.time_ada = None
else:
raise NotImplementedError
print(f'time fusion mode: {self.time_fusion}')
# context
# use a simple projection
self.use_context = False
self.context_cross = False
self.context_max_length = context_max_length
self.context_fusion = 'none'
if context_dim is not None:
self.use_context = True
self.context_embed = nn.Sequential(
nn.Linear(context_dim, embed_dim, bias=True),
nn.SiLU(),
nn.Linear(embed_dim, embed_dim, bias=True),)
self.context_fusion = context_fusion
if context_fusion == 'concat' or context_fusion == 'joint':
self.extras += context_max_length
self.context_pe = PE_wrapper(dim=embed_dim,
method=context_pe_method,
length=context_max_length)
# no cross attention layers
context_dim = None
elif context_fusion == 'cross':
self.context_pe = PE_wrapper(dim=embed_dim,
method=context_pe_method,
length=context_max_length)
self.context_cross = True
context_dim = embed_dim
else:
raise NotImplementedError
print(f'context fusion mode: {context_fusion}')
print(f'context position embedding: {context_pe_method}')
if self.context_fusion == 'joint':
Block = JointDiTBlock
self.use_skip = skip[0]
else:
Block = DiTBlock
self.use_skip = skip
# norm layers
if norm_layer == 'layernorm':
norm_layer = nn.LayerNorm
elif norm_layer == 'rmsnorm':
norm_layer = RMSNorm
else:
raise NotImplementedError
print(f'use long skip connection: {skip}')
self.in_blocks = nn.ModuleList([
Block(
dim=embed_dim, context_dim=context_dim, num_heads=num_heads,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm,
act_layer=act_layer, norm_layer=norm_layer,
time_fusion=time_fusion,
ada_lora_rank=ada_lora_rank, ada_lora_alpha=ada_lora_alpha,
skip=False, skip_norm=False,
rope_mode=self.rope,
context_norm=context_norm,
use_checkpoint=use_checkpoint)
for _ in range(depth // 2)])
self.mid_block = Block(
dim=embed_dim, context_dim=context_dim, num_heads=num_heads,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm,
act_layer=act_layer, norm_layer=norm_layer,
time_fusion=time_fusion,
ada_lora_rank=ada_lora_rank, ada_lora_alpha=ada_lora_alpha,
skip=False, skip_norm=False,
rope_mode=self.rope,
context_norm=context_norm,
use_checkpoint=use_checkpoint)
self.out_blocks = nn.ModuleList([
Block(
dim=embed_dim, context_dim=context_dim, num_heads=num_heads,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm,
act_layer=act_layer, norm_layer=norm_layer,
time_fusion=time_fusion,
ada_lora_rank=ada_lora_rank, ada_lora_alpha=ada_lora_alpha,
skip=skip, skip_norm=skip_norm,
rope_mode=self.rope,
context_norm=context_norm,
use_checkpoint=use_checkpoint)
for _ in range(depth // 2)])
# FinalLayer block
self.use_conv = use_conv
self.final_block = FinalBlock(embed_dim=embed_dim,
patch_size=patch_size,
img_size=img_size,
in_chans=out_chans,
input_type=input_type,
norm_layer=norm_layer,
use_conv=use_conv,
use_adanorm=self.use_adanorm)
self.initialize_weights()
def _init_ada(self):
if self.time_fusion == 'ada':
nn.init.constant_(self.time_ada_final.weight, 0)
nn.init.constant_(self.time_ada_final.bias, 0)
for block in self.in_blocks:
nn.init.constant_(block.adaln.time_ada.weight, 0)
nn.init.constant_(block.adaln.time_ada.bias, 0)
nn.init.constant_(self.mid_block.adaln.time_ada.weight, 0)
nn.init.constant_(self.mid_block.adaln.time_ada.bias, 0)
for block in self.out_blocks:
nn.init.constant_(block.adaln.time_ada.weight, 0)
nn.init.constant_(block.adaln.time_ada.bias, 0)
elif self.time_fusion == 'ada_single':
nn.init.constant_(self.time_ada.weight, 0)
nn.init.constant_(self.time_ada.bias, 0)
nn.init.constant_(self.time_ada_final.weight, 0)
nn.init.constant_(self.time_ada_final.bias, 0)
elif self.time_fusion in ['ada_lora', 'ada_lora_bias']:
nn.init.constant_(self.time_ada.weight, 0)
nn.init.constant_(self.time_ada.bias, 0)
nn.init.constant_(self.time_ada_final.weight, 0)
nn.init.constant_(self.time_ada_final.bias, 0)
for block in self.in_blocks:
nn.init.kaiming_uniform_(block.adaln.lora_a.weight,
a=math.sqrt(5))
nn.init.constant_(block.adaln.lora_b.weight, 0)
nn.init.kaiming_uniform_(self.mid_block.adaln.lora_a.weight,
a=math.sqrt(5))
nn.init.constant_(self.mid_block.adaln.lora_b.weight, 0)
for block in self.out_blocks:
nn.init.kaiming_uniform_(block.adaln.lora_a.weight,
a=math.sqrt(5))
nn.init.constant_(block.adaln.lora_b.weight, 0)
def initialize_weights(self):
# Basic init for all layers
def _basic_init(module):
if isinstance(module, nn.Linear):
torch.nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.constant_(module.bias, 0)
self.apply(_basic_init)
# init patch Conv like Linear
w = self.patch_embed.proj.weight.data
nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
nn.init.constant_(self.patch_embed.proj.bias, 0)
# Zero-out AdaLN
if self.use_adanorm:
self._init_ada()
# Zero-out Cross Attention
if self.context_cross:
for block in self.in_blocks:
nn.init.constant_(block.cross_attn.proj.weight, 0)
nn.init.constant_(block.cross_attn.proj.bias, 0)
nn.init.constant_(self.mid_block.cross_attn.proj.weight, 0)
nn.init.constant_(self.mid_block.cross_attn.proj.bias, 0)
for block in self.out_blocks:
nn.init.constant_(block.cross_attn.proj.weight, 0)
nn.init.constant_(block.cross_attn.proj.bias, 0)
# Zero-out cls embedding
if self.cls_embed:
if self.use_adanorm:
nn.init.constant_(self.cls_embed[-1].weight, 0)
nn.init.constant_(self.cls_embed[-1].bias, 0)
# Zero-out Output
# might not zero-out this when using v-prediction
# it could be good when using noise-prediction
# nn.init.constant_(self.final_block.linear.weight, 0)
# nn.init.constant_(self.final_block.linear.bias, 0)
# if self.use_conv:
# nn.init.constant_(self.final_block.final_layer.weight.data, 0)
# nn.init.constant_(self.final_block.final_layer.bias, 0)
# init out Conv
if self.use_conv:
nn.init.xavier_uniform_(self.final_block.final_layer.weight)
nn.init.constant_(self.final_block.final_layer.bias, 0)
def _concat_x_context(self, x, context, x_mask=None, context_mask=None):
assert context.shape[-2] == self.context_max_length
# Check if either x_mask or context_mask is provided
B = x.shape[0]
# Create default masks if they are not provided
if x_mask is None:
x_mask = torch.ones(B, x.shape[-2], device=x.device).bool()
if context_mask is None:
context_mask = torch.ones(B, context.shape[-2],
device=context.device).bool()
# Concatenate the masks along the second dimension (dim=1)
x_mask = torch.cat([context_mask, x_mask], dim=1)
# Concatenate context and x along the second dimension (dim=1)
x = torch.cat((context, x), dim=1)
return x, x_mask
def forward(self, x, timesteps, context,
x_mask=None, context_mask=None,
cls_token=None
):
# make it compatible with int time step during inference
if timesteps.dim() == 0:
timesteps = timesteps.expand(x.shape[0]).to(x.device, dtype=torch.long)
x = self.patch_embed(x)
x = self.x_pe(x)
B, L, D = x.shape
if self.use_context:
context_token = self.context_embed(context)
context_token = self.context_pe(context_token)
if self.context_fusion == 'concat' or self.context_fusion == 'joint':
x, x_mask = self._concat_x_context(x=x, context=context_token,
x_mask=x_mask,
context_mask=context_mask)
context_token, context_mask = None, None
else:
context_token, context_mask = None, None
time_token = self.time_embed(timesteps)
if self.cls_embed:
cls_token = self.cls_embed(cls_token)
time_ada = None
time_ada_final = None
if self.use_adanorm:
if self.cls_embed:
time_token = time_token + cls_token
time_token = self.time_act(time_token)
time_ada_final = self.time_ada_final(time_token)
if self.time_ada is not None:
time_ada = self.time_ada(time_token)
else:
time_token = time_token.unsqueeze(dim=1)
if self.cls_embed:
cls_token = cls_token.unsqueeze(dim=1)
time_token = torch.cat([time_token, cls_token], dim=1)
time_token = self.time_pe(time_token)
x = torch.cat((time_token, x), dim=1)
if x_mask is not None:
x_mask = torch.cat(
[torch.ones(B, time_token.shape[1], device=x_mask.device).bool(),
x_mask], dim=1)
time_token = None
skips = []
for blk in self.in_blocks:
x = blk(x=x, time_token=time_token, time_ada=time_ada,
skip=None, context=context_token,
x_mask=x_mask, context_mask=context_mask,
extras=self.extras)
if self.use_skip:
skips.append(x)
x = self.mid_block(x=x, time_token=time_token, time_ada=time_ada,
skip=None, context=context_token,
x_mask=x_mask, context_mask=context_mask,
extras=self.extras)
for blk in self.out_blocks:
skip = skips.pop() if self.use_skip else None
x = blk(x=x, time_token=time_token, time_ada=time_ada,
skip=skip, context=context_token,
x_mask=x_mask, context_mask=context_mask,
extras=self.extras)
x = self.final_block(x, time_ada=time_ada_final, extras=self.extras)
return x