Spaces:
Paused
Paused
| 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 |