import torch
import torch.nn as nn
from .utils.modules import PatchEmbed, TimestepEmbedder
from .utils.modules import PE_wrapper, RMSNorm
from .blocks import DiTBlock, JointDiTBlock
from .utils.span_mask import compute_mask_indices
class DiTControlNetEmbed(nn.Module):
def __init__(self, in_chans, out_chans, blocks,
cond_mask=False, cond_mask_prob=None,
cond_mask_ratio=None, cond_mask_span=None):
super().__init__()
self.conv_in = nn.Conv1d(in_chans, blocks[0], kernel_size=1)
self.cond_mask = cond_mask
if self.cond_mask:
self.mask_embed = nn.Parameter(torch.zeros((blocks[0])))
self.mask_prob = cond_mask_prob
self.mask_ratio = cond_mask_ratio
self.mask_span = cond_mask_span
blocks[0] = blocks[0] + 1
conv_blocks = []
for i in range(len(blocks) - 1):
channel_in = blocks[i]
channel_out = blocks[i + 1]
block = nn.Sequential(
nn.Conv1d(channel_in, channel_in, kernel_size=3, padding=1),
nn.SiLU(),
nn.Conv1d(channel_in, channel_out, kernel_size=3, padding=1, stride=2),
nn.SiLU(),)
conv_blocks.append(block)
self.blocks = nn.ModuleList(conv_blocks)
self.conv_out = nn.Conv1d(blocks[-1], out_chans, kernel_size=1)
nn.init.zeros_(self.conv_out.weight)
nn.init.zeros_(self.conv_out.bias)
def random_masking(self, gt, mask_ratios, mae_mask_infer=None):
B, D, L = gt.shape
if mae_mask_infer is None:
# mask = torch.rand(B, L).to(gt.device) < mask_ratios.unsqueeze(1)
mask_ratios = mask_ratios.cpu().numpy()
mask = compute_mask_indices(shape=[B, L],
padding_mask=None,
mask_prob=mask_ratios,
mask_length=self.mask_span,
mask_type="static",
mask_other=0.0,
min_masks=1,
no_overlap=False,
min_space=0,)
# only apply mask to some batches
mask_batch = torch.rand(B) < self.mask_prob
mask[~mask_batch] = False
mask = mask.unsqueeze(1).expand_as(gt)
else:
mask = mae_mask_infer
mask = mask.expand_as(gt)
gt[mask] = self.mask_embed.view(1, D, 1).expand_as(gt)[mask].type_as(gt)
return gt, mask.type_as(gt)
def forward(self, conditioning, cond_mask_infer=None):
embedding = self.conv_in(conditioning)
if self.cond_mask:
B, D, L = embedding.shape
if not self.training and cond_mask_infer is None:
cond_mask_infer = torch.zeros_like(embedding).bool()
mask_ratios = torch.FloatTensor(B).uniform_(*self.mask_ratio).to(embedding.device)
embedding, cond_mask = self.random_masking(embedding, mask_ratios, cond_mask_infer)
embedding = torch.cat([embedding, cond_mask[:, 0:1, :]], dim=1)
for block in self.blocks:
embedding = block(embedding)
embedding = self.conv_out(embedding)
# B, L, C
embedding = embedding.transpose(1, 2).contiguous()
return embedding
class DiTControlNet(nn.Module):
def __init__(self,
img_size=(224, 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,
# controlnet configs
cond_in=None, cond_blocks=None,
cond_mask=False, cond_mask_prob=None,
cond_mask_ratio=None, cond_mask_span=None,
**kwargs):
super().__init__()
self.num_features = self.embed_dim = embed_dim
# 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
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
else:
Block = DiTBlock
# norm layers
if norm_layer == 'layernorm':
norm_layer = nn.LayerNorm
elif norm_layer == 'rmsnorm':
norm_layer = RMSNorm
else:
raise NotImplementedError
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.controlnet_pre = DiTControlNetEmbed(in_chans=cond_in, out_chans=embed_dim,
blocks=cond_blocks,
cond_mask=cond_mask,
cond_mask_prob=cond_mask_prob,
cond_mask_ratio=cond_mask_ratio,
cond_mask_span=cond_mask_span)
controlnet_zero_blocks = []
for i in range(depth // 2):
block = nn.Linear(embed_dim, embed_dim)
nn.init.zeros_(block.weight)
nn.init.zeros_(block.bias)
controlnet_zero_blocks.append(block)
self.controlnet_zero_blocks = nn.ModuleList(controlnet_zero_blocks)
print('ControlNet ready \n')
def set_trainable(self):
for param in self.parameters():
param.requires_grad = False
# only train input_proj, blocks, and output_proj
for module_name in ['controlnet_pre', 'in_blocks', 'controlnet_zero_blocks']:
module = getattr(self, module_name, None)
if module is not None:
for param in module.parameters():
param.requires_grad = True
module.train()
else:
print(f'\n!!!warning missing trainable blocks: {module_name}!!!\n')
def forward(self, x, timesteps, context,
x_mask=None, context_mask=None,
cls_token=None,
condition=None, cond_mask_infer=None,
conditioning_scale=1.0):
# 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)
# add condition to x
condition = self.controlnet_pre(condition)
x = x + condition
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
if self.use_adanorm:
if self.cls_embed:
time_token = time_token + cls_token
time_token = self.time_act(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)
skips.append(x)
controlnet_skips = []
for skip, controlnet_block in zip(skips, self.controlnet_zero_blocks):
controlnet_skips.append(controlnet_block(skip) * conditioning_scale)
return controlnet_skips