Spaces:
Running on Zero
Running on Zero
| """ | |
| ein notation: | |
| b - batch | |
| n - sequence | |
| nt - text sequence | |
| nw - raw wave length | |
| d - dimension | |
| """ | |
| # ruff: noqa: F722 F821 | |
| from __future__ import annotations | |
| import math | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from torch.nn.utils.rnn import pad_sequence | |
| from x_transformers.x_transformers import RotaryEmbedding | |
| from wavegen_v_spec.model.modules import ( | |
| AdaLayerNorm_Final, | |
| ConvNeXtV2Block, | |
| ConvPositionEmbedding, | |
| DiTBlock, | |
| TimestepEmbedding, | |
| precompute_freqs_cis, | |
| ) | |
| # Text embedding | |
| class TextEmbedding(nn.Module): | |
| def __init__( | |
| self, text_num_embeds, text_dim, mask_padding=True, average_upsampling=False, conv_layers=0, conv_mult=2 | |
| ): | |
| super().__init__() | |
| self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token | |
| self.mask_padding = mask_padding # mask filler and batch padding tokens or not | |
| self.average_upsampling = average_upsampling # zipvoice-style text late average upsampling (after text encoder) | |
| if average_upsampling: | |
| assert mask_padding, "text_embedding_average_upsampling requires text_mask_padding to be True" | |
| if conv_layers > 0: | |
| self.extra_modeling = True | |
| self.precompute_max_pos = 8192 # 8192 waveform tokens is ~81.92s at 16k / wav_frame_len=160 | |
| 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 average_upsample_text_by_mask(self, text, text_mask): | |
| batch, seq_len, text_dim = text.shape | |
| text_lens = text_mask.sum(dim=1) # [batch] | |
| upsampled_text = torch.zeros_like(text) | |
| for i in range(batch): | |
| valid_text_len = text_lens[i].item() | |
| if valid_text_len == 0: | |
| continue | |
| valid_ind = torch.where(text_mask[i])[0] | |
| valid_data = text[i, valid_ind, :] # [valid_text_len, text_dim] | |
| base_repeat = seq_len // valid_text_len | |
| remainder = seq_len % valid_text_len | |
| indices = [] | |
| for j in range(valid_text_len): | |
| repeat_count = base_repeat + (1 if j >= valid_text_len - remainder else 0) | |
| indices.extend([j] * repeat_count) | |
| indices = torch.tensor(indices[:seq_len], device=text.device, dtype=torch.long) | |
| upsampled = valid_data[indices] # [seq_len, text_dim] | |
| upsampled_text[i, :seq_len, :] = upsampled | |
| return upsampled_text | |
| def forward(self, text: int["b nt"], seq_len, drop_text=False): | |
| 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 waveform tokens | |
| text = F.pad(text, (0, seq_len - text.shape[1]), value=0) # (opt.) if not self.average_upsampling: | |
| if self.mask_padding: | |
| text_mask = text == 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 | |
| text = text + self.freqs_cis[:seq_len, :] | |
| # convnextv2 blocks | |
| if self.mask_padding: | |
| text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) | |
| for block in self.text_blocks: | |
| text = block(text) | |
| text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) | |
| else: | |
| text = self.text_blocks(text) | |
| if self.average_upsampling: | |
| text = self.average_upsample_text_by_mask(text, ~text_mask) | |
| return text | |
| # noised waveform and context mixing embedding | |
| class MelInformationBottleneck(nn.Module): | |
| """Dropout-linear-dropout suffix for the noised mel trajectory.""" | |
| def __init__(self, bottleneck_dim: int, dropout: float = 0.5): | |
| super().__init__() | |
| if not 0.0 <= dropout < 1.0: | |
| raise ValueError(f"mel bottleneck dropout must be in [0, 1), got {dropout}") | |
| self.dropout1 = nn.Dropout(dropout) | |
| self.linear2 = nn.Linear(bottleneck_dim, bottleneck_dim, bias=True) | |
| self.dropout2 = nn.Dropout(dropout) | |
| def _inverted_dropout( | |
| value: torch.Tensor, | |
| probability: float, | |
| generator: torch.Generator | None, | |
| ) -> torch.Tensor: | |
| if probability == 0.0: | |
| return value | |
| if generator is None: | |
| raise RuntimeError( | |
| "Inference Mel bottleneck dropout requires a dedicated generator" | |
| ) | |
| keep = torch.rand( | |
| value.shape, | |
| device=value.device, | |
| dtype=torch.float32, | |
| generator=generator, | |
| ).ge(probability) | |
| return value * keep.to(dtype=value.dtype) / (1.0 - probability) | |
| def forward( | |
| self, | |
| mel: torch.Tensor, | |
| *, | |
| inference_dropout_probability: float | None = None, | |
| inference_dropout_generator: torch.Generator | None = None, | |
| ) -> torch.Tensor: | |
| if inference_dropout_probability is None: | |
| mel = self.linear2(self.dropout1(mel)) | |
| return self.dropout2(mel) | |
| probability = float(inference_dropout_probability) | |
| if not math.isfinite(probability) or not 0.0 <= probability < 1.0: | |
| raise ValueError( | |
| "inference Mel bottleneck dropout must be finite and in [0, 1), " | |
| f"got {inference_dropout_probability}" | |
| ) | |
| if self.training: | |
| raise RuntimeError( | |
| "Explicit inference Mel bottleneck dropout requires eval mode" | |
| ) | |
| mel = self._inverted_dropout( | |
| mel, probability, inference_dropout_generator | |
| ) | |
| mel = self.linear2(mel) | |
| return self._inverted_dropout( | |
| mel, probability, inference_dropout_generator | |
| ) | |
| class InputEmbedding(nn.Module): | |
| def __init__( | |
| self, | |
| wav_frame_len, | |
| text_dim, | |
| mel_dim, | |
| out_dim, | |
| use_mel_information_bottleneck: bool = False, | |
| mel_bottleneck_dim: int | None = None, | |
| mel_bottleneck_dropout: float = 0.5, | |
| mel_bottleneck_zero_init: bool = False, | |
| use_audio_proj: bool = False, | |
| audio_proj_dim: int | None = None, | |
| audio_proj_hidden: int | None = None, | |
| ): | |
| super().__init__() | |
| self.use_audio_proj = use_audio_proj | |
| self.use_mel_information_bottleneck = bool(use_mel_information_bottleneck) | |
| self.mel_bottleneck_dim = int( | |
| mel_dim if mel_bottleneck_dim is None else mel_bottleneck_dim | |
| ) | |
| if self.mel_bottleneck_dim <= 0: | |
| raise ValueError( | |
| f"mel_bottleneck_dim must be positive, got {self.mel_bottleneck_dim}" | |
| ) | |
| if self.use_mel_information_bottleneck and self.mel_bottleneck_dim > mel_dim: | |
| raise ValueError( | |
| "mel_bottleneck_dim must not exceed mel_channels when the " | |
| "bottleneck is enabled: " | |
| f"{self.mel_bottleneck_dim} > {mel_dim}" | |
| ) | |
| if self.use_mel_information_bottleneck: | |
| self.mel_proj = nn.Linear(mel_dim, self.mel_bottleneck_dim, bias=False) | |
| self.mel_bottleneck = MelInformationBottleneck( | |
| self.mel_bottleneck_dim, dropout=mel_bottleneck_dropout | |
| ) | |
| mel_xt_dim = self.mel_bottleneck_dim | |
| else: | |
| # Identity keeps the disabled path and old state dict unchanged. | |
| self.mel_proj = nn.Identity() | |
| self.mel_bottleneck = nn.Identity() | |
| mel_xt_dim = mel_dim | |
| self.mel_bottleneck_zero_init = bool(mel_bottleneck_zero_init) | |
| if self.mel_bottleneck_zero_init: | |
| if not self.use_mel_information_bottleneck: | |
| raise ValueError( | |
| "mel_bottleneck_zero_init requires " | |
| "use_mel_information_bottleneck=True" | |
| ) | |
| nn.init.zeros_(self.mel_proj.weight) | |
| nn.init.zeros_(self.mel_bottleneck.linear2.bias) | |
| if not use_audio_proj: | |
| self.proj = nn.Linear( | |
| wav_frame_len * 2 + mel_xt_dim + mel_dim + text_dim, | |
| out_dim, | |
| ) | |
| else: | |
| audio_proj_dim = out_dim if audio_proj_dim is None else audio_proj_dim | |
| audio_proj_hidden = audio_proj_dim if audio_proj_hidden is None else audio_proj_hidden | |
| self.x_proj = nn.Sequential( | |
| nn.Linear(wav_frame_len, audio_proj_hidden, bias=False), | |
| nn.Linear(audio_proj_hidden, audio_proj_dim), | |
| ) | |
| self.cond_proj = nn.Sequential( | |
| nn.Linear(wav_frame_len, audio_proj_hidden, bias=False), | |
| nn.Linear(audio_proj_hidden, audio_proj_dim), | |
| ) | |
| self.fuse = nn.Linear( | |
| audio_proj_dim * 2 + mel_xt_dim + mel_dim + 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"], | |
| mel: float["b n d"], | |
| mel_cond: float["b n d"], | |
| text_embed: float["b n d"], | |
| drop_audio_cond=False, | |
| audio_mask: bool["b n"] | None = None, | |
| inference_bottleneck_dropout: float | None = None, | |
| inference_bottleneck_generator: torch.Generator | None = None, | |
| ): | |
| if drop_audio_cond: | |
| cond = torch.zeros_like(cond) | |
| mel_cond = torch.zeros_like(mel_cond) | |
| mel_h = self.mel_proj(mel) | |
| if self.use_mel_information_bottleneck: | |
| mel_h = self.mel_bottleneck( | |
| mel_h, | |
| inference_dropout_probability=inference_bottleneck_dropout, | |
| inference_dropout_generator=inference_bottleneck_generator, | |
| ) | |
| if not self.use_audio_proj: | |
| h = self.proj(torch.cat((x, cond, mel_h, mel_cond, text_embed), dim=-1)) | |
| else: | |
| x_h = self.x_proj(x) | |
| c_h = self.cond_proj(cond) | |
| h = self.fuse(torch.cat((x_h, c_h, mel_h, mel_cond, text_embed), dim=-1)) | |
| h = self.conv_pos_embed(h, mask=audio_mask) + h | |
| return h | |
| class DiTBackbone(nn.Module): | |
| def __init__( | |
| self, | |
| *, | |
| dim, | |
| depth=8, | |
| heads=8, | |
| dim_head=64, | |
| dropout=0.1, | |
| ff_mult=4, | |
| wav_frame_len=160, | |
| mel_dim=64, | |
| mel_channels: int | None = None, | |
| use_mel_information_bottleneck: bool = False, | |
| mel_bottleneck_dim: int | None = None, | |
| mel_bottleneck_dropout: float = 0.5, | |
| mel_bottleneck_zero_init: bool = False, | |
| text_num_embeds=256, | |
| text_dim=None, | |
| text_mask_padding=True, | |
| text_embedding_average_upsampling=False, | |
| qk_norm=None, | |
| conv_layers=0, | |
| pe_attn_head=None, | |
| attn_backend="torch", # "torch" | "flash_attn" | |
| attn_mask_enabled=False, | |
| long_skip_connection=False, | |
| checkpoint_activations=False, | |
| use_audio_proj: bool = False, | |
| audio_proj_dim: int | None = None, | |
| audio_proj_hidden: int | None = None, | |
| ): | |
| super().__init__() | |
| # `mel_dim` remains as a backward-compatible API alias. Configs use | |
| # `mel_channels` as the single source for forcing input/output width. | |
| mel_dim = int(mel_dim if mel_channels is None else mel_channels) | |
| if mel_dim <= 0: | |
| raise ValueError(f"mel_channels must be positive, got {mel_dim}") | |
| self.time_embed = TimestepEmbedding(dim) | |
| if text_dim is None: | |
| text_dim = wav_frame_len | |
| self.text_embed = TextEmbedding( | |
| text_num_embeds, | |
| text_dim, | |
| mask_padding=text_mask_padding, | |
| average_upsampling=text_embedding_average_upsampling, | |
| conv_layers=conv_layers, | |
| ) | |
| self.text_cond, self.text_uncond = None, None # text cache | |
| self.input_embed = InputEmbedding( | |
| wav_frame_len, text_dim, mel_dim, dim, | |
| use_mel_information_bottleneck=use_mel_information_bottleneck, | |
| mel_bottleneck_dim=mel_bottleneck_dim, | |
| mel_bottleneck_dropout=mel_bottleneck_dropout, | |
| mel_bottleneck_zero_init=mel_bottleneck_zero_init, | |
| use_audio_proj=use_audio_proj, | |
| audio_proj_dim=audio_proj_dim, | |
| audio_proj_hidden=audio_proj_hidden, | |
| ) | |
| self.rotary_embed = RotaryEmbedding(dim_head) | |
| self.dim = dim | |
| self.depth = depth | |
| self.mel_dim = mel_dim | |
| self.mel_channels = mel_dim | |
| self.use_mel_information_bottleneck = bool(use_mel_information_bottleneck) | |
| self.mel_bottleneck_dim = self.input_embed.mel_bottleneck_dim | |
| self.mel_bottleneck_dropout = float(mel_bottleneck_dropout) | |
| self.mel_bottleneck_zero_init = bool(mel_bottleneck_zero_init) | |
| self.transformer_blocks = nn.ModuleList( | |
| [ | |
| DiTBlock( | |
| dim=dim, | |
| heads=heads, | |
| dim_head=dim_head, | |
| ff_mult=ff_mult, | |
| dropout=dropout, | |
| qk_norm=qk_norm, | |
| pe_attn_head=pe_attn_head, | |
| attn_backend=attn_backend, | |
| attn_mask_enabled=attn_mask_enabled, | |
| ) | |
| for _ in range(depth) | |
| ] | |
| ) | |
| self.wav_ddt = nn.ModuleList( | |
| [ | |
| DiTBlock( | |
| dim=dim, | |
| heads=heads, | |
| dim_head=dim_head, | |
| ff_mult=ff_mult, | |
| dropout=dropout, | |
| qk_norm=qk_norm, | |
| pe_attn_head=pe_attn_head, | |
| attn_backend=attn_backend, | |
| attn_mask_enabled=attn_mask_enabled, | |
| ) | |
| for _ in range(2) | |
| ] | |
| ) | |
| self.mel_ddt = nn.ModuleList( | |
| [ | |
| DiTBlock( | |
| dim=dim, | |
| heads=heads, | |
| dim_head=dim_head, | |
| ff_mult=ff_mult, | |
| dropout=dropout, | |
| qk_norm=qk_norm, | |
| pe_attn_head=pe_attn_head, | |
| attn_backend=attn_backend, | |
| attn_mask_enabled=attn_mask_enabled, | |
| ) | |
| for _ in range(2) | |
| ] | |
| ) | |
| self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None | |
| self.norm_out = AdaLayerNorm_Final(dim) # final modulation | |
| self.proj_out_dim = wav_frame_len | |
| self.proj_out = nn.Linear(dim, wav_frame_len) | |
| self.proj_out_output_layer = self.proj_out | |
| self.mel_norm_out = AdaLayerNorm_Final(dim) # final modulation | |
| self.mel_proj_out_dim = mel_dim | |
| self.mel_proj_out = nn.Linear(dim, mel_dim) | |
| self.mel_proj_out_output_layer = self.mel_proj_out | |
| self.checkpoint_activations = checkpoint_activations | |
| self.initialize_weights() | |
| # raw waveform tokenization. | |
| self.wav_frame_len = wav_frame_len | |
| def initialize_weights(self): | |
| # 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) | |
| # # Initialize timestep embedding MLP: | |
| # nn.init.normal_(self.t_embed.mlp[0].weight, std=0.02) | |
| # nn.init.normal_(self.t_embed.mlp[2].weight, std=0.02) | |
| # Zero-out AdaLN layers in DiT blocks: | |
| for block in self.transformer_blocks: | |
| nn.init.constant_(block.attn_norm.linear.weight, 0) | |
| nn.init.constant_(block.attn_norm.linear.bias, 0) | |
| for block in self.wav_ddt: | |
| nn.init.constant_(block.attn_norm.linear.weight, 0) | |
| nn.init.constant_(block.attn_norm.linear.bias, 0) | |
| for block in self.mel_ddt: | |
| nn.init.constant_(block.attn_norm.linear.weight, 0) | |
| nn.init.constant_(block.attn_norm.linear.bias, 0) | |
| # Zero-out output layers: | |
| nn.init.constant_(self.norm_out.linear.weight, 0) | |
| nn.init.constant_(self.norm_out.linear.bias, 0) | |
| nn.init.constant_(self.proj_out_output_layer.weight, 0) | |
| if self.proj_out_output_layer.bias is not None: | |
| nn.init.constant_(self.proj_out_output_layer.bias, 0) | |
| nn.init.constant_(self.mel_norm_out.linear.weight, 0) | |
| nn.init.constant_(self.mel_norm_out.linear.bias, 0) | |
| nn.init.constant_(self.mel_proj_out_output_layer.weight, 0) | |
| if self.mel_proj_out_output_layer.bias is not None: | |
| nn.init.constant_(self.mel_proj_out_output_layer.bias, 0) | |
| def set_wav_frame_len(self, wav_frame_len: int): | |
| self.wav_frame_len = int(wav_frame_len) | |
| if self.wav_frame_len != self.proj_out_dim: | |
| raise ValueError( | |
| f"wav_frame_len ({self.wav_frame_len}) must equal proj_out_dim ({self.proj_out_dim}) " | |
| "for reshape wav front-end." | |
| ) | |
| def _wav_to_tokens( | |
| self, | |
| wav: torch.Tensor, | |
| mask: torch.Tensor | None = None, | |
| lens: torch.Tensor | None = None, | |
| ): | |
| assert wav.ndim == 2, f"Expected [B, N] wav input, got {tuple(wav.shape)}" | |
| bsz, num_samples = wav.shape | |
| frame_len = self.wav_frame_len | |
| pad_len = (frame_len - (num_samples % frame_len)) % frame_len | |
| if pad_len > 0: | |
| wav = F.pad(wav, (0, pad_len), value=0.0) | |
| if mask is not None: | |
| mask = F.pad(mask, (0, pad_len), value=False) | |
| tokens = wav.view(bsz, -1, frame_len) | |
| token_mask = None | |
| if mask is not None: | |
| token_mask = mask.view(bsz, -1, frame_len).any(dim=-1) | |
| token_lens = None | |
| if lens is not None: | |
| token_lens = (lens.to(dtype=torch.long, device=wav.device) + frame_len - 1) // frame_len | |
| return tokens, token_mask, token_lens | |
| def _tokens_to_wav(self, tokens: torch.Tensor, target_num_samples: int): | |
| wav = tokens.reshape(tokens.shape[0], -1) | |
| return wav[:, :target_num_samples] | |
| def _align_mel_to_tokens(self, mel: torch.Tensor, seq_len: int, name: str) -> torch.Tensor: | |
| if mel.ndim != 3: | |
| raise ValueError(f"Expected {name} [B, N, M], got {tuple(mel.shape)}") | |
| if mel.shape[-1] != self.mel_channels: | |
| raise ValueError( | |
| f"Expected {name} with {self.mel_channels} mel channels, " | |
| f"got {mel.shape[-1]}" | |
| ) | |
| if mel.shape[1] < seq_len: | |
| mel = F.pad(mel, (0, 0, 0, seq_len - mel.shape[1]), value=0.0) | |
| elif mel.shape[1] > seq_len: | |
| mel = mel[:, :seq_len] | |
| return mel | |
| def get_input_embed( | |
| self, | |
| x, # b n d | |
| cond, # b n d | |
| mel, # b n d | |
| mel_cond, # b n d | |
| text, # b nt | |
| drop_audio_cond: bool = False, | |
| drop_text: bool = False, | |
| cache: bool = True, | |
| audio_mask: bool["b n"] | None = None, | |
| inference_bottleneck_dropout: float | None = None, | |
| inference_bottleneck_generator: torch.Generator | None = None, | |
| ): | |
| if self.text_uncond is None or self.text_cond is None or not cache: | |
| if audio_mask is None: | |
| text_embed = self.text_embed(text, x.shape[1], drop_text=drop_text) | |
| else: | |
| batch = x.shape[0] | |
| seq_lens = audio_mask.sum(dim=1) # Calculate the actual sequence length for each sample | |
| text_embed_list = [] | |
| for i in range(batch): | |
| text_embed_i = self.text_embed( | |
| text[i].unsqueeze(0), | |
| seq_len=seq_lens[i].item(), | |
| drop_text=drop_text, | |
| ) | |
| text_embed_list.append(text_embed_i[0]) | |
| text_embed = pad_sequence(text_embed_list, batch_first=True, padding_value=0) | |
| if cache: | |
| if drop_text: | |
| self.text_uncond = text_embed | |
| else: | |
| self.text_cond = text_embed | |
| if cache: | |
| if drop_text: | |
| text_embed = self.text_uncond | |
| else: | |
| text_embed = self.text_cond | |
| x = self.input_embed( | |
| x, | |
| cond, | |
| mel, | |
| mel_cond, | |
| text_embed, | |
| drop_audio_cond=drop_audio_cond, | |
| audio_mask=audio_mask, | |
| inference_bottleneck_dropout=inference_bottleneck_dropout, | |
| inference_bottleneck_generator=inference_bottleneck_generator, | |
| ) | |
| return x | |
| def clear_cache(self): | |
| self.text_cond, self.text_uncond = None, None | |
| def forward( | |
| self, | |
| x: float["b nw"], # noised waveform | |
| cond: float["b nw"], # masked conditioning waveform | |
| mel: float["b n m"], # noised waveform | |
| mel_cond: float["b n m"], # masked conditioning waveform | |
| text: int["b nt"], # text | |
| time: float["b"] | float[""], # time step | |
| mask: bool["b n"] | bool["b nw"] | None = None, | |
| drop_audio_cond: bool = False, # cfg for conditioning waveform | |
| drop_text: bool = False, # cfg for text | |
| cfg_infer: bool = False, # cfg inference, pack cond & uncond forward | |
| cache: bool = False, | |
| lens: int["b"] | None = None, | |
| inference_bottleneck_dropout: float | None = None, | |
| inference_bottleneck_generator: torch.Generator | None = None, | |
| cfg_null_audio: bool = True, # False preserves audio for a text-only null | |
| prompt_mask: bool["b nw"] | None = None, | |
| predict_mel: bool = True, | |
| ): | |
| if x.ndim != 2 or cond.ndim != 2: | |
| raise ValueError(f"WaveGen DiT expects raw waveform x/cond [B, N], got {x.ndim}D and {cond.ndim}D.") | |
| if prompt_mask is not None and ( | |
| not torch.is_tensor(prompt_mask) | |
| or prompt_mask.dtype != torch.bool or prompt_mask.shape != x.shape | |
| or prompt_mask.device != x.device | |
| ): | |
| raise ValueError("prompt_mask must match waveform shape/device and have boolean dtype") | |
| target_num_samples = x.shape[1] | |
| x, token_mask, token_lens = self._wav_to_tokens(x, mask=mask, lens=lens) | |
| if prompt_mask is not None: | |
| prompt_mask, _, _ = self._wav_to_tokens(prompt_mask) | |
| cond, _, _ = self._wav_to_tokens(cond, mask=None, lens=None) | |
| mask = token_mask | |
| lens = token_lens | |
| batch, seq_len = x.shape[0], x.shape[1] | |
| mel = self._align_mel_to_tokens(mel, seq_len, "mel") | |
| mel_cond = self._align_mel_to_tokens(mel_cond, seq_len, "mel_cond") | |
| if time.ndim == 0: | |
| time = time.repeat(batch) | |
| # t: conditioning time, text: text, x: noised waveform + conditioning waveform + text | |
| t = self.time_embed(time) | |
| def null_inputs(): | |
| if prompt_mask is None: | |
| return x, mel | |
| # Joint-null drops both clean conditioning and the noised prompt; | |
| # text-only null keeps the reference in both input streams. | |
| return ( | |
| x.masked_fill(prompt_mask, 0.0), | |
| mel.masked_fill(prompt_mask.any(dim=-1).unsqueeze(-1), 0.0), | |
| ) | |
| if cfg_infer: # pack cond & uncond forward: b n d -> 2b n d | |
| x_cond = self.get_input_embed( | |
| x, cond, mel, mel_cond, text, drop_audio_cond=False, drop_text=False, | |
| cache=cache, audio_mask=mask, | |
| inference_bottleneck_dropout=inference_bottleneck_dropout, | |
| inference_bottleneck_generator=inference_bottleneck_generator, | |
| ) | |
| null_x, null_mel = null_inputs() if cfg_null_audio else (x, mel) | |
| x_uncond = self.get_input_embed( | |
| null_x, cond, null_mel, mel_cond, text, drop_audio_cond=cfg_null_audio, drop_text=True, | |
| cache=cache, audio_mask=mask, | |
| inference_bottleneck_dropout=inference_bottleneck_dropout, | |
| inference_bottleneck_generator=inference_bottleneck_generator, | |
| ) | |
| x = torch.cat((x_cond, x_uncond), dim=0) | |
| t = torch.cat((t, t), dim=0) | |
| mask = torch.cat((mask, mask), dim=0) if mask is not None else None | |
| else: | |
| if drop_audio_cond and drop_text: | |
| x, mel = null_inputs() | |
| x = self.get_input_embed( | |
| x, cond, mel, mel_cond, text, drop_audio_cond=drop_audio_cond, | |
| drop_text=drop_text, cache=cache, audio_mask=mask, | |
| inference_bottleneck_dropout=inference_bottleneck_dropout, | |
| inference_bottleneck_generator=inference_bottleneck_generator, | |
| ) | |
| rope = self.rotary_embed.forward_from_seq_len(seq_len) | |
| if self.long_skip_connection is not None: | |
| residual = x | |
| for i, block in enumerate(self.transformer_blocks): | |
| if self.checkpoint_activations: | |
| # https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint | |
| x = torch.utils.checkpoint.checkpoint(block, x, t, mask, rope, use_reentrant=False) | |
| else: | |
| x = block(x, t, mask=mask, rope=rope) | |
| if self.long_skip_connection is not None: | |
| x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) | |
| mel_x = x | |
| for i, wav_block in enumerate(self.wav_ddt): | |
| if self.checkpoint_activations: | |
| # https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint | |
| x = torch.utils.checkpoint.checkpoint(wav_block, x, t, mask, rope, use_reentrant=False) | |
| else: | |
| x = wav_block(x, t, mask=mask, rope=rope) | |
| if predict_mel: | |
| for i, mel_block in enumerate(self.mel_ddt): | |
| if self.checkpoint_activations: | |
| # https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint | |
| mel_x = torch.utils.checkpoint.checkpoint(mel_block, mel_x, t, mask, rope, use_reentrant=False) | |
| else: | |
| mel_x = mel_block(mel_x, t, mask=mask, rope=rope) | |
| x = self.norm_out(x, t) | |
| output = self.proj_out(x) | |
| output = self._tokens_to_wav(output, target_num_samples=target_num_samples) | |
| mel_output = None | |
| if predict_mel: | |
| mel_x = self.mel_norm_out(mel_x, t) | |
| mel_output = self.mel_proj_out(mel_x) | |
| return (output, mel_output), None, None | |
| DiT = DiTBackbone | |
| __all__ = ["DiT", "DiTBackbone"] | |