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