mirror of
https://github.com/RVC-Boss/GPT-SoVITS.git
synced 2026-08-13 01:53:43 +08:00
365 lines
12 KiB
Python
365 lines
12 KiB
Python
"""
|
|
ein notation:
|
|
b - batch
|
|
n - sequence
|
|
nt - text sequence
|
|
nw - raw wave length
|
|
d - dimension
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
from x_transformers.x_transformers import apply_rotary_pos_emb
|
|
|
|
|
|
# sinusoidal position embedding
|
|
|
|
|
|
class SinusPositionEmbedding(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.dim = dim
|
|
|
|
def forward(self, x, scale=1000):
|
|
device = x.device
|
|
half_dim = self.dim // 2
|
|
emb = math.log(10000) / (half_dim - 1)
|
|
emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
|
|
emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
|
|
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
|
|
return emb
|
|
|
|
|
|
# convolutional position embedding
|
|
|
|
|
|
class ConvPositionEmbedding(nn.Module):
|
|
def __init__(self, dim, kernel_size=31, groups=16):
|
|
super().__init__()
|
|
assert kernel_size % 2 != 0
|
|
self.conv1d = nn.Sequential(
|
|
nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
|
|
nn.Mish(),
|
|
nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
|
|
nn.Mish(),
|
|
)
|
|
|
|
def forward(self, x: float["b n d"], mask: bool["b n"] | None = None): # noqa: F722
|
|
if mask is not None:
|
|
mask = mask[..., None]
|
|
x = x.masked_fill(~mask, 0.0)
|
|
|
|
x = x.permute(0, 2, 1)
|
|
x = self.conv1d(x)
|
|
out = x.permute(0, 2, 1)
|
|
|
|
if mask is not None:
|
|
out = out.masked_fill(~mask, 0.0)
|
|
|
|
return out
|
|
|
|
|
|
# rotary positional embedding related
|
|
|
|
|
|
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):
|
|
# proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
|
|
# has some connection to NTK literature
|
|
# https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
|
|
# https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
|
|
theta *= theta_rescale_factor ** (dim / (dim - 2))
|
|
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
|
t = torch.arange(end, device=freqs.device) # type: ignore
|
|
freqs = torch.outer(t, freqs).float() # type: ignore
|
|
freqs_cos = torch.cos(freqs) # real part
|
|
freqs_sin = torch.sin(freqs) # imaginary part
|
|
return torch.cat([freqs_cos, freqs_sin], dim=-1)
|
|
|
|
|
|
def get_pos_embed_indices(start, length, max_pos, scale=1.0):
|
|
# length = length if isinstance(length, int) else length.max()
|
|
scale = scale * torch.ones_like(start, dtype=torch.float32) # in case scale is a scalar
|
|
pos = (
|
|
start.unsqueeze(1)
|
|
+ (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
|
|
)
|
|
# avoid extra long error.
|
|
pos = torch.where(pos < max_pos, pos, max_pos - 1)
|
|
return pos
|
|
|
|
|
|
# Global Response Normalization layer (Instance Normalization ?)
|
|
|
|
|
|
class GRN(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
|
|
self.beta = nn.Parameter(torch.zeros(1, 1, dim))
|
|
|
|
def forward(self, x):
|
|
Gx = torch.norm(x, p=2, dim=1, keepdim=True)
|
|
Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
|
|
return self.gamma * (x * Nx) + self.beta + x
|
|
|
|
|
|
# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py
|
|
# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108
|
|
|
|
|
|
class ConvNeXtV2Block(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
intermediate_dim: int,
|
|
dilation: int = 1,
|
|
):
|
|
super().__init__()
|
|
padding = (dilation * (7 - 1)) // 2
|
|
self.dwconv = nn.Conv1d(
|
|
dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
|
|
) # depthwise conv
|
|
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
|
self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers
|
|
self.act = nn.GELU()
|
|
self.grn = GRN(intermediate_dim)
|
|
self.pwconv2 = nn.Linear(intermediate_dim, dim)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
residual = x
|
|
x = x.transpose(1, 2) # b n d -> b d n
|
|
x = self.dwconv(x)
|
|
x = x.transpose(1, 2) # b d n -> b n d
|
|
x = self.norm(x)
|
|
x = self.pwconv1(x)
|
|
x = self.act(x)
|
|
x = self.grn(x)
|
|
x = self.pwconv2(x)
|
|
return residual + x
|
|
|
|
|
|
# AdaLayerNormZero
|
|
# return with modulated x for attn input, and params for later mlp modulation
|
|
|
|
|
|
class AdaLayerNormZero(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
|
|
self.silu = nn.SiLU()
|
|
self.linear = nn.Linear(dim, dim * 6)
|
|
|
|
self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
|
|
|
def forward(self, x, emb=None):
|
|
emb = self.linear(self.silu(emb))
|
|
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
|
|
|
|
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
|
|
return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
|
|
|
|
|
|
# AdaLayerNormZero for final layer
|
|
# return only with modulated x for attn input, cuz no more mlp modulation
|
|
|
|
|
|
class AdaLayerNormZero_Final(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
|
|
self.silu = nn.SiLU()
|
|
self.linear = nn.Linear(dim, dim * 2)
|
|
|
|
self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
|
|
|
def forward(self, x, emb):
|
|
emb = self.linear(self.silu(emb))
|
|
scale, shift = torch.chunk(emb, 2, dim=1)
|
|
|
|
x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
|
|
return x
|
|
|
|
|
|
# FeedForward
|
|
|
|
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):
|
|
super().__init__()
|
|
inner_dim = int(dim * mult)
|
|
dim_out = dim_out if dim_out is not None else dim
|
|
|
|
activation = nn.GELU(approximate=approximate)
|
|
project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
|
|
self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
|
|
|
|
def forward(self, x):
|
|
return self.ff(x)
|
|
|
|
|
|
# Attention
|
|
# modified from diffusers/src/diffusers/models/attention_processor.py
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
processor: AttnProcessor,
|
|
dim: int,
|
|
heads: int = 8,
|
|
dim_head: int = 64,
|
|
dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
|
|
if not hasattr(F, "scaled_dot_product_attention"):
|
|
raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
|
|
|
self.processor = processor
|
|
|
|
self.dim = dim
|
|
self.heads = heads
|
|
self.inner_dim = dim_head * heads
|
|
self.dropout = dropout
|
|
|
|
self.to_q = nn.Linear(dim, self.inner_dim)
|
|
self.to_k = nn.Linear(dim, self.inner_dim)
|
|
self.to_v = nn.Linear(dim, self.inner_dim)
|
|
|
|
self.to_out = nn.ModuleList([])
|
|
self.to_out.append(nn.Linear(self.inner_dim, dim))
|
|
self.to_out.append(nn.Dropout(dropout))
|
|
|
|
def forward(
|
|
self,
|
|
x: float["b n d"], # noised input x # noqa: F722
|
|
mask: bool["b n"] | None = None, # noqa: F722
|
|
rope=None, # rotary position embedding for x
|
|
) -> torch.Tensor:
|
|
return self.processor(self, x, mask=mask, rope=rope)
|
|
|
|
|
|
# Attention processor
|
|
|
|
|
|
# from torch.nn.attention import SDPBackend
|
|
# torch.backends.cuda.enable_flash_sdp(True)
|
|
class AttnProcessor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(
|
|
self,
|
|
attn: Attention,
|
|
x: float["b n d"], # noised input x # noqa: F722
|
|
mask: bool["b n"] | None = None, # noqa: F722
|
|
rope=None, # rotary position embedding
|
|
) -> torch.FloatTensor:
|
|
batch_size = x.shape[0]
|
|
|
|
# `sample` projections.
|
|
query = attn.to_q(x)
|
|
key = attn.to_k(x)
|
|
value = attn.to_v(x)
|
|
|
|
# apply rotary position embedding
|
|
if rope is not None:
|
|
freqs, xpos_scale = rope
|
|
q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
|
|
|
|
query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
|
|
key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
|
|
|
|
# attention
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // attn.heads
|
|
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
|
|
|
# mask. e.g. inference got a batch with different target durations, mask out the padding
|
|
if mask is not None:
|
|
attn_mask = mask
|
|
attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'
|
|
# print(3433333333,attn_mask.shape)
|
|
attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
|
|
else:
|
|
attn_mask = None
|
|
# with torch.nn.attention.sdpa_kernel(backends=[SDPBackend.EFFICIENT_ATTENTION]):
|
|
# with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=True):
|
|
# with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
|
|
# print(torch.backends.cuda.flash_sdp_enabled())
|
|
# print(torch.backends.cuda.mem_efficient_sdp_enabled())
|
|
# print(torch.backends.cuda.math_sdp_enabled())
|
|
x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
|
|
x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
|
x = x.to(query.dtype)
|
|
|
|
# linear proj
|
|
x = attn.to_out[0](x)
|
|
# dropout
|
|
x = attn.to_out[1](x)
|
|
|
|
if mask is not None:
|
|
mask = mask.unsqueeze(-1)
|
|
x = x.masked_fill(~mask, 0.0)
|
|
|
|
return x
|
|
|
|
|
|
# DiT Block
|
|
|
|
|
|
class DiTBlock(nn.Module):
|
|
def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1):
|
|
super().__init__()
|
|
|
|
self.attn_norm = AdaLayerNormZero(dim)
|
|
self.attn = Attention(
|
|
processor=AttnProcessor(),
|
|
dim=dim,
|
|
heads=heads,
|
|
dim_head=dim_head,
|
|
dropout=dropout,
|
|
)
|
|
|
|
self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
|
self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
|
|
|
|
def forward(self, x, t, mask=None, rope=None): # x: noised input, t: time embedding
|
|
# pre-norm & modulation for attention input
|
|
norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
|
|
|
|
# attention
|
|
attn_output = self.attn(x=norm, mask=mask, rope=rope)
|
|
|
|
# process attention output for input x
|
|
x = x + gate_msa.unsqueeze(1) * attn_output
|
|
|
|
norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
|
ff_output = self.ff(norm)
|
|
x = x + gate_mlp.unsqueeze(1) * ff_output
|
|
|
|
return x
|
|
|
|
|
|
# time step conditioning embedding
|
|
|
|
|
|
class TimestepEmbedding(nn.Module):
|
|
def __init__(self, dim, freq_embed_dim=256):
|
|
super().__init__()
|
|
self.time_embed = SinusPositionEmbedding(freq_embed_dim)
|
|
self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
|
|
|
|
def forward(self, timestep: float["b"]): # noqa: F821
|
|
time_hidden = self.time_embed(timestep)
|
|
time_hidden = time_hidden.to(timestep.dtype)
|
|
time = self.time_mlp(time_hidden) # b d
|
|
return time
|