[Model] Add Ming-omni-tts dense 0.5B pipeline (#2906)
Signed-off-by: akshatvishu <akshatnayak197@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
akshatvishuandHongsheng Liu authored and GitHub committed
1 parent
39ae6e12f8
commit
2fa4c7a4a9
55 files changed
+6263
-549
No files matched your search
@@ -0,0 +1,193 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from x_transformers.x_transformers import apply_rotary_pos_emb
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim, eps=1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x):
|
||||
if self.weight.dtype in [torch.float16, torch.bfloat16]:
|
||||
x = x.to(self.weight.dtype)
|
||||
x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
|
||||
return x
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate="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)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
heads=8,
|
||||
dim_head=64,
|
||||
dropout=0.0,
|
||||
qk_norm=None,
|
||||
pe_attn_head=None,
|
||||
attn_mask_enabled=True,
|
||||
):
|
||||
super().__init__()
|
||||
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)
|
||||
if qk_norm is None:
|
||||
self.q_norm = None
|
||||
self.k_norm = None
|
||||
elif qk_norm == "rms_norm":
|
||||
self.q_norm = RMSNorm(dim_head)
|
||||
self.k_norm = RMSNorm(dim_head)
|
||||
else:
|
||||
raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
|
||||
|
||||
self.to_out = nn.ModuleList([])
|
||||
self.to_out.append(nn.Linear(self.inner_dim, dim))
|
||||
self.to_out.append(nn.Dropout(dropout))
|
||||
self.pe_attn_head = pe_attn_head
|
||||
self.attn_mask_enabled = attn_mask_enabled
|
||||
|
||||
def forward(self, x, mask=None, rope=None):
|
||||
batch_size = x.shape[0]
|
||||
query = self.to_q(x)
|
||||
key = self.to_k(x)
|
||||
value = self.to_v(x)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // self.heads
|
||||
query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
||||
|
||||
if self.q_norm is not None:
|
||||
query = self.q_norm(query)
|
||||
if self.k_norm is not None:
|
||||
key = self.k_norm(key)
|
||||
|
||||
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)
|
||||
if self.pe_attn_head is not None:
|
||||
on = self.pe_attn_head
|
||||
query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
|
||||
key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
|
||||
else:
|
||||
query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
|
||||
key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
|
||||
|
||||
if self.attn_mask_enabled and mask is not None:
|
||||
valid_sample_indices = mask.any(dim=1)
|
||||
final_output = torch.zeros_like(query).to(query.device)
|
||||
attn_mask = mask[valid_sample_indices]
|
||||
query = query[valid_sample_indices]
|
||||
key = key[valid_sample_indices]
|
||||
value = value[valid_sample_indices]
|
||||
attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
|
||||
attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2])
|
||||
else:
|
||||
attn_mask = None
|
||||
|
||||
x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
|
||||
if self.attn_mask_enabled and mask is not None:
|
||||
final_output[valid_sample_indices] = x
|
||||
x = final_output
|
||||
|
||||
x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
|
||||
x = x.to(query.dtype)
|
||||
x = self.to_out[0](x)
|
||||
x = self.to_out[1](x)
|
||||
|
||||
if mask is not None:
|
||||
mask = mask.unsqueeze(-1)
|
||||
x = x.masked_fill(~mask, 0.0)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class DiTBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
num_heads,
|
||||
mlp_ratio=4.0,
|
||||
dropout=0.1,
|
||||
qk_norm=None,
|
||||
pe_attn_head=None,
|
||||
attn_mask_enabled=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.norm1 = RMSNorm(hidden_size)
|
||||
self.attn = Attention(
|
||||
dim=hidden_size,
|
||||
heads=num_heads,
|
||||
dim_head=hidden_size // num_heads,
|
||||
dropout=dropout,
|
||||
qk_norm=qk_norm,
|
||||
pe_attn_head=pe_attn_head,
|
||||
attn_mask_enabled=attn_mask_enabled,
|
||||
)
|
||||
self.norm2 = RMSNorm(hidden_size)
|
||||
self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh")
|
||||
|
||||
def forward(self, x, mask, rope):
|
||||
x = x + self.attn(self.norm1(x), mask=mask, rope=rope)
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
return x
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
def __init__(self, hidden_size, out_channels):
|
||||
super().__init__()
|
||||
self.norm_final = RMSNorm(hidden_size)
|
||||
self.linear = nn.Linear(hidden_size, out_channels, bias=True)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.norm_final(x)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
class CondEmbedder(nn.Module):
|
||||
def __init__(self, input_feature_size, hidden_size):
|
||||
super().__init__()
|
||||
self.cond_embedder = nn.Linear(input_feature_size, hidden_size)
|
||||
|
||||
def forward(self, llm_cond):
|
||||
return self.cond_embedder(llm_cond)
|
||||
|
||||
|
||||
def get_epss_timesteps(n, device, dtype):
|
||||
dt = 1 / 32
|
||||
predefined_timesteps = {
|
||||
5: [0, 2, 4, 8, 16, 32],
|
||||
6: [0, 2, 4, 6, 8, 16, 32],
|
||||
7: [0, 2, 4, 6, 8, 16, 24, 32],
|
||||
10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
|
||||
12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
|
||||
16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
|
||||
}
|
||||
t = predefined_timesteps.get(n, [])
|
||||
if not t:
|
||||
return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
|
||||
return dt * torch.tensor(t, device=device, dtype=dtype)
|
||||
Reference in new issue
Block a user