Spaces:
Running
on
Zero
Running
on
Zero
| """ | |
| ein notation: | |
| b - batch | |
| n - sequence | |
| nt - text sequence | |
| nw - raw wave length | |
| d - dimension | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| from einops import repeat | |
| from x_transformers.x_transformers import RotaryEmbedding | |
| from capspeech.nar.model.modules import ( | |
| TimestepEmbedding, | |
| ConvNeXtV2Block, | |
| ConvPositionEmbedding, | |
| CrossDiTBlock, | |
| DiTBlock, | |
| AdaLayerNormZero_Final, | |
| precompute_freqs_cis, get_pos_embed_indices, | |
| ) | |
| class TextEmbedding(nn.Module): | |
| def __init__(self, text_num_embeds, text_dim, conv_layers = 0, conv_mult = 2): | |
| super().__init__() | |
| self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token | |
| if conv_layers > 0: | |
| self.extra_modeling = True | |
| self.precompute_max_pos = 4096 # ~44s of 24khz audio | |
| self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False) | |
| self.text_blocks = nn.Sequential(*[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]) | |
| else: | |
| self.extra_modeling = False | |
| def forward(self, text: int['b nt'], seq_len, drop_text=False): | |
| batch, text_len = text.shape[0], text.shape[1] | |
| text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx() | |
| text = text[:, :seq_len] # curtail if character tokens are more than the mel spec tokens | |
| text = F.pad(text, (0, seq_len - text_len), value = 0) | |
| if drop_text: # cfg for text | |
| text = torch.zeros_like(text) | |
| text = self.text_embed(text) # b n -> b n d | |
| # possible extra modeling | |
| if self.extra_modeling: | |
| # sinus pos emb | |
| batch_start = torch.zeros((batch,), dtype=torch.long) | |
| pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos) | |
| text_pos_embed = self.freqs_cis[pos_idx] | |
| text = text + text_pos_embed | |
| # convnextv2 blocks | |
| text = self.text_blocks(text) | |
| return text | |
| # noised input audio and context mixing embedding | |
| class InputEmbedding(nn.Module): | |
| def __init__(self, mel_dim, text_dim, out_dim): | |
| super().__init__() | |
| self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim) | |
| self.conv_pos_embed = ConvPositionEmbedding(dim = out_dim) | |
| def forward(self, x: float['b n d'], cond: float['b n d'], | |
| text_embed: float['b n d'], drop_audio_cond = False): | |
| if drop_audio_cond or cond is None: # cfg for cond audio | |
| cond = torch.zeros_like(x) | |
| x = self.proj(torch.cat((x, cond, text_embed), dim=-1)) | |
| x = self.conv_pos_embed(x) + x | |
| return x | |
| # Transformer backbone using DiT blocks | |
| class CrossDiT(nn.Module): | |
| def __init__(self, | |
| dim, depth=8, heads=8, dim_head=64, dropout=0.0, ff_mult=4, | |
| mel_dim=100, t5_dim=512, clap_dim=512, | |
| text_num_embeds=256, text_dim=None, conv_layers=0, | |
| skip=False, use_checkpoint=True, qk_norm=True, | |
| ): | |
| super().__init__() | |
| self.time_embed = TimestepEmbedding(dim) | |
| if text_dim is None: | |
| text_dim = mel_dim | |
| self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers) | |
| self.input_embed = InputEmbedding(mel_dim, text_dim, dim) | |
| self.caption_embedding = nn.Sequential( | |
| nn.Linear(t5_dim, dim), | |
| nn.SiLU(), | |
| nn.Linear(dim, dim) | |
| ) | |
| self.clap_embedding = nn.Sequential( | |
| nn.Linear(clap_dim, dim), | |
| nn.SiLU(), | |
| nn.Linear(dim, text_dim) | |
| ) | |
| # self.null_clap = nn.Parameters | |
| # self.null_prompt = nn.Parameters | |
| self.rotary_embed = RotaryEmbedding(dim_head) | |
| self.dim = dim | |
| self.depth = depth | |
| self.skip = skip | |
| self.in_blocks = nn.ModuleList([ | |
| CrossDiTBlock(dim=dim, | |
| heads=heads, | |
| dim_head=dim_head, | |
| ff_mult=ff_mult, | |
| dropout=dropout, | |
| use_checkpoint=use_checkpoint, | |
| qk_norm=qk_norm, | |
| skip=False | |
| ) | |
| for _ in range(depth//2) | |
| ]) | |
| self.mid_block = CrossDiTBlock(dim=dim, | |
| heads=heads, | |
| dim_head=dim_head, | |
| ff_mult=ff_mult, | |
| dropout=dropout, | |
| use_checkpoint=use_checkpoint, | |
| qk_norm=qk_norm, | |
| skip=False) | |
| self.out_blocks = nn.ModuleList([ | |
| CrossDiTBlock(dim=dim, | |
| heads=heads, | |
| dim_head=dim_head, | |
| ff_mult=ff_mult, | |
| dropout=dropout, | |
| use_checkpoint=use_checkpoint, | |
| qk_norm=qk_norm, | |
| skip=skip | |
| ) | |
| for _ in range(depth//2) | |
| ]) | |
| self.norm_out = AdaLayerNormZero_Final(dim) # final modulation | |
| self.proj_out = nn.Linear(dim, mel_dim) | |
| def forward( | |
| self, | |
| x: float['b n d'], # nosied input audio | |
| cond: float['b n d'], # masked cond audio | |
| prompt: float['b n d'], # speech caption | |
| clap: float['b n d'], # sound effects | |
| text: int['b nt'], # text | |
| time: float['b'] | float[''], # time step | |
| mask: bool['b n'] | None = None, | |
| prompt_mask: bool['b n'] | None = None, | |
| ): | |
| batch, seq_len = x.shape[0], x.shape[1] | |
| if time.ndim == 0: | |
| time = repeat(time, ' -> b', b=batch) | |
| t = self.time_embed(time) | |
| text_embed = self.text_embed(text, seq_len-1) | |
| prompt_embed = self.caption_embedding(prompt) | |
| clap_embed = self.clap_embedding(clap).unsqueeze(1) | |
| text_embed = torch.cat([clap_embed, text_embed], dim=1) | |
| x = self.input_embed(x, cond, text_embed) | |
| rope = self.rotary_embed.forward_from_seq_len(seq_len) | |
| skips = [] | |
| for i, block in enumerate(self.in_blocks): | |
| x = block(x, t, mask=mask, rope=rope, | |
| context=prompt_embed, context_mask=prompt_mask) | |
| if self.skip: | |
| skips.append(x) | |
| x = self.mid_block(x, t, mask=mask, rope=rope, | |
| context=prompt_embed, context_mask=prompt_mask) | |
| for i, block in enumerate(self.out_blocks): | |
| if self.skip: | |
| skip = skips.pop() | |
| else: | |
| skip = None | |
| x = block(x, t, mask=mask, rope=rope, | |
| context=prompt_embed, context_mask=prompt_mask, skip=skip) | |
| x = self.norm_out(x, t) | |
| output = self.proj_out(x) | |
| return output | |